[Refactor] Worker/ModelRunner runner correctness fixes (G2/N) (#5452)

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

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<rhs>[^\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."
+        )