[Bugfix] Recognize mixed-attention warmup requests in MRv2

tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py

@@ -124,18 +124,19 @@ 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)
     output = sampler.sample_step(None, batch, None, None, standard)
     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
 
 

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],

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()

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)

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")

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])

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])

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", [])))

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_"))