From a5bc5ab3420587553429420983d645202edcff1a Mon Sep 17 00:00:00 2001 From: Weiming Liao Date: Thu, 22 Oct 2026 23:49:58 +0800 Subject: [PATCH] [Bugfix] Recognize mixed-attention warmup requests in MRv2 Signed-off-by: Weiming Liao --- .../models/higgs_audio_v3/test_mrv2_state.py | 9 +++---- tests/worker_v2/test_omni_gpu_model_runner.py | 5 ++-- tests/worker_v2/test_omni_model_state.py | 24 ++++++++++++++++++- .../models/fish_omni/model_state.py | 3 ++- .../models/higgs_audio_v3/model_state.py | 5 ++-- .../models/moss_tts/local_model_state.py | 3 ++- .../model_states/omni_model_state.py | 9 +++---- vllm_omni/worker_v2/omni_model_runner.py | 3 ++- vllm_omni/worker_v2/warmup.py | 7 ++++++ 9 files changed, 52 insertions(+), 16 deletions(-) create mode 100644 vllm_omni/worker_v2/warmup.py diff --git a/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py b/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py index ccf2eb5ef3c..bb4aa185df3 100644 --- a/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py +++ b/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py @@ -124,10 +124,11 @@ def test_finalizer_filters_unsampled_prefill_and_keeps_owned_audio(): assert payload["_higgs_invalid_rows"] == () -def test_upstream_warmup_retains_generic_text_sampling(): +@pytest.mark.parametrize("req_id", ["_warmup_0_", "_v2_mixed_warmup_decode_", "_v2_mixed_warmup_prefill_"]) +def test_upstream_warmup_retains_generic_text_sampling(req_id): s = make_state() - s.add_request(0, request("_warmup_0_", 0, repetition_penalty=1.2, logprobs=1)) - batch = SimpleNamespace(req_ids=["_warmup_0_"]) + s.add_request(0, request(req_id, 0, repetition_penalty=1.2, logprobs=1)) + batch = SimpleNamespace(req_ids=[req_id]) sampler, rejection = s.custom_sampler(Mock()) expected = (object(), torch.ones(1), torch.zeros(1)) standard = Mock(return_value=expected) @@ -135,7 +136,7 @@ def test_upstream_warmup_retains_generic_text_sampling(): assert output.sampler_output is expected[0] and output.multimodal_outputs is None assert rejection is None standard.assert_called_once_with(None, batch, None) - s.remove_request("_warmup_0_") + s.remove_request(req_id) assert not s.requests diff --git a/tests/worker_v2/test_omni_gpu_model_runner.py b/tests/worker_v2/test_omni_gpu_model_runner.py index 2a6ba1709b4..c097c387a29 100644 --- a/tests/worker_v2/test_omni_gpu_model_runner.py +++ b/tests/worker_v2/test_omni_gpu_model_runner.py @@ -168,7 +168,8 @@ def parent_add_requests(_runner, _output): assert observed == ([17] if new_request else [0]) -def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip(): +@pytest.mark.parametrize("warmup_req_id", ["_warmup_0_", "_v2_mixed_warmup_decode_", "_v2_mixed_warmup_prefill_"]) +def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip(warmup_req_id): runner = _make_runner() runner.model_config = object.__new__(OmniModelConfig) runner.model_config.async_chunk = True @@ -180,7 +181,7 @@ def test_prepare_native_data_plane_terminal_abort_split_and_warmup_skip(): ) runner._omni_data_plane = plane new_req = SimpleNamespace(req_id="r1") - warmup_req = SimpleNamespace(req_id="_warmup_0_") + warmup_req = SimpleNamespace(req_id=warmup_req_id) handle = SimpleNamespace(request_id="r2", external_req_id="ext-r2") scheduler_output = SimpleNamespace( scheduled_new_reqs=[new_req, warmup_req], diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 44716113a6a..a20fc74aab3 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -151,7 +151,10 @@ def _fwd_ctx(): yield -@pytest.mark.parametrize("req_id,expect_validity", [("_warmup_0_", True), ("r1", False)]) +@pytest.mark.parametrize( + "req_id,expect_validity", + [("_warmup_0_", True), ("_v2_mixed_warmup_decode_", True), ("_v2_mixed_warmup_prefill_", True), ("r1", False)], +) def test_add_request_declared_validity_only_for_warmup(req_id, expect_validity): # Real requests must not be masked with a fabricated validity key. state = _make_state() @@ -163,6 +166,25 @@ def test_add_request_declared_validity_only_for_warmup(req_id, expect_validity): assert buf["meta"]["codec_frame_valid"].item() is False +@pytest.mark.parametrize("req_id", ["_warmup_0_", "_v2_mixed_warmup_decode_", "_v2_mixed_warmup_prefill_"]) +def test_warmup_skips_model_preprocess_and_postprocess(req_id): + state = _make_state(has_preprocess=True, has_postprocess=True) + _fill_buffers(state, req_id) + state.model.preprocess.side_effect = AssertionError("warmup has no model conditioning") + state.model.postprocess.side_effect = AssertionError("warmup must not update model state") + state.model.postprocess_batch_mrv2 = MagicMock(side_effect=AssertionError("warmup must not use batch postprocess")) + batch = _DummyInputBatch([0], num_computed_tokens_cpu=[0]) + inputs = {"input_ids": torch.tensor([1]), "inputs_embeds": torch.zeros(1, 4)} + + state.run_preprocess(batch, inputs) + state.run_postprocess(torch.zeros(1, 4), batch) + + state.model.preprocess.assert_not_called() + state.model.postprocess.assert_not_called() + state.model.postprocess_batch_mrv2.assert_not_called() + assert state.intermediate_buffer.buffers[0] == {"req_id": req_id} + + @pytest.mark.parametrize("mode", ["index", "req_id", "unknown"]) def test_remove_request_state_isolation(mode): state = _make_state() diff --git a/vllm_omni/model_executor/models/fish_omni/model_state.py b/vllm_omni/model_executor/models/fish_omni/model_state.py index f412b106a30..e0c452acfdf 100644 --- a/vllm_omni/model_executor/models/fish_omni/model_state.py +++ b/vllm_omni/model_executor/models/fish_omni/model_state.py @@ -33,6 +33,7 @@ from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot, RequestOutputSnapshot from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState +from vllm_omni.worker_v2.warmup import is_warmup_request logger = init_logger(__name__) @@ -148,7 +149,7 @@ def add_request(self, req_index: int, new_req_data) -> None: req_id = str(new_req_data.req_id) prompt_ids = list(new_req_data.prompt_token_ids or []) token_ids = list(new_req_data.prefill_token_ids or prompt_ids) - if req_id.startswith("_warmup_"): + if is_warmup_request(req_id): # Warmup requests: text rows with zero codes. rows = torch.zeros((len(prompt_ids), self.frame_width), dtype=torch.int64) rows[:, 0] = torch.tensor(prompt_ids, dtype=torch.int64) diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/model_state.py b/vllm_omni/model_executor/models/higgs_audio_v3/model_state.py index d275ea7f8f4..d6a1c510de7 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/model_state.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/model_state.py @@ -14,6 +14,7 @@ from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState from vllm_omni.worker_v2.omni_sampler import OmniSampler, OmniSamplingOutput +from vllm_omni.worker_v2.warmup import is_warmup_request logger = init_logger(__name__) @@ -72,7 +73,7 @@ def _ensure_slot_capacity(self, slots): def add_request(self, req_index, new_req_data): params = new_req_data.sampling_params - is_warmup = str(new_req_data.req_id).startswith("_warmup_") + is_warmup = is_warmup_request(new_req_data.req_id) if not is_warmup and ( params.frequency_penalty != 0 or params.presence_penalty != 0 @@ -288,7 +289,7 @@ def __init__(self, base_sampler, state: HiggsModelState): def sample_step(self, hidden_states, input_batch, req_states, grammar_output, standard_sample): # Upstream warmup exercises its text sampler, including grammar and # penalties, with synthetic non-audio prompts. Keep that path intact. - if all(str(r).startswith("_warmup_") for r in input_batch.req_ids): + if all(is_warmup_request(r) for r in input_batch.req_ids): return super().sample_step(hidden_states, input_batch, req_states, grammar_output, standard_sample) if grammar_output is not None or input_batch.num_draft_tokens: raise ValueError("Higgs MRV2 does not support grammar or speculative sampling") diff --git a/vllm_omni/model_executor/models/moss_tts/local_model_state.py b/vllm_omni/model_executor/models/moss_tts/local_model_state.py index 9b65dd5165f..02ef4c0fbe2 100644 --- a/vllm_omni/model_executor/models/moss_tts/local_model_state.py +++ b/vllm_omni/model_executor/models/moss_tts/local_model_state.py @@ -24,6 +24,7 @@ from vllm_omni.model_executor.output_snapshot import PackedOutputSnapshot, RequestOutputSnapshot from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState +from vllm_omni.worker_v2.warmup import is_warmup_request logger = init_logger(__name__) @@ -382,7 +383,7 @@ def run_preprocess(self, input_batch, model_inputs, req_states=None, mtp_batch_d for row in range(input_batch.num_reqs): slot = int(input_batch.idx_mapping_np[row]) buf = self.intermediate_buffer.buffers[slot] - if not buf or str(buf.get("req_id", "")).startswith("_warmup_"): + if not buf or is_warmup_request(buf.get("req_id", "")): continue start = int(input_batch.query_start_loc_np[row]) count = int(input_batch.num_scheduled_tokens[row]) diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index 8fe72d73c5a..9b946dd94c6 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -44,6 +44,7 @@ OmniIntermediateBuffer, ) from vllm_omni.worker_v2.streaming_audio import StreamingAudioOutput +from vllm_omni.worker_v2.warmup import is_warmup_request logger = init_logger(__name__) @@ -419,7 +420,7 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: def _initialize_upstream_warmup_buffer(self, req_index: int, req_id: str) -> None: """Satisfy model-declared output contracts for vLLM warmup requests.""" - if not str(req_id).startswith("_warmup_"): + if not is_warmup_request(req_id): return validity_key = getattr(self.model, "mtp_validity_key", None) @@ -796,7 +797,7 @@ def run_preprocess( if settled is not None and n_tok == 1 and settled[req_idx] == buf["req_id"]: settled_rows.append((i, req_idx, start, buf["req_id"])) continue - if str(buf["req_id"]).startswith("_warmup_"): + if is_warmup_request(buf["req_id"]): continue info = {key: value for key, value in buf.items() if isinstance(key, str)} @@ -1506,7 +1507,7 @@ def run_postprocess(self, hidden_states: torch.Tensor, input_batch: InputBatch) if callable(batch_postprocess) and input_batch.num_reqs: req_indices = [int(input_batch.idx_mapping_np[i]) for i in range(input_batch.num_reqs)] buffers = [self.intermediate_buffer.buffers[req_idx] for req_idx in req_indices] - if all(buf and "req_id" in buf and not str(buf["req_id"]).startswith("_warmup_") for buf in buffers): + if all(buf and "req_id" in buf and not is_warmup_request(buf["req_id"]) for buf in buffers): query_start_loc = getattr(input_batch, "query_start_loc", None) if ( isinstance(query_start_loc, torch.Tensor) @@ -1563,7 +1564,7 @@ def run_postprocess(self, hidden_states: torch.Tensor, input_batch: InputBatch) buf = self.intermediate_buffer.buffers[req_idx] if not buf or "req_id" not in buf: continue - if str(buf["req_id"]).startswith("_warmup_"): + if is_warmup_request(buf["req_id"]): continue start = int(input_batch.query_start_loc_np[i]) n_tok = int(input_batch.num_scheduled_tokens[i]) diff --git a/vllm_omni/worker_v2/omni_model_runner.py b/vllm_omni/worker_v2/omni_model_runner.py index 4ae6c97ee7d..28dd6dabe98 100644 --- a/vllm_omni/worker_v2/omni_model_runner.py +++ b/vllm_omni/worker_v2/omni_model_runner.py @@ -48,6 +48,7 @@ ) from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState from vllm_omni.worker_v2.omni_data_plane import OmniRunnerDataPlane +from vllm_omni.worker_v2.warmup import is_warmup_request if TYPE_CHECKING: from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache @@ -290,7 +291,7 @@ def _prepare_native_data_plane(self, scheduler_output: SchedulerOutput) -> None: if plane is None: return for request_data in getattr(scheduler_output, "scheduled_new_reqs", []): - if str(getattr(request_data, "req_id", "")).startswith("_warmup_"): + if is_warmup_request(request_data.req_id): continue plane.register_request(request_data) plane.register_receivers(list(getattr(scheduler_output, "pending_input_registrations", []))) diff --git a/vllm_omni/worker_v2/warmup.py b/vllm_omni/worker_v2/warmup.py new file mode 100644 index 00000000000..b9c56644061 --- /dev/null +++ b/vllm_omni/worker_v2/warmup.py @@ -0,0 +1,7 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + + +def is_warmup_request(req_id: str) -> bool: + """Recognize vLLM kernel and mixed-attention warmup requests.""" + return req_id.startswith(("_warmup_", "_v2_mixed_warmup_"))