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 new file mode 100644 index 000000000..8192a7b06 --- /dev/null +++ b/tests/core/sched/test_omni_ar_scheduler_unified_decode_deferral.py @@ -0,0 +1,114 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest +from vllm.v1.request import RequestStatus + +import vllm_omni.core.sched.omni_ar_scheduler as scheduler_mod +from vllm_omni.model_executor.models.voxcpm2.scheduler import VoxCPM2OmniARAsyncScheduler + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _MockQueue: + def __init__(self, items: list | None = None) -> None: + self._items = list(items or []) + + def __bool__(self) -> bool: + return bool(self._items) + + def __iter__(self): + return iter(self._items) + + def __len__(self) -> int: + return len(self._items) + + def add_request(self, request) -> None: + self._items.append(request) + + def prepend_requests(self, requests) -> None: + self._items = list(requests) + self._items + + +class _MockRequest: + def __init__( + self, + request_id: str, + *, + status: RequestStatus = RequestStatus.RUNNING, + num_prompt_tokens: int = 4, + num_computed_tokens: int = 4, + num_output_placeholders: int = 0, + ) -> None: + self.request_id = request_id + self.status = status + self.num_prompt_tokens = num_prompt_tokens + self.num_computed_tokens = num_computed_tokens + self.num_output_placeholders = num_output_placeholders + + def is_finished(self) -> bool: + return RequestStatus.is_finished(self.status) + + +def _make_scheduler(*, enable_unified_decode_graph: bool = True) -> VoxCPM2OmniARAsyncScheduler: + sched = VoxCPM2OmniARAsyncScheduler.__new__(VoxCPM2OmniARAsyncScheduler) + runtime_config = SimpleNamespace(enable_unified_decode_graph=enable_unified_decode_graph) + sched.vllm_config = SimpleNamespace( + model_config=SimpleNamespace( + hf_config=SimpleNamespace(voxcpm2_runtime_config=runtime_config), + ) + ) + return sched + + +def test_voxcpm2_unified_decode_graph_defers_waiting_when_decode_ready() -> None: + scheduler = _make_scheduler() + scheduler.running = [_MockRequest("decode")] + scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + + assert scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_voxcpm2_unified_decode_graph_does_not_defer_without_decode_ready() -> None: + scheduler = _make_scheduler() + scheduler.running = [_MockRequest("prefill-running", num_prompt_tokens=8, num_computed_tokens=4)] + scheduler.waiting = _MockQueue([_MockRequest("waiting", status=RequestStatus.WAITING)]) + + assert not scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_voxcpm2_unified_decode_graph_does_not_defer_when_disabled() -> None: + scheduler = _make_scheduler(enable_unified_decode_graph=False) + scheduler.running = [_MockRequest("decode")] + scheduler.waiting = _MockQueue([_MockRequest("prefill", status=RequestStatus.WAITING)]) + + assert not scheduler._should_defer_waiting_for_unified_decode_graph() + + +def test_unified_decode_graph_deferral_restores_waiting_queue(monkeypatch) -> None: + scheduler = _make_scheduler() + scheduler.running = [_MockRequest("decode")] + original_waiting_req = _MockRequest("waiting", status=RequestStatus.WAITING) + deferred_by_upstream = _MockRequest("deferred-by-upstream", status=RequestStatus.WAITING) + original_waiting = _MockQueue([original_waiting_req]) + scheduler.waiting = original_waiting + scheduler.policy = "fcfs" + scheduler.chunk_transfer_adapter = None + scheduler.input_coordinator = None + scheduler._consume_pending_connector_output = lambda model_mode: None + + monkeypatch.setattr(scheduler_mod, "create_request_queue", lambda _policy: _MockQueue()) + + def fake_upstream_schedule(self, throttle_prefills: bool = False): + assert self.waiting is not original_waiting + assert not self.waiting + self.waiting.add_request(deferred_by_upstream) + raise RuntimeError("stop before output wrapping") + + monkeypatch.setattr(scheduler_mod.VLLMScheduler, "schedule", fake_upstream_schedule) + + with pytest.raises(RuntimeError, match="stop before output wrapping"): + scheduler.schedule() + + assert scheduler.waiting._items == [deferred_by_upstream, original_waiting_req] diff --git a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py index c26e21039..315fba6ef 100644 --- a/tests/e2e/online_serving/test_qwen_image_layered_expansion.py +++ b/tests/e2e/online_serving/test_qwen_image_layered_expansion.py @@ -7,7 +7,7 @@ and are supported by Qwen-Image-Layered model. Kept cases (maximum feature coverage, no cpu-offload): sp_001 : cache_dit + Ulysses-SP 2 (2×H100) cfg_parallel_002 : cache_dit + CFG-Parallel 2 (2×H100) - layers_guard_001 : layerwise offload + layers=3 (1×H100, issue #1969 guard) + layers_guard_001 : layerwise offload + layers=2 (1×H100, issue #1969 guard) Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload. """ @@ -116,8 +116,8 @@ LAYERS_GUARD_CASES = [ model=MODEL, server_args=["--enable-layerwise-offload"], ), - 3, - id="layers_guard_001_layers3", + 2, + id="layers_guard_001_layers2", marks=SINGLE_CARD_FEATURE_MARKS, ), ] diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 31a82102a..5a7ef7bf0 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -1244,7 +1244,7 @@ def test_parameter_validation(): # Invalid layers for layered models (must stay within the backend-supported range) with pytest.raises(ValueError): - ImageGenerationRequest(prompt="test", layers=2) + ImageGenerationRequest(prompt="test", layers=1) with pytest.raises(ValueError): ImageGenerationRequest(prompt="test", layers=11) @@ -1663,13 +1663,13 @@ def test_image_edit_invalid_layers(async_omni_test_client): files=[("image", img_bytes)], data={ "prompt": "test", - "layers": 2, + "layers": 1, }, ) assert response.status_code == 400 detail = response.json()["detail"] assert "Invalid layers" in detail - assert "layers must be between 3 and 10 inclusive" in detail + assert "layers must be between 2 and 10 inclusive" in detail # Test layers above the supported range response = async_omni_test_client.post( @@ -1683,7 +1683,7 @@ def test_image_edit_invalid_layers(async_omni_test_client): assert response.status_code == 400 detail = response.json()["detail"] assert "Invalid layers" in detail - assert "layers must be between 3 and 10 inclusive" in detail + assert "layers must be between 2 and 10 inclusive" in detail def test_image_edit_resolution_and_size_conflict(async_omni_test_client): diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 5923a2500..8f2797bc7 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -39,6 +39,7 @@ from vllm_omni.entrypoints.openai.serving_speech import ( OmniOpenAIServingSpeech, _create_wav_header, ) +from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest from vllm_omni.model_executor.models.fish_speech.prompt_utils import ( FISH_TEXT_ONLY_SYSTEM_PROMPT, build_fish_voice_clone_prompt_ids, @@ -3364,6 +3365,40 @@ class TestTTSAsyncOffloading: assert stage0_params.extra_args["tts_local_seed"] == 42 assert qwen3_tts_server.engine_client.default_sampling_params_list[0].extra_args is None + def test_prepare_speech_generation_uses_adapter_model_type_label( + self, + voxtral_server, + mocker: MockerFixture, + ): + """Adapter model_type should replace the legacy _tts_model_type label ladder.""" + legacy_tts_model_type = "dummy_tts" + adapter_model_type = "adapter_dummy_tts" + + class FakeAdapter: + def validate(self, request): + return None + + async def build(self, request, sampling_params_list, has_inline_ref_audio): + return PreparedRequest( + prompt={"prompt": request.input}, + tts_params={}, + model_type=adapter_model_type, + ) + + voxtral_server._tts_model_type = legacy_tts_model_type + mocker.patch.object(voxtral_server, "_get_tts_adapter", return_value=FakeAdapter()) + log_info = mocker.patch("vllm_omni.entrypoints.openai.serving_speech.logger.info") + + asyncio.run(voxtral_server._prepare_speech_generation(OpenAICreateSpeechRequest(input="hello"))) + + assert adapter_model_type != legacy_tts_model_type + assert any( + call.args + and call.args[0] == "TTS speech request %s: text=%r, model=%s" + and call.args[3] == adapter_model_type + for call in log_info.call_args_list + ) + def test_prepare_speech_generation_treats_sse_as_streaming(self, qwen3_tts_server, mocker: MockerFixture): """stream_format=sse should request delta-style multimodal outputs.""" qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None) diff --git a/tests/helpers/stage_config.py b/tests/helpers/stage_config.py index 4351eb27d..8706bb1bc 100644 --- a/tests/helpers/stage_config.py +++ b/tests/helpers/stage_config.py @@ -633,6 +633,18 @@ def get_deploy_config_path(rel_path: str) -> str: return str(_DEPLOY_DIR / rel_path) +def get_deploy_config_stage(rel_path: str, stage_id: int) -> dict[str, Any]: + """Return one stage entry from a deploy yaml by ``stage_id``.""" + with open(get_deploy_config_path(rel_path), encoding="utf-8") as f: + cfg = yaml.safe_load(f) or {} + + stage_key = "stages" if "stages" in cfg else "stage_args" + for stage in cfg.get(stage_key, []): + if stage.get("stage_id") == stage_id: + return stage + raise KeyError(f"No stage_id={stage_id} in deploy config {rel_path!r}") + + def _get_config_value_by_path(config_dict: dict, path: str) -> Any: """Read a dot-separated path from a nested dict (e.g. ``engine_args.load_format``).""" current: Any = config_dict @@ -707,6 +719,7 @@ def stage_config_path_for_run_level(stage_config_path: str | None, run_level: st __all__ = [ "get_deploy_config_path", + "get_deploy_config_stage", "modify_stage_config", "stage_config_path_for_run_level", ] diff --git a/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py new file mode 100644 index 000000000..33ac80a6f --- /dev/null +++ b/tests/model_executor/models/voxcpm2/test_runner_assisted_unified_graph_static.py @@ -0,0 +1,204 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Static regression checks for VoxCPM2 runner-assisted unified graph hooks. + +These tests intentionally avoid importing torch/vLLM so they can run in a +lightweight local checkout. Runtime audio quality still requires a CUDA test. +""" + +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[4] +RUNNER = REPO_ROOT / "vllm_omni" / "worker" / "gpu_ar_model_runner.py" +RUNNER_ASSISTED_METADATA = REPO_ROOT / "vllm_omni" / "worker" / "runner_assisted_metadata.py" +TALKER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "voxcpm2_talker.py" +VOXCPM2_PIPELINE = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "pipeline.py" +VOXCPM2_SCHEDULER = REPO_ROOT / "vllm_omni" / "model_executor" / "models" / "voxcpm2" / "scheduler.py" +OMNI_AR_SCHEDULER = REPO_ROOT / "vllm_omni" / "core" / "sched" / "omni_ar_scheduler.py" +DEPLOY = REPO_ROOT / "vllm_omni" / "deploy" / "voxcpm2.yaml" + + +def test_ar_runner_exposes_runner_assisted_full_metadata_hook(): + source = RUNNER.read_text() + + assert "Models without this hook keep the normal runner path" in source + assert "get_runner_assisted_full_attention_metadata_request" in source + assert "set_runner_assisted_full_attention_metadata_context" in source + assert "runner_assisted_full_attn" in source + assert "runner_assisted_full_attn_capture" in source + assert "pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn" in source + assert "for_cudagraph_capture=runner_assisted_full_attn_capture" in source + assert "cudagraph_runtime_mode=(" in source + assert "CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode" in source + refresh_source = source[source.index("def _refresh_runner_assisted_full_attention_metadata_buffers") :] + refresh_source = refresh_source[: refresh_source.index("def _set_runner_assisted_full_attention_metadata_context")] + assert "num_computed_tokens_cpu" in refresh_source + assert "num_scheduled_tokens_np" in refresh_source + assert "optimistic_seq_lens_cpu[:num_reqs]=1" not in "".join(refresh_source.split()) + padding_source = source[ + source.index("runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request") : + ] + padding_source = padding_source[: padding_source.index("ubatch_slices, ubatch_slices_padded")] + assert "runner_assisted_full_attn_request.num_reqs_padded" in padding_source + assert "runner_assisted_full_attn_request.for_cudagraph_capture" in padding_source + assert "num_reqs_padded * max_num_scheduled_tokens" in padding_source + + +def test_ar_runner_without_model_hook_stays_on_normal_path(): + source = RUNNER.read_text() + contract_source = RUNNER_ASSISTED_METADATA.read_text() + request_source = source[source.index("def _get_runner_assisted_full_attention_metadata_request") :] + request_source = request_source[ + : request_source.index("def _refresh_runner_assisted_full_attention_metadata_buffers") + ] + compact_request_source = "".join(request_source.split()) + + assert "Models without this hook keep the normal runner path" in request_source + assert "class RunnerAssistedAttentionMetadataProvider(Protocol)" in contract_source + assert "class RunnerAssistedFullAttentionMetadataRequest(NamedTuple)" in contract_source + assert "tuple[int, bool]" not in contract_source + assert 'hook=getattr(self.model,"get_runner_assisted_full_attention_metadata_request",None)' in ( + compact_request_source + ) + assert "ifnotcallable(hook):returnNone" in compact_request_source + assert ".tolist()" not in compact_request_source + assert "isinstance(request,RunnerAssistedFullAttentionMetadataRequest)" in compact_request_source + assert "request.num_reqs_padded" in request_source + assert "request.for_cudagraph_capture" in request_source + assert "exceptException" not in compact_request_source + + set_context_source = source[source.index("def _set_runner_assisted_full_attention_metadata_context") :] + set_context_source = set_context_source[: set_context_source.index("def _deferred_prefix_cache_mm_keys")] + assert "except Exception" not in set_context_source + + +def test_voxcpm2_graph_paths_fail_closed_and_preserve_deterministic_noise(): + source = TALKER.read_text() + compact_source = "".join(source.split()) + + assert "_voxcpm2_compile_without_inductor_cudagraphs" in source + assert "def _compile_without_inductor_cudagraphs" not in source + assert "_voxcpm2_unwrap_torch_compile" in source + assert "_orig_mod" in source + compile_source = source[source.index("def _voxcpm2_compile_without_inductor_cudagraphs") :] + compile_source = compile_source[: compile_source.index("def _setup_torch_compile")] + assert "torch._inductor.config" not in compile_source + assert '"triton.cudagraphs": False' in compile_source + assert '"triton.cudagraph_trees": False' in compile_source + + assert "self._enable_unified_decode_graph=(use_cuda_graph" in compact_source + assert "andnotself._deterministic_cfm_noise" in compact_source + assert "decode_tail_graph" not in source + + unified_source = source[source.index("def _forward_unified_decode") :] + unified_source = unified_source[: unified_source.index("# -------------------- vllm hooks")] + assert "except Exception:" not in unified_source + assert "_failed_unified_graph_sizes" not in source + assert "self._enable_unified_decode_graph = False" not in unified_source + assert "_forward_unified_decode_fallback" not in source + assert "capture_failed" not in compact_source + assert "record_fallback" not in source + assert "fallbacks" not in source + assert "**info_dict: Unpack[VoxCPM2PreprocessInput]" in source + assert "**info: Unpack[VoxCPM2PostprocessInput]" in source + assert "reference_audio: Any" not in source + assert "prompt_audio: Any" not in source + assert "def _nullify_volatile_metadata(ctx: _ForwardContextLike) -> _ForwardContextLike" in source + assert "class _PrefillInputs(NamedTuple)" in source + assert "-> _PrefillInputs" in source + assert 'inputs["audio_feat"]' not in source + assert 'inputs["text_token"]' not in source + + capture_source = source[source.index("def _capture_unified_decode_graph") :] + capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")] + assert "_voxcpm2_compile_unified_capture_estimator" in capture_source + assert "_voxcpm2_compile_unified_capture_feat_encoder" in capture_source + graph_body = capture_source[capture_source.index("with torch.cuda.graph") :] + assert "g.cfm_noise.normal_()" not in graph_body + assert "g.cfm_noise.normal_()" in unified_source + + assert "_pre_capture_unified_graphs" not in source + assert "_unified_graph_pre_capture_sizes" not in source + assert "unified_decode_graph_pre_capture_sizes" not in source + + +def test_voxcpm2_batch_unified_graph_requires_runner_metadata_marker(): + source = TALKER.read_text() + compact_source = "".join(source.split()) + + assert "enable_runner_assisted_unified_decode_graph" not in source + assert "allow_unified_decode_graph_batch_attention" not in source + assert "get_runner_assisted_full_attention_metadata_request" in source + assert "set_runner_assisted_full_attention_metadata_context" in source + assert "RunnerAssistedFullAttentionMetadataRequest" in source + assert "_runner_assisted_unified_decode_graph_active" in source + assert "runner_full_metadata_missing" in source + assert "_build_unified_graph_bucket_sizes" in source + assert "_select_unified_graph_bucket_size" in source + + capture_source = source[source.index("def _capture_unified_decode_graph") :] + capture_source = capture_source[: capture_source.index("def _unified_decode_graph_skip_reason")] + assert "override_forward_context" in capture_source + assert "_nullify_volatile_metadata" in capture_source + assert "capture_context = override_forward_context(self._nullify_volatile_metadata(ctx))" in capture_source + assert "ifsize>1:continue" not in compact_source + forward_source = source[source.index("def _forward_unified_decode") :] + forward_source = forward_source[: forward_source.index("# -------------------- vllm hooks")] + assert "graph_size = self._select_unified_graph_bucket_size(num_reqs)" in forward_source + assert "self._unified_graphs[graph_size]" in forward_source + assert "g.input_embeds[num_reqs:graph_size].zero_()" not in forward_source + assert "padded rows as valid duplicates" in forward_source + assert "g.input_embeds[last : last + 1].expand" in forward_source + assert "ifnum_reqs>1andnotself._runner_assisted_unified_decode_graph_active" in compact_source + assert "andnotcfg.enable_runner_assisted_unified_decode_graph" not in compact_source + needs_source = source[source.index("def get_runner_assisted_full_attention_metadata_request") :] + needs_source = needs_source[: needs_source.index("def set_runner_assisted_full_attention_metadata_context")] + assert "_select_unified_graph_bucket_size(num_reqs)" in needs_source + assert "_should_use_decode_graph(num_reqs)" not in needs_source + assert "RunnerAssistedFullAttentionMetadataRequest(" in needs_source + + +def test_voxcpm2_unified_skip_preserves_segmented_decode_graphs(): + source = TALKER.read_text() + compact_source = "".join(source.split()) + forward_source = source[source.index("def forward(") :] + forward_source = forward_source[: forward_source.index("# -------------------- prefill / decode helpers")] + + assert "use_segmented_decode_graph" in forward_source + assert "unified_skip_reasonisnotNone" in compact_source + assert "notself._enable_unified_decode_graphorunified_skip_reasonisnotNone" in compact_source + assert "ifuse_segmented_decode_graph:" in compact_source + assert "anduse_segmented_decode_graph" in compact_source + + +def test_voxcpm2_scheduler_policy_stays_model_local(): + common_source = OMNI_AR_SCHEDULER.read_text() + voxcpm2_scheduler_source = VOXCPM2_SCHEDULER.read_text() + pipeline_source = VOXCPM2_PIPELINE.read_text() + + assert "VoxCPM2TalkerForConditionalGeneration" not in common_source + assert "voxcpm2_runtime_config" not in common_source + assert "pure_decode_graph" not in common_source + assert "_schedule_with_optional_waiting_deferral" not in common_source + assert "def _should_defer_waiting_admission(self) -> bool:" in common_source + + assert "class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler)" in voxcpm2_scheduler_source + assert "voxcpm2_runtime_config" in voxcpm2_scheduler_source + assert "_should_defer_waiting_for_unified_decode_graph" in voxcpm2_scheduler_source + assert "def schedule(" not in voxcpm2_scheduler_source + assert "create_request_queue" not in voxcpm2_scheduler_source + assert "original_waiting.prepend_requests(deferred_waiting)" not in voxcpm2_scheduler_source + assert "return self._should_defer_waiting_for_unified_decode_graph()" in voxcpm2_scheduler_source + assert "vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler" in pipeline_source + + +def test_voxcpm2_deploy_defaults_to_full_unified_graph_only(): + source = DEPLOY.read_text() + + assert "max_num_seqs: 8" in source + assert "enable_unified_decode_graph: true" in source + assert "unified_decode_graph_max_batch_size: 8" in source + assert "unified_decode_graph_pre_capture_sizes" not in source + assert "enable_runner_assisted_unified_decode_graph" not in source + assert "allow_unified_decode_graph_batch_attention" not in source + assert "decode_graph_capture_policy" not in source diff --git a/tests/test_config_factory.py b/tests/test_config_factory.py index 63a67d652..596a13dbb 100644 --- a/tests/test_config_factory.py +++ b/tests/test_config_factory.py @@ -14,7 +14,7 @@ from unittest.mock import patch import pytest from transformers import PretrainedConfig, Qwen3OmniMoeConfig -from tests.helpers.stage_config import get_deploy_config_path +from tests.helpers.stage_config import get_deploy_config_path, get_deploy_config_stage from vllm_omni.config.config_factory import StageConfigFactory from vllm_omni.config.pipeline_registry import OMNI_PIPELINES, register_pipeline from vllm_omni.config.stage_config import ( @@ -571,6 +571,7 @@ class TestStagePipelineConfig: assert s.final_output is False assert s.sampling_constraints == {} assert s.engine_output_type is None + assert s.scheduler_cls is None class TestPipelineConfigNew: @@ -739,6 +740,15 @@ stages: assert deploy.stages[0].compilation_config == {"pass_config": {"fuse_allreduce_rms": False}} assert "compilation_config" not in deploy.stages[0].engine_extras + def test_load_voxcpm2_deploy_config_preserves_engine_extras(self): + deploy_path = get_deploy_config_path("voxcpm2.yaml") + raw_stage = get_deploy_config_stage("voxcpm2.yaml", 0) + expected_runtime_config = raw_stage["engine_extras"]["hf_overrides"]["voxcpm2_runtime_config"] + + deploy = load_deploy_config(deploy_path) + runtime_config = deploy.stages[0].engine_extras["hf_overrides"]["voxcpm2_runtime_config"] + assert runtime_config == expected_runtime_config + def test_merge_pipeline_deploy(self): deploy_path = Path(__file__).parent.parent / "vllm_omni" / "deploy" / "qwen3_omni_moe.yaml" if not deploy_path.exists(): @@ -800,6 +810,26 @@ stages: assert stages[0].yaml_runtime["requires_multimodal_data"] is True + def test_merge_pipeline_deploy_preserves_pipeline_scheduler_cls(self): + scheduler_cls = "tests.fake.CustomScheduler" + pipeline = PipelineConfig( + model_type="test_scheduler", + model_arch="TestModel", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="ar", + execution_type=StageExecutionType.LLM_AR, + scheduler_cls=scheduler_cls, + ), + ), + ) + deploy = DeployConfig(async_chunk=False, stages=[StageDeployConfig(stage_id=0)]) + + stages = merge_pipeline_deploy(pipeline, deploy) + + assert stages[0].scheduler_cls == scheduler_cls + def test_mixed_schema_preserves_flat_fields(self): """Ensure flat fields are not dropped when engine_args are present.""" fake_config = { @@ -928,6 +958,54 @@ stages: assert "foo" in stage.engine_extras assert stage.engine_extras["foo"] == {"a": 111, "b": {"d": 9, "e": 199}, "c": 3} + def test_explicit_engine_extras_merge_with_flat_and_engine_args(self): + """Explicit engine_extras should be preserved with existing pass-through styles.""" + fake_config = { + "stages": [ + { + "stage_id": 0, + "engine_extras": { + "hf_overrides": { + "runtime_config": { + "explicit_only": True, + "shared": "explicit", + "nested": {"a": 1}, + } + } + }, + "hf_overrides": { + "runtime_config": { + "flat_only": True, + "shared": "flat", + "nested": {"b": 2}, + } + }, + "engine_args": { + "hf_overrides": { + "runtime_config": { + "engine_args_only": True, + "shared": "engine_args", + "nested": {"c": 3}, + } + } + }, + }, + ] + } + + with patch("vllm_omni.config.stage_config.resolve_deploy_yaml", return_value=fake_config): + deploy = load_deploy_config("dummy.yaml") + + assert deploy.stages[0].engine_extras["hf_overrides"] == { + "runtime_config": { + "engine_args_only": True, + "explicit_only": True, + "flat_only": True, + "nested": {"a": 1, "b": 2, "c": 3}, + "shared": "engine_args", + } + } + def test_deep_merge_does_not_mutate_inputs(self): """Merging engine_args must not mutate the base stage dict.""" fake_config = { @@ -1588,6 +1666,197 @@ class TestAuraOmniDeploy: assert stages[3].yaml_engine_args["model_arch"] == "Qwen3TTSCode2Wav" +class TestDeployCliOverrideFlow: + """Test deploy-YAML baselines overridden by CLI runtime overrides.""" + + def test_diffusion_deploy_fields_can_be_overridden_by_cli(self, tmp_path): + deploy_path = tmp_path / "diffusion_stage.yaml" + deploy_path.write_text( + """ +async_chunk: false +stages: + - stage_id: 0 + devices: "4,5" + parallel_config: + pipeline_parallel_size: 1 + data_parallel_size: 1 + tensor_parallel_size: 1 + enable_expert_parallel: false + sequence_parallel_size: 1 + ulysses_degree: 1 + ring_degree: 1 + cfg_parallel_size: 1 + vae_patch_parallel_size: 1 + use_hsdp: false + hsdp_shard_size: -1 + hsdp_replicate_size: 1 + engine_args: + cache_backend: cache_dit + diffusion_attention_backend: FLASH_ATTN + diffusion_kv_cache_dtype: auto + step_execution: false + vae_use_tiling: false + enable_cpu_offload: false + max_generated_image_size: 1048576 + tts_max_instructions_length: 1000 +""", + encoding="utf-8", + ) + + pipeline = PipelineConfig( + model_type="test_diffusion", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="dit", + execution_type=StageExecutionType.DIFFUSION, + final_output=True, + final_output_type="image", + ), + ), + ) + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + stage = stages[0] + + assert stage.yaml_engine_args["parallel_config"]["pipeline_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["data_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["tensor_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["enable_expert_parallel"] is False + assert stage.yaml_engine_args["parallel_config"]["sequence_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["ulysses_degree"] == 1 + assert stage.yaml_engine_args["parallel_config"]["ring_degree"] == 1 + assert stage.yaml_engine_args["parallel_config"]["cfg_parallel_size"] == 1 + assert stage.yaml_engine_args["parallel_config"]["vae_patch_parallel_size"] == 1 + assert stage.yaml_engine_args["cache_backend"] == "cache_dit" + assert stage.yaml_engine_args["parallel_config"]["use_hsdp"] is False + assert stage.yaml_engine_args["parallel_config"]["hsdp_shard_size"] == -1 + assert stage.yaml_engine_args["parallel_config"]["hsdp_replicate_size"] == 1 + assert stage.yaml_engine_args["diffusion_attention_backend"] == "FLASH_ATTN" + assert stage.yaml_engine_args["diffusion_kv_cache_dtype"] == "auto" + assert stage.yaml_engine_args["step_execution"] is False + assert stage.yaml_engine_args["vae_use_tiling"] is False + assert stage.yaml_engine_args["enable_cpu_offload"] is False + assert stage.yaml_engine_args["max_generated_image_size"] == 1048576 + assert stage.yaml_engine_args["tts_max_instructions_length"] == 1000 + + stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( + stage, + { + "pipeline_parallel_size": 2, + "data_parallel_size": 3, + "tensor_parallel_size": 4, + "enable_expert_parallel": True, + "sequence_parallel_size": 24, + "ulysses_degree": 2, + "ring_degree": 4, + "cfg_parallel_size": 2, + "vae_patch_parallel_size": 2, + "cache_backend": "tea_cache", + "use_hsdp": True, + "hsdp_shard_size": 8, + "hsdp_replicate_size": 2, + "diffusion_attention_backend": "SAGE_ATTN", + "diffusion_kv_cache_dtype": "fp8", + "step_execution": True, + "vae_use_tiling": True, + "enable_cpu_offload": True, + "max_generated_image_size": 2097152, + "tts_max_instructions_length": 2000, + }, + ) + + omega_config = stage.to_omegaconf() + + assert omega_config.engine_args.cache_backend == "tea_cache" + assert omega_config.engine_args.diffusion_attention_backend == "SAGE_ATTN" + assert omega_config.engine_args.diffusion_kv_cache_dtype == "fp8" + assert omega_config.engine_args.step_execution is True + assert omega_config.engine_args.vae_use_tiling is True + assert omega_config.engine_args.enable_cpu_offload is True + assert omega_config.engine_args.max_generated_image_size == 2097152 + assert omega_config.engine_args.tts_max_instructions_length == 2000 + assert omega_config.engine_args.parallel_config.pipeline_parallel_size == 2 + assert omega_config.engine_args.parallel_config.data_parallel_size == 3 + assert omega_config.engine_args.parallel_config.tensor_parallel_size == 4 + assert omega_config.engine_args.parallel_config.enable_expert_parallel is True + assert omega_config.engine_args.parallel_config.sequence_parallel_size == 24 + assert omega_config.engine_args.parallel_config.ulysses_degree == 2 + assert omega_config.engine_args.parallel_config.ring_degree == 4 + assert omega_config.engine_args.parallel_config.cfg_parallel_size == 2 + assert omega_config.engine_args.parallel_config.vae_patch_parallel_size == 2 + assert omega_config.engine_args.parallel_config.use_hsdp is True + assert omega_config.engine_args.parallel_config.hsdp_shard_size == 8 + assert omega_config.engine_args.parallel_config.hsdp_replicate_size == 2 + + def test_llm_deploy_fields_can_be_overridden_by_cli(self, tmp_path): + deploy_path = tmp_path / "llm_stage.yaml" + deploy_path.write_text( + """ +async_chunk: false +stages: + - stage_id: 0 + devices: "0,1" + engine_args: + tensor_parallel_size: 1 + enable_expert_parallel: false + gpu_memory_utilization: 0.5 + max_num_seqs: 16 + max_num_batched_tokens: 1024 + max_model_len: 4096 + enforce_eager: false +""", + encoding="utf-8", + ) + + pipeline = PipelineConfig( + model_type="test_llm", + stages=( + StagePipelineConfig( + stage_id=0, + model_stage="thinker", + execution_type=StageExecutionType.LLM_AR, + final_output=True, + final_output_type="text", + ), + ), + ) + deploy = load_deploy_config(deploy_path) + stages = merge_pipeline_deploy(pipeline, deploy) + stage = stages[0] + + assert stage.yaml_engine_args["tensor_parallel_size"] == 1 + assert stage.yaml_engine_args["enable_expert_parallel"] is False + assert stage.yaml_engine_args["gpu_memory_utilization"] == 0.5 + assert stage.yaml_engine_args["max_num_seqs"] == 16 + assert stage.yaml_engine_args["max_num_batched_tokens"] == 1024 + assert stage.yaml_engine_args["max_model_len"] == 4096 + assert stage.yaml_engine_args["enforce_eager"] is False + + stage.runtime_overrides = StageConfigFactory._merge_cli_overrides( + stage, + { + "tensor_parallel_size": 2, + "enable_expert_parallel": True, + "gpu_memory_utilization": 0.9, + "max_num_seqs": 32, + "max_num_batched_tokens": 2048, + "max_model_len": 8192, + "enforce_eager": True, + }, + ) + + omega_config = stage.to_omegaconf() + + assert omega_config.engine_args.tensor_parallel_size == 2 + assert omega_config.engine_args.enable_expert_parallel is True + assert omega_config.engine_args.gpu_memory_utilization == 0.9 + assert omega_config.engine_args.max_num_seqs == 32 + assert omega_config.engine_args.max_num_batched_tokens == 2048 + assert omega_config.engine_args.max_model_len == 8192 + assert omega_config.engine_args.enforce_eager is True + + class TestSentinelDefaultPrecedence: """Caller-typed (non-None) values win over YAML; None values fall through to YAML / dataclass defaults (#3035).""" diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index a37c9b993..0e3a39abc 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from types import SimpleNamespace import numpy as np @@ -10,6 +11,7 @@ from vllm_omni.worker.gpu_ar_model_runner import ( GPUARModelRunner, OmniAsyncGPUModelRunnerOutput, ) +from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -58,6 +60,101 @@ def test_sparse_mm_req_ids_requires_sparse_audio_marker(): assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"] +def test_runner_assisted_full_attention_metadata_request_is_opt_in(): + runner = object.__new__(GPUARModelRunner) + runner.model = object() + runner.scheduler_config = SimpleNamespace(max_num_seqs=16) + + request = runner._get_runner_assisted_full_attention_metadata_request( + req_ids=["r1", "r2"], + num_reqs=2, + num_reqs_padded=4, + num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32), + num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32), + max_num_scheduled_tokens=1, + ) + + assert request is None + + +def test_runner_assisted_full_attention_metadata_request_and_context_hooks(): + calls = [] + + class Model: + def get_runner_assisted_full_attention_metadata_request( + self, + *, + req_ids: Sequence[str], + num_reqs: int, + num_scheduled_tokens: Sequence[int], + num_computed_tokens: Sequence[int], + max_num_scheduled_tokens: int, + ) -> RunnerAssistedFullAttentionMetadataRequest: + calls.append( + ( + "request", + { + "req_ids": list(req_ids), + "num_reqs": num_reqs, + "num_scheduled_tokens": [int(n) for n in num_scheduled_tokens], + "num_computed_tokens": [int(n) for n in num_computed_tokens], + "max_num_scheduled_tokens": max_num_scheduled_tokens, + }, + ) + ) + return RunnerAssistedFullAttentionMetadataRequest( + num_reqs_padded=12, + for_cudagraph_capture=True, + ) + + def set_runner_assisted_full_attention_metadata_context( + self, + *, + enabled: bool, + num_reqs: int = 0, + ) -> None: + calls.append(("context", {"enabled": enabled, "num_reqs": num_reqs})) + + runner = object.__new__(GPUARModelRunner) + runner.model = Model() + runner.scheduler_config = SimpleNamespace(max_num_seqs=8) + + request = runner._get_runner_assisted_full_attention_metadata_request( + req_ids=["r1", "r2"], + num_reqs=2, + num_reqs_padded=4, + num_scheduled_tokens_np=np.array([1, 1], dtype=np.int32), + num_computed_tokens_cpu=np.array([5, 6], dtype=np.int32), + max_num_scheduled_tokens=1, + ) + context_enabled = runner._set_runner_assisted_full_attention_metadata_context( + enabled=True, + num_reqs=2, + ) + context_disabled = runner._set_runner_assisted_full_attention_metadata_context(enabled=False) + + assert request == RunnerAssistedFullAttentionMetadataRequest( + num_reqs_padded=8, + for_cudagraph_capture=True, + ) + assert context_enabled + assert context_disabled + assert calls == [ + ( + "request", + { + "req_ids": ["r1", "r2"], + "num_reqs": 2, + "num_scheduled_tokens": [1, 1], + "num_computed_tokens": [5, 6], + "max_num_scheduled_tokens": 1, + }, + ), + ("context", {"enabled": True, "num_reqs": 2}), + ("context", {"enabled": False, "num_reqs": 0}), + ] + + def test_omni_async_gpu_model_runner_output_builds_lazily_once(): async_output = object.__new__(OmniAsyncGPUModelRunnerOutput) calls = [] @@ -357,6 +454,47 @@ def test_async_snapshot_payload_omits_hidden_when_model_opts_out(): assert payload["multimodal_outputs"]["codes"]["audio"].tolist() == [[1]] +def test_runner_assisted_full_attention_metadata_refresh_pads_buffers(): + class QueryStartLoc: + def __init__(self): + self.np = np.full(5, -1, dtype=np.int32) + self.copied = False + + def copy_to_gpu(self): + self.copied = True + + class BlockTable: + def __init__(self): + self.commits = [] + + def commit_block_table(self, num_reqs_padded): + self.commits.append(num_reqs_padded) + + runner = object.__new__(GPUARModelRunner) + block_table = BlockTable() + runner.input_batch = SimpleNamespace( + num_computed_tokens_cpu=torch.tensor([10, 20, 99, 99], dtype=torch.int32), + block_table=block_table, + ) + runner.optimistic_seq_lens_cpu = torch.zeros(4, dtype=torch.int32) + runner.seq_lens = torch.empty(4, dtype=torch.int32) + runner.query_pos = SimpleNamespace(np=np.empty(4, dtype=np.int32)) + runner.query_start_loc = QueryStartLoc() + runner._get_cumsum_and_arange = lambda scheduled, _query_pos: np.cumsum(scheduled, dtype=np.int32) + + runner._refresh_runner_assisted_full_attention_metadata_buffers( + num_reqs=2, + num_reqs_padded=4, + num_scheduled_tokens_np=np.array([1, 2], dtype=np.int32), + ) + + assert runner.optimistic_seq_lens_cpu.tolist() == [11, 22, 0, 0] + assert runner.seq_lens.tolist() == [11, 22, 0, 0] + assert runner.query_start_loc.np.tolist() == [0, 1, 3, 3, 3] + assert runner.query_start_loc.copied + assert block_table.commits == [4] + + @pytest.mark.parametrize("query_start_loc_attr", ["method", "tensor_attr"]) def test_sample_tokens_tail_only_prefix_cache_uses_staged_cpu_hidden_states(monkeypatch, query_start_loc_attr): runner = object.__new__(GPUARModelRunner) diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index a4eb683e6..e4f11b1e7 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -1264,7 +1264,7 @@ if "daily-omni" not in OPENAI_COMPATIBLE_BACKENDS: # Prevent import order from causing patch failures from vllm.benchmarks import serve from vllm.benchmarks.lib.ready_checker import wait_for_endpoint -from vllm.benchmarks.serve import TaskType, _merge_overrides, calculate_metrics_for_embeddings, get_request +from vllm.benchmarks.serve import TaskType, calculate_metrics_for_embeddings, get_request from vllm_omni.benchmarks.metrics.metrics import ( MultiModalsBenchmarkMetrics, @@ -1276,6 +1276,19 @@ from vllm_omni.benchmarks.metrics.metrics import ( benchmark_old = serve.benchmark +def _merge_overrides(base: dict | None, overrides: dict | None) -> dict | None: + """Merge benchmark extra_body with per-request overrides. + + vLLM 0.24 removed the private helper from ``vllm.benchmarks.serve``. + Keep the same shallow-merge behavior here, with request overrides winning. + """ + if not base and not overrides: + return None + merged = dict(base or {}) + merged.update(overrides or {}) + return merged + + async def benchmark( task_type: TaskType, endpoint_type: str, diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index a269216da..5447fde7a 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -63,7 +63,12 @@ def build_stage_runtime_overrides( if internal_keys is None: from vllm_omni.engine.arg_utils import SHARED_FIELDS, internal_blacklist_keys - internal_keys = internal_blacklist_keys() | SHARED_FIELDS + # Some fields are modeled as orchestrator-owned for top-level CLI + # parsing, but are also legitimate deploy-time stage overrides. Keep + # the default blacklist for true orchestrator/shared fields while + # allowing any field explicitly represented by the deploy schema to + # continue flowing into per-stage overrides. + internal_keys = (internal_blacklist_keys() | SHARED_FIELDS) - deploy_runtime_override_keys() result: dict[str, Any] = {} @@ -223,6 +228,7 @@ class StagePipelineConfig: prompt_expand_func: str | None = None cfg_kv_collect_func: str | None = None omni_kv_config: dict[str, Any] | None = None + scheduler_cls: str | None = None # Model subdirectory indirections: for multi-component HF repos where the # stage's config/tokenizer lives in a subdirectory (e.g. GLM-Image's AR # config is in ``vision_language_encoder/``). Consumed at stage-init time @@ -303,7 +309,7 @@ class StageDeployConfig: the top level of ``DeployConfig`` and propagated to every stage. """ - # === Omni fields === + # === Omni stage wrapper fields === # Stage identity and Omni runtime placement. stage_id: int devices: str | None = None @@ -316,22 +322,33 @@ class StageDeployConfig: default_sampling_params: dict[str, Any] | None = None subtalker_sampling_params: dict[str, Any] | None = None - # === vLLM EngineArgs fields === - # Parallelism and scheduler/memory capacity. + # === Generic stage engine fields === + # Parallelism, scheduler, and memory-capacity controls. tensor_parallel_size: int | None = None + enable_expert_parallel: bool | None = None gpu_memory_utilization: float | None = None max_num_seqs: int | None = None max_num_batched_tokens: int | None = None max_model_len: int | None = None - # Execution, scheduling, and KV/cache behavior. + # Generic execution, scheduling, and KV/cache behavior. enforce_eager: bool | None = None async_scheduling: bool | None = None disable_hybrid_kv_cache_manager: bool | None = None mm_processor_cache_gb: float | None = None - # Diffusion parallel_config deploy override fields. - enable_expert_parallel: bool | None = None + # Generic compilation, profiling, tokenizer/config parsing, and model + # loading controls. + compilation_config: dict[str, Any] | None = None + profiler_config: dict[str, Any] | None = None + skip_mm_profiling: bool | None = None + enable_flashinfer_autotune: bool | None = None + config_format: str | None = None + load_format: str | None = None + tokenizer_mode: str | None = None + + # === Diffusion stage runtime fields === + # Diffusion parallel_config deploy/runtime override fields. ulysses_degree: int | None = None ulysses_mode: str | None = None ring_degree: int | None = None @@ -343,16 +360,44 @@ class StageDeployConfig: hsdp_shard_size: int | None = None hsdp_replicate_size: int | None = None - # Compilation, profiling, tokenizer/config parsing, and model loading. - compilation_config: dict[str, Any] | None = None - profiler_config: dict[str, Any] | None = None - skip_mm_profiling: bool | None = None - enable_flashinfer_autotune: bool | None = None - config_format: str | None = None - load_format: str | None = None - tokenizer_mode: str | None = None + # Diffusion model loading and adapter construction. + model_class_name: str | None = None + diffusion_load_format: str | None = None + diffusers_load_kwargs: dict[str, Any] | None = None + diffusers_call_kwargs: dict[str, Any] | None = None + diffusion_quantization_config: str | None = None + diffusion_attention_backend: str | None = None + diffusion_attention_config: dict[str, Any] | None = None - # Pass-through vLLM EngineArgs fields that are not represented above. + # Diffusion execution, cache, and VAE behavior. + cache_backend: str | None = None + cache_config: dict[str, Any] | None = None + enable_cache_dit_summary: bool | None = None + step_execution: bool | None = None + vae_use_slicing: bool | None = None + vae_use_tiling: bool | None = None + boundary_ratio: float | None = None + flow_shift: float | None = None + diffusion_kv_cache_dtype: str | None = None + diffusion_kv_cache_skip_steps: str | None = None + diffusion_kv_cache_skip_layers: str | None = None + auxiliary_text_encoder: str | None = None + + # Runtime optimizations used by diffusion loading/execution. + enable_multithread_weight_load: bool | None = None + num_weight_load_threads: int | None = None + enable_cpu_offload: bool | None = None + enable_layerwise_offload: bool | None = None + + # Diffusion-specific debug and observability knobs. + enable_diffusion_pipeline_profiler: bool | None = None + + # Modality/service constraints consumed outside the core engine config. + max_generated_image_size: int | None = None + tts_max_instructions_length: int | None = None + + # === Pass-through stage engine fields === + # Pass-through stage engine args that are not represented above. engine_extras: dict[str, Any] = field(default_factory=dict) @@ -409,10 +454,22 @@ _STAGE_RESERVED_KEYS = frozenset( _STAGE_DEPLOY_FIELDS = {f.name: f for f in fields(StageDeployConfig) if f.name not in _STAGE_RESERVED_KEYS} +def deploy_runtime_override_keys() -> frozenset[str]: + """Return deploy-schema fields that are valid CLI/runtime overrides. + + These keys form the positive contract for stage override propagation: + stage-scoped deploy knobs plus top-level pipeline-wide engine settings. + They must remain overridable even if they are also modeled on + ``OrchestratorArgs`` for top-level CLI parsing. + """ + return frozenset(_STAGE_DEPLOY_FIELDS) | frozenset(_PIPELINE_WIDE_ENGINE_FIELDS) + + def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: """Parse a single stage entry from deploy YAML into StageDeployConfig.""" # Get the non-reserved keys for this stage flat_args = {k: v for k, v in stage_data.items() if k not in _STAGE_RESERVED_KEYS} + explicit_engine_extras = dict(stage_data.get("engine_extras") or {}) runtime_cfg = dict(stage_data.get("runtime", {})) devices = runtime_cfg.get("devices", stage_data.get("devices")) num_replicas = runtime_cfg.get("num_replicas", stage_data.get("num_replicas", 1)) @@ -440,7 +497,7 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: kwargs["output_connectors"] = stage_data.get("output_connectors") kwargs["input_connectors"] = stage_data.get("input_connectors") kwargs["default_sampling_params"] = stage_data.get("default_sampling_params") - kwargs["engine_extras"] = flat_args + kwargs["engine_extras"] = _get_recursively_merged_dict(explicit_engine_extras, flat_args) return StageDeployConfig(**kwargs) @@ -840,7 +897,7 @@ def merge_pipeline_deploy( final_output=ps.final_output, final_output_type=ps.final_output_type, worker_type=worker_type, - scheduler_cls=_scheduler_path(sched_cls), + scheduler_cls=ps.scheduler_cls or _scheduler_path(sched_cls), hf_config_name=ps.hf_config_name, is_comprehension=ps.owns_tokenizer, yaml_engine_args=engine_args, diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 178fedee3..217c81543 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -11,6 +11,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStat from vllm.logger import init_logger from vllm.v1.core.sched.async_scheduler import AsyncScheduler as AsyncVLLMScheduler 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, FinishReason @@ -133,6 +134,9 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self._omits_kv_transfer_cache[rid] = result return result + def _should_defer_waiting_admission(self) -> bool: + return False + def _process_kv_transfer_trigger(self, request: Request, new_token_ids: list[int]) -> bool: """ Check triggers and process side effects (marking transfer). @@ -222,9 +226,19 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self.waiting, self.running, scheduler_requests=self.requests ) + original_waiting = None + if self._should_defer_waiting_admission(): + original_waiting = self.waiting + self.waiting = create_request_queue(self.policy) + try: scheduler_output = super().schedule(throttle_prefills) finally: + if original_waiting is not None: + deferred_waiting = list(self.waiting) + if deferred_waiting: + original_waiting.prepend_requests(deferred_waiting) + self.waiting = original_waiting if self.chunk_transfer_adapter: # Add request waiting for chunk to the waiting and running queue self.chunk_transfer_adapter.restore_queues( diff --git a/vllm_omni/deploy/voxcpm2.yaml b/vllm_omni/deploy/voxcpm2.yaml index 5d366521b..d534a41c1 100644 --- a/vllm_omni/deploy/voxcpm2.yaml +++ b/vllm_omni/deploy/voxcpm2.yaml @@ -8,16 +8,17 @@ dtype: bfloat16 stages: - stage_id: 0 - max_num_seqs: 32 + max_num_seqs: 8 gpu_memory_utilization: 0.9 # Right-size the KV cache instead of letting it fill gpu_memory_utilization. - # max_num_seqs(32) * max_model_len(4096) ~= 131k tokens; 6 GiB holds ~175k - # (42x concurrency, above the 32-seq scheduler cap, so no throughput loss). + # max_num_seqs(8) * max_model_len(4096) ~= 32k tokens; 6 GiB leaves headroom + # for longer live requests while keeping VoxCPM2 TTS decode batches small + # enough to avoid the high ITL observed with large admission caps. # The remaining VRAM stays free for the VoxCPM2 diffusion side-path # (LocDiT / AudioVAE / code2wav + batch CUDA-graph pools), which allocates # outside vLLM's KV accounting during inference. Keeps peak ~13 GiB on any # card and avoids per-rebase OOM tuning on small CI GPUs. - # Assumes prefix caching off and max_num_seqs*max_model_len <= ~175k tokens; + # Assumes prefix caching off and typical live decode concurrency <= ~8; # raise this cap if either grows or enable_prefix_caching is turned on. kv_cache_memory_bytes: 6442450944 # 6 GiB enforce_eager: true @@ -27,18 +28,20 @@ stages: max_num_batched_tokens: 4096 max_model_len: 4096 devices: "0" - engine_extras: - hf_overrides: - voxcpm2_runtime_config: - audio_emit_every: 1 - vae_decode_every: 3 - enable_batched_cfm: true - enable_batched_vae_decode: true - coalesce_audio_d2h: true - enable_loc_dit_fused_qkv: true - enable_loc_dit_fused_mlp: true - enable_loc_dit_zero_dt_cache: true - enable_loc_dit_skip_qkv_contig: true + hf_overrides: + voxcpm2_runtime_config: + audio_emit_every: 1 + vae_decode_every: 3 + enable_batched_cfm: true + enable_batched_vae_decode: true + coalesce_audio_d2h: true + enable_loc_dit_fused_qkv: true + enable_loc_dit_fused_mlp: true + enable_loc_dit_zero_dt_cache: true + enable_loc_dit_skip_qkv_contig: true + enable_loc_dit_reduce_overhead_no_cg: true + enable_unified_decode_graph: true + unified_decode_graph_max_batch_size: 8 default_sampling_params: temperature: 0.0 top_p: 1.0 diff --git a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py index f3e0f3357..c67f8cb24 100644 --- a/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py +++ b/vllm_omni/diffusion/models/ming_flash_omni/pipeline_ming_imagegen.py @@ -24,7 +24,6 @@ from __future__ import annotations import logging import os -from dataclasses import dataclass, field from pathlib import Path from typing import Any @@ -49,7 +48,9 @@ from vllm_omni.diffusion.models.ming_flash_omni.ming_zimage_transformer import ( MingZImageTransformer2DModel, ) from vllm_omni.diffusion.models.z_image.pipeline_z_image import ZImagePipeline +from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch +from vllm_omni.inputs.data import OmniDiffusionSamplingParams from vllm_omni.model_executor.model_loader.weight_utils import ( download_weights_from_hf_specific, ) @@ -58,38 +59,6 @@ from vllm_omni.transformers_utils.configs.ming_flash_omni import MingImageGenCon logger = logging.getLogger(__name__) -@dataclass -class _ZPipelineSamplingParams: - """Typed shim satisfying the attributes ZImagePipeline.forward reads - from ``req.sampling_params``. Unlike SimpleNamespace this fails loudly - at construction if a required field is missing.""" - - height: int - width: int - num_inference_steps: int - guidance_scale: float - generator: torch.Generator | None = None - strength: float | None = None - sigmas: list[float] | None = None - max_sequence_length: int = 512 - guidance_rescale: float | None = None - num_outputs_per_prompt: int = 1 - - -@dataclass -class _ZPipelineRequest: - """Typed shim satisfying the attributes ZImagePipeline.forward reads - from ``req``.""" - - request_id: str - sampling_params: _ZPipelineSamplingParams - prompts: list[dict[str, Any]] = field(default_factory=lambda: [{"prompt": "", "negative_prompt": ""}]) - - @property - def num_reqs(self) -> int: - return 1 - - class MingImagePipeline(ZImagePipeline): """Ming-flash-omni-2.0 text-to-image diffusion pipeline. @@ -123,6 +92,10 @@ class MingImagePipeline(ZImagePipeline): self.device = self._execution_device # Ming convention alias self._dtype = dtype + # Request-scoped conditioning handed to the inherited encode_prompt override + self._pending_prompt_embeds: list[torch.Tensor] | None = None + self._pending_negative_prompt_embeds: list[torch.Tensor] | None = None + # Ming's per-checkpoint image-gen configuration. We cannot rely on # ``od_config.hf_config.image_gen_config`` because the diffusion # stage is started with ``hf_config_name: thinker_config`` (the @@ -273,6 +246,21 @@ class MingImagePipeline(ZImagePipeline): latent = self.vae.encode(ref).latent_dist.mode() return (latent - self.vae.config.shift_factor) * self.vae.config.scaling_factor + def encode_prompt(self, *args, **kwargs): # noqa: ARG002 + """Return Ming's precomputed conditioning instead of encoding text. + + NOTE: Ming has no Z-Image text_encoder; its conditioning (cap_feats, optionally ByT5-augmented) + is computed in forward and stashed on `self._pending_*` immediately before + the `super().forward` call, so we simply hand it back here. + """ + if self._pending_prompt_embeds is None: + raise RuntimeError( + "MingImagePipeline.encode_prompt called without pending " + "conditioning; it must run within MingImagePipeline.forward." + ) + + return self._pending_prompt_embeds, self._pending_negative_prompt_embeds + # ------------------------------------------------------------------ # Forward # ------------------------------------------------------------------ @@ -416,17 +404,27 @@ class MingImagePipeline(ZImagePipeline): negative_prompt_embeds = [negative_cap_feats[i] for i in range(negative_cap_feats.shape[0])] else: negative_prompt_embeds = [self.condition_encoder.zero_negative(e) for e in prompt_embeds] + self._pending_prompt_embeds = prompt_embeds + self._pending_negative_prompt_embeds = negative_prompt_embeds - z_sp = _ZPipelineSamplingParams( + # Build a real single-request DiffusionRequestBatch; + # the prompt text itself is a neutral placeholder here. + z_sp = OmniDiffusionSamplingParams( height=height, width=width, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, generator=generator, + output_type="pt", ) - z_req = _ZPipelineRequest( - request_id=req.request_id or "ming-imagegen", - sampling_params=z_sp, + z_req = DiffusionRequestBatch( + requests=[ + OmniDiffusionRequest( + prompt={"prompt": ""}, + sampling_params=z_sp, + request_id=req.request_id or "ming-imagegen", + ) + ] ) # Reference image (img2img) → VAE-encoded latent published on the @@ -446,21 +444,12 @@ class MingImagePipeline(ZImagePipeline): None if ref_latent is None else tuple(ref_latent.shape), ) try: - outputs: DiffusionOutput = super().forward( - z_req, - prompt=None, - height=height, - width=width, - num_inference_steps=num_inference_steps, - guidance_scale=guidance_scale, - prompt_embeds=prompt_embeds, - negative_prompt_embeds=negative_prompt_embeds, - generator=generator, - output_type="pt", - return_dict=True, - ) + outputs: DiffusionOutput = super().forward(z_req) finally: set_forward_context_ref_latent(None) + # Drop request-scoped conditioning so we don't retain GPU tensors. + self._pending_prompt_embeds = None + self._pending_negative_prompt_embeds = None raw = outputs.output if not isinstance(raw, torch.Tensor): diff --git a/vllm_omni/entrypoints/openai/image_api_utils.py b/vllm_omni/entrypoints/openai/image_api_utils.py index b958d786a..5bd79e936 100644 --- a/vllm_omni/entrypoints/openai/image_api_utils.py +++ b/vllm_omni/entrypoints/openai/image_api_utils.py @@ -14,7 +14,7 @@ import io import PIL.Image SUPPORTED_LAYERED_RESOLUTIONS = (640, 1024) -SUPPORTED_LAYERED_LAYERS_RANGE = range(3, 11) +SUPPORTED_LAYERED_LAYERS_RANGE = range(2, 11) def parse_size(size_str: str) -> tuple[int, int]: diff --git a/vllm_omni/entrypoints/openai/protocol/images.py b/vllm_omni/entrypoints/openai/protocol/images.py index 78f63d387..b7247962c 100644 --- a/vllm_omni/entrypoints/openai/protocol/images.py +++ b/vllm_omni/entrypoints/openai/protocol/images.py @@ -60,7 +60,7 @@ class ImageGenerationRequest(BaseModel): user: str | None = Field(default=None, description="User identifier for tracking") layers: int | None = Field( default=None, - description="Number of output layers for layered image models. Supported range: 3-10.", + description="Number of output layers for layered image models. Supported range: 2-10.", ) @field_validator("size") diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 4fa9c7bed..ae1a7c1fd 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -3513,11 +3513,13 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): # apply uploaded speakers inside validate(), which sets request.ref_audio # in place. The builders need to know whether the caller supplied audio # inline vs. via an uploaded voice. + model_type: str | None = None has_inline_ref_audio = request.ref_audio is not None if self._tts_model_type == "ming_flash_omni_tts": # ming_flash_omni is intentionally NOT migrated onto the adapter # framework in this PR (it has no registered adapter); keep it on the # legacy inline dispatch so serving still works. + model_type = "ming_flash_omni_tts" validation_error = self._validate_ming_flash_omni_tts_request(request) if validation_error: raise ValueError(validation_error) @@ -3531,6 +3533,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): prepared = await adapter.build(request, sampling_params_list, has_inline_ref_audio) prompt = prepared.prompt tts_params = prepared.tts_params + model_type = prepared.model_type qwen3_ref_audio_warmup_artifact_key = prepared.warmup_artifact_key else: # Qwen omni models (Qwen3-Omni, Qwen2.5-Omni) use a "talker" @@ -3554,36 +3557,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): tts_params = {} prompt = {"prompt": request.input} - if self._tts_model_type == "step_audio2": - model_type = "step_audio2" - elif self._is_fish_speech: - model_type = "fish_speech" - elif self._tts_model_type == "covo_audio": - model_type = "covo_audio" - elif self._tts_model_type == "voxtral_tts": - model_type = "voxtral_tts" - elif self._tts_model_type == "cosyvoice3": - model_type = "cosyvoice3" - elif self._tts_model_type == "voxcpm2": - model_type = "voxcpm2" - elif self._tts_model_type == "ming_flash_omni_tts": - model_type = "ming_flash_omni_tts" - elif self._tts_model_type == "ming_tts": - model_type = "ming_tts" - elif self._tts_model_type == "moss_tts_nano": - model_type = "moss_tts_nano" - elif self._tts_model_type == "moss_tts": - model_type = "moss_tts" - elif self._tts_model_type == "higgs_audio_v2": - model_type = "higgs_audio_v2" - elif self._tts_model_type == "glm_tts": - model_type = "glm_tts" - elif self._tts_model_type == "indextts2": - model_type = "indextts2" - elif self._is_tts: - model_type = tts_params.get("task_type", ["unknown"])[0] - else: - model_type = "generic" + if model_type is None: + if self._is_tts: + model_type = tts_params.get("task_type", ["unknown"])[0] + else: + model_type = "generic" logger.info( "TTS speech request %s: text=%r, model=%s", request_id, diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index 282abed23..a3dfd7bf1 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -587,9 +587,9 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): input_ids: torch.Tensor | None, positions: torch.Tensor | None, inputs_embeds: torch.Tensor | None, - req_ids: list[str], - num_computed_tokens: list[int], - num_scheduled_tokens: list[int], + req_ids: Sequence[str], + num_computed_tokens: Sequence[int], + num_scheduled_tokens: Sequence[int], input_ids_buffer: torch.Tensor | None = None, ) -> tuple[torch.Tensor | None, torch.Tensor | None]: """Restore input_ids so _adjust_positions_for_img2img can locate @@ -599,7 +599,7 @@ class OmniBagelForConditionalGeneration(BagelForConditionalGeneration): input_ids = input_ids_buffer return input_ids, positions - def flush_pending_metadata(self, req_ids: list[str]) -> None: + def flush_pending_metadata(self, req_ids: Sequence[str]) -> None: """Map pending metadata (batch order) to req_ids after forward(). Guard: if a request already has metadata with ``image_shape`` diff --git a/vllm_omni/model_executor/models/voxcpm2/pipeline.py b/vllm_omni/model_executor/models/voxcpm2/pipeline.py index 347fce177..0579ebc8c 100644 --- a/vllm_omni/model_executor/models/voxcpm2/pipeline.py +++ b/vllm_omni/model_executor/models/voxcpm2/pipeline.py @@ -26,6 +26,7 @@ VOXCPM2_PIPELINE = PipelineConfig( final_output_type="audio", owns_tokenizer=True, engine_output_type="audio", + scheduler_cls="vllm_omni.model_executor.models.voxcpm2.scheduler.VoxCPM2OmniARAsyncScheduler", sampling_constraints={ "detokenize": False, "stop_token_ids": [1], diff --git a/vllm_omni/model_executor/models/voxcpm2/scheduler.py b/vllm_omni/model_executor/models/voxcpm2/scheduler.py new file mode 100644 index 000000000..8aade3e74 --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm2/scheduler.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from vllm.v1.request import RequestStatus + +from vllm_omni.core.sched.omni_ar_scheduler import OmniARAsyncScheduler + + +class VoxCPM2OmniARAsyncScheduler(OmniARAsyncScheduler): + """VoxCPM2 scheduler variant for full unified decode graph serving. + + VoxCPM2's full unified decode graph only applies to pure decode batches. + When a decode-ready request is already running, this scheduler defers new + waiting admissions for the current tick so the decode batch can stay on the + VoxCPM2 graph path. This is a model-local serving policy, not a generic AR + scheduler rule. + """ + + def _unified_decode_graph_enabled(self) -> bool: + hf_config = getattr(self.vllm_config.model_config, "hf_config", None) + runtime_config = getattr(hf_config, "voxcpm2_runtime_config", None) + if isinstance(runtime_config, dict): + return bool(runtime_config.get("enable_unified_decode_graph", False)) + return bool(getattr(runtime_config, "enable_unified_decode_graph", False)) + + def _should_defer_waiting_for_unified_decode_graph(self) -> bool: + if not self._unified_decode_graph_enabled(): + return False + if not self.waiting or not self.running: + return False + + for request in self.running: + if getattr(request, "status", None) != RequestStatus.RUNNING or request.is_finished(): + continue + if self._get_confirmed_num_computed_tokens(request) >= request.num_prompt_tokens: + return True + return False + + def _should_defer_waiting_admission(self) -> bool: + return self._should_defer_waiting_for_unified_decode_graph() diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 9505bfbee..555fdf6a6 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -16,12 +16,13 @@ import hashlib import logging import math import time -from collections.abc import Iterable +from collections.abc import Callable, Iterable, Sequence from types import MethodType -from typing import Any +from typing import Any, NamedTuple, Protocol, TypedDict import torch import torch.nn as nn +from typing_extensions import Unpack from vllm.config import VllmConfig from vllm.forward_context import get_forward_context, override_forward_context from vllm.inputs import tokens_input @@ -42,6 +43,7 @@ from vllm_omni.utils.speaker_cache import ( iter_custom_voice_profiles, load_validated_profile_tensors, ) +from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest from .minicpm4_paged import MiniCPM4PagedForVoxCPM2, MiniCPM4PagedResidualLM from .voxcpm2_import_utils import import_voxcpm2_core @@ -56,6 +58,52 @@ _ENABLE_NVTX_PROFILE = False _ACTIVE_STATE_LEAK_WARN_MIN = 512 +class _VoxCPM2PromptConfigLike(Protocol): + audio_vae_config: dict[str, Any] + patch_size: int + + +class _AttentionMetadataLike(Protocol): + scheduler_metadata: Any + + +class _ForwardContextLike(Protocol): + attn_metadata: dict[str, _AttentionMetadataLike] | Any + + +class VoxCPM2PreprocessInput(TypedDict, total=False): + additional_information: dict[str, Any] + request_id: str + text_token_ids: list[list[int]] + reference_audio: object + ref_audio: object + prompt_audio: object + prompt_text: str | list[str] | None + voice_profile: dict[str, Any] | list[dict[str, Any]] | None + voice_name: str | list[str] | None + voice_created_at: int | str | None + + +class VoxCPM2PostprocessInput(TypedDict, total=False): + request_id: str + + +class _PrefillInputs(NamedTuple): + text_token: torch.Tensor + audio_feat: torch.Tensor + text_mask: torch.Tensor + audio_mask: torch.Tensor + + +class _PrefillResidualMeta(NamedTuple): + lm_hidden: torch.Tensor + prefix_feat_cond: torch.Tensor + + +class _DecodeResidualMeta(NamedTuple): + new_lm_hidden: torch.Tensor + + @dataclasses.dataclass(frozen=True) class _VoxCPM2RuntimeConfig: enable_profiling: bool = False @@ -84,6 +132,9 @@ class _VoxCPM2RuntimeConfig: enable_batched_vae_decode: bool = True enable_batched_fsq_fusion: bool = False batched_fsq_fusion_max_batch: int = 32 + enable_batched_prefill_tail: bool = False + enable_unified_decode_graph: bool = True + unified_decode_graph_max_batch_size: int = 64 @classmethod def from_vllm_config(cls, vllm_config: VllmConfig) -> _VoxCPM2RuntimeConfig: @@ -133,6 +184,13 @@ class _VoxCPM2RuntimeConfig: "audio_emit_every > 1, and VAE CUDA Graph; disabling batched VAE decode." ) cfg = dataclasses.replace(cfg, enable_batched_vae_decode=False) + if cfg.enable_unified_decode_graph: + if cfg.enable_cfm_cuda_graph: + logger.info("VoxCPM2 unified decode graph supersedes independent CFM CUDA Graph; disabling it.") + cfg = dataclasses.replace(cfg, enable_cfm_cuda_graph=False, enable_cfm_prealloc_output=False) + if not cfg.enable_batched_cfm: + logger.info("VoxCPM2 unified decode graph requires batched CFM; enabling it.") + cfg = dataclasses.replace(cfg, enable_batched_cfm=True) return cfg @staticmethod @@ -195,11 +253,11 @@ def split_multichar_chinese(token_ids: list[int], split_map: dict[int, list[int] def build_voxcpm2_prompt( - hf_config: Any, + hf_config: _VoxCPM2PromptConfigLike, tokenizer: Any, split_map: dict[int, list[int]], text: str, - ref_audio: Any | None = None, + ref_audio: Sequence[float] | torch.Tensor | None = None, ref_sr: int | None = None, ref_text: str | None = None, voice_profile: dict[str, Any] | None = None, @@ -234,6 +292,8 @@ def build_voxcpm2_prompt( prefill_len += len(prompt_ids) prefill_len += audio_feat_len elif ref_audio is not None: + if ref_sr is None: + raise ValueError("VoxCPM2 ref_sr is required when ref_audio is provided.") vae = hf_config.audio_vae_config patch_samples = hf_config.patch_size * math.prod(vae["encoder_rates"]) ref_len = math.ceil(math.ceil(len(ref_audio) * vae["sample_rate"] / ref_sr) / patch_samples) @@ -343,6 +403,46 @@ class _CapturedCFMGraph: buffers: _CFMBufferManager +@dataclasses.dataclass +class _CapturedUnifiedDecodeGraph: + graph: torch.cuda.CUDAGraph + batch_size: int + # Static input buffers (copied before replay) + input_embeds: torch.Tensor + positions: torch.Tensor + prev_feat_embed: torch.Tensor + prefix_feat_cond: torch.Tensor + # Static output buffers (read after replay) + next_feat_embed: torch.Tensor + cfm_output: torch.Tensor + lm_hidden: torch.Tensor + # Internal buffers for CFM solver + cfm_buffers: _CFMBufferManager + cfm_noise: torch.Tensor + + +@dataclasses.dataclass +class _UnifiedDecodeGraphStats: + captures: int = 0 + replays: int = 0 + skips: dict[str, int] = dataclasses.field(default_factory=dict) + real_batch_sizes: dict[int, int] = dataclasses.field(default_factory=dict) + graph_bucket_sizes: dict[int, int] = dataclasses.field(default_factory=dict) + logged_replays: int = 0 + logged_skips: int = 0 + + def record_skip(self, reason: str) -> None: + self.skips[reason] = self.skips.get(reason, 0) + 1 + + def record_replay_batch(self, *, real_batch_size: int, graph_bucket_size: int) -> None: + self.real_batch_sizes[real_batch_size] = self.real_batch_sizes.get(real_batch_size, 0) + 1 + self.graph_bucket_sizes[graph_bucket_size] = self.graph_bucket_sizes.get(graph_bucket_size, 0) + 1 + + @property + def total_skips(self) -> int: + return sum(self.skips.values()) + + @dataclasses.dataclass class _PendingAudioCopy: host: torch.Tensor @@ -927,6 +1027,17 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._max_cached_graphs = self._max_batch_size self._cuda_graph_warmup_steps = 0 self._cuda_graph_warmup_threshold = 3 + self._unified_graphs: dict[int, _CapturedUnifiedDecodeGraph] = {} + self._enable_unified_decode_graph = ( + use_cuda_graph and self._runtime_config.enable_unified_decode_graph and not self._deterministic_cfm_noise + ) + self._unified_decode_graph_max_batch_size = self._runtime_config.unified_decode_graph_max_batch_size + self._unified_graph_bucket_sizes: frozenset[int] = self._build_unified_graph_bucket_sizes( + min(self._max_batch_size, self._unified_decode_graph_max_batch_size) + ) + self._unified_graph_stats = _UnifiedDecodeGraphStats() + self._runner_assisted_unified_decode_graph_active = False + self._runner_assisted_unified_decode_graph_batch_size = 0 self._multichar_zh_split: dict[int, list[int]] | None = None @@ -1035,10 +1146,10 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): def _build_prompt_cache( self, - ref_audio: Any = None, - prompt_audio: Any = None, + ref_audio: object = None, + prompt_audio: object = None, prompt_text: str | None = None, - ) -> dict | None: + ) -> dict[str, Any] | None: """Build prompt cache, handling both file paths and raw audio data. The OpenAI speech API sends decoded audio as [samples_list, sr] @@ -1106,6 +1217,84 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): max_batch_size=self._max_batch_size, ) + @staticmethod + def _voxcpm2_unwrap_torch_compile(module: nn.Module | Callable) -> nn.Module | Callable: + """Return the eager module behind torch.compile wrappers.""" + seen: set[int] = set() + while hasattr(module, "_orig_mod") and id(module) not in seen: + seen.add(id(module)) + module = getattr(module, "_orig_mod") + return module + + @staticmethod + def _voxcpm2_compile_without_inductor_cudagraphs( + module: nn.Module | Callable, + *, + mode: str | None, + fullgraph: bool, + ) -> Callable: + """Compile VoxCPM2 LocDiT for capture inside VoxCPM2-owned CUDA graphs.""" + kwargs: dict[str, Any] = { + "fullgraph": fullgraph, + "options": { + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + }, + } + if mode is not None: + kwargs["mode"] = mode + return torch.compile(module, **kwargs) + + @staticmethod + def _raise_if_cuda_runtime_error(error: BaseException) -> None: + if isinstance(error, torch.cuda.OutOfMemoryError): + raise error + if isinstance(error, RuntimeError): + message = str(error).lower() + cuda_markers = ( + "cuda error", + "device-side assert", + "illegal memory access", + "out of memory", + "cublas", + "cudnn", + "cufft", + ) + if any(marker in message for marker in cuda_markers): + raise error + + def _voxcpm2_compile_unified_capture_estimator(self, module: nn.Module | Callable) -> Callable: + eager_module = self._voxcpm2_unwrap_torch_compile(module) + if getattr(self, "_estimator_for_unified_capture_source", None) is eager_module and hasattr( + self, "_estimator_for_unified_capture" + ): + return self._estimator_for_unified_capture + compiled = self._voxcpm2_compile_without_inductor_cudagraphs( + eager_module, + mode=None, + fullgraph=False, + ) + compiled._compiled = True + self._estimator_for_unified_capture_source = eager_module + self._estimator_for_unified_capture = compiled + return compiled + + def _voxcpm2_compile_unified_capture_feat_encoder(self, module: nn.Module | Callable) -> Callable: + eager_module = self._voxcpm2_unwrap_torch_compile(module) + if getattr(self, "_feat_encoder_for_unified_capture_source", None) is eager_module and hasattr( + self, "_feat_encoder_for_unified_capture" + ): + return self._feat_encoder_for_unified_capture + compiled = self._voxcpm2_compile_without_inductor_cudagraphs( + eager_module, + mode=None, + fullgraph=False, + ) + compiled._compiled = True + self._feat_encoder_for_unified_capture_source = eager_module + self._feat_encoder_for_unified_capture = compiled + return compiled + def _setup_torch_compile(self) -> None: if not self._enable_torch_compile: return @@ -1133,30 +1322,42 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if _install_locdit_zero_dt_cache(estimator): targets.append("LocDiT zero-dt embedding cache") + external_cfm_capture = self._enable_cfm_cuda_graph or self._enable_unified_decode_graph if cfg.enable_loc_dit_layer_nvtx: patched = _install_locdit_layer_nvtx(estimator) if patched: targets.append(f"LocDiT layer NVTX ({patched} layers, compile skipped)") estimator._compiled = True - elif self._enable_cfm_cuda_graph: + elif self._enable_unified_decode_graph: + try: + compiled_ro = torch.compile(estimator, mode="reduce-overhead", fullgraph=False) + compiled_ro._compiled = True + tts.feat_decoder.estimator = compiled_ro + targets.append("LocDiT (reduce-overhead serving + lazy no-cg unified capture)") + except Exception as e: + self._raise_if_cuda_runtime_error(e) + logger.warning("torch.compile LocDiT dual-mode failed: %s", e) + try: + tts.feat_decoder.estimator = self._voxcpm2_compile_without_inductor_cudagraphs( + estimator, + mode="reduce-overhead", + fullgraph=False, + ) + tts.feat_decoder.estimator._compiled = True + self._estimator_for_unified_capture = tts.feat_decoder.estimator + targets.append("LocDiT (no-cg only, unified capture)") + except Exception as inner_e: + self._raise_if_cuda_runtime_error(inner_e) + logger.warning("torch.compile LocDiT failed completely: %s", inner_e) + elif external_cfm_capture: try: if cfg.enable_loc_dit_reduce_overhead_no_cg: - import torch._inductor.config as inductor_config - - old_cudagraphs = inductor_config.triton.cudagraphs - old_cudagraph_trees = inductor_config.triton.cudagraph_trees - inductor_config.triton.cudagraphs = False - inductor_config.triton.cudagraph_trees = False - try: - tts.feat_decoder.estimator = torch.compile( - estimator, - mode="reduce-overhead", - fullgraph=False, - ) - finally: - inductor_config.triton.cudagraphs = old_cudagraphs - inductor_config.triton.cudagraph_trees = old_cudagraph_trees - targets.append("LocDiT (reduce-overhead no-inductor-cudagraph + CUDA Graph)") + tts.feat_decoder.estimator = self._voxcpm2_compile_without_inductor_cudagraphs( + estimator, + mode="reduce-overhead", + fullgraph=False, + ) + targets.append("LocDiT (reduce-overhead no-inductor-cudagraph + external CUDA Graph)") else: tts.feat_decoder.estimator = torch.compile( estimator, @@ -1166,25 +1367,11 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): "triton.cudagraph_trees": False, }, ) - graph_mode = "fullgraph" if cfg.enable_loc_dit_fullgraph_no_cg else "partitioned" - targets.append(f"LocDiT ({graph_mode} compile no-inductor-cudagraph + CUDA Graph)") + targets.append("LocDiT (no-inductor-cudagraph + external CUDA Graph)") tts.feat_decoder.estimator._compiled = True except Exception as e: - logger.warning("torch.compile LocDiT reduce-overhead for CUDA Graph failed: %s", e) - try: - tts.feat_decoder.estimator = torch.compile( - estimator, - fullgraph=cfg.enable_loc_dit_fullgraph_no_cg, - options={ - "triton.cudagraphs": False, - "triton.cudagraph_trees": False, - }, - ) - tts.feat_decoder.estimator._compiled = True - graph_mode = "fullgraph" if cfg.enable_loc_dit_fullgraph_no_cg else "partitioned" - targets.append(f"LocDiT ({graph_mode} compile no-inductor-cudagraph + CUDA Graph)") - except Exception as inner_e: - logger.warning("torch.compile LocDiT for CUDA Graph failed: %s", inner_e) + self._raise_if_cuda_runtime_error(e) + logger.warning("torch.compile LocDiT for external CUDA Graph failed: %s", e) else: try: tts.feat_decoder.estimator = torch.compile( @@ -1195,14 +1382,34 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): tts.feat_decoder.estimator._compiled = True targets.append("LocDiT") except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("torch.compile LocDiT failed: %s", e) try: if not hasattr(tts.feat_encoder, "_compiled"): - tts.feat_encoder = torch.compile(tts.feat_encoder, mode="reduce-overhead", fullgraph=False) - tts.feat_encoder._compiled = True - targets.append("feat_encoder") + if self._enable_unified_decode_graph: + feat_encoder = tts.feat_encoder + compiled_ro = torch.compile(feat_encoder, mode="reduce-overhead", fullgraph=False) + compiled_ro._compiled = True + tts.feat_encoder = compiled_ro + targets.append("feat_encoder (reduce-overhead serving + lazy no-cg unified capture)") + elif external_cfm_capture: + tts.feat_encoder = torch.compile( + tts.feat_encoder, + fullgraph=False, + options={ + "triton.cudagraphs": False, + "triton.cudagraph_trees": False, + }, + ) + tts.feat_encoder._compiled = True + targets.append("feat_encoder (no-inductor-cudagraph)") + else: + tts.feat_encoder = torch.compile(tts.feat_encoder, mode="reduce-overhead", fullgraph=False) + tts.feat_encoder._compiled = True + targets.append("feat_encoder") except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("torch.compile feat_encoder failed: %s", e) if self._compile_vae and not self._enable_vae_cuda_graph: @@ -1212,6 +1419,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): tts.audio_vae._compiled = True targets.append("AudioVAE") except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("torch.compile AudioVAE failed: %s", e) if not self._enable_cuda_graph: @@ -1220,6 +1428,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): targets.extend(f"scaffold.{t}" for t in self.model.compile_selective()) self.model._selective_compiled = True except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("scaffold compile failed: %s", e) if not getattr(self.residual_model, "_selective_compiled", False): @@ -1227,6 +1436,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): targets.extend(f"residual.{t}" for t in self.residual_model.compile_selective()) self.residual_model._selective_compiled = True except Exception as e: + self._raise_if_cuda_runtime_error(e) logger.warning("residual compile failed: %s", e) else: self.model.precompute_fused_qkv() @@ -1240,6 +1450,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._projections_compiled = True targets.append("projections") except Exception as e: + self._raise_if_cuda_runtime_error(e) self._compiled_dit_proj = self._compiled_stop_fn = None logger.warning("projections compile failed: %s", e) @@ -1255,7 +1466,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): return tts.stop_head(tts.stop_actn(tts.stop_proj(lm_h))) @staticmethod - def _nullify_volatile_metadata(ctx: Any) -> Any: + def _nullify_volatile_metadata(ctx: _ForwardContextLike) -> _ForwardContextLike: """Set ``scheduler_metadata`` to None on all attention layers. This is the only tensor FA3 reallocates each step (variable shape). @@ -1351,6 +1562,31 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): ) return True + @staticmethod + def _build_unified_graph_bucket_sizes(max_size: int) -> frozenset[int]: + if max_size <= 0: + return frozenset() + sizes = {1} + size = 2 + while size <= min(max_size, 16): + sizes.add(size) + size <<= 1 + bucket = 24 + while bucket <= max_size: + sizes.add(bucket) + bucket += 8 + sizes.add(max_size) + return frozenset(size for size in sizes if 0 < size <= max_size) + + def _select_unified_graph_bucket_size(self, batch_size: int) -> int | None: + if ( + batch_size <= 0 + or batch_size > self._unified_decode_graph_max_batch_size + or batch_size > self._max_cached_graphs + ): + return None + return min((size for size in self._unified_graph_bucket_sizes if size >= batch_size), default=None) + def _capture_vae_graph(self, feat: torch.Tensor) -> _CapturedVAEGraph: batch_size = feat.shape[0] num_frames = feat.shape[-1] @@ -1535,8 +1771,297 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state.cfm_output_gpu.copy_(graph.output) return state.cfm_output_gpu.transpose(1, 2) + # -------------------- unified decode graph -------------------- + + def _capture_unified_decode_graph(self, batch_size: int) -> _CapturedUnifiedDecodeGraph: + H = self.config.hidden_size + D = self._feat_dim + P = self._patch_size + dev = torch.device(self._device) + dtype = self._side_dtype + + self._setup_cfm_buffers() + if self._enable_torch_compile: + self._setup_torch_compile() + self.model.precompute_fused_qkv() + self.residual_model.precompute_fused_qkv() + + tts = self.tts + dit_hidden = tts.lm_to_dit_proj.out_features + tts.res_to_dit_proj.out_features + cfm_buffers = _CFMBufferManager( + device=dev, + dtype=dtype, + feat_dim=D, + patch_size=P, + dit_hidden_size=dit_hidden, + max_batch_size=batch_size, + ) + + g = _CapturedUnifiedDecodeGraph( + graph=torch.cuda.CUDAGraph(), + batch_size=batch_size, + input_embeds=torch.zeros(batch_size, H, device=dev, dtype=dtype), + positions=torch.zeros(batch_size, device=dev, dtype=torch.long), + prev_feat_embed=torch.zeros(batch_size, H, device=dev, dtype=dtype), + prefix_feat_cond=torch.zeros(batch_size, P, D, device=dev, dtype=dtype), + next_feat_embed=torch.zeros(batch_size, H, device=dev, dtype=dtype), + cfm_output=torch.zeros(batch_size, D, P, device=dev, dtype=dtype), + lm_hidden=torch.zeros(batch_size, H, device=dev, dtype=dtype), + cfm_buffers=cfm_buffers, + cfm_noise=torch.zeros(batch_size, D, P, device=dev, dtype=dtype), + ) + + dit_proj = self._dit_proj_fn + + def unified_fwd(): + scaffold_h = self.model(input_ids=None, positions=g.positions, inputs_embeds=g.input_embeds) + if isinstance(scaffold_h, tuple): + scaffold_h = scaffold_h[0] + lm_h = tts.fsq_layer(scaffold_h) + res_input = tts.fusion_concat_proj(torch.cat([lm_h, g.prev_feat_embed], dim=-1)) + res_h = self.residual_model(positions=g.positions, inputs_embeds=res_input) + dit_h = dit_proj(lm_h, res_h) + cond = g.prefix_feat_cond.transpose(1, 2).contiguous() + cfm_out = _optimized_solve_euler_with_noise( + tts.feat_decoder, + dit_h, + P, + cond, + g.cfm_noise, + self._inference_timesteps, + self._cfg_value, + cfm_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + ) + feat_enc = tts.feat_encoder(cfm_out.transpose(1, 2).unsqueeze(1)).squeeze(1) + next_embed = tts.enc_to_lm_proj(feat_enc) + return next_embed, cfm_out, lm_h + + ctx = get_forward_context() + capture_context = override_forward_context(self._nullify_volatile_metadata(ctx)) + + original_estimator = tts.feat_decoder.estimator + capture_estimator = self._voxcpm2_compile_unified_capture_estimator(original_estimator) + tts.feat_decoder.estimator = capture_estimator + + original_feat_encoder = tts.feat_encoder + capture_feat_encoder = self._voxcpm2_compile_unified_capture_feat_encoder(original_feat_encoder) + tts.feat_encoder = capture_feat_encoder + + try: + with capture_context: + with torch.no_grad(): + for _ in range(3): + g.cfm_noise.normal_() + unified_fwd() + + g.cfm_noise.normal_() + with torch.cuda.graph(g.graph, pool=current_platform.get_global_graph_pool()): + g.next_feat_embed, g.cfm_output, g.lm_hidden = unified_fwd() + finally: + tts.feat_decoder.estimator = original_estimator + tts.feat_encoder = original_feat_encoder + self._unified_graph_stats.captures += 1 + logger.info( + "CUDA Graph captured for unified decode (batch_size=%d, captures=%d)", + batch_size, + self._unified_graph_stats.captures, + ) + return g + + def _unified_decode_graph_skip_reason( + self, + *, + can_use_graph: bool, + is_all_decode: bool, + num_reqs: int, + ) -> str | None: + if not self._enable_unified_decode_graph: + return "disabled" + if not can_use_graph: + return "graph_not_ready" + if not is_all_decode: + return "not_all_decode" + if num_reqs > self._unified_decode_graph_max_batch_size: + return "batch_exceeds_unified_max" + if num_reqs > self._max_cached_graphs: + return "batch_too_large" + graph_size = self._select_unified_graph_bucket_size(num_reqs) + if graph_size is None: + return "capture_policy" + if num_reqs > 1 and not self._runner_assisted_unified_decode_graph_active: + return "runner_full_metadata_missing" + for req_id, _, _, _ in self._pending_requests: + state = self._active_states[req_id] + if not state.prefill_completed: + return "prefill_incomplete" + if state.prev_feat_embed is None or state.curr_prefix_feat_cond is None: + return "state_not_ready" + return None + + def _maybe_log_unified_graph_stats(self) -> None: + if not self._enable_profiling: + return + stats = self._unified_graph_stats + total_skips = stats.total_skips + should_log = (stats.replays - stats.logged_replays >= 100) or (total_skips - stats.logged_skips >= 50) + if not should_log: + return + stats.logged_replays = stats.replays + stats.logged_skips = total_skips + logger.info( + "VoxCPM2 unified graph stats: captures=%d replays=%d skips=%s real_batch_sizes=%s graph_bucket_sizes=%s", + stats.captures, + stats.replays, + stats.skips, + dict(sorted(stats.real_batch_sizes.items())), + dict(sorted(stats.graph_bucket_sizes.items())), + ) + + def _forward_unified_decode( + self, + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + num_reqs: int, + ) -> torch.Tensor: + graph_size = self._select_unified_graph_bucket_size(num_reqs) + if graph_size is None: + raise RuntimeError(f"No unified decode graph bucket for batch size {num_reqs}") + g = self._unified_graphs.get(graph_size) + if g is None: + g = self._capture_unified_decode_graph(graph_size) + self._unified_graphs[graph_size] = g + + self._unified_graph_stats.replays += 1 + if self._enable_profiling: + self._unified_graph_stats.record_replay_batch(real_batch_size=num_reqs, graph_bucket_size=graph_size) + self._maybe_log_unified_graph_stats() + states: list[_RequestState] = [] + commit_mask: list[bool] = [] + for req_id, _is_prefill, _embeds, _n in self._pending_requests: + state = self._active_states[req_id] + states.append(state) + already_stopping = state.is_stopping + commit_mask.append(not already_stopping) + if not already_stopping: + state.decode_step_count += 1 + if state.decode_step_count >= self._max_decode_steps: + state.is_stopping = True + + self._perf.start("unified.copy_inputs") + for i, state in enumerate(states): + pfe = state.prev_feat_embed + g.prev_feat_embed[i].copy_(pfe.squeeze(0) if pfe.ndim > 1 else pfe) + pfc = state.curr_prefix_feat_cond + if pfc.ndim == 2: + g.prefix_feat_cond[i].copy_(pfc) + else: + g.prefix_feat_cond[i].copy_(pfc.squeeze(0)) + + g.input_embeds[:num_reqs].copy_(inputs_embeds[:num_reqs]) + g.positions[:num_reqs].copy_(positions[:num_reqs]) + if graph_size > num_reqs: + # Runner-assisted FULL metadata pads attention to graph_size. Keep + # padded rows as valid duplicates instead of position-0 zero rows; + # some attention backends may still touch padded slots before + # honoring seq_lens=0. + last = num_reqs - 1 + g.input_embeds[num_reqs:graph_size].copy_(g.input_embeds[last : last + 1].expand(graph_size - num_reqs, -1)) + g.positions[num_reqs:graph_size].copy_(g.positions[last : last + 1].expand(graph_size - num_reqs)) + g.prev_feat_embed[num_reqs:graph_size].copy_( + g.prev_feat_embed[last : last + 1].expand(graph_size - num_reqs, -1) + ) + g.prefix_feat_cond[num_reqs:graph_size].copy_( + g.prefix_feat_cond[last : last + 1].expand(graph_size - num_reqs, -1, -1) + ) + g.cfm_noise.normal_() + self._perf.stop("unified.copy_inputs") + + self._perf.start("unified.replay") + g.graph.replay() + self._perf.stop("unified.replay") + + self._perf.start("unified.commit") + with torch.no_grad(): + all_stop_logits = self._stop_fn(g.lm_hidden[:num_reqs]) + for i, state in enumerate(states): + stop_logits_i = all_stop_logits[i : i + 1] + if not commit_mask[i]: + state.precomputed_stop_logits = stop_logits_i + continue + next_embed_i = g.next_feat_embed[i : i + 1].clone() + cfm_out_i = g.cfm_output[i : i + 1].transpose(1, 2) + self._commit_decode_state(state, stop_logits_i, next_embed_i, cfm_out_i) + self._perf.stop("unified.commit") + + self._perf.start("unified.audio") + self._precompute_stop_flags_for_audio_collect(states) + ready_audio = self._drain_ready_audio_copies_for_states(states) + audio_by_req = self._collect_audio_batch(states, initial_delayed_chunks_by_req=ready_audio) + for state in states: + self._results_queue.append((state.request_id, state.precomputed_stop_logits)) + self._audio_queue.append((state.request_id, audio_by_req.get(state.request_id))) + + self._pending_requests.clear() + self._flush_deferred_cleanup() + self._perf.stop("unified.audio") + return g.next_feat_embed[:num_reqs] + # -------------------- vllm hooks -------------------- + def get_runner_assisted_full_attention_metadata_request( + self, + *, + req_ids: Sequence[str], + num_reqs: int, + num_scheduled_tokens: Sequence[int], + num_computed_tokens: Sequence[int], + max_num_scheduled_tokens: int, + ) -> RunnerAssistedFullAttentionMetadataRequest | None: + if not self._enable_unified_decode_graph or not self._enable_cuda_graph: + return None + if num_reqs <= 1 or num_reqs > self._unified_decode_graph_max_batch_size: + return None + if num_reqs > self._max_cached_graphs or self._select_unified_graph_bucket_size(num_reqs) is None: + return None + if max_num_scheduled_tokens != 1 or len(num_scheduled_tokens) != num_reqs: + return None + if any(int(n) != 1 for n in num_scheduled_tokens): + return None + if len(req_ids) != num_reqs or len(num_computed_tokens) != num_reqs: + return None + if any(int(n) <= 0 for n in num_computed_tokens): + return None + tts_compiled = getattr(self.tts.feat_decoder.estimator, "_compiled", False) if self._tts is not None else False + if not tts_compiled or self._cuda_graph_warmup_steps < self._cuda_graph_warmup_threshold: + return None + for req_id in req_ids: + state = self._active_states.get(req_id) + if ( + state is None + or not state.prefill_completed + or state.prev_feat_embed is None + or state.curr_prefix_feat_cond is None + ): + return None + bucket_size = self._select_unified_graph_bucket_size(num_reqs) + if bucket_size is None: + return None + return RunnerAssistedFullAttentionMetadataRequest( + num_reqs_padded=bucket_size, + for_cudagraph_capture=bucket_size not in self._unified_graphs, + ) + + def set_runner_assisted_full_attention_metadata_context( + self, + *, + enabled: bool, + num_reqs: int = 0, + ) -> None: + enabled = bool(enabled) + self._runner_assisted_unified_decode_graph_active = enabled + self._runner_assisted_unified_decode_graph_batch_size = int(num_reqs) if enabled else 0 + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: return self.model.embed_input_ids(input_ids) @@ -1565,12 +2090,30 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._enable_cuda_graph and graph_ready and intermediate_tensors is None and inputs_embeds is not None ) - if can_use_graph and is_all_decode and num_reqs <= self._max_cached_graphs: - use_scaffold_graph = self._should_use_decode_graph(num_reqs) - else: - use_scaffold_graph = False + unified_skip_reason = self._unified_decode_graph_skip_reason( + can_use_graph=can_use_graph, + is_all_decode=is_all_decode, + num_reqs=num_reqs, + ) + if unified_skip_reason is None: + self._perf.start("unified_decode") + result = self._forward_unified_decode(inputs_embeds, positions, num_reqs) + self._perf.stop("unified_decode") + self._perf.stop("forward_total") + return result + if self._enable_unified_decode_graph and num_reqs > 0 and unified_skip_reason is not None: + self._unified_graph_stats.record_skip(unified_skip_reason) + self._maybe_log_unified_graph_stats() - if use_scaffold_graph: + use_segmented_decode_graph = ( + can_use_graph + and is_all_decode + and num_reqs <= self._max_cached_graphs + and (not self._enable_unified_decode_graph or unified_skip_reason is not None) + and self._should_use_decode_graph(num_reqs) + ) + + if use_segmented_decode_graph: self._perf.start("scaffold_fwd") if num_reqs not in self._scaffold_graphs: self._scaffold_graphs[num_reqs] = self._capture_graph(self.model, num_reqs, "scaffold") @@ -1590,7 +2133,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): token_offset = 0 residual_inputs: list[torch.Tensor] = [] residual_positions: list[torch.Tensor] = [] - req_metas: list[tuple[_RequestState, bool, dict]] = [] + req_metas: list[tuple[_RequestState, bool, _PrefillResidualMeta | _DecodeResidualMeta]] = [] pending_decode_fsq: list[tuple[_RequestState, torch.Tensor, torch.Tensor]] = [] def flush_decode_fsq_batch() -> None: @@ -1656,6 +2199,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): residual_batch_size = batch_in.shape[0] use_residual_graph = ( self._enable_cuda_graph + and use_segmented_decode_graph and is_all_decode and graph_ready and residual_batch_size == num_reqs # 1 token per request @@ -1679,12 +2223,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): batch_out = self.residual_model(batch_pos, batch_in) self._perf.stop("residual_fwd") - can_batch_decode_tail = ( + can_finish_decode_batch = ( bool(req_metas) and all(not is_prefill for _, is_prefill, _ in req_metas) and all(x.shape[0] == 1 for x in residual_inputs) ) - if can_batch_decode_tail: + if can_finish_decode_batch: self._finish_decode_batch(req_metas, batch_out) self._precompute_stop_flags_for_audio_collect([state for state, _, _ in req_metas]) ready_audio_by_req = self._drain_ready_audio_copies_for_states([state for state, _, _ in req_metas]) @@ -1698,17 +2242,24 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): else: offset = 0 decoded_states: list[_RequestState] = [] + prefill_batch: list[tuple[_RequestState, _PrefillResidualMeta, torch.Tensor]] = [] for idx, (state, is_prefill, meta) in enumerate(req_metas): n = residual_inputs[idx].shape[0] res_out = batch_out[offset : offset + n] offset += n if is_prefill: - self._finish_prefill(state, meta, res_out, dev) + prefill_batch.append((state, meta, res_out)) else: - self._finish_decode(state, meta, res_out, dev) + self._finish_decode(state, meta, res_out) decoded_states.append(state) + if prefill_batch and self._runtime_config.enable_batched_prefill_tail: + self._finish_prefill_batch(prefill_batch) + else: + for state, meta, res_out in prefill_batch: + self._finish_prefill(state, meta, res_out, dev) + collect_states = [state for state, _, _ in req_metas] self._precompute_stop_flags_for_audio_collect(collect_states) ready_audio_by_req = self._drain_ready_audio_copies_for_states(collect_states) @@ -1731,7 +2282,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): # -------------------- prefill / decode helpers -------------------- - def _prepare_residual_prefill(self, state: _RequestState, base_lm_out: torch.Tensor, dev: Any): + def _prepare_residual_prefill( + self, + state: _RequestState, + base_lm_out: torch.Tensor, + dev: torch.device, + ) -> tuple[torch.Tensor, _PrefillResidualMeta]: tts = self.tts text_mask, feat_mask, feat, feat_embed = state.prefill_masks state.prefill_masks = None @@ -1754,10 +2310,15 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): lm_hidden = enc_outputs[:, -1, :] residual_input = tts.fusion_concat_proj(torch.cat([enc_outputs, feat_mask.unsqueeze(-1) * feat_embed], dim=-1)) - meta = {"lm_hidden": lm_hidden, "prefix_feat_cond": prefix_feat_cond} + meta = _PrefillResidualMeta(lm_hidden=lm_hidden, prefix_feat_cond=prefix_feat_cond) return residual_input.squeeze(0), meta - def _prepare_residual_decode(self, state: _RequestState, base_lm_out: torch.Tensor, dev: Any): + def _prepare_residual_decode( + self, + state: _RequestState, + base_lm_out: torch.Tensor, + dev: torch.device, + ) -> tuple[torch.Tensor, _DecodeResidualMeta]: tts = self.tts state.decode_step_count += 1 @@ -1774,14 +2335,14 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if prev.ndim == 1: prev = prev.unsqueeze(0) res_input = tts.fusion_concat_proj(torch.cat([lm_h, prev], dim=-1)) - return res_input, {"new_lm_hidden": lm_h} + return res_input, _DecodeResidualMeta(new_lm_hidden=lm_h) def _prepare_residual_decode_batch( self, states: list[_RequestState], base_lm_outs: list[torch.Tensor], - dev: Any, - ) -> list[tuple[torch.Tensor, dict]]: + dev: torch.device, + ) -> list[tuple[torch.Tensor, _DecodeResidualMeta]]: tts = self.tts for state in states: @@ -1811,7 +2372,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): return [ ( res_batch[i : i + 1], - {"new_lm_hidden": lm_h_batch[i : i + 1]}, + _DecodeResidualMeta(new_lm_hidden=lm_h_batch[i : i + 1]), ) for i in range(len(states)) ] @@ -1906,14 +2467,96 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): out.normal_(generator=gen) state.cfm_noise_step += 1 - def _finish_prefill(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any): + def _finish_prefill_batch(self, batch: list[tuple[_RequestState, _PrefillResidualMeta, torch.Tensor]]) -> None: + if len(batch) == 1: + state, meta, res_out = batch[0] + self._finish_prefill(state, meta, res_out, None) + return + + self._perf.start("prefill_tail_batch") tts = self.tts - lm_hidden = meta["lm_hidden"] - prefix_feat_cond = meta["prefix_feat_cond"] + self._setup_cfm_buffers() + if self._enable_torch_compile: + self._setup_torch_compile() + + dit_hs: list[torch.Tensor] = [] + conds: list[torch.Tensor] = [] + stop_logits_list: list[torch.Tensor] = [] + + for state, meta, res_out in batch: + lm_hidden = meta.lm_hidden + prefix_feat_cond = meta.prefix_feat_cond + residual_hidden = res_out[-1:, :] + stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach() + dit_h = torch.cat([tts.lm_to_dit_proj(lm_hidden), tts.res_to_dit_proj(residual_hidden)], dim=-1) + cond = prefix_feat_cond.transpose(1, 2).contiguous() + dit_hs.append(dit_h) + conds.append(cond) + stop_logits_list.append(stop_logits) + + batched_dit_h = torch.cat(dit_hs, dim=0) + batched_cond = torch.cat(conds, dim=0) + b = batched_dit_h.shape[0] + + if self._deterministic_cfm_noise and self._cfm_buffers is not None: + noise = self._cfm_buffers.noise[:b] + for i, (state, _, _) in enumerate(batch): + self._fill_deterministic_cfm_noise(state, noise[i : i + 1]) + else: + noise = torch.randn( + b, + self._feat_dim, + self._patch_size, + device=batched_dit_h.device, + dtype=batched_dit_h.dtype, + ) + pred_feats = _optimized_solve_euler_with_noise( + tts.feat_decoder, + batched_dit_h, + self._patch_size, + batched_cond, + noise, + self._inference_timesteps, + self._cfg_value, + self._cfm_buffers, + cfg_cutoff_ratio=self._cfg_cutoff_ratio, + ).transpose(1, 2) + + with torch.no_grad(): + curr_embeds = tts.enc_to_lm_proj(tts.feat_encoder(pred_feats.unsqueeze(1))).squeeze(1) + + for i, (state, meta, res_out) in enumerate(batch): + pred_feat_i = pred_feats[i : i + 1] + curr_embed_i = curr_embeds[i : i + 1] + self._commit_decode_state(state, stop_logits_list[i], curr_embed_i, pred_feat_i) + state.decode_step_count = 0 + state.request_start_time = time.perf_counter() + state.prefill_completed = True + + self._perf.stop("prefill_tail_batch") + if self._enable_profiling: + logger.info("PREFILL_BATCH[%d requests] tail breakdown:\n%s", len(batch), self._perf.breakdown()) + self._perf.reset() + + def _finish_prefill( + self, + state: _RequestState, + meta: _PrefillResidualMeta, + res_out: torch.Tensor, + dev: torch.device | None, + ) -> None: + self._perf.start("prefill_tail") + tts = self.tts + lm_hidden = meta.lm_hidden + prefix_feat_cond = meta.prefix_feat_cond residual_hidden = res_out[-1:, :] + self._perf.start("prefill.stop_fn") stop_logits = tts.stop_head(tts.stop_actn(tts.stop_proj(lm_hidden))).detach() + self._perf.stop("prefill.stop_fn") + self._perf.start("prefill.dit_proj") dit_h = torch.cat([tts.lm_to_dit_proj(lm_hidden), tts.res_to_dit_proj(residual_hidden)], dim=-1) + self._perf.stop("prefill.dit_proj") self._setup_cfm_buffers() if self._enable_torch_compile: @@ -1921,24 +2564,29 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): pred_feat = self._run_cfm_for_state(state, dit_h, prefix_feat_cond.transpose(1, 2).contiguous()) + self._perf.start("prefill.feat_encoder") with torch.no_grad(): curr_embed = tts.enc_to_lm_proj(tts.feat_encoder(pred_feat.unsqueeze(1))).squeeze(1) + self._perf.stop("prefill.feat_encoder") self._commit_decode_state(state, stop_logits, curr_embed, pred_feat) state.decode_step_count = 0 state.request_start_time = time.perf_counter() state.prefill_completed = True + self._perf.stop("prefill_tail") if logger.isEnabledFor(logging.DEBUG): # Only compute the norm (which forces a GPU->CPU sync) if we will log it. logger.debug("PREFILL[%s]: patch norm=%.4f", state.request_id, pred_feat.norm().item()) + if self._enable_profiling: + logger.info("PREFILL[%s] tail breakdown:\n%s", state.request_id, self._perf.breakdown()) self._perf.reset() - def _finish_decode(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any): + def _finish_decode(self, state: _RequestState, meta: _DecodeResidualMeta, res_out: torch.Tensor) -> None: self._perf.start("decode_step") tts = self.tts - lm_h = meta["new_lm_hidden"] + lm_h = meta.new_lm_hidden res_h = res_out.unsqueeze(0) if res_out.ndim == 1 else res_out dit_proj = getattr(self, "_compiled_dit_proj", None) or self._dit_proj_fn @@ -1959,7 +2607,11 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if self._enable_profiling and state.decode_step_count % 20 == 0: logger.info("Step %d[%s]:\n%s", state.decode_step_count, state.request_id, self._perf.breakdown()) - def _finish_decode_batch(self, req_metas: list[tuple], batch_out: torch.Tensor): + def _finish_decode_batch( + self, + req_metas: list[tuple[_RequestState, bool, _DecodeResidualMeta]], + batch_out: torch.Tensor, + ) -> None: self._perf.start("decode_step") tts = self.tts @@ -1968,7 +2620,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if self._enable_batched_cfm and not self._enable_cfm_cuda_graph: states = [state for state, _, _ in req_metas] - lm_h = torch.cat([meta["new_lm_hidden"] for _, _, meta in req_metas], dim=0) + lm_h = torch.cat([meta.new_lm_hidden for _, _, meta in req_metas], dim=0) pfc = torch.cat( [ state.curr_prefix_feat_cond.to(self._side_dtype).unsqueeze(0) @@ -2016,7 +2668,7 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): return for i, (state, _, meta) in enumerate(req_metas): - lm_h = meta["new_lm_hidden"] + lm_h = meta.new_lm_hidden res_h = batch_out[i : i + 1] pfc = state.curr_prefix_feat_cond.to(self._side_dtype) if pfc.ndim == 2: @@ -2474,7 +3126,10 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): # -------------------- preprocess / postprocess -------------------- def preprocess( - self, input_ids: torch.Tensor, input_embeds: torch.Tensor | None, **info_dict: Any + self, + input_ids: torch.Tensor, + input_embeds: torch.Tensor | None, + **info_dict: Unpack[VoxCPM2PreprocessInput], ) -> tuple[torch.Tensor, torch.Tensor, dict[str, Any]]: additional = info_dict.get("additional_information") if isinstance(additional, dict): @@ -2570,34 +3225,28 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): ) if state.prompt_cache is None and (ref_audio or (prompt_audio and prompt_text)): - try: - state.prompt_cache = self._build_prompt_cache( - ref_audio=ref_audio, - prompt_audio=prompt_audio, - prompt_text=prompt_text, - ) - if ( - voice_name - and state.prompt_cache is not None - and state.prompt_cache.get("mode") == "reference" - and "ref_audio_feat" in state.prompt_cache - ): - _key = self._speaker_cache.make_cache_key( - voice_name, model_type="voxcpm2", created_at=_created_at - ) - self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()}) - logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name) - except Exception as e: - logger.warning("build_prompt_cache failed: %s; falling back to zero-shot", e) - state.prompt_cache = None + state.prompt_cache = self._build_prompt_cache( + ref_audio=ref_audio, + prompt_audio=prompt_audio, + prompt_text=prompt_text, + ) + if ( + voice_name + and state.prompt_cache is not None + and state.prompt_cache.get("mode") == "reference" + and "ref_audio_feat" in state.prompt_cache + ): + _key = self._speaker_cache.make_cache_key(voice_name, model_type="voxcpm2", created_at=_created_at) + self._speaker_cache.put(_key, {"ref_audio_feat": state.prompt_cache["ref_audio_feat"].cpu()}) + logger.debug("Speaker cache STORE for VoxCPM2 speaker '%s'", voice_name) inputs = self._build_prefill_inputs(token_ids, dev, req_id) tts = self.tts - feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs["audio_feat"])) - text_embed = self.model.embed_input_ids(inputs["text_token"].to(dev)) - text_mask, feat_mask = inputs["text_mask"], inputs["audio_mask"] + feat_embed = tts.enc_to_lm_proj(tts.feat_encoder(inputs.audio_feat)) + text_embed = self.model.embed_input_ids(inputs.text_token.to(dev)) + text_mask, feat_mask = inputs.text_mask, inputs.audio_mask embeds = (text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed).squeeze(0) - state.prefill_masks = (text_mask, feat_mask, inputs["audio_feat"], feat_embed) + state.prefill_masks = (text_mask, feat_mask, inputs.audio_feat, feat_embed) else: state = self._active_states.get(req_id) curr = state.curr_embed_for_next if state else None @@ -2609,23 +3258,31 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._pending_requests.append((req_id, is_prefill, embeds, span_len)) return input_ids, embeds, {} - def postprocess(self, hidden_states: torch.Tensor, **info: Any) -> dict[str, Any]: + def postprocess(self, hidden_states: torch.Tensor, **info: Unpack[VoxCPM2PostprocessInput]) -> dict[str, Any]: req_id = info.get("request_id", self._current_request_id or "default") if self._enable_profiling: state = self._active_states.get(req_id) if state and state.decode_step_count > 0: logger.info( - "REQUEST DONE[%s]: %d steps, %.2fs\n%s", + "REQUEST DONE[%s]: %d steps, %.2fs\n%s\nUnified graph: captures=%d replays=%d skips=%s", req_id, state.decode_step_count, time.perf_counter() - state.request_start_time, self._perf.breakdown(), + self._unified_graph_stats.captures, + self._unified_graph_stats.replays, + self._unified_graph_stats.skips, ) return {} # -------------------- build prefill inputs -------------------- - def _build_prefill_inputs(self, token_ids: list[int], dev: Any, req_id: str = "default") -> dict: + def _build_prefill_inputs( + self, + token_ids: list[int], + dev: torch.device, + req_id: str = "default", + ) -> _PrefillInputs: tts = self.tts dtype = self._side_dtype state = self._active_states.get(req_id) @@ -2673,12 +3330,12 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): text_mask = torch.cat([rtm.cpu(), ones_t, zeros_p]) audio_mask = torch.cat([ram.cpu(), zeros_t, ones_p]) - return { - "text_token": text_token.unsqueeze(0).to(dev), - "audio_feat": audio_feat.unsqueeze(0).to(dev).to(dtype), - "text_mask": text_mask.unsqueeze(0).to(dev), - "audio_mask": audio_mask.unsqueeze(0).to(dev), - } + return _PrefillInputs( + text_token=text_token.unsqueeze(0).to(dev), + audio_feat=audio_feat.unsqueeze(0).to(dev).to(dtype), + text_mask=text_mask.unsqueeze(0).to(dev), + audio_mask=audio_mask.unsqueeze(0).to(dev), + ) # -------------------- weight loading -------------------- diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index b469234c1..43d2c9d19 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -621,8 +621,8 @@ class NPUARModelRunner(OmniNPUModelRunner): positions=positions, inputs_embeds=inputs_embeds, req_ids=req_ids[:num_reqs], - num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)], - num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)], + num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs], + num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs], input_ids_buffer=self.input_ids.gpu[:num_tokens_padded], ) # -------------------------------------- Omni-new ------------------------------------------------- @@ -690,8 +690,9 @@ class NPUARModelRunner(OmniNPUModelRunner): with record_function_or_nullcontext("post process"): # -------------------------------------- Omni-new ------------------------------------------------- # [Omni] Map pending ropes metadata to req_ids. - if hasattr(self.model, "flush_pending_metadata"): - self.model.flush_pending_metadata(list(req_ids)) + flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None) + if callable(flush_pending_metadata): + flush_pending_metadata(req_ids[:num_reqs]) hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index c1812b8c9..816246410 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -8,7 +8,7 @@ from __future__ import annotations import gc import threading -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from contextlib import nullcontext from copy import copy from dataclasses import replace @@ -46,6 +46,7 @@ from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin +from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest logger = init_logger(__name__) @@ -588,6 +589,93 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): model = getattr(self, "model", None) return bool(getattr(model, "requires_full_prefix_cached_hidden_states", True)) + def _get_runner_assisted_full_attention_metadata_request( + self, + *, + req_ids: Sequence[str], + num_reqs: int, + num_reqs_padded: int, + num_scheduled_tokens_np: np.ndarray, + num_computed_tokens_cpu: torch.Tensor | np.ndarray, + max_num_scheduled_tokens: int, + ) -> RunnerAssistedFullAttentionMetadataRequest | None: + # Models without this hook keep the normal runner path. VoxCPM2 uses + # it to ask the runner for padded FULL attention metadata while keeping + # graph policy and lifecycle in the model layer. + hook = getattr(self.model, "get_runner_assisted_full_attention_metadata_request", None) + if not callable(hook): + return None + num_computed_tokens = num_computed_tokens_cpu + if hasattr(num_computed_tokens, "numpy"): + num_computed_tokens = num_computed_tokens.numpy() + request = hook( + req_ids=req_ids, + num_reqs=num_reqs, + num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs], + num_computed_tokens=num_computed_tokens, + max_num_scheduled_tokens=max_num_scheduled_tokens, + ) + if request is None: + return None + if not isinstance(request, RunnerAssistedFullAttentionMetadataRequest): + raise TypeError( + "runner-assisted full attention metadata hook must return " + "RunnerAssistedFullAttentionMetadataRequest or None, got " + f"{type(request).__name__}" + ) + padded_num_reqs = max( + num_reqs_padded, + min(int(request.num_reqs_padded), self.scheduler_config.max_num_seqs), + ) + return RunnerAssistedFullAttentionMetadataRequest( + num_reqs_padded=padded_num_reqs, + for_cudagraph_capture=bool(request.for_cudagraph_capture), + ) + + def _refresh_runner_assisted_full_attention_metadata_buffers( + self, + *, + num_reqs: int, + num_reqs_padded: int, + num_scheduled_tokens_np: np.ndarray, + ) -> None: + num_computed_tokens = self.input_batch.num_computed_tokens_cpu[:num_reqs] + if hasattr(num_computed_tokens, "numpy"): + num_computed_tokens = num_computed_tokens.numpy() + np.add( + num_computed_tokens, + num_scheduled_tokens_np, + out=self.optimistic_seq_lens_cpu[:num_reqs].numpy(), + ) + self.optimistic_seq_lens_cpu[num_reqs:num_reqs_padded].fill_(0) + self.seq_lens.copy_(self.optimistic_seq_lens_cpu, non_blocking=True) + + cum_num_tokens = self._get_cumsum_and_arange( + num_scheduled_tokens_np, + self.query_pos.np, + ) + self.query_start_loc.np[0] = 0 + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1]) + self.query_start_loc.copy_to_gpu() + + self.input_batch.block_table.commit_block_table(num_reqs_padded) + + def _set_runner_assisted_full_attention_metadata_context( + self, + *, + enabled: bool, + num_reqs: int = 0, + ) -> bool: + hook = getattr(self.model, "set_runner_assisted_full_attention_metadata_context", None) + if not callable(hook): + return False + hook( + enabled=enabled, + num_reqs=num_reqs, + ) + return True + def _deferred_prefix_cache_mm_keys(self) -> set[str]: """Model-declared multimodal keys whose prefix-cache writes are deferred.""" model = getattr(self, "model", None) @@ -1028,6 +1116,26 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): ) num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + num_computed_tokens_cpu = self.input_batch.num_computed_tokens_cpu[:num_reqs] + + runner_assisted_full_attn_request = self._get_runner_assisted_full_attention_metadata_request( + req_ids=req_ids[:num_reqs], + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + num_scheduled_tokens_np=num_scheduled_tokens_np, + num_computed_tokens_cpu=num_computed_tokens_cpu, + max_num_scheduled_tokens=max_num_scheduled_tokens, + ) + runner_assisted_full_attn_capture = False + if runner_assisted_full_attn_request is not None: + num_reqs_padded = runner_assisted_full_attn_request.num_reqs_padded + runner_assisted_full_attn_capture = runner_assisted_full_attn_request.for_cudagraph_capture + num_tokens_padded = max( + num_tokens_padded, + num_reqs_padded * max_num_scheduled_tokens, + ) + runner_assisted_full_attn = runner_assisted_full_attn_request is not None + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( should_ubatch, num_scheduled_tokens_np, @@ -1035,8 +1143,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): num_reqs_padded, self.parallel_config.num_ubatches, ) - - pad_attn = cudagraph_mode == CUDAGraphMode.FULL + pad_attn = cudagraph_mode == CUDAGraphMode.FULL or runner_assisted_full_attn use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices @@ -1058,11 +1165,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): num_tokens_unpadded=num_tokens_unpadded, ubatch_slices=ubatch_slices_padded, ) + if runner_assisted_full_attn: + self._refresh_runner_assisted_full_attention_metadata_buffers( + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded, + num_scheduled_tokens_np=num_scheduled_tokens_np, + ) + attn_num_reqs = num_reqs_padded if runner_assisted_full_attn else num_reqs attn_metadata, spec_decode_common_attn_metadata = self._build_attention_metadata( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, - num_reqs=num_reqs, + num_reqs=attn_num_reqs, num_reqs_padded=num_reqs_padded if pad_attn else None, max_query_len=max_num_scheduled_tokens, ubatch_slices=ubatch_slices_attn, @@ -1071,6 +1185,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): num_scheduled_tokens=scheduler_output.num_scheduled_tokens, cascade_attn_prefix_lens=cascade_attn_prefix_lens, slot_mappings=slot_mappings_by_group, + for_cudagraph_capture=runner_assisted_full_attn_capture, ) self._maybe_attach_attention_metadata_extensions( attn_metadata=attn_metadata, @@ -1079,6 +1194,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): max_query_len=max_num_scheduled_tokens, pad_attn=pad_attn, num_scheduled_tokens_np=num_scheduled_tokens_np, + for_cudagraph_capture=runner_assisted_full_attn_capture, ) ( @@ -1092,14 +1208,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # Let the model adjust inputs before forward (e.g. restore input_ids # for multimodal position detection, fix decode position offsets). - if hasattr(self.model, "prepare_runner_inputs"): - input_ids, positions = self.model.prepare_runner_inputs( + prepare_runner_inputs = getattr(self.model, "prepare_runner_inputs", None) + if callable(prepare_runner_inputs): + input_ids, positions = prepare_runner_inputs( input_ids=input_ids, positions=positions, inputs_embeds=inputs_embeds, req_ids=req_ids[:num_reqs], - num_computed_tokens=[int(self.input_batch.num_computed_tokens_cpu[i]) for i in range(num_reqs)], - num_scheduled_tokens=[int(num_scheduled_tokens_np[i]) for i in range(num_reqs)], + num_computed_tokens=self.input_batch.num_computed_tokens_cpu[:num_reqs], + num_scheduled_tokens=num_scheduled_tokens_np[:num_reqs], input_ids_buffer=self.input_ids.gpu[:num_tokens_padded], ) @@ -1108,46 +1225,60 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): # with CUDA graph capture. if self.calculate_kv_scales: cudagraph_mode = CUDAGraphMode.NONE + runner_assisted_full_attn = False # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False + runner_assisted_context_enabled = False + if runner_assisted_full_attn: + runner_assisted_context_enabled = self._set_runner_assisted_full_attention_metadata_context( + enabled=True, + num_reqs=num_reqs, + ) + runner_assisted_full_attn = runner_assisted_context_enabled + # Run the model. # Use persistent buffers for CUDA graphs. # When spec decode is enabled, defer connector finalization # (wait_for_save + clear metadata) until after draft model runs. defer_kv_connector_finalize = self.speculative_config is not None - with ( - nullcontext(), - set_forward_context( - attn_metadata, - self.vllm_config, - num_tokens=num_tokens_padded, - num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_mode, - batch_descriptor=batch_desc, - ubatch_slices=ubatch_slices_padded, - slot_mapping=slot_mappings, # OMNI: required for KV cache operations - ), - record_function_or_nullcontext("gpu_model_runner: forward"), - self.maybe_get_kv_connector_output( - scheduler_output, - defer_finalize=defer_kv_connector_finalize, - ) as kv_connector_output, - ): - model_output = self._model_forward( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - sampling_metadata=self.input_batch.sampling_metadata, - logits_index=logits_indices, - sampler=self.sampler, - ) + try: + with ( + nullcontext(), + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_padded, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=(CUDAGraphMode.FULL if runner_assisted_full_attn else cudagraph_mode), + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices_padded, + slot_mapping=slot_mappings, # OMNI: required for KV cache operations + ), + record_function_or_nullcontext("gpu_model_runner: forward"), + self.maybe_get_kv_connector_output( + scheduler_output, + defer_finalize=defer_kv_connector_finalize, + ) as kv_connector_output, + ): + model_output = self._model_forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + sampling_metadata=self.input_batch.sampling_metadata, + logits_index=logits_indices, + sampler=self.sampler, + ) - # [Omni] Map pending ropes metadata to req_ids. - if hasattr(self.model, "flush_pending_metadata"): - self.model.flush_pending_metadata(list(req_ids)) + # [Omni] Map pending ropes metadata to req_ids. + flush_pending_metadata = getattr(self.model, "flush_pending_metadata", None) + if callable(flush_pending_metadata): + flush_pending_metadata(req_ids[:num_reqs]) + finally: + if runner_assisted_context_enabled: + self._set_runner_assisted_full_attention_metadata_context(enabled=False) with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: diff --git a/vllm_omni/worker/runner_assisted_metadata.py b/vllm_omni/worker/runner_assisted_metadata.py new file mode 100644 index 000000000..a71180a25 --- /dev/null +++ b/vllm_omni/worker/runner_assisted_metadata.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +from collections.abc import Sequence +from typing import NamedTuple, Protocol + + +class RunnerAssistedFullAttentionMetadataRequest(NamedTuple): + num_reqs_padded: int + for_cudagraph_capture: bool + + +class RunnerAssistedAttentionMetadataProvider(Protocol): + """Optional model hook for runner-built full attention metadata.""" + + def get_runner_assisted_full_attention_metadata_request( + self, + *, + req_ids: Sequence[str], + num_reqs: int, + num_scheduled_tokens: Sequence[int], + num_computed_tokens: Sequence[int], + max_num_scheduled_tokens: int, + ) -> RunnerAssistedFullAttentionMetadataRequest | None: ... + + def set_runner_assisted_full_attention_metadata_context( + self, + *, + enabled: bool, + num_reqs: int = 0, + ) -> None: ...