[MR V2] Migrate to Model Runner V2

Squashed migration rebased onto current main. Collapses the 46-commit
dev/migrate-MR-v2 history (10 merge commits + superseded pre-rebase work)
into a single commit, so the diff reflects only the genuine MR-V2 feature
(worker_v2/ runners + model_states, omni schedulers, stage engine pipeline,
pooling_output serialization bridge). Content already merged to main
separately (step_audio2, tts_adapters, npu/_310p, etc.) drops out.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

tests/worker_v2/test_omni_gpu_model_runner.py

@@ -0,0 +1,137 @@
+"""Unit tests for OmniGPUModelRunner v2 overrides."""
+
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner, _needs_capture_tensor_unwrap
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _DummyInputBatch:
+    def __init__(self, idx_mapping=None):
+        self.idx_mapping_np = idx_mapping or []
+
+
+def _make_runner():
+    """Create an OmniGPUModelRunner without calling __init__."""
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.model = MagicMock()
+    runner.req_states = SimpleNamespace(req_id_to_index={"r1": 0, "r2": 1})
+    runner.execute_model_state = None
+    return runner
+
+
+def test_finish_requests_calls_remove_for_finished():
+    runner = _make_runner()
+    mock_state = MagicMock()
+    runner.model_state = mock_state
+
+    sched_output = SimpleNamespace(
+        finished_req_ids={"r1"},
+        preempted_req_ids=set(),
+    )
+
+    with patch.object(type(runner).__bases__[0], "finish_requests", return_value=None):
+        runner.finish_requests(sched_output)
+
+    mock_state.remove_request.assert_called_once_with(0)
+
+
+def test_finish_requests_calls_remove_for_preempted():
+    runner = _make_runner()
+    mock_state = MagicMock()
+    runner.model_state = mock_state
+
+    sched_output = SimpleNamespace(
+        finished_req_ids=set(),
+        preempted_req_ids={"r2"},
+    )
+
+    with patch.object(type(runner).__bases__[0], "finish_requests", return_value=None):
+        runner.finish_requests(sched_output)
+
+    mock_state.remove_request.assert_called_once_with(1)
+
+
+def test_finish_requests_ignores_unknown_req_ids():
+    runner = _make_runner()
+    mock_state = MagicMock()
+    runner.model_state = mock_state
+
+    sched_output = SimpleNamespace(
+        finished_req_ids={"unknown"},
+        preempted_req_ids=set(),
+    )
+
+    with patch.object(type(runner).__bases__[0], "finish_requests", return_value=None):
+        runner.finish_requests(sched_output)
+
+    mock_state.remove_request.assert_not_called()
+
+
+def test_finish_requests_handles_both_finished_and_preempted():
+    runner = _make_runner()
+    mock_state = MagicMock()
+    runner.model_state = mock_state
+
+    sched_output = SimpleNamespace(
+        finished_req_ids={"r1"},
+        preempted_req_ids={"r2"},
+    )
+
+    with patch.object(type(runner).__bases__[0], "finish_requests", return_value=None):
+        runner.finish_requests(sched_output)
+
+    assert mock_state.remove_request.call_count == 2
+
+
+def test_thinker_stage_needs_capture_tensor_unwrap():
+    assert _needs_capture_tensor_unwrap(SimpleNamespace(model_stage="thinker"))
+    assert not _needs_capture_tensor_unwrap(SimpleNamespace(model_stage="talker"))
+
+
+def test_capture_model_unwraps_tuple_outputs():
+    runner = object.__new__(OmniGPUModelRunner)
+    hidden = torch.ones(1, 2)
+
+    def original_forward():
+        return hidden, {"layers": {}}
+
+    runner.model = SimpleNamespace(forward=original_forward)
+    runner._model_returns_tuple = True
+    runner._exclude_full_graph = False
+
+    def capture_model(_self):
+        assert torch.equal(runner.model.forward(), hidden)
+        return 3
+
+    with patch.object(type(runner).__bases__[0], "capture_model", capture_model):
+        assert runner.capture_model() == 3
+
+    assert runner.model.forward is original_forward
+
+
+def test_capture_model_unwraps_omni_outputs():
+    runner = object.__new__(OmniGPUModelRunner)
+    hidden = torch.ones(1, 2)
+
+    def original_forward():
+        return OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
+
+    runner.model = SimpleNamespace(forward=original_forward)
+    runner._model_returns_tuple = True
+    runner._exclude_full_graph = False
+
+    def capture_model(_self):
+        assert torch.equal(runner.model.forward(), hidden)
+        return 5
+
+    with patch.object(type(runner).__bases__[0], "capture_model", capture_model):
+        assert runner.capture_model() == 5
+
+    assert runner.model.forward is original_forward