fix: align MRV2 profiling and capture with vLLM 0.31

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/worker_v2/test_omni_generation_model_runner.py

@@ -20,6 +20,17 @@ from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModel
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+@pytest.mark.parametrize("randomize_inputs", [False, True])
+def test_profile_run_accepts_release_flag_without_running_codec(monkeypatch, randomize_inputs):
+    runner = object.__new__(OmniGenerationModelRunner)
+    runner.model = MagicMock(side_effect=AssertionError("profiling ran codec model"))
+    synchronize = MagicMock()
+    monkeypatch.setattr(torch.accelerator, "synchronize", synchronize)
+    runner.profile_run(randomize_inputs=randomize_inputs)
+    synchronize.assert_called_once_with()
+    runner.model.assert_not_called()
+
+
 class _FakeStagedField:
     def __init__(self, data):
         self.np = data

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -95,7 +95,8 @@ def test_full_payload_receive_is_polled_without_scheduled_tokens(mocker):
 
 
 @pytest.mark.parametrize("output_form", ["tuple", "omni"])
-def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form):
+@pytest.mark.parametrize("profile_only", [False, True])
+def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form, profile_only):
     runner = object.__new__(OmniGPUModelRunner)
     hidden = torch.ones(1, 2)
 
@@ -117,16 +118,21 @@ def test_capture_model_unwraps_exclude_full_and_capture_mtp(output_form):
     runner.model_state = SimpleNamespace(capture_mtp_graphs=MagicMock())
     runner._dispatch_mtp_batch_descriptor = MagicMock(return_value="desc")
 
-    def assert_unwrapped(_self):
+    def assert_unwrapped(_self, *, profile_only: bool = False):
         assert torch.equal(runner.model.forward(), hidden)  # forward unwrapped during capture
+        assert profile_only is expected_profile_only
         return 3
 
+    expected_profile_only = profile_only
     with patch.object(GPUModelRunner, "capture_model", assert_unwrapped):
-        assert runner.capture_model() == 3
+        assert runner.capture_model(profile_only=profile_only) == 3
 
     assert runner.model.forward is original_forward  # restored after capture
     assert runner.cudagraph_manager._capture_descs == {CUDAGraphMode.PIECEWISE: [piecewise]}
-    runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor)
+    if profile_only:
+        runner.model_state.capture_mtp_graphs.assert_not_called()
+    else:
+        runner.model_state.capture_mtp_graphs.assert_called_once_with(runner._dispatch_mtp_batch_descriptor)
 
 
 @pytest.mark.parametrize(
@@ -209,7 +215,7 @@ def test_full_graph_aux_contract_keeps_full_and_round_trips_leaves():
     runner.model_state = SimpleNamespace()
     captured = {}
 
-    def capture(_self):
+    def capture(_self, *, profile_only=False):
         # The graph manager stores (hidden, leaves) in its aux buffers.
         assert runner.use_aux_hidden_state_outputs is True
         hidden, leaves = runner.model.forward(4)
@@ -262,7 +268,9 @@ def test_full_graph_output_without_contract_is_hidden_only():
 
 
 @pytest.mark.parametrize("runner_kind", ["gpu", "ar", "generation"])
-def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch):
+@pytest.mark.parametrize("randomize_inputs", [False, True])
+@pytest.mark.parametrize("is_profile", [False, True])
+def test_dummy_forward_uses_upstream_execution_state(runner_kind, randomize_inputs, is_profile, monkeypatch):
     from vllm.v1.worker.gpu.input_batch import InputBatch
     from vllm.v1.worker.gpu.model_runner import ExecuteModelState
 
@@ -288,6 +296,7 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch):
     runner.model_state.prepare_inputs.return_value = {}
     runner._omni_data_plane = object()
     runner.supports_mm_inputs = False
+    runner.vocab_size = 32
     runner.lora_config = None
     runner.is_encoder_decoder = False
     runner.eplb = MagicMock()
