diff --git a/.buildkite/npu/test-npu-nightly.yml b/.buildkite/npu/test-npu-nightly.yml index b89f61a39..e049783ce 100644 --- a/.buildkite/npu/test-npu-nightly.yml +++ b/.buildkite/npu/test-npu-nightly.yml @@ -81,6 +81,55 @@ steps: EXIT=$$? buildkite-agent artifact upload "tests/e2e/accuracy/minicpmo_4_5/results/*.json" exit $$EXIT + - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test" + key: nightly-omni-performance-minicpmo-4-5-npu + timeout_in_minutes: 180 + mirror_hardwares: a3_npu_2 + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + HF_TOKEN: "${HF_TOKEN}" + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "npu and A3" --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + - label: ":full_moon: Omni · MiniCPM-o 4.5 · Perf Test (A2)" + key: nightly-omni-performance-minicpmo-4-5-npu-a2 + timeout_in_minutes: 180 + mirror_hardwares: a2b3_npu_1 + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + HF_TOKEN: "${HF_TOKEN}" + HF_HUB_OFFLINE: "1" + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "npu and A2" --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + - label: ":full_moon: Omni · MiniCPM-o 4.5 · Duplex Seed-TTS Perf Test" + key: nightly-omni-performance-minicpmo-4-5-duplex-seed-tts-npu + timeout_in_minutes: 180 + mirror_hardwares: a3_npu_2 + env: + VLLM_WORKER_MULTIPROC_METHOD: spawn + VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1" + HF_TOKEN: "${HF_TOKEN}" + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m "npu and A3" --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT - group: ":card_index_dividers: TTS Model Test" key: nightly-tts-test-group depends_on: upload-nightly-pipeline diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 41fc15df5..cacc8142f 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1127,9 +1127,11 @@ class TestDeployConfigLoading: assert all("Async" not in (stage.scheduler_cls or "") for stage in stages) assert [stage.devices for stage in deploy.stages] == ["0", "0", "0"] assert deploy.stages[1].enforce_eager is False - assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 2048 + assert stages[1].yaml_extras["default_sampling_params"]["max_tokens"] == 4096 assert stages[1].yaml_extras["default_sampling_params"]["min_tokens"] == 0 - assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [1] + assert stages[1].yaml_extras["default_sampling_params"]["temperature"] == 0.8 + assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [6561] + assert "codec_sampling_params" not in stages[1].yaml_engine_args @pytest.mark.parametrize( ("filename", "stage0_devices", "stage1_devices", "stage2_devices", "stage1_replicas"), @@ -2111,6 +2113,23 @@ class TestMingFlashOmniPipeline: class TestBaseConfigInheritance: """Test deploy YAML base_config inheritance.""" + def test_minicpmo_overlays_inherit_talker_sampling_params(self): + """Overlays must keep the base Talker codec Sampler knobs.""" + for filename in ( + "minicpmo_4_5_duplex.yaml", + "minicpmo_4_5_3gpu_stage1_replicas.yaml", + "minicpmo_4_5_4gpu_stage1_replicas.yaml", + "minicpmo_4_5_8x4090_stage1_replicas.yaml", + ): + path = Path(get_deploy_config_path(filename)) + if not path.exists(): + pytest.skip(f"{filename} not found") + deploy = load_deploy_config(path) + sampling = deploy.stages[1].default_sampling_params + assert sampling is not None, f"{filename} stage 1 lost default_sampling_params" + assert sampling["temperature"] == 0.8, filename + assert sampling["top_k"] == 100, filename + def test_ci_inherits_from_main(self): ci_path = Path(get_deploy_config_path("ci/qwen3_omni_moe.yaml")) if not ci_path.exists(): diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5.json b/tests/dfx/perf/tests/test_minicpmo_4_5.json index cfb57169b..9296cf5b3 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5.json @@ -39,6 +39,7 @@ 8 ], "no_oversample": true, + "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", "extra_body": { @@ -79,9 +80,92 @@ 0.8568, 1.5425 ] + }, + "A3": { + "request_throughput": [ + 0.5383, + 0.6042, + 0.7547 + ], + "mean_ttft_ms": [ + 333.2633, + 533.8656, + 547.3388 + ], + "mean_e2el_ms": [ + 1857.2154, + 6600.5101, + 10450.6811 + ], + "mean_audio_ttfp_ms": [ + 986.4666, + 3411.0754, + 3352.7515 + ], + "mean_audio_rtf": [ + 0.4423, + 1.5734, + 2.3024 + ] } } } ] + }, + { + "test_name": "test_minicpmo_4_5_a2", + "mark": [ + { + "hardware_marks": { + "res": { + "npu": "A2" + }, + "num_cards": 1 + } + }, + "full_model", + "omni" + ], + "server_params": { + "model": "openbmb/MiniCPM-o-4_5", + "extra_cli_args": [ + "--deploy-config", + "vllm_omni/deploy/minicpmo_4_5.yaml", + "--trust-remote-code" + ] + }, + "benchmark_params": [ + { + "dataset_name": "seed-tts", + "dataset_path": "zhaochenyang20/seed-tts-eval", + "backend": "openai-chat-omni", + "endpoint": "/v1/chat/completions", + "num_prompts": [ + 32, + 64, + 128 + ], + "max_concurrency": [ + 1, + 4, + 8 + ], + "no_oversample": true, + "disable_shuffle": true, + "trust_remote_code": true, + "seed_tts_locale": "en", + "extra_body": { + "modalities": [ + "text", + "audio" + ], + "chat_template_kwargs": { + "enable_thinking": false, + "use_tts_template": true + } + }, + "percentile-metrics": "ttft,e2el,audio_ttfp,audio_rtf,audio_duration" + } + ] } ] diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json index cfbcc3abb..6ef27680c 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_duplex_seed_tts.json @@ -35,6 +35,7 @@ 1 ], "no_oversample": true, + "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", "seed_tts_turns_per_session": 4, @@ -60,6 +61,23 @@ "mean_audio_rtf": [ 0.46163 ] + }, + "A3": { + "request_throughput": [ + 0.1594 + ], + "mean_ttft_ms": [ + 274.164 + ], + "mean_e2el_ms": [ + 6250.7696 + ], + "mean_audio_ttfp_ms": [ + 994.422 + ], + "mean_audio_rtf": [ + 0.46163 + ] } } } diff --git a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json index e327be510..15d52e6ad 100644 --- a/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json +++ b/tests/dfx/perf/tests/test_minicpmo_4_5_ready.json @@ -37,6 +37,7 @@ 4 ], "no_oversample": true, + "disable_shuffle": true, "trust_remote_code": true, "seed_tts_locale": "en", "extra_body": { diff --git a/tests/e2e/online_serving/test_minicpmo_4_5.py b/tests/e2e/online_serving/test_minicpmo_4_5.py index 6953ef291..9c8bcd908 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5.py @@ -65,6 +65,15 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) +# Close and emit <|tts_bos|> so Talker speaks the answer, not reasoning. +_TTS_EXTRA_BODY = { + "chat_template_kwargs": { + "use_tts_template": True, + "enable_thinking": False, + } +} + + @pytest.mark.core_model @pytest.mark.advanced_model @pytest.mark.omni @@ -111,6 +120,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None: "messages": messages, "stream": True, "key_words": {"audio": ["Beijing"]}, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config) @@ -141,12 +151,7 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None: "stream": True, "key_words": {"text": ["Beijing"]}, "modalities": ["text", "audio"], - "extra_body": { - "chat_template_kwargs": { - "use_tts_template": True, - "enable_thinking": False, - } - }, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @@ -175,6 +180,7 @@ def test_image_to_text_audio_001(omni_server, openai_client) -> None: "model": omni_server.model, "messages": messages, "stream": True, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @@ -203,6 +209,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None: "model": omni_server.model, "messages": messages, "stream": True, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) @@ -236,6 +243,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: "model": omni_server.model, "messages": messages, "stream": True, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) diff --git a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py index e84985e3e..5a21e8692 100644 --- a/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py +++ b/tests/e2e/online_serving/test_minicpmo_4_5_expansion.py @@ -1,7 +1,7 @@ """E2E online expansion tests for MiniCPM-o 4.5. These cover modality combinations, separate Code2Wav behavior, sequential -request isolation, longer audio output, and Chinese speech. +request isolation, longer audio output, long-form generation, and Chinese speech. Like Qwen3-Omni expansion, ``default`` exercises ``--no-async-chunk`` and ``async_chunk`` exercises ``--async-chunk``. @@ -75,6 +75,15 @@ def get_max_batch_size(size_type="few"): return batch_sizes.get(size_type, 5) +# Close and emit <|tts_bos|> so Talker speaks the answer, not reasoning. +_TTS_EXTRA_BODY = { + "chat_template_kwargs": { + "use_tts_template": True, + "enable_thinking": False, + } +} + + @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_video_to_text_001(omni_server, openai_client) -> None: @@ -127,6 +136,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None: "model": omni_server.model, "messages": messages_1, "stream": True, + "extra_body": _TTS_EXTRA_BODY, }, request_num=1, ) @@ -138,6 +148,7 @@ def test_sequential_requests_independent(omni_server, openai_client) -> None: "messages": messages_2, "stream": True, "key_words": {"text": ["Beijing"]}, + "extra_body": _TTS_EXTRA_BODY, }, request_num=1, ) @@ -164,20 +175,39 @@ def test_text_to_audio_long_output_001(omni_server, openai_client) -> None: "messages": messages, "modalities": ["text", "audio"], "stream": True, - # Delimit the assistant answer with the TTS template so the Talker - # does not spend its codec-token budget speaking hidden reasoning. - "extra_body": { - "chat_template_kwargs": { - "use_tts_template": True, - "enable_thinking": False, - } - }, + "extra_body": _TTS_EXTRA_BODY, "key_words": {"audio": ["Beijing"]}, } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) +@hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_audio_long_form_001(omni_server, openai_client) -> None: + """ + Input Modal: text only (long-form generation prompt). + Output Modal: text, audio (default ``modalities``); + Input Setting: stream=True + Datasets: single request + """ + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + content_text="帮我讲一个300字的故事.", + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "extra_body": _TTS_EXTRA_BODY, + } + responses = openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + text = responses[0].text_content if responses else "" + word_count = len(text.split()) + assert word_count >= 200, f"Expected at least 200 words in long output, got {word_count}" + + @hardware_test(res={"cuda": "H100", "npu": "A3"}, num_cards=1) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_chinese_text_to_audio(omni_server, openai_client) -> None: @@ -193,5 +223,6 @@ def test_chinese_text_to_audio(omni_server, openai_client) -> None: "messages": messages, "stream": True, "key_words": {"text": ["北京"]}, + "extra_body": _TTS_EXTRA_BODY, } openai_client.send_omni_request(request_config) diff --git a/tests/engine/test_stage_request_token_budget.py b/tests/engine/test_stage_request_token_budget.py new file mode 100644 index 000000000..045dfd31f --- /dev/null +++ b/tests/engine/test_stage_request_token_budget.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Token-budget clamping when the orchestrator builds a stage request.""" + +from __future__ import annotations + +import pytest +from vllm.sampling_params import SamplingParams + +from vllm_omni.engine.orchestrator import build_engine_core_request_from_tokens + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _StageModelConfig: + """Only the field ``build_engine_core_request_from_tokens`` reads.""" + + max_model_len = 4096 + + +class TestStageRequestTokenBudget: + def test_min_tokens_is_clamped_to_the_remaining_context(self): + # A min_tokens wider than the leftover window makes ``check_stop`` + # return early forever: the request can neither stop nor be granted + # more tokens, so it stalls instead of being length-capped. + params = SamplingParams(max_tokens=4096, min_tokens=50) + + request = build_engine_core_request_from_tokens( + request_id="req-clamped", + prompt={"prompt_token_ids": [0] * 4050}, + params=params, + model_config=_StageModelConfig(), + ) + + assert request.sampling_params.min_tokens == 46 + assert params.min_tokens == 50 + + def test_min_tokens_survives_when_the_prompt_leaves_room(self): + params = SamplingParams(max_tokens=4096, min_tokens=50) + + request = build_engine_core_request_from_tokens( + request_id="req-untouched", + prompt={"prompt_token_ids": [0] * 68}, + params=params, + model_config=_StageModelConfig(), + ) + + assert request.sampling_params.min_tokens == 50 + assert request.sampling_params.max_tokens == 4096 + + def test_min_tokens_floors_at_zero_when_the_prompt_fills_the_context(self): + params = SamplingParams(max_tokens=16, min_tokens=8) + + request = build_engine_core_request_from_tokens( + request_id="req-full", + prompt={"prompt_token_ids": [0] * 4096}, + params=params, + model_config=_StageModelConfig(), + ) + + assert request.sampling_params.min_tokens == 0 + + def test_absent_max_tokens_still_defaults_to_the_remaining_context(self): + params = SamplingParams(min_tokens=0) + params.max_tokens = None + + request = build_engine_core_request_from_tokens( + request_id="req-default", + prompt={"prompt_token_ids": [0] * 96}, + params=params, + model_config=_StageModelConfig(), + ) + + assert request.sampling_params.max_tokens == 4000 diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index ba3c7292b..a1ef3b2ec 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py @@ -7,6 +7,8 @@ import torch.nn as nn from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( BatchedToken2Wav, + plan_token2wav_encode_slices, + relpos_encode_token_budget, ) from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import ( MiniCPMO45Code2Wav, @@ -217,6 +219,43 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches(): assert cache1[0, 0, 0, 0, 0].item() == 20 +def test_relpos_encode_token_budget_leaves_room_for_upsample_and_cache(): + # The NPU crash was 6968 vs 985: 3333 tokens * 2 + cache overflowed max_pos=5000. + assert relpos_encode_token_budget(max_pos=5000, stride=2, cache_offset=150, lookahead=3) <= 1024 + assert relpos_encode_token_budget(max_pos=5000, stride=2, cache_offset=150, lookahead=3, cap=3000) == 2346 + assert plan_token2wav_encode_slices(3333, max_frames=1024, min_nonfinal=4, last_chunk=True) == [ + (0, 1024), + (1024, 2048), + (2048, 3072), + (3072, 3333), + ] + assert plan_token2wav_encode_slices(50, max_frames=40, min_nonfinal=20, last_chunk=False) == [ + (0, 30), + (30, 50), + ] + + +def test_decode_batch_splits_overlong_prefill_inside_relpos_budget(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + states = adapter.setup_batch(prompt, 1) + setup_encodes = len(token2wav.flow.encoder.calls) + adapter._max_encode_token_frames = lambda _states: 4 + + audios, _ = adapter.decode_batch( + torch.arange(10, dtype=torch.long).reshape(1, -1), + prompt, + states, + last_chunk=True, + ) + + assert token2wav.flow.encoder.calls[setup_encodes:] == [1, 1, 1] + assert token2wav.flow.encoder.last_chunk_calls[setup_encodes:] == [False, False, True] + assert len(audios) == 1 + assert audios[0].numel() > 0 + + def test_fade_in_out_limits_overlap_to_available_previous_audio(): speech = torch.arange(6, dtype=torch.float32).reshape(1, -1) previous = torch.full((1, 3), 2.0) diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index bbc011284..42bc364cb 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -105,6 +105,10 @@ class TestPipelineTopology: assert talker.engine_output_type == "latent" # scope KV cache / mrope sizing to talker sub-config assert talker.hf_config_name == "tts_config" + assert talker.sampling_constraints == { + "detokenize": False, + "stop_token_ids": [6561], + } assert talker.custom_process_next_stage_input_func == ( "vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.tts2code2wav_full_payload" ) diff --git a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py index 8a7593623..4115c1a29 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_talker_batching.py @@ -4,6 +4,8 @@ from __future__ import annotations +from dataclasses import dataclass + import pytest import torch import torch.nn as nn @@ -11,18 +13,30 @@ import torch.nn as nn from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni import ( MiniCPMO45OmniForConditionalGeneration, ) +from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_llm import ( + ConditionalChatTTSConfig, +) from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts import ( + _DUPLEX_CODEC_TOKENS_PER_CHUNK, _REPETITION_PENALTY_CHUNK_SIZE, MiniCPMO45OmniTTSForConditionalGeneration, _apply_batched_repetition_penalty, - _max_audio_tokens, + _native_duplex_chunk_budget, _restore_weight_norm_weight, + blank_scheduler_prompt_for_penalties, ) from vllm_omni.utils.mm_outputs import to_payload_element pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +@dataclass +class _PenaltyMetadata: + """Stand-in for the penalty-relevant slice of vLLM's SamplingMetadata.""" + + prompt_token_ids: torch.Tensor | None + + class _FakeNativeTalker(nn.Module): has_preprocess = True @@ -35,6 +49,12 @@ class _FakeNativeTalker(nn.Module): return torch.ones(2, 4) +def test_tts_config_exposes_codec_vocab_to_vllm() -> None: + cfg = ConditionalChatTTSConfig(num_audio_tokens=6562) + assert cfg.vocab_size == 6562 + assert cfg.eos_token_id == 6561 + + def test_wrapper_always_delegates_talker_to_native_ar_path() -> None: model = MiniCPMO45OmniForConditionalGeneration.__new__(MiniCPMO45OmniForConditionalGeneration) nn.Module.__init__(model) @@ -55,12 +75,14 @@ def _make_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: talker = MiniCPMO45OmniTTSForConditionalGeneration.__new__(MiniCPMO45OmniTTSForConditionalGeneration) nn.Module.__init__(talker) talker._num_audio_tokens = 8 - talker._batch_stop_logits = None - talker._request_generators = {} + talker._codec_eos_id = 7 + talker._force_eos_rows = None + talker._mask_eos_rows = None talker._request_audio_states = {} talker._deferred_cleanup_ids = set() - talker._codec_min_tokens = 50 - talker._codec_seed = 42 + talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) + with torch.no_grad(): + talker.head_code[0].weight.copy_(torch.eye(8, 2)) return talker @@ -90,46 +112,6 @@ def _reference_repetition_penalty( return torch.where(logits < 0, logits * alpha, logits / alpha) -def _make_sampling_talker() -> MiniCPMO45OmniTTSForConditionalGeneration: - talker = _make_talker() - talker._codec_temperature = 0.8 - talker._codec_top_k = 5 - talker._codec_top_p = 0.85 - talker._codec_repetition_penalty = 1.05 - talker._request_audio_states = { - "req-a": {"min_tokens": 2}, - "req-b": {"min_tokens": 0}, - } - talker.head_code = nn.ModuleList([nn.Linear(2, 8, bias=False)]) - with torch.no_grad(): - talker.head_code[0].weight.copy_( - torch.tensor( - [ - [0.1, -0.3], - [0.2, 0.4], - [-0.5, 0.2], - [0.7, -0.1], - [-0.2, 0.8], - [0.6, 0.3], - [-0.4, -0.7], - [0.3, 0.5], - ] - ) - ) - return talker - - -@pytest.mark.parametrize( - ("condition_tokens", "expected"), - [(3, 64), (100, 1000), (1000, 2048)], -) -def test_audio_token_limit_scales_with_condition_length( - condition_tokens: int, - expected: int, -) -> None: - assert _max_audio_tokens(condition_tokens) == expected - - def test_weight_norm_restore_matches_checkpoint_parametrization_in_bfloat16() -> None: generator = torch.Generator().manual_seed(42) weight_v = torch.randn(8, 16, generator=generator, dtype=torch.bfloat16) @@ -220,133 +202,212 @@ def test_batched_repetition_penalty_matches_rows_across_chunks(mocker) -> None: ] -def test_batched_codec_sampling_is_request_order_independent() -> None: - talker = _make_sampling_talker() - histories = [torch.tensor([1, 1, 2]), torch.tensor([3, 4, 4])] - hidden = torch.tensor([[0.25, -0.75], [-0.5, 0.5]]) +def test_scheduler_prompt_is_fully_blanked_for_penalties() -> None: + # The handoff path fills the prompt with 0s and the non-handoff path passes + # thinker token ids; neither is codec history, so every position is padded. + prompt = torch.tensor([[0, 0, 0], [12, 0, 340]]) - forward_order = talker._sample_audio_codes( - hidden, - histories, - ["req-a", "req-b"], - [0, 0], - ) - talker._request_generators = {} - reverse_order = talker._sample_audio_codes( - hidden.flip(0), - list(reversed(histories)), - ["req-b", "req-a"], - [0, 0], + blanked = blank_scheduler_prompt_for_penalties(prompt, vocab_size=8) + + assert blanked.tolist() == [[8, 8, 8], [8, 8, 8]] + assert prompt.tolist() == [[0, 0, 0], [12, 0, 340]] + + +def test_blanked_prompt_leaves_every_codec_logit_unpenalized() -> None: + from vllm.model_executor.layers.utils import apply_penalties + + vocab_size = 8 + logits = torch.arange(1.0, vocab_size + 1.0).unsqueeze(0) + prompt = blank_scheduler_prompt_for_penalties(torch.tensor([[0, 0, 3]]), vocab_size) + + penalized = apply_penalties( + logits.clone(), + prompt, + torch.full((1, 1), vocab_size, dtype=torch.long), + torch.zeros(1), + torch.zeros(1), + torch.full((1,), 1.02), ) - assert torch.equal(forward_order[0], reverse_order[1]) - assert torch.equal(forward_order[1], reverse_order[0]) + torch.testing.assert_close(penalized, logits) -def test_batched_codec_sampling_preserves_request_rng_after_compaction() -> None: - talker = _make_sampling_talker() - history_a = torch.tensor([1, 1, 2]) - history_b = torch.tensor([3, 4, 4]) - hidden_a = torch.tensor([[0.25, -0.75]]) - hidden_b = torch.tensor([[-0.5, 0.5]]) - - first_batch = talker._sample_audio_codes( - torch.cat([hidden_a, hidden_b]), - [history_a, history_b], - ["req-a", "req-b"], - [0, 0], - ) - compacted_b = talker._sample_audio_codes( - hidden_b, - [history_b], - ["req-b"], - [1], - ) - - talker._request_generators = {} - standalone_b_first = talker._sample_audio_codes( - hidden_b, - [history_b], - ["req-b"], - [0], - )[0] - standalone_b_second = talker._sample_audio_codes( - hidden_b, - [history_b], - ["req-b"], - [1], - )[0] - - assert torch.equal(first_batch[1], standalone_b_first) - assert torch.equal(compacted_b[0], standalone_b_second) - - -def test_talker_emits_request_aligned_codec_deltas_after_compaction(mocker) -> None: +def test_talker_sample_blanks_penalty_prompt_without_touching_runner_state(mocker) -> None: talker = _make_talker() - seen: list[tuple[str, list[float], list[int]]] = [] + metadata = _PenaltyMetadata(prompt_token_ids=torch.tensor([[0, 2, 0]])) + captured: dict[str, torch.Tensor] = {} - def sample(hidden, histories, request_ids, steps): - assert steps == [0, 0] - seen.extend( - (request_id, hidden[row].reshape(-1).tolist(), histories[row].tolist()) - for row, request_id in enumerate(request_ids) - ) - return torch.tensor([2, 3]) + class _Recorder: + def __call__(self, logits, sampling_metadata): + captured["prompt"] = sampling_metadata.prompt_token_ids + return None - batched_sample = mocker.patch.object(talker, "_sample_audio_codes", side_effect=sample) + mocker.patch( + "vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_omni_tts.Sampler", + return_value=_Recorder(), + ) + talker.sample(torch.zeros(1, 8), metadata) + + assert captured["prompt"].tolist() == [[8, 8, 8]] + assert metadata.prompt_token_ids.tolist() == [[0, 2, 0]] + + +def test_compute_logits_returns_codec_vocab() -> None: + talker = _make_talker() + hidden = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + + logits = talker.compute_logits(hidden) + + assert logits.shape == (2, 8) + assert torch.allclose(logits[0], talker.head_code[0](hidden[0])) + + +def test_compute_logits_forces_eos_for_empty_speech() -> None: + talker = _make_talker() + talker._request_audio_states["req-empty"] = {"finished": True} + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-empty", "codes": {"audio": torch.empty(0)}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert output.multimodal_outputs["codes"]["audio"][0].numel() == 0 + assert output.multimodal_outputs["meta"]["finished"][0].item() is True + assert int(logits.argmax(dim=-1)) == 7 + + +def test_native_duplex_chunk_budget_masks_eos_until_generate_chunk_is_full() -> None: + assert _native_duplex_chunk_budget({}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, _DUPLEX_CODEC_TOKENS_PER_CHUNK) + assert _native_duplex_chunk_budget({"turn_start": True}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, 0) + assert _native_duplex_chunk_budget({"turn_end": True}) == (_DUPLEX_CODEC_TOKENS_PER_CHUNK, 0) + + talker = _make_talker() + talker._request_audio_states["req-chunk"] = { + "finished": False, + "step": 0, + "max_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK, + "min_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-chunk", "codes": {"audio": torch.empty(0)}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert logits[0, 7].item() == float("-inf") + assert torch.isfinite(logits[0, 0]) + assert output.multimodal_outputs["meta"]["finished"][0].item() is False + + +def test_native_duplex_chunk_budget_forces_eos_after_twenty_five_frames() -> None: + talker = _make_talker() + talker._request_audio_states["req-full"] = { + "finished": False, + "step": 24, + "max_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK, + "min_tokens": _DUPLEX_CODEC_TOKENS_PER_CHUNK, + } + + output = talker.make_omni_output( + torch.ones(1, 2), + model_intermediate_buffer=[{"request_id": "req-full", "codes": {"audio": torch.tensor([[3]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert talker._request_audio_states["req-full"]["step"] == 25 + assert talker._request_audio_states["req-full"]["finished"] is True + assert int(logits.argmax(dim=-1)) == 7 + assert output.multimodal_outputs["meta"]["finished"][0].item() is True + + +def test_simplex_talker_does_not_force_eos_after_a_duplex_sized_chunk() -> None: + talker = _make_talker() + talker._request_audio_states["req-simplex"] = {"finished": False, "step": 0} + + output = talker.make_omni_output( + torch.tensor([[1.0, 0.0]]), + model_intermediate_buffer=[{"request_id": "req-simplex", "codes": {"audio": torch.tensor([[3]])}}], + request_token_spans=[(0, 1)], + ) + logits = talker.compute_logits(output.text_hidden_states) + + assert talker._request_audio_states["req-simplex"]["step"] == 0 + assert torch.allclose(logits[0], talker.head_code[0](output.text_hidden_states[0]).float()) + + +def test_native_duplex_prefill_records_generate_chunk_budget() -> None: + talker = _make_talker() + talker.emb_text = nn.Embedding(1, 2) + talker.emb_code = nn.ModuleList([nn.Embedding(8, 2)]) + talker._text_eos_id = 0 + talker._tts_bos_id = 0 + + _, _, mid_turn = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + request_id="req-mid", + native_duplex=True, + tts_token_ids=torch.empty(0, dtype=torch.long), + tts_hidden_states=torch.empty(0, 2), + ) + _, _, boundary = talker.preprocess( + torch.zeros(2, dtype=torch.long), + None, + _omni_is_prefill=True, + request_id="req-end", + native_duplex=True, + meta={"turn_end": True}, + tts_token_ids=torch.empty(0, dtype=torch.long), + tts_hidden_states=torch.empty(0, 2), + ) + + assert mid_turn["audio_state"]["max_tokens"] == _DUPLEX_CODEC_TOKENS_PER_CHUNK + assert mid_turn["audio_state"]["min_tokens"] == _DUPLEX_CODEC_TOKENS_PER_CHUNK + assert boundary["audio_state"]["min_tokens"] == 0 + + +def test_make_omni_output_packs_the_previous_codec_id() -> None: + talker = _make_talker() infos = [ - {"request_id": "req-a", "audio_codes": {"accumulated": torch.tensor([1])}}, - {"request_id": "req-b", "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}}, + {"request_id": "req-a", "codes": {"audio": torch.tensor([[2]])}}, + {"request_id": "req-b", "codes": {"audio": torch.tensor([[3]])}}, ] output = talker.make_omni_output( - torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]), + torch.ones(2, 2), model_intermediate_buffer=infos, - request_token_spans=[(0, 2), (2, 3)], + request_token_spans=[(0, 1), (1, 2)], ) - assert seen == [ - ("req-a", [2.0, 0.0], [1]), - ("req-b", [3.0, 0.0], []), - ] - batched_sample.assert_called_once() - assert infos[0]["audio_codes"]["accumulated"].tolist() == [1, 2] - assert infos[1]["audio_codes"]["accumulated"].tolist() == [3] - assert set(output.multimodal_outputs) == {"codes", "meta"} - assert "model_outputs" not in output.multimodal_outputs - assert "sr" not in output.multimodal_outputs assert _routed(output, 0)["codes"]["audio"].tolist() == [[2]] assert _routed(output, 1)["codes"]["audio"].tolist() == [[3]] assert _routed(output, 0)["meta"]["finished"].item() is False - assert set(output.multimodal_outputs["meta"]) == {"finished"} - assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [0, 0] -def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None: +def test_talker_projects_request_aligned_duplex_metadata() -> None: talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2, 2])) infos = [ { "request_id": "req-a", "native_duplex": True, "duplex": {"epoch": 3, "turn_id": 7}, "ids": {"tts": [41]}, - "meta": { - "native_duplex_segment_text": "first", - "turn_eos_token_id": 99, - }, - "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, + "meta": {"native_duplex_segment_text": "first", "turn_eos_token_id": 99}, + "codes": {"audio": torch.empty(0)}, }, { "request_id": "req-b", "native_duplex": True, "duplex": {"epoch": 4, "turn_id": 8}, "ids": {"tts": [42, 99]}, - "meta": { - "native_duplex_segment_text": "second", - "turn_eos_token_id": 99, - }, - "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, + "meta": {"native_duplex_segment_text": "second", "turn_eos_token_id": 99}, + "codes": {"audio": torch.empty(0)}, }, ] @@ -360,7 +421,6 @@ def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None: assert [value.item() for value in meta["native_duplex"]] == [True, True] assert [value.item() for value in meta["duplex_epoch"]] == [3, 4] assert [value.item() for value in meta["duplex_turn_id"]] == [7, 8] - assert "native_duplex_segment_text" not in meta assert [bytes(value.tolist()).decode("utf-8") for value in meta["llm_output_text_utf8"]] == [ "first", "second", @@ -368,14 +428,9 @@ def test_talker_projects_request_aligned_duplex_metadata(mocker) -> None: assert [value.item() for value in meta["turn_end"]] == [False, True] -def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None: +def test_talker_rejects_native_duplex_without_fence_identity() -> None: talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2])) - info = { - "request_id": "req-missing-fence", - "native_duplex": True, - "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, - } + info = {"request_id": "req-missing-fence", "native_duplex": True, "codes": {"audio": torch.empty(0)}} with pytest.raises(RuntimeError, match="requires non-negative integer epoch and turn_id"): talker.make_omni_output( @@ -385,113 +440,35 @@ def test_talker_rejects_native_duplex_without_fence_identity(mocker) -> None: ) -def test_incomplete_prefill_emits_no_code_and_does_not_advance_state(mocker) -> None: +def test_decode_preprocess_embeds_the_sampled_codec_id() -> None: talker = _make_talker() - sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([2])) - infos = [ - { - "request_id": "req-prefill", - "audio_state": {"step": 0}, - "audio_codes": {"accumulated": torch.empty(0, dtype=torch.long)}, - }, - { - "request_id": "req-decode", - "audio_state": {"step": 4}, - "audio_codes": {"accumulated": torch.tensor([1])}, - }, - ] + talker.emb_code = nn.ModuleList([nn.Embedding(8, 4)]) + with torch.no_grad(): + talker.emb_code[0].weight.copy_(torch.arange(32, dtype=torch.float32).reshape(8, 4)) - output = talker.make_omni_output( - torch.tensor([[1.0, 0.0], [2.0, 0.0], [3.0, 0.0]]), - model_intermediate_buffer=infos, - request_token_spans=[(0, 2), (2, 3)], - request_sample_eligible=[False, True], + _, embeds, updates = talker.preprocess( + torch.tensor([3]), + None, + request_id="req-decode", + audio_state={"finished": False}, ) - sample.assert_called_once() - assert sample.call_args.args[2] == ["req-decode"] - assert infos[0]["audio_state"]["step"] == 0 - assert infos[0]["audio_codes"]["accumulated"].numel() == 0 - assert infos[1]["audio_state"]["step"] == 5 - assert _routed(output, 0)["codes"]["audio"].shape == (0, 1) - assert _routed(output, 1)["codes"]["audio"].tolist() == [[2]] + assert torch.equal(embeds, talker.emb_code[0](torch.tensor([3]))) + assert updates["codes"]["audio"].tolist() == [[3]] -def test_eos_is_terminal_once_and_never_enters_codec_history(mocker) -> None: +def test_decode_preprocess_drops_codec_eos() -> None: talker = _make_talker() - sample = mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([7])) - info = { - "request_id": "req-stop", - "audio_state": {"step": 3}, - "audio_codes": {"accumulated": torch.tensor([4, 5])}, - } + talker.emb_code = nn.ModuleList([nn.Embedding(8, 4)]) - first = talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - first_logits = talker.compute_logits(first.text_hidden_states) - second = talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], + _, _, updates = talker.preprocess( + torch.tensor([7]), + None, + request_id="req-eos", + audio_state={"finished": False}, ) - sample.assert_called_once() - assert info["audio_codes"]["accumulated"].tolist() == [4, 5] - assert first.multimodal_outputs["codes"]["audio"][0].shape == (0, 1) - assert first.multimodal_outputs["meta"]["finished"][0].item() is True - assert second.multimodal_outputs["meta"]["finished"][0].item() is False - assert first_logits.argmax(dim=-1).tolist() == [1] - assert talker.compute_logits(second.text_hidden_states).argmax(dim=-1).tolist() == [1] - - -def test_max_token_terminal_drops_unconsumed_codec_delta(mocker) -> None: - talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3])) - info = { - "request_id": "req-limit", - "audio_state": {"step": 1, "max_tokens": 2}, - "audio_codes": {"accumulated": torch.tensor([4, 5])}, - } - - output = talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[info], - request_token_spans=[(0, 1)], - ) - - # MiniCPMTTS.generate_chunk samples once at the max-token boundary to - # advance RNG state, but the sampled code is not fed into KV or returned. - assert info["audio_codes"]["accumulated"].tolist() == [4, 5] - assert output.multimodal_outputs["codes"]["audio"][0].shape == (0, 1) - assert output.multimodal_outputs["meta"]["finished"][0].item() is True - assert talker.compute_logits(output.text_hidden_states).argmax(dim=-1).tolist() == [1] - - -def test_request_local_state_survives_missing_runner_buffer_update(mocker) -> None: - talker = _make_talker() - mocker.patch.object(talker, "_sample_audio_codes", return_value=torch.tensor([3])) - first_info = { - "request_id": "req-local-state", - "audio_state": {"step": 1, "max_tokens": 3}, - "audio_codes": {"accumulated": torch.tensor([4])}, - } - - talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[first_info], - request_token_spans=[(0, 1)], - ) - second = talker.make_omni_output( - torch.ones(1, 2), - model_intermediate_buffer=[{"request_id": "req-local-state"}], - request_token_spans=[(0, 1)], - ) - - assert second.multimodal_outputs["meta"]["finished"][0].item() is True - assert talker._request_audio_states["req-local-state"]["step"] == 3 + assert updates["codes"]["audio"].numel() == 0 def test_missing_conditioning_fails_clearly() -> None: @@ -524,18 +501,19 @@ def test_empty_speech_segment_finishes_without_sampling_codes() -> None: assert torch.equal(embeds, talker.emb_text(torch.tensor([5, 6]))) assert updates["audio_state"]["finished"] is True + assert updates["codes"]["audio"].numel() == 0 - # Stage 1's sampling min_tokens keeps scheduling decode steps until the stop - # token becomes eligible, and those steps have no previous code to embed. - _, decode_embeds, _ = talker.preprocess( + # min_tokens can keep scheduling leftover decode rows after an empty + # speech segment. Those rows have no previous codec id to embed. + _, decode_embeds, decode_updates = talker.preprocess( torch.zeros(1, dtype=torch.long), None, request_id="req-empty", audio_state=updates["audio_state"], - audio_codes=updates["audio_codes"], ) assert decode_embeds.shape == (1, 4) + assert decode_updates["codes"]["audio"].numel() == 0 def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> None: @@ -544,7 +522,7 @@ def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> Non condition = torch.arange(18, dtype=torch.float32).reshape(9, 2) mocker.patch.object(talker, "_build_condition_embeddings", return_value=condition) - _, embeds, _ = talker.preprocess( + _, embeds, updates = talker.preprocess( torch.zeros(9, dtype=torch.long), None, _omni_is_prefill=True, @@ -556,46 +534,9 @@ def test_chunked_prefill_tail_aligns_condition_with_prompt_length(mocker) -> Non ) assert torch.equal(embeds, condition) + assert updates["codes"]["audio"].numel() == 0 state = talker._request_audio_states["req-chunked-prefill"] - assert state["min_tokens"] == 50 - assert state["max_tokens"] == 64 - - -@pytest.mark.parametrize( - ("meta", "expected_min_tokens"), - [ - ({"turn_start": True}, 0), - ({}, 26), - ({"turn_end": True}, 0), - ], -) -def test_native_duplex_prefill_uses_official_chunk_limits( - mocker, - meta, - expected_min_tokens, -) -> None: - talker = _make_talker() - talker.emb_text = nn.Embedding(1, 2) - mocker.patch.object( - talker, - "_build_condition_embeddings", - return_value=torch.ones(3, 2), - ) - - talker.preprocess( - torch.zeros(3, dtype=torch.long), - None, - _omni_is_prefill=True, - request_id="req-duplex-chunk", - native_duplex=True, - meta=meta, - tts_token_ids=torch.tensor([1]), - tts_hidden_states=torch.ones(1, 2), - ) - - state = talker._request_audio_states["req-duplex-chunk"] - assert state["min_tokens"] == expected_min_tokens - assert state["max_tokens"] == 26 + assert state["finished"] is False def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None: @@ -626,13 +567,11 @@ def test_native_duplex_condition_matches_official_text_plus_audio_bos() -> None: assert condition.shape[0] == token_ids.shape[0] + 1 -def test_request_cleanup_evicts_ar_rng_and_decode_state() -> None: +def test_request_cleanup_evicts_decode_state() -> None: talker = _make_talker() - talker._request_generators["req-done"] = torch.Generator() - talker._request_audio_states["req-done"] = {"step": 1} + talker._request_audio_states["req-done"] = {"finished": False} talker.on_requests_finished(["req-done"]) talker._flush_deferred_cleanup() - assert "req-done" not in talker._request_generators assert "req-done" not in talker._request_audio_states diff --git a/tests/worker/test_native_duplex_hooks.py b/tests/worker/test_native_duplex_hooks.py index 383379202..ba50be72c 100644 --- a/tests/worker/test_native_duplex_hooks.py +++ b/tests/worker/test_native_duplex_hooks.py @@ -1,5 +1,6 @@ from __future__ import annotations +import ast import inspect from types import SimpleNamespace @@ -2466,3 +2467,128 @@ def test_minicpmo_stage0_session_context_includes_resolved_ref_audio(): assert state.context_token_ids == [1, 2, 3, 151683, 151683, 4, 5] assert len(state.context_embeds) == 6 + + +# --------------------------------------------------------------------------- +# Cross-platform wiring: every AR model runner must reach the duplex hook. +# +# GPUARModelRunner and NPUARModelRunner are siblings -- neither inherits the +# other's _sample -- so the NPU runner once copied GPU _sample without the +# prepare_duplex_sampling call. Parse source instead of importing NPUARModelRunner: +# that class needs vllm_ascend, which CPU test images do not have. +# --------------------------------------------------------------------------- + +_REQUIRED_DUPLEX_CALLS = { + "__init__": "_init_duplex_sampling_state", + "load_model": "_resolve_duplex_sampling_hook", + "_update_states": "_update_duplex_sampling_states", + "_sample": "_apply_duplex_sampling", +} + +_MODEL_SAMPLER_CALL = "model_sample" + + +def _repo_root(): + from pathlib import Path + + return Path(__file__).parents[2] + + +def _ar_runner_sources(): + root = _repo_root() + paths = [root / "vllm_omni/worker/gpu_ar_model_runner.py"] + paths.extend(sorted(root.glob("vllm_omni/platforms/*/worker/*_ar_model_runner.py"))) + return paths + + +def _ar_runner_classdefs(): + root = _repo_root() + found = {} + for path in _ar_runner_sources(): + module = ast.parse(path.read_text()) + for node in module.body: + if isinstance(node, ast.ClassDef) and node.name.endswith("ARModelRunner"): + found[node.name] = (node, str(path.relative_to(root))) + assert found, "no AR model runners discovered; the glob in _ar_runner_sources() is stale" + return found + + +def _method_calls(method): + calls = {} + for node in ast.walk(method): + if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute): + calls.setdefault(node.func.attr, node.lineno) + return calls + + +def _wires_duplex_sampling_directly(classdef) -> bool: + bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)} + if "DuplexSamplingRunnerMixin" not in bases: + return False + methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)} + return all( + method_name in methods and required_call in _method_calls(methods[method_name]) + for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items() + ) + + +def test_every_ar_runner_reaches_the_duplex_sampling_hook(): + classdefs = _ar_runner_classdefs() + wired = {name for name, (classdef, _path) in classdefs.items() if _wires_duplex_sampling_directly(classdef)} + + for class_name, (classdef, path) in sorted(classdefs.items()): + if class_name in wired: + continue + base_names = {base.id for base in classdef.bases if isinstance(base, ast.Name)} + assert base_names & wired, ( + f"{path}::{class_name} neither wires the duplex sampling hook nor derives from a " + f"runner that does (wired runners: {sorted(wired)}). Inherit DuplexSamplingRunnerMixin " + f"and call {sorted(_REQUIRED_DUPLEX_CALLS.values())} from " + f"{sorted(_REQUIRED_DUPLEX_CALLS)} respectively." + ) + + +@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"]) +def test_ar_runner_calls_every_duplex_sampling_hook_site(class_name: str): + classdefs = _ar_runner_classdefs() + assert class_name in classdefs, f"{class_name} not found by _ar_runner_sources()" + classdef, path = classdefs[class_name] + + bases = {base.id for base in classdef.bases if isinstance(base, ast.Name)} + assert "DuplexSamplingRunnerMixin" in bases, ( + f"{path}::{class_name} must inherit DuplexSamplingRunnerMixin so the duplex " + f"sampling hook is wired the same way on every platform" + ) + + methods = {node.name: node for node in classdef.body if isinstance(node, ast.FunctionDef)} + for method_name, required_call in _REQUIRED_DUPLEX_CALLS.items(): + method = methods.get(method_name) + assert method is not None, f"{path}::{class_name}.{method_name} is missing; it must call {required_call}" + calls = _method_calls(method) + assert required_call in calls, f"{path}::{class_name}.{method_name} must call self.{required_call}" + + +@pytest.mark.parametrize("class_name", ["GPUARModelRunner", "NPUARModelRunner"]) +def test_ar_runner_applies_duplex_sampling_before_the_model_sampler(class_name: str): + classdef, path = _ar_runner_classdefs()[class_name] + sample = next(node for node in classdef.body if isinstance(node, ast.FunctionDef) and node.name == "_sample") + + apply_lineno = _method_calls(sample).get("_apply_duplex_sampling") + sampler_lineno = next( + ( + node.lineno + for node in ast.walk(sample) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) and node.func.id == _MODEL_SAMPLER_CALL + ), + None, + ) + assert apply_lineno is not None, f"{path}::{class_name}._sample must call self._apply_duplex_sampling" + assert sampler_lineno is not None, ( + f"{path}::{class_name}._sample no longer calls {_MODEL_SAMPLER_CALL}(); update _MODEL_SAMPLER_CALL " + f"so the ordering check keeps testing something" + ) + assert apply_lineno < sampler_lineno, ( + f"{path}::{class_name}._sample calls self._apply_duplex_sampling at line {apply_lineno}, after " + f"{_MODEL_SAMPLER_CALL}() at line {sampler_lineno}. The hook masks the logits and publishes the " + f"row -> session map the sampler reads, so running it afterwards is a silent no-op." + ) diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index cdfd35ee2..0e4457733 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -532,7 +532,13 @@ def _parse_stage_deploy(stage_data: dict[str, Any]) -> StageDeployConfig: _DEEP_MERGE_KEYS = frozenset( - {"default_sampling_params", "default_pooling_params", "subtalker_sampling_params", "engine_extras", "engine_args"} + { + "default_sampling_params", + "default_pooling_params", + "subtalker_sampling_params", + "engine_extras", + "engine_args", + } ) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 532de5e79..cd7705f52 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -64,15 +64,15 @@ stages: devices: "0" output_connectors: to_stage_2: connector_of_shared_memory - # Codec sampling comes from the checkpoint's tts_config. These parameters - # govern only vLLM's binary continue/stop token. + # Talker is a single-vocab codec LM. These knobs go to vLLM's Sampler + # (head_code logits). EOS is tts_config.eos_token_id (num_audio_tokens-1). default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 + temperature: 0.8 + top_p: 0.8 + top_k: 100 + repetition_penalty: 1.02 min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] + max_tokens: 4096 seed: 42 detokenize: false diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 5b43206b8..fd56b376f 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -52,16 +52,13 @@ stages: devices: "1" output_connectors: to_stage_2: connector_of_shared_memory - # Codec sampling is model-owned by the checkpoint's tts_config (defaults: - # T=0.8, top-k=25, top-p=0.85, repetition penalty=1.05, seed=42). - # These parameters only govern vLLM's binary continue/stop token. default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 + temperature: 0.8 + top_p: 0.8 + top_k: 100 + repetition_penalty: 1.02 min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] + max_tokens: 4096 seed: 42 detokenize: false diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index edbde0c5d..0bb21b177 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -57,15 +57,13 @@ stages: devices: "1" output_connectors: to_stage_2: connector_of_shared_memory - # vLLM samples the binary continue/stop row; codec sampling comes from - # the checkpoint's tts_config and currently uses a deterministic seed. default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 + temperature: 0.8 + top_p: 0.8 + top_k: 100 + repetition_penalty: 1.02 min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] + max_tokens: 4096 seed: 42 detokenize: false diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 813a2e2eb..0734e2d96 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -69,15 +69,13 @@ stages: devices: "4" output_connectors: to_stage_2: connector_of_shared_memory - # vLLM samples the binary continue/stop row; codec sampling comes from - # the checkpoint's tts_config and currently uses a deterministic seed. default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 + temperature: 0.8 + top_p: 0.8 + top_k: 100 + repetition_penalty: 1.02 min_tokens: 50 - max_tokens: 2048 - stop_token_ids: [1] + max_tokens: 4096 seed: 42 detokenize: false diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml index df0e40be9..b51c46a9d 100644 --- a/vllm_omni/deploy/minicpmo_4_5_duplex.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_duplex.yaml @@ -29,12 +29,13 @@ stages: async_scheduling: false # Talker context is 4096 (tts_config.max_position_embeddings). max_model_len: 4096 - # The split Talker exposes a binary continue/stop vocabulary; native - # segment and turn boundaries travel as explicit output metadata. + # Overlay only the knobs duplex changes. Codec temperature / top_k stay + # on the base YAML. Native duplex still ends each Talker request after + # one generate_chunk (26 samples, last is codec EOS) inside the model; + # YAML max_tokens is only the hard ceiling. default_sampling_params: - max_tokens: 2048 + max_tokens: 4096 min_tokens: 0 - stop_token_ids: [1] - stage_id: 2 max_num_seqs: 2 diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index f112983e1..0fb3fb14c 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -127,8 +127,16 @@ def build_engine_core_request_from_tokens( pooling_params = None if isinstance(params, SamplingParams): sampling_params = params.clone() - if sampling_params.max_tokens is None and model_config is not None: - sampling_params.max_tokens = model_config.max_model_len - len(prompt_token_ids) + if model_config is not None: + remaining = model_config.max_model_len - len(prompt_token_ids) + if sampling_params.max_tokens is None: + sampling_params.max_tokens = remaining + # ``check_stop`` returns early while ``min_tokens`` is unmet, so a + # min_tokens wider than the leftover context window never reaches + # the length cap: the request stalls once the scheduler can grant + # no further tokens. Stage prompts are built here, so clamp now. + if sampling_params.min_tokens > remaining: + sampling_params.min_tokens = max(0, remaining) else: pooling_params = params.clone() diff --git a/vllm_omni/experimental/fullduplex/model_executor.py b/vllm_omni/experimental/fullduplex/model_executor.py index eebe01876..58ba597df 100644 --- a/vllm_omni/experimental/fullduplex/model_executor.py +++ b/vllm_omni/experimental/fullduplex/model_executor.py @@ -4,7 +4,12 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + import torch + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.sample.metadata import SamplingMetadata @dataclass(frozen=True, slots=True) @@ -103,4 +108,55 @@ class DuplexSamplingHelper: self.hook_active = False -__all__ = ["DuplexSamplingHelper", "DuplexSamplingRow"] +class DuplexSamplingRunnerMixin: + """Runner-side wiring for the experimental duplex sampling hook. + + GPU and NPU AR runners are siblings, not parent and child, so neither + inherits the other's ``_sample``. Keep the hook here and have each runner + opt in, instead of copying the GPU ``_sample`` path and dropping the + ``prepare_duplex_sampling`` call. + """ + + def _init_duplex_sampling_state(self) -> None: + """Reset hook resolution. Call from ``__init__``, before ``load_model``.""" + self._duplex_sampling_hook = None + self._duplex_sampling_hook_resolved = False + + def _resolve_duplex_sampling_hook(self, *, force: bool = False): + """Bind ``model.prepare_duplex_sampling``. Call with ``force`` after load.""" + if not force and getattr(self, "_duplex_sampling_hook_resolved", False): + return self._duplex_sampling_hook + candidate = getattr(getattr(self, "model", None), "prepare_duplex_sampling", None) + self._duplex_sampling_hook = candidate if callable(candidate) else None + self._duplex_sampling_hook_resolved = True + if self._duplex_sampling_hook is not None and not hasattr(self, "_duplex_sampling_helper"): + self._duplex_sampling_helper = DuplexSamplingHelper() + return self._duplex_sampling_hook + + def _update_duplex_sampling_states(self, scheduler_output: SchedulerOutput) -> None: + """Refresh which batch rows are duplex rows. Call from ``_update_states``.""" + if self._resolve_duplex_sampling_hook() is None: + return + helper = getattr(self, "_duplex_sampling_helper", None) + if helper is not None: + helper.update_states(self, scheduler_output) + + def _apply_duplex_sampling(self, logits: torch.Tensor, prepared_sampling_metadata: SamplingMetadata) -> None: + """Hand the model its duplex rows. Call from ``_sample``, before the sampler.""" + prepare_duplex_sampling = self._resolve_duplex_sampling_hook() + if prepare_duplex_sampling is None: + return + helper = getattr(self, "_duplex_sampling_helper", None) + rows = helper.rows(self) if helper is not None and helper.active_request_ids else () + if rows or (helper is not None and helper.hook_active): + prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) + if helper is not None: + helper.hook_active = bool(rows) + + def _clear_duplex_sampling(self) -> None: + helper = getattr(self, "_duplex_sampling_helper", None) + if helper is not None: + helper.clear() + + +__all__ = ["DuplexSamplingHelper", "DuplexSamplingRow", "DuplexSamplingRunnerMixin"] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index fe5c8a57e..1553a5921 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -19,6 +19,65 @@ from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper logger = init_logger(__name__) _SILENCE_TOKEN = 4218 +# CosyVoice2 RelPos PE is built for max_len=5000 and `forward_chunk` never +# calls ``extend_pe``. After the 2x upsample a single Code2Wav prefill longer +# than this many codec tokens overflows the PE slice (matrix_ac vs matrix_bd). +_DEFAULT_RELPOS_MAX_POS = 5000 +_MAX_ENCODE_TOKEN_CAP = 1024 + + +def relpos_encode_token_budget( + *, + max_pos: int, + stride: int, + cache_offset: int, + lookahead: int, + cap: int = _MAX_ENCODE_TOKEN_CAP, +) -> int: + """How many codec tokens ``forward_chunk`` can take before RelPos PE wraps. + + ``position_encoding(size)`` needs ``size <= max_pos``. After upsample, + ``size = cache_offset * stride + ~stride * token_frames`` (plus last-chunk + lookahead pad). + """ + stride = max(1, int(stride)) + lookahead = max(0, int(lookahead)) + room = int(max_pos) // stride - max(0, int(cache_offset)) - lookahead - 1 + return max(lookahead + 1, min(int(cap), room)) + + +def plan_token2wav_encode_slices( + num_frames: int, + *, + max_frames: int, + min_nonfinal: int, + last_chunk: bool, +) -> list[tuple[int, int]]: + """Split a Code2Wav token span so every non-final piece has a lookahead window.""" + if num_frames <= 0: + return [] + min_nonfinal = max(1, int(min_nonfinal)) + max_frames = max(int(max_frames), min_nonfinal) + last_min = 1 if last_chunk else min_nonfinal + if num_frames <= max_frames: + return [(0, num_frames)] + + slices: list[tuple[int, int]] = [] + start = 0 + while start < num_frames: + remaining = num_frames - start + if remaining <= max_frames: + slices.append((start, num_frames)) + break + take = max_frames + tail = remaining - take + if tail <= max_frames and tail < last_min: + take = remaining - last_min + if take < min_nonfinal: + take = remaining + slices.append((start, start + take)) + start += take + return slices def _autocast_disabled(device: torch.device): @@ -237,6 +296,41 @@ class BatchedToken2Wav(nn.Module): width = getattr(layer, "pre_lookahead_len", None) return int(width) if width is not None else None + def _relpos_max_pos(self) -> int: + embed = getattr(self.flow.encoder, "embed", None) + pe = getattr(embed, "pe", None) + if isinstance(pe, torch.Tensor) and pe.numel() > 0: + return max(1, int(pe.size(1) // 2)) + return _DEFAULT_RELPOS_MAX_POS + + def _upsample_stride(self) -> int: + stride = getattr(getattr(self.flow.encoder, "up_layer", None), "stride", None) + return max(1, int(stride)) if stride is not None else 2 + + def _max_encode_token_frames(self, states: list[BatchedToken2WavState]) -> int: + att = None + if states: + att = self._stack_flow_cache(states).get("conformer_att_cache") + offset1 = int(att.shape[3] // 2) if att is not None else 0 + lookahead = self._pre_lookahead_len() or 0 + return relpos_encode_token_budget( + max_pos=self._relpos_max_pos(), + stride=self._upsample_stride(), + cache_offset=offset1, + lookahead=lookahead, + ) + + def _ensure_relpos_pe(self, tokens: torch.Tensor, att_cache: torch.Tensor | None) -> None: + """Grow CosyVoice RelPos PE before ``forward_chunk``, which never calls extend_pe.""" + embed = getattr(self.flow.encoder, "embed", None) + extend_pe = getattr(embed, "extend_pe", None) + if not callable(extend_pe): + return + offset1 = int(att_cache.shape[3] // 2) if att_cache is not None else 0 + lookahead = self._pre_lookahead_len() or 0 + needed = offset1 * self._upsample_stride() + self._upsample_stride() * (int(tokens.shape[1]) + lookahead + 1) + extend_pe(tokens.new_zeros((1, max(needed, 1)))) + def _encode_chunk( self, tokens: torch.Tensor, @@ -245,6 +339,7 @@ class BatchedToken2Wav(nn.Module): cnn_cache: torch.Tensor | None, att_cache: torch.Tensor | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + self._ensure_relpos_pe(tokens, att_cache) embedded = self.flow.input_embedding(tokens) hidden, new_cnn, new_att = self.flow.encoder.forward_chunk( xs=embedded, @@ -508,14 +603,68 @@ class BatchedToken2Wav(nn.Module): # must carry at least one full kernel. Only the final chunk is allowed # to be shorter: ``forward_chunk`` zero-pads it by the lookahead width. lookahead = self._pre_lookahead_len() + num_frames = int(tokens.shape[1]) if lookahead is not None and not last_chunk: - num_frames = int(tokens.shape[1]) if num_frames <= lookahead: raise RuntimeError( "MiniCPMO45Code2WavBatchError " f'{{"reason":"chunk_below_lookahead_window","frames":{num_frames},' f'"minimum":{lookahead + 1}}}' ) + # A non-async Talker dump can land thousands of codec tokens in one + # Code2Wav prefill. CosyVoice RelPos PE (max_len=5000) plus 2x upsample + # cannot score that in one ``forward_chunk`` (6968 vs 985 on NPU). + max_frames = self._max_encode_token_frames(states) + slices = plan_token2wav_encode_slices( + num_frames, + max_frames=max_frames, + min_nonfinal=(lookahead + 1) if lookahead is not None else 1, + last_chunk=last_chunk, + ) + if len(slices) > 1: + logger.info( + "MiniCPM-o Code2Wav splitting %d codec tokens into %d encoder windows " + "(max_frames=%d) to stay inside RelPos PE.", + num_frames, + len(slices), + max_frames, + ) + parts: list[list[torch.Tensor]] = [[] for _ in range(batch_size)] + current = states + for index, (start, end) in enumerate(slices): + is_last_piece = index == len(slices) - 1 + audios, current = self._decode_batch_once( + tokens[:, start:end], + features, + current, + last_chunk=last_chunk and is_last_piece, + flush_encoder=flush_encoder and is_last_piece, + ) + for row, audio in enumerate(audios): + parts[row].append(audio) + merged = [ + torch.cat(row_parts) if row_parts else tokens.new_zeros((0,), dtype=torch.float32) + for row_parts in parts + ] + return merged, current + return self._decode_batch_once( + tokens, + features, + states, + last_chunk=last_chunk, + flush_encoder=flush_encoder, + ) + + def _decode_batch_once( + self, + tokens: torch.Tensor, + features: PromptFeatures, + states: list[BatchedToken2WavState], + *, + last_chunk: bool, + flush_encoder: bool = False, + ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: + batch_size = int(tokens.shape[0]) flow_cache = self._stack_flow_cache(states) speakers = features.speaker_embedding.expand(batch_size, -1) with self._autocast(tokens.device): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index 490d8e0cf..af9401872 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -301,6 +301,10 @@ class ConditionalChatTTSConfig(PretrainedConfig): self.top_p = top_p self.top_k = top_k self.repetition_penalty = repetition_penalty + # Talker stage is a single-vocab codec LM: vLLM's Sampler and + # embed_tokens both address this table, not the Thinker text vocab. + self.vocab_size = int(num_audio_tokens) + self.eos_token_id = int(num_audio_tokens) - 1 class MiniCPMOConfig(Qwen2Config): diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py index 9e9dfe952..0480176b7 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_tts.py @@ -8,11 +8,12 @@ Pipeline: 1. Receive thinker hidden_states + full token IDs via additional_information 2. Extract tts_bos..tts_eos region 3. Build condition: emb_text(tokens) + projector_semantic(hidden) (hidden_text_merge) - 4. Continuously generate request-aligned discrete audio-code deltas + 4. Project last hidden through head_code; vLLM Sampler picks the codec id + 5. Next decode embeds that id with emb_code and emits it to Code2Wav """ -from collections.abc import Iterable, Sequence -from dataclasses import dataclass +from collections.abc import Iterable, Mapping, Sequence +from dataclasses import replace from typing import Any import torch @@ -32,45 +33,34 @@ from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) -_REPETITION_WINDOW = 16 _REPETITION_PENALTY_CHUNK_SIZE = 16 -_MIN_AUDIO_TOKENS = 64 -_MAX_AUDIO_TOKENS = 2048 -_AUDIO_TOKENS_PER_TEXT_TOKEN = 10 -# Codec-token sampling happens inside the model; vLLM sampling parameters -# only choose the Talker's binary continue/stop row. -_CODEC_SEED = 42 -_CODEC_TEMPERATURE = 0.8 -_CODEC_TOP_K = 25 -_CODEC_TOP_P = 0.85 -_CODEC_REPETITION_PENALTY = 1.05 -_CODEC_MIN_TOKENS = 50 +# Native duplex Talker must finish after one MiniCPMTTS.generate_chunk: +# 25 codec frames (``codec_chunk_frames``) plus the terminating sample. +# Without this, the single-vocab Sampler keeps the stage-1 request alive +# until codec EOS / 4096 and Thinker never starts the next model turn. _DUPLEX_CODEC_TOKENS_PER_CHUNK = 26 -@dataclass(slots=True) -class _PendingCodecSample: - output_index: int - hidden_row: torch.Tensor - codes: torch.Tensor - request_id: str - step: int - state: dict[str, Any] - info: dict[str, Any] +def _native_duplex_chunk_budget(meta: Mapping[str, Any] | None) -> tuple[int, int]: + """Return ``(max_tokens, min_tokens)`` for one native-duplex Talker request.""" + boundary = isinstance(meta, Mapping) and (bool(meta.get("turn_start")) or bool(meta.get("turn_end"))) + ceiling = _DUPLEX_CODEC_TOKENS_PER_CHUNK + return ceiling, 0 if boundary else ceiling -def _max_audio_tokens(condition_tokens: int) -> int: - """Bound codec generation with a conservative text-length estimate. +def blank_scheduler_prompt_for_penalties( + prompt_token_ids: torch.Tensor, + vocab_size: int, +) -> torch.Tensor: + """Return a penalty prompt whose every position is the pad id (``vocab_size``). - EOS is masked for the first 50 steps, so a direct ``text_tokens * 10`` - limit can terminate short responses before EOS is eligible. The 2048 - ceiling matches the checkpoint's native generation default and keeps the - sequence within the Talker's 4096-position context. + No Talker prompt position is codec history: prefill conditioning arrives as + embeddings from ``preprocess`` and decode embeds sampled ids with + ``emb_code``. The scheduler ids are placeholders (``llm2tts`` fills them + with ``0``, or with thinker token ids on the non-handoff path), so scoring + them would tax unrelated codec tokens. """ - return max( - _MIN_AUDIO_TOKENS, - min(_MAX_AUDIO_TOKENS, condition_tokens * _AUDIO_TOKENS_PER_TEXT_TOKEN), - ) + return torch.full_like(prompt_token_ids, int(vocab_size)) def _restore_weight_norm_weight(weight_g: torch.Tensor, weight_v: torch.Tensor) -> torch.Tensor: @@ -122,32 +112,6 @@ def _apply_batched_repetition_penalty( return penalized -def _apply_top_k_top_p( - logits: torch.Tensor, - *, - top_k: int | None, - top_p: float | None, - min_tokens_to_keep: int = 3, -) -> torch.Tensor: - """Apply the same candidate floors as the upstream Transformers warpers.""" - filtered = logits.clone() - vocab_size = filtered.shape[-1] - # MiniCPM-o's gen_logits() appends TopPLogitsWarper before - # TopKLogitsWarper. The order is observable for fixed-seed sampling. - if top_p is not None and 0.0 < top_p < 1.0: - sorted_logits, sorted_indices = torch.sort(filtered, descending=False, dim=-1) - cumulative_probs = torch.softmax(sorted_logits, dim=-1).cumsum(dim=-1) - remove = cumulative_probs <= (1.0 - float(top_p)) - remove[..., -min_tokens_to_keep:] = False - remove = remove.scatter(-1, sorted_indices, remove) - filtered.masked_fill_(remove, float("-inf")) - if top_k is not None and top_k > 0: - keep = min(vocab_size, max(int(top_k), min_tokens_to_keep)) - threshold = torch.topk(filtered, keep, dim=-1).values[..., -1, None] - filtered.masked_fill_(filtered < threshold, float("-inf")) - return filtered - - class _MiniCPMTTSProjector(nn.Module): """Checkpoint-compatible hidden-state projector used by MiniCPMTTS.""" @@ -173,8 +137,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): config: MiniCPMOConfig = vllm_config.model_config.hf_config self.config = config self.vllm_config = vllm_config - self._batch_stop_logits: torch.Tensor | None = None - self._request_generators: dict[str, torch.Generator] = {} + self._force_eos_rows: list[bool] | None = None + self._mask_eos_rows: list[bool] | None = None self._request_audio_states: dict[str, dict[str, Any]] = {} self._deferred_cleanup_ids: set[str] = set() @@ -186,23 +150,18 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): self._tts_bos_id = getattr(tts_config, "audio_bos_token_id", 151687) self._text_eos_id = getattr(tts_config, "text_eos_token_id", 151692) self._num_audio_tokens = getattr(tts_config, "num_audio_tokens", 6562) + self._codec_eos_id = int(getattr(tts_config, "eos_token_id", self._num_audio_tokens - 1)) self._hidden_size = getattr(tts_config, "hidden_size", 768) self._normalize = getattr(tts_config, "normalize_projected_hidden", True) - self._codec_seed = int(getattr(tts_config, "seed", _CODEC_SEED)) - self._codec_temperature = float(getattr(tts_config, "temperature", _CODEC_TEMPERATURE)) - self._codec_top_k = int(getattr(tts_config, "top_k", _CODEC_TOP_K)) - self._codec_top_p = float(getattr(tts_config, "top_p", _CODEC_TOP_P)) - self._codec_repetition_penalty = float(getattr(tts_config, "repetition_penalty", _CODEC_REPETITION_PENALTY)) - self._codec_min_tokens = int(getattr(tts_config, "min_new_tokens", _CODEC_MIN_TOKENS)) else: self._tts_config = None + self._codec_eos_id = 0 self.has_preprocess = True self.has_postprocess = False - self.gpu_resident_buffer_keys: set[tuple[str, str]] = { - ("audio_codes", "current"), - ("audio_codes", "accumulated"), - } + # Same-step codes travel through make_omni_output from the previous + # sampled id (decode preprocess embeds that id via emb_code). + self.gpu_resident_buffer_keys: set[tuple[str, str]] = {("codes", "audio")} self._init_native_talker(prefix) def _init_native_talker(self, prefix: str) -> None: @@ -326,6 +285,7 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): info_dict.get("request_id"), ) native_duplex = bool(info_dict.get("native_duplex", False)) + meta = info_dict.get("meta") full_embeds = self._build_condition_embeddings( token_ids, hidden_states, @@ -333,7 +293,6 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): ) offset = int(info_dict.get("_omni_num_computed_tokens", 0)) request_id = str(info_dict.get("request_id", "0")) - meta = info_dict.get("meta") # The handoff rebuilds only the tail-aligned Talker condition. # Materialize zero-token embeddings for any scheduler prompt # prefix so chunked prefill can slice from a non-zero offset. @@ -356,21 +315,11 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): f"tts_ids={token_ids.shape[0]} tts_hidden={hidden_states.shape[0]} " f"prompt_len={info_dict.get('_omni_prompt_len')}" ) - duplex_boundary = isinstance(meta, dict) and ( - bool(meta.get("turn_start", False)) or bool(meta.get("turn_end", False)) - ) + state: dict[str, Any] = {"finished": empty_condition, "step": 0} if native_duplex: - max_tokens = _DUPLEX_CODEC_TOKENS_PER_CHUNK - min_tokens = 0 if duplex_boundary else _DUPLEX_CODEC_TOKENS_PER_CHUNK - else: - max_tokens = _max_audio_tokens(int(token_ids.numel())) - min_tokens = self._codec_min_tokens - state = { - "step": 0, - "max_tokens": max_tokens, - "min_tokens": min_tokens, - "finished": empty_condition, - } + max_tokens, min_tokens = _native_duplex_chunk_budget(meta if isinstance(meta, Mapping) else None) + state["max_tokens"] = max_tokens + state["min_tokens"] = min_tokens request_states = getattr(self, "_request_audio_states", None) if request_states is None: request_states = {} @@ -382,98 +331,38 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): embeds, { "audio_state": state, - "audio_codes": { - "current": empty_codes, - "accumulated": empty_codes, - }, + # Prefill has no previous codec id. vLLM samples the first + # code after this forward; the next decode emits it. + "codes": {"audio": empty_codes}, }, ) - current = (info_dict.get("audio_codes", {}) or {}).get("current") - if not isinstance(current, torch.Tensor) or current.numel() != 1: - if state.get("finished"): - # A request that finished before sampling any code can still be - # scheduled for decode steps while sampling min_tokens masks the - # stop token. make_omni_output ignores its hidden states, so any - # shape-correct embedding will do. - weight = self.emb_code[0].weight - return input_ids, weight.new_zeros((span_len, weight.shape[1])), {} - raise RuntimeError("MiniCPM-o Talker decode is missing the previous request-local audio code") - code = current.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(1) + request_id = str(info_dict.get("request_id", "0")) + stored = self._request_audio_states.get(request_id) + if isinstance(stored, dict): + state = stored + if isinstance(state, dict) and state.get("finished"): + # An empty speech segment can still be scheduled until EOS is + # eligible. The sampler is forced to EOS; any shape-correct + # embedding is enough for these leftover decode rows. + weight = self.emb_code[0].weight + empty_codes = torch.empty(0, dtype=torch.long, device=weight.device) + return input_ids, weight.new_zeros((span_len, weight.shape[1])), {"codes": {"audio": empty_codes}} + + # Decode: vLLM's previous sampled codec id is this step's input. + # Embed it with the codec table (not the 32k Llama embed_tokens) and + # hand the same id to make_omni_output so Code2Wav sees it this step. + code = input_ids.to(device=self.emb_code[0].weight.device, dtype=torch.long).reshape(-1)[-1:] embeds = self.emb_code[0](code) - return input_ids, embeds, {} - - def _request_generator(self, request_id: str, device: torch.device) -> torch.Generator: - generator = self._request_generators.get(request_id) - if generator is None: - generator = torch.Generator(device=device) - generator.manual_seed(self._codec_seed) - self._request_generators[request_id] = generator - return generator - - def _sample_audio_codes( - self, - hidden_states: torch.Tensor, - histories: Sequence[torch.Tensor], - request_ids: Sequence[str], - steps: Sequence[int], - ) -> torch.Tensor: - batch_size = int(hidden_states.shape[0]) - if not (len(histories) == len(request_ids) == len(steps) == batch_size): - raise ValueError( - "MiniCPM-o batched codec sampling requires one history, request id, " - f"and step per hidden row, got batch={batch_size}, histories={len(histories)}, " - f"request_ids={len(request_ids)}, steps={len(steps)}" - ) - if batch_size == 0: - return torch.empty(0, dtype=torch.long, device=hidden_states.device) - - logits = self.head_code[0](hidden_states).float() / self._codec_temperature - eos_id = self._num_audio_tokens - 1 - logits = _apply_batched_repetition_penalty( - logits, - histories, - penalty=self._codec_repetition_penalty, - window_size=_REPETITION_WINDOW, - ) - request_states = getattr(self, "_request_audio_states", {}) - mask_eos_values: list[bool] = [] - for request_id, step in zip(request_ids, steps, strict=True): - state = request_states.get(request_id) - min_tokens = ( - int(state.get("min_tokens", self._codec_min_tokens)) - if isinstance(state, dict) - else self._codec_min_tokens - ) - mask_eos_values.append(step < min_tokens) - mask_eos = torch.tensor( - mask_eos_values, - dtype=torch.bool, - device=logits.device, - ) - logits[:, eos_id].masked_fill_(mask_eos, float("-inf")) - logits = _apply_top_k_top_p( - logits, - top_k=self._codec_top_k, - top_p=self._codec_top_p, - min_tokens_to_keep=3, - ) - probabilities = torch.softmax(logits, dim=-1) - # torch.multinomial accepts one generator for the entire call. Keep it - # per row so compaction, request ordering, or another request finishing - # cannot perturb request-local RNG streams; the expensive projection - # and logit transforms above remain fully batched. - return torch.cat( - [ - torch.multinomial( - probabilities[row : row + 1], - num_samples=1, - generator=self._request_generator(request_id, probabilities.device), - ) - for row, request_id in enumerate(request_ids) - ], - dim=0, - ).reshape(-1) + if int(code.item()) == int(self._codec_eos_id): + if isinstance(state, dict): + state["finished"] = True + elif stored is None: + self._request_audio_states[request_id] = {"finished": True, "step": 0} + delta = torch.empty(0, dtype=torch.long, device=code.device) + else: + delta = code.reshape(1, 1) + return input_ids, embeds, {"codes": {"audio": delta}} def make_omni_output( self, @@ -487,18 +376,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): spans = kwargs.get("request_token_spans") if spans is None or len(spans) != len(infos): raise RuntimeError("MiniCPM-o continuous Talker requires one request_token_span per request") - sample_eligible = kwargs.get("request_sample_eligible") - if sample_eligible is None: - sample_eligible = [True] * len(infos) - if len(sample_eligible) != len(infos): - raise RuntimeError( - f"MiniCPM-o continuous Talker received {len(sample_eligible)} sampling flags for {len(infos)} requests" - ) emit_duplex_metadata = any(isinstance(info, dict) and info.get("native_duplex") is True for info in infos) - # Rows default to continue. Only previously finished or newly terminal - # requests stop; prefill/ineligible rows stay aligned as False. - stop_flags = [False] * len(infos) native_duplex_flags: list[torch.Tensor] = [] duplex_epochs: list[torch.Tensor] = [] duplex_turn_ids: list[torch.Tensor] = [] @@ -507,7 +386,8 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): empty_delta = hidden.new_empty((0, 1), dtype=torch.long) codec_deltas = [empty_delta for _ in infos] terminal_flags = [torch.tensor(False, dtype=torch.bool) for _ in infos] - pending_samples: list[_PendingCodecSample] = [] + force_eos_rows = [False] * len(infos) + mask_eos_rows = [False] * len(infos) for index, info in enumerate(infos): info_dict = info if isinstance(info, dict) else {} native_duplex = info_dict.get("native_duplex") is True @@ -554,100 +434,39 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): if not isinstance(info, dict): continue - start, end = spans[index] - end = min(int(end), int(hidden.shape[0])) - if int(start) >= end: - continue request_id = str(info.get("request_id", index)) - request_states = getattr(self, "_request_audio_states", None) - if request_states is None: - request_states = {} - self._request_audio_states = request_states - state = request_states.get(request_id) + state = self._request_audio_states.get(request_id) if not isinstance(state, dict): state = dict(info.get("audio_state", {}) or {}) - request_states[request_id] = state - if state.get("finished"): - stop_flags[index] = True - continue - if not sample_eligible[index]: - # vLLM computes a logit row for incomplete chunked prefills but - # discards its sampled token. Advancing codec/RNG state here - # would make output depend on prefill chunking and compaction. - continue - codes = state.get("codes") - if not isinstance(codes, torch.Tensor): - codes = (info.get("audio_codes", {}) or {}).get("accumulated") - if not isinstance(codes, torch.Tensor): - codes = torch.empty(0, dtype=torch.long, device=hidden.device) - else: - codes = codes.to(device=hidden.device, dtype=torch.long).reshape(-1) + self._request_audio_states[request_id] = state + empty_speech = bool(state.get("finished")) + codes = info.get("codes", {}) + audio = codes.get("audio") if isinstance(codes, Mapping) else None + if isinstance(audio, torch.Tensor) and audio.numel() > 0: + codec_deltas[index] = audio.to(device=hidden.device, dtype=torch.long).reshape(-1, 1) + if state.get("max_tokens") is not None: + state["step"] = int(state.get("step", 0)) + 1 + max_tokens = state.get("max_tokens") + min_tokens = state.get("min_tokens") step = int(state.get("step", 0)) - pending_samples.append( - _PendingCodecSample( - output_index=index, - hidden_row=hidden[end - 1 : end], - codes=codes, - request_id=request_id, - step=step, - state=state, - info=info, - ) - ) + # 26 samples include the terminating EOS, so force it after 25 + # forwarded frames — the same cadence as generate_chunk. + hit_chunk_limit = max_tokens is not None and step >= int(max_tokens) - 1 + chunk_done = empty_speech or hit_chunk_limit + if hit_chunk_limit: + # The next sampled id is forced EOS and will finish the + # request; mark the chunk done on this step so the data + # plane does not wait for a follow-up empty decode. + state["finished"] = True + force_eos_rows[index] = chunk_done + mask_eos_rows[index] = not force_eos_rows[index] and min_tokens is not None and step < int(min_tokens) + terminal_flags[index] = torch.tensor(chunk_done, dtype=torch.bool) - if pending_samples: - active_hidden_rows = [pending.hidden_row for pending in pending_samples] - active_hidden = ( - active_hidden_rows[0] if len(active_hidden_rows) == 1 else torch.cat(active_hidden_rows, dim=0) - ) - sampled_batch = self._sample_audio_codes( - active_hidden, - [pending.codes for pending in pending_samples], - [pending.request_id for pending in pending_samples], - [pending.step for pending in pending_samples], - ) - # One batched device-to-host synchronization replaces one .item() - # synchronization per request. - sampled_ids = sampled_batch.detach().to(device="cpu").tolist() - else: - sampled_batch = hidden.new_empty((0,), dtype=torch.long) - sampled_ids = [] - - for row, pending in enumerate(pending_samples): - sampled = sampled_batch[row].reshape(()) - sampled_id = int(sampled_ids[row]) - codes = pending.codes - state = pending.state - info = pending.info - is_eos = sampled_id == self._num_audio_tokens - 1 - state["step"] = int(state.get("step", 0)) + 1 - reached_limit = int(state["step"]) >= int(state.get("max_tokens", 2048)) - finished = is_eos or reached_limit - state["finished"] = finished - # MiniCPMTTS.generate_chunk consumes the boundary sample but - # returns only codes that were fed into the retained KV state. - if not is_eos and not reached_limit: - codes = torch.cat([codes[-(_REPETITION_WINDOW - 1) :], sampled.reshape(1)]) - delta = sampled.reshape(1, 1) - else: - delta = empty_delta - state["codes"] = codes - info["audio_state"] = state - info["audio_codes"] = { - "current": sampled.reshape(1), - "accumulated": codes, - } - codec_deltas[pending.output_index] = delta - terminal_flags[pending.output_index] = torch.tensor(finished, dtype=torch.bool) - stop_flags[pending.output_index] = finished - - self._batch_stop_logits = ( - hidden.new_tensor([[float("-inf"), 0.0] if finished else [0.0, float("-inf")] for finished in stop_flags]) - if stop_flags - else hidden.new_empty((0, 2)) - ) - # Lists are deliberate: the runner routes element i to request i, - # preserving compaction alignment while emitting only this step's code. + # Empty-speech and finished duplex chunks must sample codec EOS so + # the scheduler releases the request; mid-chunk rows mask EOS until + # the 26-token generate_chunk budget is met. + self._force_eos_rows = force_eos_rows + self._mask_eos_rows = mask_eos_rows meta_outputs = {"finished": terminal_flags} if emit_duplex_metadata: meta_outputs.update( @@ -659,13 +478,12 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): "turn_end": turn_end_flags, } ) - multimodal_outputs: dict[str, Any] = { - "codes": {"audio": codec_deltas}, - "meta": meta_outputs, - } return OmniOutput( text_hidden_states=hidden, - multimodal_outputs=multimodal_outputs, + multimodal_outputs={ + "codes": {"audio": codec_deltas}, + "meta": meta_outputs, + }, ) def on_requests_finished(self, finished_req_ids: set[str] | list[str]) -> None: @@ -674,7 +492,6 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): def _flush_deferred_cleanup(self) -> None: request_audio_states = getattr(self, "_request_audio_states", {}) for request_id in self._deferred_cleanup_ids: - self._request_generators.pop(request_id, None) request_audio_states.pop(request_id, None) self._deferred_cleanup_ids.clear() @@ -723,18 +540,38 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): def compute_logits(self, hidden_states, *args, **kwargs): if not isinstance(hidden_states, torch.Tensor): return None - if self._batch_stop_logits is None: - return torch.zeros( - hidden_states.shape[0], - 2, - device=hidden_states.device, - dtype=torch.float32, - ) - logits = self._batch_stop_logits - self._batch_stop_logits = None + if hidden_states.numel() == 0: + return hidden_states.new_empty((0, int(self._num_audio_tokens))) + logits = self.head_code[0](hidden_states).float() + force_eos = self._force_eos_rows + mask_eos = self._mask_eos_rows + self._force_eos_rows = None + self._mask_eos_rows = None + need_force = bool(force_eos and len(force_eos) == logits.shape[0] and any(force_eos)) + need_mask = bool(mask_eos and len(mask_eos) == logits.shape[0] and any(mask_eos)) + if need_force or need_mask: + logits = logits.clone() + eos_id = int(self._codec_eos_id) + if need_force: + assert force_eos is not None + forced = torch.tensor(force_eos, dtype=torch.bool, device=logits.device) + logits[forced] = float("-inf") + logits[forced, eos_id] = 0.0 + if need_mask: + assert mask_eos is not None + masked = torch.tensor(mask_eos, dtype=torch.bool, device=logits.device) + logits[masked, eos_id] = float("-inf") return logits def sample(self, logits, sampling_metadata): + prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None) + if isinstance(logits, torch.Tensor) and isinstance(prompt_ids, torch.Tensor): + # Copy rather than mutate: the runner may hand us the input batch's + # own persistent SamplingMetadata. + sampling_metadata = replace( + sampling_metadata, + prompt_token_ids=blank_scheduler_prompt_for_penalties(prompt_ids, logits.shape[-1]), + ) return Sampler()(logits, sampling_metadata) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): @@ -782,9 +619,10 @@ class MiniCPMO45OmniTTSForConditionalGeneration(nn.Module, SupportsPP): return loaded def get_input_embeddings(self, input_ids, multimodal_embeddings=None, **kwargs): - if hasattr(self, "emb_text") and self.emb_text is not None: - return self.emb_text(input_ids) - return torch.zeros(input_ids.shape[0], 1) + del multimodal_embeddings + # Decode tokens live in the codec table. Prefill overwrites these + # embeddings in preprocess with emb_text + projected thinker hidden. + return self.emb_code[0](input_ids) def embed_input_ids(self, input_ids, **kwargs): return self.get_input_embeddings(input_ids, **kwargs) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index b8fed05f1..f4e8e4c32 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py @@ -62,7 +62,12 @@ MINICPMO_4_5_PIPELINE = PipelineConfig( custom_process_input_func=f"{_PROC}.llm2tts", custom_process_next_stage_input_func=f"{_PROC}.tts2code2wav_full_payload", async_chunk_process_next_stage_input_func=f"{_PROC}.tts2code2wav_async_chunk", - sampling_constraints={"detokenize": False}, + sampling_constraints={ + "detokenize": False, + # MiniCPM-o 4.5 codec EOS is tts_config.num_audio_tokens - 1. + # Same pattern as Qwen3 talker's stop_token_ids: [2150]. + "stop_token_ids": [6561], + }, ), StagePipelineConfig( stage_id=2, diff --git a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py index d12a3b382..cd5860786 100644 --- a/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/minicpmo_4_5_omni.py @@ -945,6 +945,9 @@ def llm2tts( if handoff_ids is not None and handoff_hidden is not None: condition_suffix_length = 1 if is_native_duplex_handoff else 2 condition_length = max(len(handoff_ids), len(handoff_hidden)) + condition_suffix_length + # Dummy ids only reserve scheduler slots; prefill overwrites the + # embeddings. Talker.sample() blanks the whole prompt out of the + # repetition penalty, so codec id 0 is not taxed from step one. scheduler_prompt_token_ids = [0] * condition_length handoff_meta = model_intermediate_buffer.setdefault("meta", {}) handoff_meta["next_stage_prompt_len"] = condition_length 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 63f3912f7..1829fafa2 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -48,6 +48,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import ( get_stage_connector_role, stage_sends_async_output, ) +from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRunnerMixin from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.platforms.npu.worker.npu_model_runner import OmniNPUModelRunner from vllm_omni.utils.mm_outputs import build_mm_cpu, partition_payload_list, to_payload_element @@ -106,7 +107,7 @@ class ExecuteModelState(NamedTuple): batch_desc: BatchDescriptor multimodal_outputs: Any # Omni-Specific -class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): +class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin): """Autoregressive NPU model runner that returns hidden states per request.""" def __init__(self, *args, **kwargs): @@ -140,8 +141,18 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): kv_transfer_manager=self.kv_transfer_manager, ) self._downstream_payload_cache: dict[str, bool] = {} + self._init_duplex_sampling_state() # -------------------------------------- Omni-new ------------------------------------------------- + def load_model(self, *args, **kwargs) -> None: + super().load_model(*args, **kwargs) + self._resolve_duplex_sampling_hook(force=True) + + def _update_states(self, scheduler_output: SchedulerOutput): + deferred_state_corrections_fn = super()._update_states(scheduler_output) + self._update_duplex_sampling_states(scheduler_output) + return deferred_state_corrections_fn + def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size from vllm_omni.distributed.ray_utils.utils import ( @@ -899,10 +910,9 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): self.input_batch.idx_mapping_np, self.input_batch.positions[self.input_batch.logits_indices], ) - sampler_output = model_sample( - logits, - self._sampling_metadata_for_model_sampler(sampling_metadata), - ) + prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata) + self._apply_duplex_sampling(logits, prepared_sampling_metadata) + sampler_output = model_sample(logits, prepared_sampling_metadata) if sampler_output is not None: return sampler_output return self.sampler( diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index b3c388e3a..e66f759f8 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -45,6 +45,7 @@ from vllm_omni.distributed.omni_connectors.utils.config import ( get_stage_connector_role, stage_sends_async_output, ) +from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingRunnerMixin from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.utils.mm_outputs import ( build_mm_cpu, @@ -294,7 +295,7 @@ class ExecuteModelState(NamedTuple): slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None = None -class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): +class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, DuplexSamplingRunnerMixin): """Autoregressive GPU model runner that returns hidden states per request. Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and @@ -362,25 +363,12 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): kv_transfer_manager=self.kv_transfer_manager, ) self._downstream_payload_cache: dict[str, bool] = {} - self._duplex_sampling_hook = None - self._duplex_sampling_hook_resolved = False + self._init_duplex_sampling_state() def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) self._resolve_duplex_sampling_hook(force=True) - def _resolve_duplex_sampling_hook(self, *, force: bool = False): - if not force and getattr(self, "_duplex_sampling_hook_resolved", False): - return self._duplex_sampling_hook - candidate = getattr(getattr(self, "model", None), "prepare_duplex_sampling", None) - self._duplex_sampling_hook = candidate if callable(candidate) else None - self._duplex_sampling_hook_resolved = True - if self._duplex_sampling_hook is not None and not hasattr(self, "_duplex_sampling_helper"): - from vllm_omni.experimental.fullduplex.model_executor import DuplexSamplingHelper - - self._duplex_sampling_helper = DuplexSamplingHelper() - return self._duplex_sampling_hook - def _make_buffer(self, *size, dtype, numpy=True): # Prevent ray from pinning the buffer due to large size from vllm_omni.distributed.ray_utils.utils import ( @@ -449,11 +437,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): def _update_states(self, scheduler_output: SchedulerOutput) -> Callable | None: deferred_state_corrections_fn = super()._update_states(scheduler_output) - if self._resolve_duplex_sampling_hook() is None: - return deferred_state_corrections_fn - helper = getattr(self, "_duplex_sampling_helper", None) - if helper is not None: - helper.update_states(self, scheduler_output) + self._update_duplex_sampling_states(scheduler_output) return deferred_state_corrections_fn def _request_final_stage_id(self, req_id: str) -> int | None: @@ -581,9 +565,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self._downstream_payload_cache.clear() if hasattr(self, "model_intermediate_buffer"): self.model_intermediate_buffer.clear() - duplex_helper = getattr(self, "_duplex_sampling_helper", None) - if duplex_helper is not None: - duplex_helper.clear() + self._clear_duplex_sampling() # 5. Release all CUDA graphs unconditionally (upstream only does this # on ROCm; on CUDA the graphs are only freed by Python GC during @@ -1501,14 +1483,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self.input_batch.positions[self.input_batch.logits_indices], ) prepared_sampling_metadata = self._sampling_metadata_for_model_sampler(sampling_metadata) - prepare_duplex_sampling = self._resolve_duplex_sampling_hook() - if prepare_duplex_sampling is not None: - helper = getattr(self, "_duplex_sampling_helper", None) - rows = helper.rows(self) if helper is not None and helper.active_request_ids else () - if rows or (helper is not None and helper.hook_active): - prepare_duplex_sampling(logits, prepared_sampling_metadata, rows) - if helper is not None: - helper.hook_active = bool(rows) + self._apply_duplex_sampling(logits, prepared_sampling_metadata) sampler_output = model_sample(logits, prepared_sampling_metadata) if sampler_output is not None: return sampler_output