[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/test_gpu_ar_model_runner_bookkeeping_compat.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
+
+
+def test_gpu_ar_model_runner_calls_bookkeeping_with_current_vllm_signature():
+    runner = object.__new__(GPUARModelRunner)
+    calls = []
+
+    def bookkeeping_sync(*args):
+        calls.append(args)
+        return (0, None, [], {}, [], {}, [])
+
+    runner._bookkeeping_sync = bookkeeping_sync
+    runner._bookkeeping_accepts_spec_decode_metadata = False
+    scheduler_output = SimpleNamespace(total_num_scheduled_tokens=3)
+    sampler_output = SimpleNamespace()
+    logits = torch.empty(1, 1)
+    hidden_states = torch.empty(1, 1)
+
+    result = runner._bookkeeping_sync_compat(
+        scheduler_output,
+        sampler_output,
+        logits,
+        hidden_states,
+        scheduler_output.total_num_scheduled_tokens,
+        spec_decode_metadata=object(),
+    )
+
+    assert result == (0, None, [], {}, [], {}, [])
+    assert calls == [
+        (
+            scheduler_output,
+            sampler_output,
+            logits,
+            hidden_states,
+            3,
+        )
+    ]
+
+
+if __name__ == "__main__":
+    test_gpu_ar_model_runner_calls_bookkeeping_with_current_vllm_signature()
+
+
+def test_bookkeeping_compat_does_not_reflect_per_call():
+    import inspect
+
+    source = inspect.getsource(GPUARModelRunner._bookkeeping_sync_compat)
+    assert "inspect.signature" not in source
+
+
+def test_sample_tokens_passes_spec_decode_metadata_to_bookkeeping_compat():
+    import ast
+    import inspect
+
+    source = inspect.getsource(GPUARModelRunner.sample_tokens)
+    tree = ast.parse(source)
+
+    calls = [
+        node
+        for node in ast.walk(tree)
+        if isinstance(node, ast.Call)
+        and isinstance(node.func, ast.Attribute)
+        and node.func.attr == "_bookkeeping_sync_compat"
+    ]
+    assert len(calls) == 1
+    assert any(isinstance(arg, ast.Name) and arg.id == "spec_decode_metadata" for arg in calls[0].args)

tests/worker/test_omni_generation_scheduler_vllm20_compat.py

@@ -0,0 +1,42 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+from vllm_omni.core.sched.omni_generation_scheduler import (
+    _get_request_num_cached_tokens,
+    _make_engine_core_output,
+)
+
+
+def test_get_request_num_cached_tokens_defaults_to_zero_when_field_is_missing():
+    request = SimpleNamespace(request_id="rid", num_computed_tokens=5)
+
+    assert _get_request_num_cached_tokens(request) == 0
+
+
+def test_make_engine_core_output_ignores_fields_missing_from_current_vllm():
+    output = _make_engine_core_output(
+        request_id="rid",
+        new_token_ids=[1],
+        finish_reason=None,
+        new_logprobs=None,
+        new_prompt_logprobs_tensors=None,
+        pooling_output=None,
+        stop_reason=None,
+        events=[],
+        kv_transfer_params=None,
+        trace_headers={},
+        num_cached_tokens=3,
+        num_external_computed_tokens=4,
+        routed_experts=None,
+        num_nans_in_logits=0,
+    )
+
+    assert output.request_id == "rid"
+    assert output.new_token_ids == [1]
+
+
+if __name__ == "__main__":
+    test_get_request_num_cached_tokens_defaults_to_zero_when_field_is_missing()
+    test_make_engine_core_output_ignores_fields_missing_from_current_vllm()

tests/worker/test_omni_generation_worker_compile_time.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+
+import vllm_omni.worker.gpu_generation_worker as worker
+from vllm_omni.worker.gpu_generation_worker import _make_compilation_times
+
+
+def test_make_compilation_times_matches_current_vllm_shape():
+    result = _make_compilation_times(0.0)
+
+    if isinstance(result, float):
+        assert result == 0.0
+    else:
+        assert result.language_model == 0.0
+        assert result.encoder == 0.0
+
+
+if __name__ == "__main__":
+    test_make_compilation_times_matches_current_vllm_shape()
+
+
+def test_make_compilation_times_filters_to_current_fields():
+    result = _make_compilation_times(1.5, speculative_model=2.0)
+
+    if isinstance(result, float):
+        assert result == 1.5
+    else:
+        assert result.language_model == 1.5
+        assert result.encoder == 0.0
+
+
+def test_compilation_times_import_is_lazy():
+    source_before_helper = inspect.getsource(worker).split("def _make_compilation_times", 1)[0]
+    assert "CompilationTimes" not in source_before_helper

tests/worker/test_omni_scheduler_engine_core_output_compat.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+from vllm_omni.core.sched.omni_ar_scheduler import (
+    _get_request_num_cached_tokens,
+    _get_request_num_external_computed_tokens,
+    _make_engine_core_output,
+)
+
+
+def test_get_request_num_cached_tokens_defaults_to_zero_when_field_is_missing():
+    request = SimpleNamespace(request_id="rid", num_computed_tokens=5)
+
+    assert _get_request_num_cached_tokens(request) == 0
+
+
+def test_get_request_num_external_computed_tokens_defaults_to_zero_when_field_is_missing():
+    request = SimpleNamespace(request_id="rid")
+
+    assert _get_request_num_external_computed_tokens(request) == 0
+
+
+def test_make_engine_core_output_ignores_fields_missing_from_current_vllm():
+    output = _make_engine_core_output(
+        request_id="rid",
+        new_token_ids=[1],
+        finish_reason=None,
+        new_logprobs=None,
+        new_prompt_logprobs_tensors=None,
+        pooling_output=None,
+        stop_reason=None,
+        events=[],
+        kv_transfer_params=None,
+        trace_headers={},
+        num_cached_tokens=3,
+        num_external_computed_tokens=4,
+        routed_experts=None,
+        num_nans_in_logits=0,
+        is_segment_finished=None,
+        new_prompt_len_snapshot=None,
+    )
+
+    assert output.request_id == "rid"
+    assert output.new_token_ids == [1]
+
+
+if __name__ == "__main__":
+    test_get_request_num_cached_tokens_defaults_to_zero_when_field_is_missing()
+    test_get_request_num_external_computed_tokens_defaults_to_zero_when_field_is_missing()
+    test_make_engine_core_output_ignores_fields_missing_from_current_vllm()

tests/worker/test_omni_v2_async_chunk_updates.py

@@ -0,0 +1,164 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+import torch
+
+from vllm_omni.core.sched.output import OmniCachedRequestData
+from vllm_omni.model_executor.models.output_templates import OmniOutput
+from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner
+
+
+class _Array:
+    def __init__(self):
+        self.np = [0]
+        self.writes = []
+        self.applied = False
+
+    def stage_write_elem(self, idx, value):
+        self.writes.append((idx, value))
+
+    def apply_write(self):
+        self.applied = True
+
+
+class _TokenIds:
+    def __init__(self):
+        self.writes = []
+
+    def stage_write(self, idx, start, values):
+        self.writes.append((idx, start, list(values)))
+
+
+class _ReqStates:
+    def __init__(self, has_request=True):
+        self.req_id_to_index = {"r1": 0} if has_request else {}
+        self.prompt_len = SimpleNamespace(np=[0])
+        self.prefill_len = SimpleNamespace(np=[0])
+        self.total_len = _Array()
+        self.all_token_ids = _TokenIds()
+        self.num_computed_tokens = _Array()
+        self.num_computed_prefill_tokens = [7]
+        self.applied = False
+        self.removed = []
+
+    def apply_staged_writes(self):
+        self.applied = True
+
+    def remove_request(self, req_id):
+        self.removed.append(req_id)
+        self.req_id_to_index.pop(req_id, None)
+
+
+class _IntermediateBuffer:
+    def __init__(self):
+        self.buffers = [{"stale": "value", "req_id": "r1"}]
+
+    def remove_request(self, idx):
+        self.buffers[idx] = {}
+
+
+def test_async_chunk_update_clears_stale_intermediate_buffer():
+    runner = object.__new__(OmniGenerationModelRunner)
+    runner.req_states = _ReqStates()
+    runner.model_state = SimpleNamespace(intermediate_buffer=_IntermediateBuffer())
+    cached = OmniCachedRequestData(
+        req_ids=["r1"],
+        resumed_req_ids=set(),
+        new_token_ids=[],
+        all_token_ids=[],
+        new_block_ids=[],
+        num_computed_tokens=[],
+        num_output_tokens=[],
+        prompt_token_ids={"r1": [1, 2, 3]},
+        additional_information={"r1": {"fresh": "value"}},
+    )
+
+    runner._handle_async_chunk_updates(SimpleNamespace(scheduled_cached_reqs=cached))
+
+    assert runner.model_state.intermediate_buffer.buffers[0] == {}
+    assert runner.req_states.prompt_len.np[0] == 3
+    assert runner.req_states.prefill_len.np[0] == 3
+    assert runner.req_states.all_token_ids.writes == [(0, 0, [1, 2, 3])]
+    assert runner.req_states.num_computed_prefill_tokens[0] == 0
+    assert runner.req_states.applied is True
+
+
+def test_async_chunk_update_readds_released_cached_request():
+    runner = object.__new__(OmniGenerationModelRunner)
+    runner.req_states = _ReqStates(has_request=False)
+    runner.model_state = SimpleNamespace(intermediate_buffer=_IntermediateBuffer())
+    added = []
+    runner.add_requests = lambda scheduler_output: added.extend(scheduler_output.scheduled_new_reqs)
+    cached = OmniCachedRequestData(
+        req_ids=["r1"],
+        resumed_req_ids=set(),
+        new_token_ids=[],
+        all_token_ids={"r1": [1, 2, 3]},
+        new_block_ids=[([7],)],
+        num_computed_tokens=[0],
+        num_output_tokens=[0],
+        prompt_token_ids={"r1": [1, 2, 3]},
+        additional_information={"r1": {"fresh": "value"}},
+    )
+
+    runner._handle_async_chunk_updates(SimpleNamespace(scheduled_cached_reqs=cached))
+
+    assert len(added) == 1
+    assert added[0].req_id == "r1"
+    assert added[0].prompt_token_ids == [1, 2, 3]
+    assert added[0].prefill_token_ids == [1, 2, 3]
+    assert added[0].block_ids == ([7],)
+    assert added[0].additional_information == {"fresh": "value"}
+
+
+def _make_runner_for_sample(async_chunk):
+    runner = object.__new__(OmniGenerationModelRunner)
+    runner._gen_model_output = OmniOutput(
+        text_hidden_states=torch.zeros(1, 1),
+        multimodal_outputs={"audio": [torch.zeros(1)]},
+    )
+    runner._gen_input_batch = SimpleNamespace(
+        num_reqs=1,
+        idx_mapping_np=[0],
+        req_ids=["r1"],
+    )
+    runner._gen_kv_connector_output = None
+    runner.execute_model_state = object()
+    runner.model_config = SimpleNamespace(async_chunk=async_chunk)
+    runner.req_states = _ReqStates()
+    runner.model_state = SimpleNamespace(
+        intermediate_buffer=_IntermediateBuffer(),
+        remove_request=lambda idx: runner.model_state.intermediate_buffer.remove_request(idx),
+    )
+    removed = []
+    runner._remove_request = lambda req_id: removed.append(req_id) or True
+    return runner, removed
+
+
+def test_generation_sample_keeps_runner_slot_without_async_chunk():
+    runner, removed = _make_runner_for_sample(async_chunk=False)
+
+    output = runner.sample_tokens()
+
+    assert output.req_ids == ["r1"]
+    assert removed == []
+    assert runner.model_state.intermediate_buffer.buffers[0] == {"stale": "value", "req_id": "r1"}
+
+
+def test_generation_sample_releases_runner_slot_after_chunk_output():
+    runner, removed = _make_runner_for_sample(async_chunk=True)
+
+    output = runner.sample_tokens()
+
+    assert output.req_ids == ["r1"]
+    assert removed == ["r1"]
+    assert runner.model_state.intermediate_buffer.buffers[0] == {}
+
+
+if __name__ == "__main__":
+    test_async_chunk_update_clears_stale_intermediate_buffer()
+    test_async_chunk_update_readds_released_cached_request()
+    test_generation_sample_keeps_runner_slot_without_async_chunk()
+    test_generation_sample_releases_runner_slot_after_chunk_output()

tests/worker/test_omni_v2_async_output_signature.py

@@ -0,0 +1,16 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+
+from vllm_omni.worker_v2.omni_ar_model_runner import OmniAsyncOutput
+
+
+def test_omni_async_output_does_not_require_external_copy_event():
+    param = inspect.signature(OmniAsyncOutput.__init__).parameters["copy_event"]
+
+    assert param.default is None
+
+
+if __name__ == "__main__":
+    test_omni_async_output_does_not_require_external_copy_event()

tests/worker/test_omni_v2_execute_model_signature.py

@@ -0,0 +1,17 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import inspect
+
+from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+
+def test_v2_model_runners_accept_vllm_020_is_profile_keyword():
+    for runner_cls in (OmniARModelRunner, OmniGenerationModelRunner, OmniGPUModelRunner):
+        assert "is_profile" in inspect.signature(runner_cls.execute_model).parameters
+
+
+if __name__ == "__main__":
+    test_v2_model_runners_accept_vllm_020_is_profile_keyword()

tests/worker/test_omni_v2_execute_model_state_compat.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.worker_v2.omni_model_runner import _make_execute_model_state
+
+
+def test_make_execute_model_state_ignores_fields_missing_from_current_vllm():
+    state = _make_execute_model_state(
+        input_batch="input",
+        attn_metadata=None,
+        slot_mappings_by_layer=None,
+        hidden_states="hidden",
+        aux_hidden_states=None,
+        # finished_req_ids is a real field on vLLM 0.23.0's ExecuteModelState.
+        finished_req_ids=set(),
+        # These are fields the omni runners still pass but current vLLM no longer
+        # declares; they must be filtered out rather than crash the constructor.
+        kv_connector_output=None,
+        num_tokens_across_dp=8,
+    )
+
+    assert state.input_batch == "input"
+    assert state.hidden_states == "hidden"
+    assert not hasattr(state, "num_tokens_across_dp")
+    assert not hasattr(state, "kv_connector_output")
+
+
+if __name__ == "__main__":
+    test_make_execute_model_state_ignores_fields_missing_from_current_vllm()

tests/worker/test_omni_v2_profile_dispatch.py

@@ -0,0 +1,95 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from types import SimpleNamespace
+
+from vllm.compilation.cuda_graph import CUDAGraphMode
+
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+
+def test_eager_dispatch_uses_eager_without_cudagraph_manager():
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.cudagraph_manager = None
+    runner.dp_size = 1
+    runner.dp_rank = 0
+
+    batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
+        num_reqs=1,
+        num_toks=8,
+        uniform_tok_count=8,
+        use_eager=True,
+    )
+
+    assert batch_desc.cg_mode == CUDAGraphMode.NONE
+    assert batch_desc.num_tokens == 8
+    assert batch_desc.num_reqs == 1
+    assert num_tokens_across_dp is None
+
+
+def test_eager_dispatch_syncs_dp_padding(monkeypatch):
+    runner = object.__new__(OmniGPUModelRunner)
+    runner.cudagraph_manager = None
+    runner.dp_size = 2
+    runner.dp_rank = 1
+    expected_tokens_across_dp = object()
+    calls = []
+
+    def sync_dp(cudagraph_manager, batch_desc, num_tokens, num_reqs, uniform_token_count, dp_size, dp_rank):
+        calls.append((cudagraph_manager, batch_desc, num_tokens, num_reqs, uniform_token_count, dp_size, dp_rank))
+        return batch_desc, expected_tokens_across_dp
+
+    import vllm.v1.worker.gpu.dp_utils as dp_utils
+
+    monkeypatch.setattr(dp_utils, "sync_cudagraph_and_dp_padding", sync_dp)
+
+    batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
+        num_reqs=1,
+        num_toks=8,
+        uniform_tok_count=8,
+        use_eager=True,
+    )
+
+    assert batch_desc.cg_mode == CUDAGraphMode.NONE
+    assert batch_desc.num_tokens == 8
+    assert batch_desc.num_reqs == 1
+    assert num_tokens_across_dp is expected_tokens_across_dp
+    assert calls == [(None, batch_desc, 8, 1, 8, 2, 1)]
+
+
+def test_non_eager_dispatch_uses_cudagraph_manager():
+    runner = object.__new__(OmniGPUModelRunner)
+    expected = SimpleNamespace(cg_mode=CUDAGraphMode.PIECEWISE, num_tokens=8, num_reqs=1)
+    calls = []
+    runner.cudagraph_manager = SimpleNamespace(dispatch=lambda *args: calls.append(args) or expected)
+    runner.dp_size = 1
+    runner.dp_rank = 0
+
+    batch_desc, num_tokens_across_dp = runner._dispatch_batch_descriptor(
+        num_reqs=1,
+        num_toks=8,
+        uniform_tok_count=8,
+        use_eager=False,
+    )
+
+    assert batch_desc is expected
+    assert calls == [(1, 8, 8)]
+    assert num_tokens_across_dp is None
+
+
+if __name__ == "__main__":
+    test_eager_dispatch_uses_eager_without_cudagraph_manager()
+    test_non_eager_dispatch_uses_cudagraph_manager()
+
+
+def test_skip_compiled_uses_eager_dispatch_without_overwriting_batch_descriptor():
+    source = __import__("inspect").getsource(OmniGPUModelRunner.execute_model)
+
+    assert "use_eager=is_profile or skip_compiled" in source
+    assert "if self.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs" not in source
+
+
+def test_fullgraph_requires_cudagraph_manager_contract():
+    source = __import__("inspect").getsource(OmniGPUModelRunner.execute_model)
+
+    assert "assert self.cudagraph_manager is not None" in source

tests/worker/test_omni_v2_shutdown_signature.py

@@ -0,0 +1,22 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm_omni.worker_v2.omni_ar_model_runner import OmniARModelRunner
+from vllm_omni.worker_v2.omni_generation_model_runner import OmniGenerationModelRunner
+from vllm_omni.worker_v2.omni_model_runner import OmniGPUModelRunner
+
+
+def test_v2_model_runners_have_shutdown_for_vllm_020_worker_exit():
+    for runner_cls in (OmniARModelRunner, OmniGenerationModelRunner, OmniGPUModelRunner):
+        assert hasattr(runner_cls, "shutdown")
+
+
+if __name__ == "__main__":
+    test_v2_model_runners_have_shutdown_for_vllm_020_worker_exit()
+
+
+def test_load_model_uses_class_patch_lock():
+    import inspect
+
+    source = inspect.getsource(OmniGPUModelRunner.load_model)
+    assert "_model_state_patch_lock" in source