@@ -308,7 +317,15 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch):
         cg_mode=CUDAGraphMode.NONE, num_reqs=1, num_tokens=1, num_active_loras=0, max_query_len=1
     )
     runner._dispatch_batch_descriptor = MagicMock(return_value=(batch_desc, None))
-    monkeypatch.setattr(InputBatch, "make_dummy", lambda *args, **kwargs: input_batch)
+    make_dummy = MagicMock(return_value=input_batch)
+    monkeypatch.setattr(InputBatch, "make_dummy", make_dummy)
+    randomized = []
+
+    def randomize(tensor, low, high):
+        randomized.append((low, high))
+        return tensor.fill_(7)
+
+    monkeypatch.setattr(torch.Tensor, "random_", randomize)
     monkeypatch.setattr("vllm_omni.worker_v2.omni_model_runner.build_slot_mappings_by_layer", lambda *args: {})
     for module in ("omni_model_runner", "omni_generation_model_runner"):
         monkeypatch.setattr(f"vllm_omni.worker_v2.{module}.set_forward_context", lambda *args, **kwargs: nullcontext())
@@ -318,7 +335,20 @@ def test_dummy_forward_uses_upstream_execution_state(runner_kind, monkeypatch):
 
     # Upstream _dummy_run always supplies valid_dummy_state_slots, and the
     # result must use the real upstream constructor rather than a mocked state.
-    assert runner.execute_model(scheduled, dummy_run=True, valid_dummy_state_slots=True) is None
+    assert (
+        runner.execute_model(
+            scheduled,
+            dummy_run=True,
+            valid_dummy_state_slots=True,
+            randomize_inputs=randomize_inputs,
+            is_profile=is_profile,
+        )
+        is None
+    )
+    assert make_dummy.call_args.kwargs["is_padding"] is not is_profile
+    should_randomize = randomize_inputs and runner_kind != "generation"
+    assert randomized == ([(0, 32)] if should_randomize else [])
+    assert input_batch.input_ids.tolist() == ([7] if should_randomize else [1])
     assert isinstance(runner.execute_model_state, ExecuteModelState)
     assert runner.execute_model_state.input_batch is input_batch
     assert runner.execute_model_state.cudagraph_stats is None

tests/worker_v2/test_omni_model_state.py

@@ -22,6 +22,16 @@ from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _m
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_prepare_attn_forwards_release_model_specific_metadata():
+    state = object.__new__(OmniModelState)
+    metadata = object()
+    output = object()
+    batch = SimpleNamespace()
+    with patch.object(DefaultModelState, "prepare_attn", return_value=output) as prepare:
+        assert state.prepare_attn(batch, None, (), None, [], None, model_specific_attn_metadata=metadata) is output
+    assert prepare.call_args.kwargs["model_specific_attn_metadata"] is metadata
+
+
 class _DummyInputBatch:
     is_prefilling_np: np.ndarray
 

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -31,6 +31,7 @@ from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu.input_batch import InputBatch
 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
+from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
 from vllm.v1.worker.gpu.states import RequestState
 from vllm.v1.worker.utils import AttentionGroup
 
@@ -447,6 +448,7 @@ class OmniModelState(DefaultModelState):
         kv_cache_config: KVCacheConfig,
         for_capture: bool = False,
         ubatch_idx: int = 0,
+        model_specific_attn_metadata: ModelSpecificAttnMetadata | None = None,
     ) -> dict[str, Any]:
         if (
             for_capture
@@ -471,6 +473,7 @@ class OmniModelState(DefaultModelState):
             kv_cache_config,
             for_capture=for_capture,
             ubatch_idx=ubatch_idx,
+            model_specific_attn_metadata=model_specific_attn_metadata,
         )
 
     def prepare_inputs_embeds(

vllm_omni/worker_v2/omni_ar_model_runner.py

@@ -121,6 +121,7 @@ class OmniARModelRunner(OmniGPUModelRunner):
         is_profile: bool = False,
         context_len: int = 0,
         valid_dummy_state_slots: bool = False,
+        randomize_inputs: bool = False,
     ) -> Any:
         if not dummy_run:
             self._handle_kv_transfer_pre(scheduler_output)
