rebase: keep the omni scheduler on the v1 model runner

Every omni worker already forces the v1 model runner because the v2 runner
carries no omni hooks (four worker classes do this). The scheduler was left
reading the config value, and vLLM 0.29 changed that default: 0.28 resolved
this config to v1, 0.29 ends its new fallback chain with v2.

The two then disagreed. The scheduler took its v2 fast path, which
deliberately omits resumed-request token ids ("No resumed reqs in fast
path"), while the worker-side v1 runner still read them, so resuming a
request raised KeyError: <req_id> in _update_states and killed the engine
core. That is the Multi-Replica Startup Test failure: it passes on main
6/6 and on the branch before the rebase (3005), and fails on every build
since.

Forcing it once in OmniSchedulerMixin covers both OmniARScheduler and
OmniGenerationScheduler, mirroring the wording the workers already use.

Adds tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py, since a
default-sensitive regression like this would otherwise return silently: it
pins Omni's side of the contract rather than the upstream default, and was
negative-controlled by removing the fix (2 of 3 cases fail without it).

The 4x H100 multi-replica path itself is not reproducible here, so that part
is verified only by proxy: 165 sched tests, plus a 599-test
core/diffusion/ar_diffusion sweep, and the change is inert whenever the
config already resolves to v1.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0146Xu5Sgj4vXxLxvLjV5YPj
Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py

@@ -0,0 +1,59 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""The omni scheduler must stay on the v1 model runner, like every omni worker.
+
+vLLM 0.28 resolved this config to v1; 0.29 ends its fallback chain with v2. When
+only the workers forced v1, the scheduler took its v2 fast path -- which omits
+resumed-request token ids -- while the v1 runner still read them, so resuming a
+request raised ``KeyError: <req_id>`` in ``_update_states`` and killed the engine
+core (Multi-Replica Startup Test).
+
+The default is upstream's to change, so this pins Omni's side of the contract
+rather than the default itself.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class _Scheduler(OmniSchedulerMixin):
+    """Minimal carrier for the shared init hook both omni schedulers call."""
+
+    def __init__(self, *, use_v2_model_runner: bool) -> None:
+        self.use_v2_model_runner = use_v2_model_runner
+        self.vllm_config = SimpleNamespace(
+            model_config=SimpleNamespace(
+                async_chunk=False,
+                stage_id=0,
+                pooling_output_decoder=None,
+            )
+        )
+        self._init_omni_io_scheduling_state()
+
+
+@pytest.mark.parametrize("initial", [True, False])
+def test_scheduler_ends_on_v1_model_runner(initial: bool) -> None:
+    """Whatever upstream resolves, the omni scheduler runs v1."""
+    assert _Scheduler(use_v2_model_runner=initial).use_v2_model_runner is False
+
+
+def test_v2_default_does_not_reach_the_resumed_request_fast_path() -> None:
+    """The v2 fast path is what dropped resumed-request token ids.
+
+    ``OmniGenerationScheduler`` branches on ``use_v2_model_runner`` when it
+    assembles SchedulerOutput; forcing it false keeps the v1 assembly that
+    carries ``all_token_ids`` for resumed requests, which the v1 runner reads.
+    """
+    scheduler = _Scheduler(use_v2_model_runner=True)
+
+    assert not scheduler.use_v2_model_runner, (
+        "scheduler would take the v2 fast path while omni workers run v1; "
+        "resumed requests then KeyError in _update_states"
+    )

vllm_omni/core/sched/omni_scheduler_mixin.py

@@ -107,6 +107,17 @@ class OmniSchedulerMixin:
 
     def _init_omni_io_scheduling_state(self) -> None:
         """Initialize scheduler state shared by AR and generation stages."""
+        # Every omni worker forces the v1 model runner (the v2 runner carries no
+        # omni hooks), so the scheduler has to agree or the two disagree about
+        # what the SchedulerOutput carries. vLLM 0.29 defaults this to v2 where
+        # 0.28 did not, so the scheduler took its v2 fast path -- the one that
+        # deliberately omits resumed-request token ids -- while the v1 runner
+        # still read them, and resuming a request raised ``KeyError: <req_id>``
+        # in ``_update_states``.
+        if getattr(self, "use_v2_model_runner", False):
+            logger.warning("OMNI scheduler forces v1 model runner for omni hooks.")
+            self.use_v2_model_runner = False
+
         model_config = self.vllm_config.model_config
         self.chunk_transfer_adapter = (
             OmniChunkTransferAdapter(self.vllm_config) if getattr(model_config, "async_chunk", False) else None