diff --git a/tests/model_executor/models/test_gepard_wiring.py b/tests/model_executor/models/test_gepard_wiring.py index b38b65639..9a08ddf0a 100644 --- a/tests/model_executor/models/test_gepard_wiring.py +++ b/tests/model_executor/models/test_gepard_wiring.py @@ -10,9 +10,10 @@ survives the runner's sparse-audio routing. These drive the real ``preprocess`` -> ``forward`` -> ``make_omni_output`` sequence through a stand-in for ``GPUARModelRunner.execute_model``, and use the -runner's own ``_resolve_sparse_mm_routing`` to decide what actually reaches the -caller. Only the LM backbone and the frame sampler are stubbed out; every -method under test is the real one, so a test here goes red when the engine +shared ``sparse_audio.resolve_sparse_mm_routing`` helper used by the runner to +decide what actually reaches the caller. Only the LM backbone and the frame +sampler are stubbed out; every method under test is the real one, so a test +here goes red when the engine path breaks — which a test that calls the flush by hand cannot do. """ @@ -33,7 +34,7 @@ from vllm_omni.model_executor.models.gepard.gepard_talker import ( ) from vllm_omni.model_executor.models.gepard.nanocodec import NanoCodec from vllm_omni.model_executor.models.gepard.prompt import build_gepard_prompt_ids -from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner +from vllm_omni.worker import sparse_audio pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -219,7 +220,7 @@ class _MiniEngine: self.committed_frames[req_id] = state.frame_count # --- runner: sparse-audio routing decides what reaches the caller --- - downstream, sparse_index, is_sparse = GPUARModelRunner._resolve_sparse_mm_routing( + downstream, sparse_index, is_sparse = sparse_audio.resolve_sparse_mm_routing( engine_output_type="audio", req_ids_output_copy=list(live), downstream_req_ids=list(live), diff --git a/tests/model_executor/models/voxcpm2/test_talker_output_marker.py b/tests/model_executor/models/voxcpm2/test_talker_output_marker.py new file mode 100644 index 000000000..1a76525fe --- /dev/null +++ b/tests/model_executor/models/voxcpm2/test_talker_output_marker.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""VoxCPM2 dense-mode audio outputs must carry the sparse-alignment marker. + +Dense mode (`_uses_sparse_audio_outputs()` False) still yields a strict +SUBSET of the batch whenever some requests emit no audio in a step (e.g. +prefill phase). The runner routes per-request `model_outputs` lists by batch +position unless `meta.sparse_audio` declares `meta.req_id` alignment — so a +subset without the marker misroutes audio across requests (RFC #5450 C3 +review follow-up). +""" + +from __future__ import annotations + +import functools + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +torch = pytest.importorskip("torch") + + +@functools.lru_cache(maxsize=1) +def _talker_cls(): + """Defer talker import (pulls vLLM model_executor) until first use.""" + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + VoxCPM2TalkerForConditionalGeneration, + ) + + return VoxCPM2TalkerForConditionalGeneration + + +def _make_dense_talker(): + cls = _talker_cls() + talker = cls.__new__(cls) + # Dense mode: emit every step, no delayed copy, no batched VAE decode. + talker._audio_emit_every = 1 + talker._vae_decode_every = 1 + talker._enable_delayed_audio_copy = False + talker._coalesce_audio_d2h = False + talker._sample_rate = 16000 + talker._audio_queue = [] + return talker + + +def test_dense_subset_audio_carries_sparse_alignment_marker(): + talker = _make_dense_talker() + # Batch had [r0, r1, r2]; only r1/r2 produced audio this step. + talker._audio_queue = [("r1", torch.zeros(4)), ("r2", torch.ones(4))] + + out = talker.make_omni_output(torch.zeros(1)) + + mm = out.multimodal_outputs + assert mm["meta"] == {"req_id": ["r1", "r2"], "sparse_audio": ["1"]} + assert len(mm["model_outputs"]) == 2 + assert len(mm["sr"]) == 2 + # (The coalesce-D2H sibling branch carries the same marker; it needs + # device-resident chunks, so its marker contract is pinned here via the + # shared dense-subset shape rather than a GPU-only test.) + + +def test_dense_full_batch_audio_still_carries_marker_and_order(): + # Full coverage is the degenerate subset: the marker stays truthful and + # `meta.req_id` order matches the emitted list order. + talker = _make_dense_talker() + talker._audio_queue = [("r0", torch.zeros(2)), ("r1", torch.ones(2))] + + out = talker.make_omni_output(torch.zeros(1)) + + mm = out.multimodal_outputs + assert mm["meta"]["req_id"] == ["r0", "r1"] + assert mm["meta"]["sparse_audio"] == ["1"] diff --git a/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py b/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py index 1c35c37aa..9a069fb8f 100644 --- a/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py +++ b/tests/platforms/npu/test_minicpmo_code2wav_npugraph.py @@ -32,6 +32,7 @@ pytestmark = [ class _Flow(nn.Module): def __init__(self, estimator: nn.Module): super().__init__() + self.encoder = nn.Identity() self.decoder = SimpleNamespace(estimator=estimator) diff --git a/tests/worker/test_capture_replay_contract.py b/tests/worker/test_capture_replay_contract.py index 504aaa7b0..a0ca095f3 100644 --- a/tests/worker/test_capture_replay_contract.py +++ b/tests/worker/test_capture_replay_contract.py @@ -14,9 +14,9 @@ cannot be masked by an adjacent ``else``/``elif`` that happens to read it. It matches by attribute chain (``self.input_ids.gpu``), so local variable renames do not break it. -Scope: the BASE runner (``OmniGPUModelRunner``) only. The generation runner is -KNOWN to lack the ``has_preprocess`` branch; extend this -check to the generation file when it adds the branch. +Scope: the BASE runner (``OmniGPUModelRunner``) and the generation runner's +``_dummy_run`` override (which gained the branch in RFC #5450 C6; its +``_preprocess`` is inherited from the base). """ import ast @@ -25,6 +25,7 @@ import textwrap import pytest +from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -69,11 +70,24 @@ def _has_preprocess_branch_reads_both_buffers(method) -> bool: return any(_REQUIRED_BUFFERS <= _attr_chains_in(node.body) for node in guarded) -@pytest.mark.parametrize("method_name", ["_dummy_run", "_preprocess"]) -def test_has_preprocess_branch_reads_shared_buffers(method_name): - method = getattr(OmniGPUModelRunner, method_name) +@pytest.mark.parametrize( + ("runner_cls", "method_name"), + [ + (OmniGPUModelRunner, "_dummy_run"), + (OmniGPUModelRunner, "_preprocess"), + (GPUGenerationModelRunner, "_dummy_run"), + ], +) +def test_has_preprocess_branch_reads_shared_buffers(runner_cls, method_name): + method = getattr(runner_cls, method_name) assert _has_preprocess_branch_reads_both_buffers(method), ( - f"{OmniGPUModelRunner.__name__}.{method_name}: the has_preprocess branch no " + f"{runner_cls.__name__}.{method_name}: the has_preprocess branch no " "longer reads BOTH self.input_ids.gpu and self.inputs_embeds.gpu — capture " "and replay would diverge." ) + + +def test_generation_dummy_run_overrides_with_branch(): + # Guard against the override being deleted (which would silently fall back + # to the base implementation and vacuously pass the parametrized check). + assert "_dummy_run" in GPUGenerationModelRunner.__dict__ diff --git a/tests/worker/test_gpu_ar_model_runner.py b/tests/worker/test_gpu_ar_model_runner.py index 967407169..99f609ea5 100644 --- a/tests/worker/test_gpu_ar_model_runner.py +++ b/tests/worker/test_gpu_ar_model_runner.py @@ -1,8 +1,12 @@ import asyncio import copy +import re from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from threading import Barrier from types import SimpleNamespace from typing import Any +from unittest.mock import patch import numpy as np import pytest @@ -15,12 +19,15 @@ from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechReques from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech from vllm_omni.entrypoints.openai.tts_adapters.base import PreparedRequest from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker import sparse_audio from vllm_omni.worker.gpu_ar_model_runner import ( ExecuteModelState, GPUARModelRunner, OmniAsyncGPUModelRunnerOutput, ) +from vllm_omni.worker.output import payload_build from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest +from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -301,11 +308,13 @@ def test_resolve_pooler_payload_req_ids_downstream_stage_uses_filtered_requests( def test_sparse_mm_req_ids_requires_sparse_audio_marker(): - assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None - assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None + assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"]}}) is None + assert sparse_audio.resolve_sparse_mm_req_ids({"meta.req_id": ["r1"]}) is None - assert GPUARModelRunner._sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"] - assert GPUARModelRunner._sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"] + assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"], "sparse_audio": ["1"]}}) == ["r1"] + assert sparse_audio.resolve_sparse_mm_req_ids({"meta.req_id": ["r1"], "meta.sparse_audio": ["1"]}) == ["r1"] + assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"req_id": ["r1"]}, "meta.sparse_audio": ["1"]}) == ["r1"] + assert sparse_audio.resolve_sparse_mm_req_ids({"meta": {"sparse_audio": ["1"]}, "meta.req_id": ["r1"]}) == ["r1"] def test_runner_assisted_full_attention_metadata_request_is_opt_in(): @@ -1198,3 +1207,642 @@ def test_build_omni_output_uses_combined_prefix_cache_mm_payload_for_partial_dow assert torch.equal(output.inter_stage_outputs[0]["codes.ref"], ref_codes[0]) assert torch.equal(output.inter_stage_outputs[2]["codes.ref"], ref_codes[2]) + + +class TestPromptIdsClampPaddingConvention: + """Pin `sampling_utils.clamp_prompt_ids_to_penalty_padding` (the clamp + `sample_tokens` applies before penalties). + + The audit (RFC #5450 C2) originally flagged the clamp as an off-by-one + (OOB index). It is not: upstream penalties allocate `vocab_size + 1` bins + and drop the last column, so `vocab_size` is the designed padding value + (vllm/model_executor/layers/utils.py::get_token_bin_counts_and_mask). + These tests call the runner's own helper, so a future "fix" to the clamp + fails here. + """ + + LOGITS_VOCAB = 8 + BATCH_VOCAB = 10 # wider input_batch vocab; its pad value is BATCH_VOCAB + + def _padded_prompt(self): + # Two real tokens + two padding slots padded with the batch vocab size. + return torch.tensor([[1, 2, self.BATCH_VOCAB, self.BATCH_VOCAB]]) + + def test_clamp_to_logits_vocab_keeps_padding_inert(self): + from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask + + # The runner's own helper, not a re-implementation. + clamped = clamp_prompt_ids_to_penalty_padding(self._padded_prompt(), self.LOGITS_VOCAB) + + _, mask = get_token_bin_counts_and_mask(clamped, self.LOGITS_VOCAB, 1) + + expected = torch.zeros(1, self.LOGITS_VOCAB, dtype=torch.bool) + expected[0, 1] = True + expected[0, 2] = True + # Pad ids landed in the dropped (vocab_size-th) bin: no penalty effect. + assert torch.equal(mask, expected) + + def test_clamp_one_lower_would_corrupt_last_token_penalties(self): + from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask + + # The "fix" the audit first proposed - shown here to be the actual bug: + # padding would count as occurrences of the last real vocab token. + clamped = self._padded_prompt().clamp(max=self.LOGITS_VOCAB - 1) + + _, mask = get_token_bin_counts_and_mask(clamped, self.LOGITS_VOCAB, 1) + + assert bool(mask[0, self.LOGITS_VOCAB - 1]) + + def test_unclamped_padding_overflows_penalty_bins(self): + from vllm.model_executor.layers.utils import get_token_bin_counts_and_mask + + # Why the clamp exists: the batch-level pad value (BATCH_VOCAB) is out + # of range even for the vocab_size+1 penalty bins of a narrower head. + with pytest.raises(RuntimeError): + get_token_bin_counts_and_mask(self._padded_prompt(), self.LOGITS_VOCAB, 1) + + +class TestSparseAudioMarkerRobustness: + """RFC #5450 C3: `sparse_audio.is_sparse_audio_marker` must not crash on tensor/array + markers, and the combined prefix-cache payload builder must never ship one + request's data as another's.""" + + def test_marker_list_and_str_forms(self): + marker = sparse_audio.is_sparse_audio_marker + assert marker(["1"]) is True + assert marker(["0"]) is False + assert marker("true") is True + assert marker("off") is False + + def test_marker_enforced_canonical_forms_no_error(self): + # The designed carrier (list-of-str; bare str accepted) and marker + # absence (None) are the ONLY silent forms. + marker = sparse_audio.is_sparse_audio_marker + with patch.object(sparse_audio.logger, "error") as err: + assert marker(["1"]) is True + assert marker(["0"]) is False + assert marker(["0", "1"]) is True + assert marker("true") is True + assert marker("off") is False + assert marker(None) is False + err.assert_not_called() + + def test_marker_illegal_carriers_raise_one_error_per_type(self): + # Enforcement, not normalization: any carrier outside the design + # contract is a producer bug - it raises InvalidSparseDeclarationError + # (caught inside the module by resolve_sparse_mm_routing, which + # fails closed) + # and logs at error level once per offending type. Interpreting + # deviants would need a second truth predicate, and predicate + # divergence is exactly how the np.array(["0"])-routes-as-sparse bug + # happened. + marker = sparse_audio.is_sparse_audio_marker + for value in (torch.tensor(1), torch.tensor([1, 1]), np.array(["1"]), 1, 1.0, True, ["1", 1]): + with ( + patch.object(sparse_audio, "_logged_offenders", set()), + patch.object(sparse_audio.logger, "error") as err, + ): + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + marker(value) + assert err.call_count == 1, f"expected one error for {value!r}" + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + marker(value) + assert err.call_count == 1 # deduped per type, not per step + + def test_marker_device_tensor_rejected_without_read(self): + # Deviant carriers are never read element-wise, so a device tensor + # can never introduce a D2H sync here. Meta tensors hold no data - + # any read (`.item()`/`bool()`) raises a RuntimeError, so seeing OUR + # exception type (not RuntimeError) proves the data is never touched. + marker = sparse_audio.is_sparse_audio_marker + with patch.object(sparse_audio, "_logged_offenders", set()): + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + marker(torch.ones((), device="meta")) + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + marker(torch.ones(3, device="meta")) + + def test_invalid_marker_fails_closed_at_routing(self): + # A present-but-invalid sparse declaration most likely belongs to a + # producer that INTENDED sparse output; falling back to dense would + # batch-index-assign its payloads to the WRONG requests. Routing + # must fail closed: no payload routed for the step (empty sparse + # set), never dense misrouting. + with patch.object(sparse_audio, "_logged_offenders", set()): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0", "r1"], + downstream_req_ids=["r0", "r1"], + multimodal_outputs={ + "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}, + "model_outputs": [torch.zeros(1)], + }, + ) + assert downstream == [] + assert index == {} + assert is_sparse is True + + def test_fail_closed_routing_error_remains_observable(self): + def route_invalid_marker(): + return sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0"], + downstream_req_ids=["r0"], + multimodal_outputs={ + "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}, + "model_outputs": [torch.zeros(1)], + }, + ) + + with ( + patch.object(sparse_audio, "_logged_offenders", set()), + patch.object(sparse_audio, "_fail_closed_last_log_at", None), + patch.object(sparse_audio, "_fail_closed_suppressed", 0), + patch.object(sparse_audio.time, "monotonic", side_effect=[100.0, 120.0, 161.0]), + patch.object(sparse_audio.logger, "error") as err, + ): + for _ in range(3): + assert route_invalid_marker() == ([], {}, True) + + routing_errors = [call for call in err.call_args_list if "no multimodal payload" in call.args[0]] + assert len(routing_errors) == 2 + assert [call.args[1] for call in routing_errors] == [0, 1] + + def test_fail_closed_logging_is_thread_safe(self): + workers = 8 + ready = Barrier(workers) + + def route_invalid_marker(): + return sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0"], + downstream_req_ids=["r0"], + multimodal_outputs={ + "meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}, + }, + ) + + def route_concurrently(_): + ready.wait(timeout=10) + return route_invalid_marker() + + with ( + patch.object(sparse_audio, "_logged_offenders", set()), + patch.object(sparse_audio, "_fail_closed_last_log_at", None), + patch.object(sparse_audio, "_fail_closed_suppressed", 0), + patch.object(sparse_audio.time, "monotonic", return_value=100.0) as clock, + patch.object(sparse_audio.logger, "error") as err, + ): + with ThreadPoolExecutor(max_workers=workers) as executor: + results = list(executor.map(route_concurrently, range(workers))) + assert results == [([], {}, True)] * workers + + clock.return_value = 161.0 + assert route_invalid_marker() == ([], {}, True) + + routing_errors = [call for call in err.call_args_list if "no multimodal payload" in call.args[0]] + assert len(routing_errors) == 2 + assert [call.args[1] for call in routing_errors] == [0, workers - 1] + classifier_errors = [call for call in err.call_args_list if "Illegal sparse-audio marker" in call.args[0]] + assert len(classifier_errors) == 1 + + def test_duplicate_sparse_metadata_must_be_consistent(self): + equal = { + "meta": {"req_id": ["r1"], "sparse_audio": ["1"]}, + "meta.req_id": ["r1"], + "meta.sparse_audio": ["true"], + } + assert sparse_audio.resolve_sparse_mm_req_ids(equal) == ["r1"] + + invalid_declarations = ( + { + "meta": {"req_id": ["r1"], "sparse_audio": ["1"]}, + "meta.req_id": ["r2"], + }, + { + "meta": {"req_id": ["r1"], "sparse_audio": ["1"]}, + "meta.req_id": "r1", + }, + { + "meta": {"req_id": ["r1"], "sparse_audio": ["0"]}, + "meta.sparse_audio": ["1"], + }, + {"meta": {"req_id": ["r1", "r1"], "sparse_audio": ["1"]}}, + ) + with patch.object(sparse_audio, "_logged_offenders", set()): + for declaration in invalid_declarations: + with pytest.raises(sparse_audio.InvalidSparseDeclarationError): + sparse_audio.resolve_sparse_mm_req_ids(declaration) + + def test_conflicting_sparse_metadata_fails_closed_at_public_boundary(self): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r1", "r2"], + downstream_req_ids=["r1", "r2"], + multimodal_outputs={ + "meta": {"req_id": ["r1"], "sparse_audio": ["1"]}, + "meta.req_id": ["r2"], + }, + ) + assert downstream == [] + assert index == {} + assert is_sparse is True + + def test_flat_marker_with_nested_req_ids_is_not_erased(self): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0", "r1"], + downstream_req_ids=["r0", "r1"], + multimodal_outputs={"meta": {"req_id": ["r1"]}, "meta.sparse_audio": ["1"]}, + ) + assert downstream == ["r1"] + assert index == {"r1": 0} + assert is_sparse is True + + def test_sparse_marker_with_unusable_req_ids_fails_closed(self): + # The marker says "sparse" but the alignment list is unusable - the + # declaration as a whole is out of contract; same fail-closed shape. + with patch.object(sparse_audio, "_logged_offenders", set()): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0"], + downstream_req_ids=["r0"], + multimodal_outputs={"meta": {"req_id": "r0", "sparse_audio": ["1"]}}, + ) + assert downstream == [] + assert is_sparse is True + + def test_nested_marker_without_req_id_fails_closed(self): + # {"meta": {"sparse_audio": ["1"]}} - marker set, req_id missing. + # The flattened-encoding fallback must not ERASE the nested marker + # (which used to silently demote the declaration to dense); this is + # an unusable-req_id invalid declaration -> fail closed. + with patch.object(sparse_audio, "_logged_offenders", set()): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0"], + downstream_req_ids=["r0"], + multimodal_outputs={"meta": {"sparse_audio": ["1"]}}, + ) + assert downstream == [] + assert is_sparse is True + + def test_invalid_marker_on_non_audio_pipeline_keeps_dense_routing(self): + # Non-audio pipelines never consult sparse routing; the carrier + # error is logged by the classifier but routing is unaffected. + with patch.object(sparse_audio, "_logged_offenders", set()): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="latent", + req_ids_output_copy=["r0", "r1"], + downstream_req_ids=["r0", "r1"], + multimodal_outputs={"meta": {"req_id": ["r0"], "sparse_audio": torch.tensor([1, 1])}}, + ) + assert downstream == ["r0", "r1"] + assert is_sparse is False + + def test_legal_marker_still_routes_sparse(self): + downstream, index, is_sparse = sparse_audio.resolve_sparse_mm_routing( + engine_output_type="audio", + req_ids_output_copy=["r0", "r1", "r2"], + downstream_req_ids=["r0", "r1", "r2"], + multimodal_outputs={"meta": {"req_id": ["r1", "r2"], "sparse_audio": ["1"]}}, + ) + assert downstream == ["r1", "r2"] + assert index == {"r1": 0, "r2": 1} + assert is_sparse is True + + def test_combined_payload_unwraps_aligned_list(self): + combined = {"codes.audio": {"req-1": [torch.zeros(1), torch.ones(1)]}} + payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-1", list_idx=1) + assert torch.equal(payload["codes.audio"], torch.ones(1)) + + def test_combined_payload_broadcasts_qwen_request_invariant_singletons(self): + req_ids = ("r0", "r1") + tts_embeddings = { + "embed.tts_bos": torch.full((1, 1, 2), 1.0), + "embed.tts_eos": torch.full((1, 1, 2), 2.0), + "embed.tts_pad": torch.full((1, 1, 2), 3.0), + } + combined = {key: {rid: [value.clone()] for rid in req_ids} for key, value in tts_embeddings.items()} + + for idx, rid in enumerate(req_ids): + payload = payload_build.build_omni_mm_payload( + combined_multimodal_outputs=combined, + mm_cpu=None, + rid=rid, + idx=idx, + start=idx, + end=idx + 1, + audio_sparse_output=False, + sparse_mm_index={}, + hidden_seq_len=len(req_ids), + scheduled_seq_len=len(req_ids), + ) + assert payload.keys() == tts_embeddings.keys() + for key, expected in tts_embeddings.items(): + assert torch.equal(payload[key], expected) + + def test_combined_payload_drops_sparse_protocol_metadata(self): + combined = { + "meta.sparse_audio": {"req-1": ["1"]}, + "meta.req_id": {"req-1": ["req-1"]}, + "sr": {"req-1": 22050}, + } + payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-1", list_idx=2) + assert payload == {"sr": 22050} + + def test_combined_payload_misalignment_drops_key_not_request_zero(self): + # A short multi-element list used to silently return `v[0]` - request + # 0's payload shipped under request `rid`. Never substitute another + # request's data; the key is dropped with an error log (same policy + # as the non-combined sparse mismatch) rather than raised - v1 + # runners don't catch per-request exceptions, so a raise would turn a + # per-request protocol break into a whole-stage outage. + combined = { + "codes.audio": {"req-3": [torch.zeros(1), torch.ones(1)]}, + "sr": {"req-3": 22050}, + } + payload = payload_build.build_combined_prefix_cache_mm_payload(combined, rid="req-3", list_idx=2) + assert "codes.audio" not in payload + assert payload["sr"] == 22050 + + def test_direct_path_sparse_mismatch_drops_key_and_keeps_payload_shape(self): + # Mirror of the combined-path policy on the non-prefix-cache path: + # a sparse index beyond the per-request list drops that key (error + # log) and the rest of the request's payload survives. + payload = payload_build.build_omni_mm_payload( + combined_multimodal_outputs=None, + mm_cpu={"codes.audio": [torch.zeros(1)]}, + rid="r2", + idx=1, + start=0, + end=1, + audio_sparse_output=True, + sparse_mm_index={"r2": 1}, + hidden_seq_len=1, + scheduled_seq_len=1, + ) + assert payload == {} + + def test_combined_payload_uses_sparse_index_under_sparse_routing(self): + # Batch [r0, r1, r2], sparse outputs only for [r1, r2]: producers + # align per-request lists to the SPARSE order, so r1 (batch idx 1, + # sparse idx 0) must get element 0 and r2 (batch idx 2, sparse idx 1) + # element 1. Indexing by batch idx shipped r2's payload to r1 and + # went out of range for r2. The sparse->list_idx resolution lives in + # `_build_omni_mm_payload` (the combined helper itself is + # sparse-agnostic), so this exercises the caller. + sparse_mm_index = {"r1": 0, "r2": 1} + combined = { + "codes.audio": { + "r1": [torch.zeros(1), torch.ones(1)], + "r2": [torch.zeros(1), torch.ones(1)], + } + } + payloads = {} + for rid, idx in (("r1", 1), ("r2", 2)): + payloads[rid] = payload_build.build_omni_mm_payload( + combined_multimodal_outputs=combined, + mm_cpu=None, + rid=rid, + idx=idx, + start=0, + end=1, + audio_sparse_output=True, + sparse_mm_index=sparse_mm_index, + hidden_seq_len=1, + scheduled_seq_len=1, + ) + assert torch.equal(payloads["r1"]["codes.audio"], torch.zeros(1)) + assert torch.equal(payloads["r2"]["codes.audio"], torch.ones(1)) + + +class TestMergeModelKvTransferMetadata: + """RFC #5450 C4: model KV-transfer metadata must merge into a COPY of the + scheduler's finished-request data - `scheduler_output` can be shared with + the engine-core process and must never be mutated by the runner.""" + + def _runner_with_model(self, model): + runner = object.__new__(GPUARModelRunner) + runner.model = model + return runner + + def test_merge_does_not_mutate_scheduler_data(self): + class _Model: + def get_kv_transfer_metadata(self, req_id, num_computed_tokens): + return {"talker_codes": [7], "seen_tokens": num_computed_tokens} + + original_meta = {"a": 1} + original = {"r1": {"seq_len": 3, "block_ids": [4], "custom_metadata": original_meta}} + runner = self._runner_with_model(_Model()) + + merged = runner._merge_model_kv_transfer_metadata(original) + + # The manager sees the merged view... + assert merged["r1"]["custom_metadata"] == {"a": 1, "talker_codes": [7], "seen_tokens": 3} + assert merged["r1"]["seq_len"] == 3 + # ...while the engine-shared originals are untouched. + assert original["r1"]["custom_metadata"] is original_meta + assert original_meta == {"a": 1} + assert "talker_codes" not in original["r1"]["custom_metadata"] + + def test_merge_without_model_meta_passes_entry_through(self): + class _Model: + def get_kv_transfer_metadata(self, req_id, num_computed_tokens): + return None + + data = {"seq_len": 2, "block_ids": [1]} + runner = self._runner_with_model(_Model()) + + merged = runner._merge_model_kv_transfer_metadata({"r1": data}) + + assert merged["r1"] is data + + def test_merge_failure_raises_instead_of_shipping_partial_transfer(self): + # RFC #5450 C7: the bare `except Exception -> logger.warning` used to + # drop the model's custom metadata and let the KV transfer proceed + # incomplete; the failure then surfaced much later on the consumer + # stage. It must fail at the defect instead. + class _Model: + def get_kv_transfer_metadata(self, req_id, num_computed_tokens): + raise RuntimeError("boom") + + data = {"seq_len": 2, "block_ids": [1], "custom_metadata": {"a": 1}} + runner = self._runner_with_model(_Model()) + + with pytest.raises(RuntimeError, match="boom"): + runner._merge_model_kv_transfer_metadata({"r1": data}) + assert data["custom_metadata"] == {"a": 1} + + +class TestDownstreamPayloadMemoization: + """RFC #5450 C7: `_request_needs_downstream_stage_payload` must not + memoize before the deciding marker exists - the final-stage id arrives via + `model_intermediate_buffer`, which can be unpopulated on the first call.""" + + def _runner(self, stages): + runner = object.__new__(GPUARModelRunner) + runner._downstream_payload_cache = {} + runner._request_final_stage_id = lambda rid: stages[rid] + return runner + + def test_missing_marker_defaults_true_without_memoizing(self): + stages = {"r1": None} + runner = self._runner(stages) + + assert runner._request_needs_downstream_stage_payload("r1") is True + # Not cached: the answer must refresh once the marker arrives. + assert "r1" not in runner._downstream_payload_cache + + stages["r1"] = 0 # marker lands: this stage IS the final stage + assert runner._request_needs_downstream_stage_payload("r1") is False + assert runner._downstream_payload_cache["r1"] is False + + def test_memoizes_once_marker_known(self): + stages = {"r1": 2} + runner = self._runner(stages) + + assert runner._request_needs_downstream_stage_payload("r1") is True + assert runner._downstream_payload_cache["r1"] is True + # Cached value now authoritative for the request's lifetime. + stages["r1"] = 0 + assert runner._request_needs_downstream_stage_payload("r1") is True + + +# Matches a real ASSIGNMENT of the capability (optionally `self.`-qualified +# and/or annotated), not a bare mention: a comment, a docstring, or a read +# such as `getattr(self.model, "prefer_model_sampler", False)` must not count +# as declaring it. `=(?!=)` keeps `==` comparisons out. +_PREFER_MODEL_SAMPLER_ASSIGNMENT = re.compile( + r"^[ \t]*(?:self\.)?prefer_model_sampler[ \t]*(?::[^=\n]*)?=(?!=)[ \t]*(?P[^\n]*)", + re.MULTILINE, +) + + +def _declares_prefer_model_sampler(source: str) -> bool: + """Whether `source` opts INTO the model-sampler contract. + + An explicit `prefer_model_sampler = False` is an opt-OUT and does not + count; any other right-hand side does, including runtime-conditional + forms (minicpmo_4_5 assigns `self.model_stage in {...}`), which can + evaluate True. + """ + for match in _PREFER_MODEL_SAMPLER_ASSIGNMENT.finditer(source): + if match.group("rhs").split("#")[0].strip() != "False": + return True + return False + + +class TestPreferModelSamplerNoneFallback: + """RFC #5450 C7: a `prefer_model_sampler` model returning None from + `sample()` DECLINES the step and the runner falls back to the default + sampler. The fallback is load-bearing (cosyvoice3/glm_tts decline for + empty-logits / inapplicable-stage steps), so it must keep working - now + as a documented contract with a one-time visibility log.""" + + def _runner(self, model_sample_result, default_result): + runner = object.__new__(GPUARModelRunner) + runner.input_batch = SimpleNamespace( + sampling_metadata="smd", + update_async_output_token_ids=lambda: None, + ) + runner.model = SimpleNamespace( + prefer_model_sampler=True, + sample=lambda logits, metadata: model_sample_result, + ) + runner._sampling_metadata_for_model_sampler = lambda smd: smd + runner._resolve_duplex_sampling_hook = lambda: None + runner.sampler = lambda logits, sampling_metadata: default_result + return runner + + def test_none_falls_back_to_default_sampler_and_warns_once(self): + import logging + + from vllm_omni.worker import gpu_ar_model_runner as ar_module + + class _UniqueDecliningSampler: + prefer_model_sampler = True + + @staticmethod + def sample(logits, metadata): + return None + + sentinel = object() + runner = self._runner(model_sample_result=None, default_result=sentinel) + runner.model = _UniqueDecliningSampler() + + records: list[logging.LogRecord] = [] + + class _Capture(logging.Handler): + def emit(self, record): + records.append(record) + + handler = _Capture(level=logging.WARNING) + ar_module.logger.addHandler(handler) + try: + for _ in range(3): + out = GPUARModelRunner._sample(runner, torch.zeros(1, 4), None) + assert out is sentinel + finally: + ar_module.logger.removeHandler(handler) + + fallback_records = [r for r in records if "falling back to the default sampler" in r.getMessage()] + assert len(fallback_records) == 1 + assert "_UniqueDecliningSampler" in fallback_records[0].getMessage() + + def test_model_sampler_output_short_circuits_default(self): + model_out = object() + runner = self._runner(model_sample_result=model_out, default_result=object()) + + out = GPUARModelRunner._sample(runner, torch.zeros(1, 4), None) + + assert out is model_out + + def test_declaration_matcher_ignores_mentions_and_opt_outs(self): + # Guards the guard: the inventory below is only meaningful if + # "declares" means an actual opt-in assignment. + assert _declares_prefer_model_sampler(" prefer_model_sampler = True") + assert _declares_prefer_model_sampler(" prefer_model_sampler: bool = True") + assert _declares_prefer_model_sampler(' self.prefer_model_sampler = stage in {"llm"}') + assert not _declares_prefer_model_sampler(" prefer_model_sampler = False") + assert not _declares_prefer_model_sampler(" prefer_model_sampler = False # opted out") + assert not _declares_prefer_model_sampler(" # prefer_model_sampler = True (not anymore)") + assert not _declares_prefer_model_sampler(' """Does not use prefer_model_sampler."""') + assert not _declares_prefer_model_sampler(' x = getattr(model, "prefer_model_sampler", False)') + assert not _declares_prefer_model_sampler(" if prefer_model_sampler == True:") + + def test_in_tree_declarer_inventory_is_exact(self): + # The contract's audience: an exact inventory, so ADDING a declarer + # fails this test and the new model owner consciously signs up for the + # documented None-fallback semantics (update the set when they do). + import pathlib + + import vllm_omni + + models_dir = pathlib.Path(vllm_omni.__file__).parent / "model_executor" / "models" + declarers = { + p.parent.name + for p in models_dir.rglob("*.py") + if _declares_prefer_model_sampler(p.read_text(encoding="utf-8", errors="ignore")) + } + expected = { + "cosyvoice3", + "glm_tts", + "higgs_audio_v2", + "higgs_audio_v3", + "hunyuan_image3", + "minicpmo_4_5", + "minimax_music3", + } + assert declarers == expected, ( + "The set of models declaring `prefer_model_sampler` changed:\n" + f" added: {sorted(declarers - expected) or 'none'}\n" + f" removed: {sorted(expected - declarers) or 'none'}\n\n" + "This is a deliberate consent gate, not a broken test. Declaring " + "`prefer_model_sampler` opts a model into the runner's documented " + "None-fallback contract: `model.sample()` may return None to decline " + "a step, and the runner then falls back to the DEFAULT sampler " + "(see `GPUARModelRunner._sample`; RFC #5450 C7.3).\n\n" + "If you ADDED a declarer: confirm your model relies on -- or at " + "least tolerates -- that fallback, then add its directory name to " + "`expected` above. If you REMOVED one, drop its name." + ) diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index a6df4732f..b97168161 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -16,14 +16,14 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class _DummyInputBatch: - def __init__(self): - self.req_ids = ["req-1"] - self.req_id_to_index = {"req-1": 0} - self.num_reqs = 1 + def __init__(self, num_reqs: int = 1): + self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)] + self.req_id_to_index = {rid: i for i, rid in enumerate(self.req_ids)} + self.num_reqs = num_reqs self.vocab_size = 10 -def _make_runner(multimodal_outputs): +def _make_runner(multimodal_outputs, num_reqs: int = 1): runner = object.__new__(GPUGenerationModelRunner) runner.execute_model_state = ExecuteModelState( None, @@ -40,7 +40,7 @@ def _make_runner(multimodal_outputs): None, ) runner.kv_connector_output = None - runner.input_batch = _DummyInputBatch() + runner.input_batch = _DummyInputBatch(num_reqs) runner.use_async_scheduling = False runner.device = torch.device("cpu") runner.supports_mm_inputs = False @@ -111,6 +111,12 @@ def _make_guard_runner(): runner.cache_config = CacheConfig() runner.input_batch = _DummyInputBatch() runner.model = object() + runner.parallel_config = SimpleNamespace( + distributed_executor_backend=None, + data_parallel_size=1, + ) + runner._dummy_run_calls = [] + runner._dummy_run = lambda n: runner._dummy_run_calls.append(n) runner._update_states = lambda scheduler_output: None runner.synchronize_input_prep = contextlib.nullcontext runner.attach_omni_connector_output = lambda result: result @@ -121,9 +127,112 @@ def _make_guard_runner(): def test_execute_model_skips_non_positive_scheduled_span(monkeypatch, total): """#5196: a negative span is truthy, so it used to reach `_prepare_inputs`, whose `assert total_num_scheduled_tokens > 0` killed the stage EngineCore.""" + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False) monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False) runner = _make_guard_runner() output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total)) assert output is EMPTY_MODEL_RUNNER_OUTPUT + assert runner._dummy_run_calls == [] + + +@pytest.mark.parametrize("total", [-1, 0]) +def test_execute_model_zero_tokens_runs_dp_dummy_run(monkeypatch, total): + """A 0-token step under external_launcher + DP>1 must `_dummy_run(1)` so + `coordinate_batch_across_dp` stays in sync (it used to be shadowed by an + unconditional early return and never ran -> DP hang risk).""" + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False) + monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False) + runner = _make_guard_runner() + runner.parallel_config = SimpleNamespace( + distributed_executor_backend="external_launcher", + data_parallel_size=2, + ) + + output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total)) + + assert output is EMPTY_MODEL_RUNNER_OUTPUT + assert runner._dummy_run_calls == [1] + + +def test_execute_model_zero_tokens_kv_connector_no_forward(monkeypatch): + """With a KV-transfer group, a 0-token step must go through + `kv_connector_no_forward` (also shadowed by the old early return).""" + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False) + monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: True) + runner = _make_guard_runner() + sentinel = object() + calls = [] + runner.kv_connector_no_forward = lambda scheduler_output, vllm_config: ( + calls.append(scheduler_output), + sentinel, + )[1] + runner.vllm_config = object() + + output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0)) + + assert output is sentinel + assert len(calls) == 1 + + +def test_execute_model_zero_tokens_ec_producer_takes_encoder_path(monkeypatch): + """AR lock-step: the EC-producer block sits ABOVE the zero-token block + (matching GPUARModelRunner), so a 0-token step on an EC producer returns + the empty ENCODER output. Before the C1 reorder, the unconditional early + return won and produced EMPTY_MODEL_RUNNER_OUTPUT instead — this is the + one cell whose behaviour the reorder changed.""" + monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: True) + monkeypatch.setattr(gen_runner_module, "get_ec_transfer", lambda: SimpleNamespace(is_consumer=False)) + monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False) + sentinel = object() + monkeypatch.setattr(gen_runner_module, "make_empty_encoder_model_runner_output", lambda so: sentinel) + runner = _make_guard_runner() + runner.encoder_cache = {} + runner.maybe_get_ec_connector_output = lambda scheduler_output, encoder_cache: contextlib.nullcontext() + encoder_calls = [] + runner._execute_mm_encoder = lambda so: encoder_calls.append(so) + + output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0)) + + assert output is sentinel + assert len(encoder_calls) == 1 + assert runner._dummy_run_calls == [] + + +class TestSampleTokensBatched: + """RFC #5450 C5: `sample_tokens` must slice per request for num_reqs > 1. + The tensor branch's asserts used to force num_reqs == 1 and the list + branch built a length-1 payload misaligned with `req_ids`.""" + + def test_tensor_output_sliced_per_request(self): + multimodal_outputs = torch.arange(12, dtype=torch.float32).reshape(2, 2, 3) + runner = _make_runner(multimodal_outputs, num_reqs=2) + + output = GPUGenerationModelRunner.sample_tokens(runner) + + assert len(output.multimodal_outputs) == 2 + assert torch.equal(output.multimodal_outputs[0]["model_outputs"], multimodal_outputs[0]) + assert torch.equal(output.multimodal_outputs[1]["model_outputs"], multimodal_outputs[1]) + + def test_list_output_one_entry_per_request(self): + entries = [torch.zeros(2, 1), torch.ones(2, 1)] + runner = _make_runner(list(entries), num_reqs=2) + + output = GPUGenerationModelRunner.sample_tokens(runner) + + assert len(output.multimodal_outputs) == 2 + assert torch.equal(output.multimodal_outputs[0]["model_outputs"], entries[0]) + assert torch.equal(output.multimodal_outputs[1]["model_outputs"], entries[1]) + + def test_tensor_output_batch_mismatch_raises(self): + runner = _make_runner(torch.randn(1, 2, 3), num_reqs=2) + + with pytest.raises(ValueError, match="leading dim 1 .* 2 requests"): + GPUGenerationModelRunner.sample_tokens(runner) + + def test_list_output_batch_mismatch_raises(self): + runner = _make_runner([torch.randn(2, 1)], num_reqs=2) + + with pytest.raises(ValueError, match="length 1 .* 2 requests"): + GPUGenerationModelRunner.sample_tokens(runner) diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 3a1677d49..7dda87e96 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -3013,9 +3013,18 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): merged_cpu = merged.detach().cpu().float() mm["model_outputs"] = list(merged_cpu.split(sizes)) mm["sr"] = [sr for _ in ready_req_ids] + # Dense mode still yields a strict SUBSET of the batch when + # some requests emit no audio this step (e.g. prefill + # phase). Without the marker the runner indexes these + # per-request lists by batch position and misroutes audio + # across requests; the marker declares meta.req_id + # alignment. + mm["meta"] = {"req_id": ready_req_ids, "sparse_audio": ["1"]} else: mm["model_outputs"] = list(audio_by_req.values()) mm["sr"] = [sr for _ in audio_by_req] + # Same subset hazard as the coalesce branch above. + mm["meta"] = {"req_id": list(audio_by_req), "sparse_audio": ["1"]} elif self._uses_sparse_audio_outputs(): mm["model_outputs"] = [] mm["sr"] = [] diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index e66f759f8..c87d17cc2 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -51,12 +51,13 @@ from vllm_omni.utils.mm_outputs import ( build_mm_cpu, partition_payload_list, snapshot_mm_payload, - to_payload_element, ) from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin +from vllm_omni.worker.output.payload_build import build_omni_mm_payload from vllm_omni.worker.runner_assisted_metadata import RunnerAssistedFullAttentionMetadataRequest -from vllm_omni.worker.sampling_utils import sanitize_min_tokens_stop_ids +from vllm_omni.worker.sampling_utils import clamp_prompt_ids_to_penalty_padding, sanitize_min_tokens_stop_ids +from vllm_omni.worker.sparse_audio import resolve_sparse_mm_routing logger = init_logger(__name__) @@ -457,9 +458,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex cached = self._downstream_payload_cache.get(req_id) if cached is not None: return cached - # Conservative default: keep payload if marker is missing. final_stage_id = self._request_final_stage_id(req_id) - needs_payload = final_stage_id is None or final_stage_id > 0 + if final_stage_id is None: + # Conservative default while the marker is missing: keep the + # payload, but do NOT memoize - the marker arrives via + # `model_intermediate_buffer`, which may be unpopulated on the + # first call (memoizing here pinned the request to True forever, + # never refreshing once the marker landed). + return True + needs_payload = final_stage_id > 0 self._downstream_payload_cache[req_id] = needs_payload return needs_payload @@ -472,50 +479,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex downstream_req_ids = req_ids_output_copy return engine_output_type, downstream_req_ids - @staticmethod - def _sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None: - if not isinstance(multimodal_outputs, dict): - return None - meta = multimodal_outputs.get("meta") - req_ids = None - sparse_audio = False - if isinstance(meta, dict): - req_ids = meta.get("req_id") - sparse_audio = GPUARModelRunner._is_sparse_audio_marker(meta.get("sparse_audio")) - if req_ids is None: - req_ids = multimodal_outputs.get("meta.req_id") - sparse_audio = GPUARModelRunner._is_sparse_audio_marker(multimodal_outputs.get("meta.sparse_audio")) - if not sparse_audio: - return None - if not isinstance(req_ids, list): - return None - return [rid for rid in req_ids if isinstance(rid, str)] - - @staticmethod - def _resolve_sparse_mm_routing( - *, - engine_output_type: str, - req_ids_output_copy: list[str], - downstream_req_ids: list[str], - multimodal_outputs: Any, - ) -> tuple[list[str], dict[str, int], bool]: - sparse_mm_req_ids = GPUARModelRunner._sparse_mm_req_ids(multimodal_outputs) - sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])} - if engine_output_type != "audio" or sparse_mm_req_ids is None: - return downstream_req_ids, sparse_mm_index, False - - sparse_req_id_set = set(sparse_mm_req_ids) - sparse_downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set] - return sparse_downstream_req_ids, sparse_mm_index, True - - @staticmethod - def _is_sparse_audio_marker(value: Any) -> bool: - if isinstance(value, list): - return any(str(item).lower() in ("1", "true", "yes", "on") for item in value) - if isinstance(value, str): - return value.lower() in ("1", "true", "yes", "on") - return bool(value) - def capture_model(self) -> int: result = super().capture_model() self._capture_talker_mtp_graphs() @@ -862,79 +825,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex ) return hidden_states_cpu, combined_hidden_states, combined_multimodal_outputs - @staticmethod - def _build_combined_prefix_cache_mm_payload( - combined_multimodal_outputs: dict, - *, - rid: str, - idx: int, - ) -> dict[str, object]: - def _unwrap_lists(v): - if isinstance(v, list): - return v[idx] if idx < len(v) else v[0] - if isinstance(v, dict): - return {k: _unwrap_lists(sv) for k, sv in v.items()} - return v - - return { - mm_key: _unwrap_lists(combined_multimodal_outputs[mm_key][rid]) - for mm_key in combined_multimodal_outputs.keys() - } - - def _build_omni_mm_payload( - self, - *, - combined_multimodal_outputs: dict | None, - mm_cpu: dict[str, object] | None, - rid: str, - idx: int, - start: int, - end: int, - audio_sparse_output: bool, - sparse_mm_index: dict[str, int], - hidden_seq_len: int, - scheduled_seq_len: int, - ) -> dict[str, object]: - if combined_multimodal_outputs: - return self._build_combined_prefix_cache_mm_payload( - combined_multimodal_outputs, - rid=rid, - idx=idx, - ) - - mm_payload: dict[str, object] = {} - if not mm_cpu: - return mm_payload - - for mm_key, mm_val in mm_cpu.items(): - if mm_key in {"meta.req_id", "meta.sparse_audio"}: - continue - if audio_sparse_output and isinstance(mm_val, list): - sparse_idx = sparse_mm_index.get(rid) - if sparse_idx is None: - continue - if sparse_idx >= len(mm_val): - logger.warning( - "Sparse multimodal payload mismatch for request %s: index %d >= %d.", - rid, - sparse_idx, - len(mm_val), - ) - continue - sparse_val = mm_val[sparse_idx] - mm_payload[mm_key] = sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val - continue - mm_payload[mm_key] = to_payload_element( - element=mm_val, - idx=idx, - start=start, - end=end, - pass_lists_through=False, - seq_len=hidden_seq_len, - scheduled_seq_len=scheduled_seq_len, - ) - return mm_payload - def _build_omni_pooler_payload( self, *, @@ -967,7 +857,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if req_hidden_states is not None: payload["hidden"] = req_hidden_states - mm_payload = self._build_omni_mm_payload( + mm_payload = build_omni_mm_payload( combined_multimodal_outputs=combined_multimodal_outputs, mm_cpu=mm_cpu, rid=rid, @@ -982,6 +872,47 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex payload.update(mm_payload) return payload + def _merge_model_kv_transfer_metadata( + self, + finished_reqs: dict[str, dict[str, Any]], + ) -> dict[str, dict[str, Any]]: + """Merge model-provided KV-transfer metadata copy-on-write. + + The outer mapping is always new; a per-request dict is shallow-copied + (with a fresh merged ``custom_metadata``) only when the model supplies + metadata for it, and passes through unchanged otherwise. + `scheduler_output` (and its per-request dicts) can be shared with the + engine-core process, so nothing here mutates them — the ngram block + in `execute_model` `replace()`-copies for exactly the same reason. The + merged mapping is only consumed by + `kv_transfer_manager.handle_finished_requests_kv_transfer`. + """ + merged: dict[str, dict[str, Any]] = {} + for req_id, data in finished_reqs.items(): + try: + # NOTE: seq_len is the same as num_computed_tokens_cpu in current + # async scheduling, since both exclude async placeholders. We use + # seq_len since we control it, just in case upstream async scheduler + # semantics change in the future. + model_meta = self.model.get_kv_transfer_metadata( + req_id, + num_computed_tokens=data.get("seq_len"), + ) + except Exception: + # A swallowed failure here used to ship the KV transfer WITHOUT + # the model's custom metadata and surface much later as a + # corrupt/incomplete transfer on the consumer stage. Fail at + # the point of the defect instead. + logger.exception("Failed to get custom KV-transfer metadata from model for %s", req_id) + raise + if model_meta: + data = { + **data, + "custom_metadata": {**(data.get("custom_metadata") or {}), **model_meta}, + } + merged[req_id] = data + return merged + @torch.inference_mode() def execute_model( self, @@ -1012,24 +943,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex # [Omni] Handle KV transfer BEFORE updating states (which removes finished requests) finished_reqs = getattr(scheduler_output, "finished_requests_needing_kv_transfer", {}) if finished_reqs and hasattr(self.model, "get_kv_transfer_metadata"): - for req_id, data in finished_reqs.items(): - try: - # NOTE: seq_len is the same as num_computed_tokens_cpu in current - # async scheduling, since both exclude async placeholders. We use - # seq_len since we control it, just in case upstream async scheduler - # semantics change in the future. - num_computed = data.get("seq_len") - - model_meta = self.model.get_kv_transfer_metadata( - req_id, - num_computed_tokens=num_computed, - ) - if model_meta: - existing = data.get("custom_metadata") or {} - existing.update(model_meta) - data["custom_metadata"] = existing - except Exception as e: - logger.warning(f"Failed to get custom metadata from model for {req_id}: {e}") + finished_reqs = self._merge_model_kv_transfer_metadata(finished_reqs) self.kv_extracted_req_ids = self.kv_transfer_manager.handle_finished_requests_kv_transfer( finished_reqs=finished_reqs, kv_caches=self.kv_caches, @@ -1467,6 +1381,16 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex logits: torch.Tensor | None, spec_decode_metadata: Any, ): + """Sample tokens, preferring the model's own sampler when declared. + + Model-interface contract for `prefer_model_sampler` models: a custom + `model.sample(logits, sampling_metadata)` may return ``None`` to + DECLINE the step, and the runner then falls back to the default + sampler. Declarers rely on this deliberately for empty-logits or + inapplicable-stage steps (e.g. cosyvoice3, glm_tts), so the + fallthrough is load-bearing, not an error path; it is logged once per + model for visibility. + """ sampling_metadata = self.input_batch.sampling_metadata if spec_decode_metadata is None: model_sample = getattr(self.model, "sample", None) @@ -1487,6 +1411,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex sampler_output = model_sample(logits, prepared_sampling_metadata) if sampler_output is not None: return sampler_output + # Contract: None => fall back to the default sampler (see + # docstring). If a custom sampler returns None by accident, + # its tokens silently come from the default sampler - the + # one-time log is the visibility hook for that case. + logger.warning_once( + "prefer_model_sampler model %s returned None from sample(); falling back to the default sampler.", + type(self.model).__name__, + ) return self.sampler( logits=logits, sampling_metadata=sampling_metadata, @@ -1889,7 +1821,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex combined_multimodal_outputs = None engine_output_type, downstream_req_ids = self._resolve_pooler_payload_req_ids(req_ids_output_copy) - downstream_req_ids, sparse_mm_index, audio_sparse_output = self._resolve_sparse_mm_routing( + downstream_req_ids, sparse_mm_index, audio_sparse_output = resolve_sparse_mm_routing( engine_output_type=engine_output_type, req_ids_output_copy=req_ids_output_copy, downstream_req_ids=downstream_req_ids, @@ -2089,13 +2021,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if grammar_output is not None: apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits) - # Correct padding values of prompt_token_ids to match the logits vocabulary size + # Correct padding values of prompt_token_ids to match the logits vocabulary size. if logits is not None and not self.input_batch.sampling_metadata.no_penalties: smd = self.input_batch.sampling_metadata if smd.prompt_token_ids is not None: logits_vocab = logits.shape[-1] if self.input_batch.vocab_size > logits_vocab: - smd.prompt_token_ids = smd.prompt_token_ids.clamp(max=logits_vocab) + smd.prompt_token_ids = clamp_prompt_ids_to_penalty_padding(smd.prompt_token_ids, logits_vocab) # Drop min-tokens stop ids the head cannot emit (e.g. the text # tokenizer EOS folded into all_stop_token_ids on a narrow codec diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 7cf26d005..a397c2e9d 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -167,10 +167,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"): self.model.on_requests_finished(scheduler_output.finished_req_ids) - # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). - if scheduler_output.total_num_scheduled_tokens <= 0: - return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) - if has_ec_transfer() and not get_ec_transfer().is_consumer: with self.maybe_get_ec_connector_output( scheduler_output, @@ -179,7 +175,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin self._execute_mm_encoder(scheduler_output) return self.attach_omni_connector_output(make_empty_encoder_model_runner_output(scheduler_output)) - if not num_scheduled_tokens: + # OMNI: keep this block in lock-step with the same block in + # GPUARModelRunner.execute_model. + # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). + if num_scheduled_tokens <= 0: if ( self.parallel_config.distributed_executor_backend == "external_launcher" and self.parallel_config.data_parallel_size > 1 @@ -439,16 +438,29 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin # pooler_output is no longer used for multimodal data. per_req_payloads: list[dict[str, object]] = [] if isinstance(multimodal_outputs_raw, torch.Tensor): - assert multimodal_outputs_raw.shape[0] == 1, ( - "model should return a single tensor, to return multiple tensors, use a dict" - ) - assert multimodal_outputs_raw.shape[0] == self.input_batch.num_reqs - for i in range(self.input_batch.num_reqs): + # One row per request. The old asserts (`shape[0] == 1` AND + # `shape[0] == num_reqs`) jointly forced num_reqs == 1, silently + # rejecting batched steps the generation scheduler does admit. + num_reqs = self.input_batch.num_reqs + if multimodal_outputs_raw.shape[0] != num_reqs: + raise ValueError( + f"Multimodal output tensor has leading dim {multimodal_outputs_raw.shape[0]} " + f"but the batch has {num_reqs} requests (one row per request; " + "to return multiple tensors per request, use a dict)." + ) + for i in range(num_reqs): per_req_payloads.append({"model_outputs": multimodal_outputs_raw[i].detach().to("cpu").contiguous()}) elif isinstance(multimodal_outputs_raw, list): - assert len(multimodal_outputs_raw) == 1, ( - "model should return a single list, to return multiple lists, use a dict" - ) + # One entry per request. The old `len == 1` assert did not check + # num_reqs, so a batched step built a length-1 payload list that + # misaligned with `req_ids` downstream. + num_reqs = self.input_batch.num_reqs + if len(multimodal_outputs_raw) != num_reqs: + raise ValueError( + f"Multimodal output list has length {len(multimodal_outputs_raw)} " + f"but the batch has {num_reqs} requests (one entry per request; " + "to return multiple lists per request, use a dict)." + ) for out in multimodal_outputs_raw: per_req_payloads.append( {"model_outputs": out.detach().to("cpu").contiguous() if out is not None else None} @@ -794,6 +806,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] model_kwargs = self._init_model_kwargs() + elif getattr(getattr(self, "model", None), "has_preprocess", False): + # Capture CUDA graph with inputs_embeds path so replay reads + # from the same buffer that _preprocess writes into. + input_ids = self.input_ids.gpu[:num_tokens_padded] + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] else: input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None diff --git a/vllm_omni/worker/output/__init__.py b/vllm_omni/worker/output/__init__.py new file mode 100644 index 000000000..33ea418a3 --- /dev/null +++ b/vllm_omni/worker/output/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Worker output helpers. + +Package seeded ahead of the async-output refactor (RFC #5450 series, G5-C2), +which relocates the remaining payload-copy helpers, snapshot types, and +`ExecuteModelState` here. Currently hosts the per-request multimodal payload +builders. +""" + +from vllm_omni.worker.output.payload_build import ( + build_combined_prefix_cache_mm_payload, + build_omni_mm_payload, + unwrap_combined_payload_value, +) + +__all__ = [ + "build_combined_prefix_cache_mm_payload", + "build_omni_mm_payload", + "unwrap_combined_payload_value", +] diff --git a/vllm_omni/worker/output/payload_build.py b/vllm_omni/worker/output/payload_build.py new file mode 100644 index 000000000..0f65ac293 --- /dev/null +++ b/vllm_omni/worker/output/payload_build.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Per-request multimodal payload builders. + +State-free extraction helpers that turn a step's `multimodal_outputs` +(combined prefix-cache-merged form, or the direct `mm_cpu` form) into one +request's payload dict. Sparse-routing context (`audio_sparse_output`, +`sparse_mm_index`) is resolved upstream by +`vllm_omni.worker.sparse_audio.resolve_sparse_mm_routing` and arrives here +as plain arguments. +""" + +from typing import Any + +import torch +from vllm.logger import init_logger + +from vllm_omni.utils.mm_outputs import to_payload_element + +logger = init_logger(__name__) + +_SPARSE_AUDIO_PROTOCOL_KEYS = frozenset({"meta.req_id", "meta.sparse_audio"}) + + +def unwrap_combined_payload_value( + value: Any, + *, + rid: str, + list_idx: int, + mm_key: str, +) -> tuple[bool, Any]: + """Unwrap one value from the prefix-cache-merged payload. + + Returns ``(keep, unwrapped)``: ``keep`` is False when a non-singleton + per-request list is misaligned and must be dropped. A singleton list is + request-invariant passthrough data and is shared across the batch. + """ + if isinstance(value, list): + if len(value) == 1: + # Prefix-cache passthrough intentionally preserves list + # containers and copies the whole list under every request id. + # Qwen3-Omni uses this shape for its request-invariant TTS + # BOS/EOS/PAD embeddings, so every batch index unwraps the same + # singleton. Sparse one-request lists are safe here too: routing + # limits them to the declared request at sparse index zero. + return True, value[0] + if list_idx < len(value): + return True, value[list_idx] + # A non-singleton list shorter than the request's index is per-request + # data that lost alignment. Never substitute another request's + # element (`v[0]` used to ship request 0's payload as `rid`'s); drop + # the key loudly instead — the SAME protocol-break policy as the + # sparse mismatch in + # `build_omni_mm_payload`. Not a raise: v1 runners do not catch + # per-request exceptions, so raising would turn a per-request + # protocol break into a whole-stage outage. + logger.error( + "Combined prefix-cache mm payload misaligned for request %s: index %d out of " + "range for per-request list of length %d; dropping key %s.", + rid, + list_idx, + len(value), + mm_key, + ) + return False, None + if isinstance(value, dict): + nested = {} + for k, sv in value.items(): + keep, unwrapped = unwrap_combined_payload_value(sv, rid=rid, list_idx=list_idx, mm_key=mm_key) + if keep: + nested[k] = unwrapped + return True, nested + return True, value + + +def build_combined_prefix_cache_mm_payload( + combined_multimodal_outputs: dict, + *, + rid: str, + list_idx: int, +) -> dict[str, object]: + """Unwrap one request's entry from the prefix-cache-merged payload. + + Sparse-agnostic by design: ``list_idx`` arrives pre-resolved by the + caller (`build_omni_mm_payload`, which owns the sparse-routing context) + — batch index normally, sparse index under sparse routing. + """ + payload: dict[str, object] = {} + for mm_key in combined_multimodal_outputs.keys(): + if mm_key in _SPARSE_AUDIO_PROTOCOL_KEYS: + continue + keep, unwrapped = unwrap_combined_payload_value( + combined_multimodal_outputs[mm_key][rid], + rid=rid, + list_idx=list_idx, + mm_key=mm_key, + ) + if keep: + payload[mm_key] = unwrapped + return payload + + +def build_omni_mm_payload( + *, + combined_multimodal_outputs: dict | None, + mm_cpu: dict[str, object] | None, + rid: str, + idx: int, + start: int, + end: int, + audio_sparse_output: bool, + sparse_mm_index: dict[str, int], + hidden_seq_len: int, + scheduled_seq_len: int, +) -> dict[str, object]: + """Build one request's multimodal payload for the step.""" + if combined_multimodal_outputs: + # Sparse-audio producers emit per-request lists aligned to the + # sparse request order (`meta.req_id`), not the batch order; under + # sparse routing the request's sparse index selects its element. + list_idx = sparse_mm_index.get(rid, idx) if audio_sparse_output else idx + return build_combined_prefix_cache_mm_payload( + combined_multimodal_outputs, + rid=rid, + list_idx=list_idx, + ) + + mm_payload: dict[str, object] = {} + if not mm_cpu: + return mm_payload + + for mm_key, mm_val in mm_cpu.items(): + if mm_key in _SPARSE_AUDIO_PROTOCOL_KEYS: + continue + if audio_sparse_output and isinstance(mm_val, list): + sparse_idx = sparse_mm_index.get(rid) + if sparse_idx is None: + continue + if sparse_idx >= len(mm_val): + # Misalignment means the sparse-marker protocol broke; the + # request's output for this key is dropped, so be loud. + logger.error( + "Sparse multimodal payload mismatch for request %s: index %d >= %d; dropping key %s.", + rid, + sparse_idx, + len(mm_val), + mm_key, + ) + continue + sparse_val = mm_val[sparse_idx] + mm_payload[mm_key] = sparse_val.clone() if isinstance(sparse_val, torch.Tensor) else sparse_val + continue + mm_payload[mm_key] = to_payload_element( + element=mm_val, + idx=idx, + start=start, + end=end, + pass_lists_through=False, + seq_len=hidden_seq_len, + scheduled_seq_len=scheduled_seq_len, + ) + return mm_payload diff --git a/vllm_omni/worker/sampling_utils.py b/vllm_omni/worker/sampling_utils.py index c38a9bbdc..845881138 100644 --- a/vllm_omni/worker/sampling_utils.py +++ b/vllm_omni/worker/sampling_utils.py @@ -8,7 +8,21 @@ from vllm.v1.sample.logits_processor import LogitsProcessors, MinTokensLogitsPro logger = init_logger(__name__) -__all__ = ["sanitize_min_tokens_stop_ids"] +__all__ = ["clamp_prompt_ids_to_penalty_padding", "sanitize_min_tokens_stop_ids"] + + +def clamp_prompt_ids_to_penalty_padding(prompt_token_ids: torch.Tensor, logits_vocab: int) -> torch.Tensor: + """Clamp batch-level pad ids down to ``logits_vocab`` — upstream's + designed penalty padding value. + + ``max=logits_vocab`` (NOT ``logits_vocab - 1``) is deliberate: upstream + penalty computation allocates ``vocab_size + 1`` bins and drops the last + column, so ``vocab_size`` is the padding value that never affects + penalties (vllm/model_executor/layers/utils.py:: + get_token_bin_counts_and_mask). Clamping one lower would count padding + as real occurrences of the last vocab token. + """ + return prompt_token_ids.clamp(max=logits_vocab) def sanitize_min_tokens_stop_ids(logitsprocs: LogitsProcessors, logits_vocab: int) -> None: diff --git a/vllm_omni/worker/sparse_audio.py b/vllm_omni/worker/sparse_audio.py new file mode 100644 index 000000000..f82c7a59e --- /dev/null +++ b/vllm_omni/worker/sparse_audio.py @@ -0,0 +1,232 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Sparse-audio output protocol: marker classification and payload routing. + +A step's per-request `multimodal_outputs` lists may cover only a SUBSET of +the batch (requests that produced audio this step). The producer declares +this with two metadata keys (nested ``meta`` dict or flattened ``meta.*``): + +* ``meta.req_id`` — the request ids that DO have output, in list order; +* ``meta.sparse_audio`` — the marker saying "per-request lists are aligned + to ``meta.req_id`` order, not batch order". + +The marker's carrier form is ENFORCED, not normalized. The design contract +is a list of strings (e.g. ``["1"]`` — what every in-tree producer emits); a +bare string is also accepted. Interpreting deviant carriers would need a +second truth predicate, and predicate divergence is exactly how +``np.array(["0"])`` came to route as sparse while ``["0"]`` routed as dense; +per the refactor series' Policy A the runner<->model interface owes no +tolerance shim. + +A present-but-invalid declaration (illegal marker carrier, or marker set +with an unusable ``meta.req_id``) FAILS CLOSED: it most likely belongs to a +producer that INTENDED sparse output, and falling back to dense would +batch-index-assign its payloads to the WRONG requests. `InvalidSparseDeclarationError` +is raised by the classifiers and caught INSIDE this module by +`resolve_sparse_mm_routing`, which routes no payload for the affected step — +the exception never crosses the module boundary, so it can never escape a +v1 runner (which does not catch per-request exceptions). + +This module is the single home of the contract until the model-state +refactor series (G9-C2.6) turns the marker into a typed bool resolved at +one adapter site and unifies the nested/flattened encodings. The NPU runner +still carries a tolerant pre-enforcement copy — NPU-track handoff item. +""" + +import threading +import time +from typing import Any + +from vllm.logger import init_logger + +logger = init_logger(__name__) + +# The DESIGNED marker carrier is a list of strings (bare str also accepted); +# truthiness of a string item: +_SPARSE_MARKER_TRUTHY = ("1", "true", "yes", "on") + +# Per-process dedup so a broken producer logs once per offending shape, not +# once per step: +_logged_offenders: set[str] = set() +_error_state_lock = threading.Lock() + +# A persistent invalid producer must remain observable without flooding the +# per-step output path. Classifier errors above stay deduplicated per carrier +# type; the routing consequence is reported at most once per interval, with a +# count of repeats suppressed since the previous report. +_FAIL_CLOSED_LOG_INTERVAL_S = 60.0 +_fail_closed_last_log_at: float | None = None +_fail_closed_suppressed = 0 + + +class InvalidSparseDeclarationError(ValueError): + """Sparse-audio metadata is PRESENT but out of contract. + + Internal control signal: raised by `is_sparse_audio_marker` / + `resolve_sparse_mm_req_ids` and caught by `resolve_sparse_mm_routing` + in this same module, which fails closed. Callers outside this module + only ever see the routing result, never the exception. + """ + + +def _error_once(key: str, msg: str, *args: Any) -> None: + with _error_state_lock: + if key in _logged_offenders: + return + _logged_offenders.add(key) + logger.error(msg, *args) + + +def _error_fail_closed_rate_limited(cause: str) -> None: + global _fail_closed_last_log_at, _fail_closed_suppressed + + now = time.monotonic() + with _error_state_lock: + if _fail_closed_last_log_at is not None and now - _fail_closed_last_log_at < _FAIL_CLOSED_LOG_INTERVAL_S: + _fail_closed_suppressed += 1 + return + suppressed = _fail_closed_suppressed + _fail_closed_last_log_at = now + _fail_closed_suppressed = 0 + + logger.error( + "Invalid sparse-audio declaration: failing closed - no multimodal payload is routed " + "for the affected step(s). Fix the emitting model. Suppressed repeats before this " + "report: %d. Cause: %s", + suppressed, + cause, + ) + + +def is_sparse_audio_marker(value: Any) -> bool: + """Classify a ``meta.sparse_audio`` marker. + + Legal carriers (list of strings, bare string) evaluate against the + literal truthy set; ``None`` means "marker absent" and returns ``False``. + An ILLEGAL carrier (tensors, arrays, numbers, lists with non-string + items) is a producer bug: logged at error level once per offending type, + then raised as `InvalidSparseDeclarationError`. Deviant carriers are never + read element-wise, so a device tensor can never introduce a D2H sync + here. + """ + if value is None: + return False + if isinstance(value, str): + return value.lower() in _SPARSE_MARKER_TRUTHY + if isinstance(value, list) and all(isinstance(item, str) for item in value): + return any(item.lower() in _SPARSE_MARKER_TRUTHY for item in value) + vtype = type(value).__name__ + _error_once( + vtype, + 'Illegal sparse-audio marker of type %s: the marker MUST be a list of strings (e.g. ["1"]). ' + "The emitting model is out of contract; routing fails closed for the affected step(s).", + vtype, + ) + raise InvalidSparseDeclarationError(f"illegal sparse-audio marker carrier: {vtype}") + + +def resolve_sparse_mm_req_ids(multimodal_outputs: Any) -> list[str] | None: + """Resolve the sparse request-id list. + + Returns the request-id list when a legal sparse declaration is present, + and ``None`` when there is no sparse declaration (dense). Raises + `InvalidSparseDeclarationError` when the declaration is present but out of + contract (illegal marker carrier, or marker set with an unusable + ``meta.req_id``). + """ + if not isinstance(multimodal_outputs, dict): + return None + meta = multimodal_outputs.get("meta") + nested_req_ids = meta.get("req_id") if isinstance(meta, dict) else None + flat_req_ids = multimodal_outputs.get("meta.req_id") + + # Resolve both marker encodings independently, even when the request ids + # came from the other encoding. Previously a nested `req_id` prevented the + # flattened marker from being read, while the mirror combination worked. + # Evaluating both also preserves the strict contract: an illegal carrier + # on either side wins over an otherwise-valid declaration and fails closed. + nested_marker_value = meta.get("sparse_audio") if isinstance(meta, dict) else None + flat_marker_value = multimodal_outputs.get("meta.sparse_audio") + nested_marker = is_sparse_audio_marker(nested_marker_value) + flat_marker = is_sparse_audio_marker(flat_marker_value) + if nested_marker_value is not None and flat_marker_value is not None and nested_marker != flat_marker: + _error_once( + "marker:conflict", + "Nested and flattened sparse-audio markers disagree. Failing closed.", + ) + raise InvalidSparseDeclarationError("nested and flattened sparse markers disagree") + sparse_audio = nested_marker or flat_marker + if not sparse_audio: + return None + + req_id_carriers = [ + (name, value) for name, value in (("nested", nested_req_ids), ("flattened", flat_req_ids)) if value is not None + ] + if not req_id_carriers: + req_id_carriers = [("missing", None)] + + validated_req_ids: list[list[str]] = [] + for carrier_name, req_ids in req_id_carriers: + if isinstance(req_ids, list) and all(isinstance(rid, str) for rid in req_ids): + validated_req_ids.append(req_ids) + continue + rtype = type(req_ids).__name__ + _error_once( + f"req_id:{carrier_name}:{rtype}", + "Sparse-audio marker is set but the %s meta.req_id carrier is unusable (%s; " + "must be a list of request-id strings). Failing closed.", + carrier_name, + rtype, + ) + raise InvalidSparseDeclarationError(f"sparse marker set but {carrier_name} meta.req_id is unusable: {rtype}") + + req_ids = validated_req_ids[0] + if len(validated_req_ids) == 2 and validated_req_ids[1] != req_ids: + _error_once( + "req_id:conflict", + "Nested and flattened sparse-audio request-id lists disagree. Failing closed.", + ) + raise InvalidSparseDeclarationError("nested and flattened sparse request-id lists disagree") + if len(req_ids) != len(set(req_ids)): + _error_once( + "req_id:duplicate", + "Sparse-audio request-id list contains duplicates. Failing closed.", + ) + raise InvalidSparseDeclarationError("sparse request-id list contains duplicates") + return list(req_ids) + + +def resolve_sparse_mm_routing( + *, + engine_output_type: str, + req_ids_output_copy: list[str], + downstream_req_ids: list[str], + multimodal_outputs: Any, +) -> tuple[list[str], dict[str, int], bool]: + """Resolve `(downstream_req_ids, sparse_mm_index, audio_sparse_output)`. + + The only entry point runners use; all protocol errors are absorbed here. + """ + try: + sparse_mm_req_ids = resolve_sparse_mm_req_ids(multimodal_outputs) + except InvalidSparseDeclarationError as exc: + if engine_output_type != "audio": + # Non-audio pipelines never consult sparse routing; the carrier + # error is already logged by the classifier. + return downstream_req_ids, {}, False + # FAIL CLOSED: a present-but-invalid sparse declaration most likely + # belongs to a producer that INTENDED sparse output; routing it as + # dense would batch-index-assign its payloads to the WRONG requests. + # Route no payload for this step instead - a visible, bounded gap + # rather than silent cross-request corruption. (Empty-sparse is an + # already-exercised path: the zero-audio sparse step ships + # req_id=[] through the same shape.) + _error_fail_closed_rate_limited(str(exc)) + return [], {}, True + sparse_mm_index = {rid: i for i, rid in enumerate(sparse_mm_req_ids or [])} + if engine_output_type != "audio" or sparse_mm_req_ids is None: + return downstream_req_ids, sparse_mm_index, False + + sparse_req_id_set = set(sparse_mm_req_ids) + sparse_downstream_req_ids = [rid for rid in req_ids_output_copy if rid in sparse_req_id_set] + return sparse_downstream_req_ids, sparse_mm_index, True