@@ -132,6 +133,7 @@ class OmniARModelRunner(OmniGPUModelRunner):
             is_profile=is_profile,
             context_len=context_len,
             valid_dummy_state_slots=valid_dummy_state_slots,
+            randomize_inputs=randomize_inputs,
         )
 
     # ------------------------------------------------------------------

vllm_omni/worker_v2/omni_generation_model_runner.py

@@ -292,7 +292,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
     # ------------------------------------------------------------------
 
     @torch.inference_mode()
-    def profile_run(self) -> None:
+    def profile_run(self, randomize_inputs: bool = False) -> None:
         """Generation models have no KV cache — skip profiling.
 
         Code2Wav shares GPU memory with the Talker stage (same device);
@@ -317,6 +317,7 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
         is_profile: bool = False,
         context_len: int = 0,
         valid_dummy_state_slots: bool = False,
+        randomize_inputs: bool = False,
     ) -> ModelRunnerOutput | IntermediateTensors | None:
         if not dummy_run:
             self._prepare_native_data_plane(scheduler_output)
@@ -388,7 +389,10 @@ class OmniGenerationModelRunner(OmniGPUModelRunner):
                 batch_desc.num_tokens,
                 self.input_buffers,
                 max_query_len=batch_desc.max_query_len,
+                is_padding=not is_profile,
             )
+            # Codec IDs have model-specific bounds; keep valid dummy tokens
+            # even when upstream requests randomized language-model inputs.
 
         attn_metadata = None
         slot_mappings_by_layer = None

vllm_omni/worker_v2/omni_model_runner.py

@@ -391,7 +391,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         }
         logger.info("Excluded FULL CUDA graph capture for Omni model. PIECEWISE graphs will still be captured.")
 
-    def capture_model(self) -> int:
+    def capture_model(self, *, profile_only: bool = False) -> int:
         """Handle CUDA graph capture for Omni models.
 
         Tuple-returning models use PIECEWISE graphs; FULL replay requires
@@ -424,12 +424,17 @@ class OmniGPUModelRunner(GPUModelRunner):
             use_aux = self.use_aux_hidden_state_outputs
             self.use_aux_hidden_state_outputs = use_aux or aux_outputs
             try:
-                result = super().capture_model()
+                result = super().capture_model(profile_only=profile_only)
             finally:
                 self.model.forward = original_forward  # type: ignore[assignment]
                 self.use_aux_hidden_state_outputs = use_aux
         else:
-            result = super().capture_model()
+            result = super().capture_model(profile_only=profile_only)
+
+        if profile_only:
+            # Upstream discards captures in its temporary profiling pool.
+            # Model-owned graphs must be recorded during the real capture.
+            return result
 
         capture_mtp = getattr(getattr(self, "model_state", None), "capture_mtp_graphs", None)
         if callable(capture_mtp):
@@ -494,6 +499,7 @@ class OmniGPUModelRunner(GPUModelRunner):
         is_profile: bool = False,
         context_len: int = 0,
         valid_dummy_state_slots: bool = False,
+        randomize_inputs: bool = False,
     ) -> Any:
         if not dummy_run:
             self._prepare_native_data_plane(scheduler_output)
@@ -570,7 +576,10 @@ class OmniGPUModelRunner(GPUModelRunner):
                 batch_desc.num_tokens,
                 self.input_buffers,
                 max_query_len=batch_desc.max_query_len,
+                is_padding=not is_profile,
             )
+            if randomize_inputs:
+                input_batch.input_ids.random_(0, self.vocab_size)
             if not skip_attn_for_dummy_run:
                 block_tables, slot_mappings = self.prepare_dummy_attn(input_batch, valid_dummy_state_slots)
                 if context_len: