From 94914a127087c87e4807ea5b125b0ee28f5ccbad Mon Sep 17 00:00:00 2001 From: tzhouam Date: Thu, 3 Sep 2026 16:10:14 +0800 Subject: [PATCH] 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: 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 Claude-Session: https://claude.ai/code/session_0146Xu5Sgj4vXxLxvLjV5YPj Signed-off-by: tzhouam --- ...t_omni_scheduler_forces_v1_model_runner.py | 59 +++++++++++++++++++ vllm_omni/core/sched/omni_scheduler_mixin.py | 11 ++++ 2 files changed, 70 insertions(+) create mode 100644 tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py diff --git a/tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py b/tests/core/sched/test_omni_scheduler_forces_v1_model_runner.py new file mode 100644 index 000000000..fb8546e2d --- /dev/null +++ b/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: `` 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" + ) diff --git a/vllm_omni/core/sched/omni_scheduler_mixin.py b/vllm_omni/core/sched/omni_scheduler_mixin.py index aa71245a0..9123b794d 100644 --- a/vllm_omni/core/sched/omni_scheduler_mixin.py +++ b/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: `` + # 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