From cc1939120f3dc06c9456aa1165833545feccab5d Mon Sep 17 00:00:00 2001 From: vOv <124694036+cr-gao@users.noreply.github.com> Date: Thu, 23 Jul 2026 00:01:45 +0800 Subject: [PATCH] [Bugfix] Skip non-positive scheduled token spans in omni stage runners (#5269) Signed-off-by: cr-gao Co-authored-by: Claude Opus 4.8 --- .../test_gpu_generation_model_runner.py | 43 +++++++++++++++++++ .../npu/worker/npu_ar_model_runner.py | 3 +- .../npu/worker/npu_generation_model_runner.py | 3 +- vllm_omni/worker/gpu_ar_model_runner.py | 3 +- .../worker/gpu_generation_model_runner.py | 3 +- 5 files changed, 51 insertions(+), 4 deletions(-) diff --git a/tests/worker/test_gpu_generation_model_runner.py b/tests/worker/test_gpu_generation_model_runner.py index 0cd9aeb26..e64eecb57 100644 --- a/tests/worker/test_gpu_generation_model_runner.py +++ b/tests/worker/test_gpu_generation_model_runner.py @@ -1,6 +1,12 @@ +import contextlib +from types import SimpleNamespace + import pytest import torch +from vllm.config import CacheConfig +from vllm.v1.worker.gpu_model_runner import EMPTY_MODEL_RUNNER_OUTPUT +import vllm_omni.worker.gpu_generation_model_runner as gen_runner_module from vllm_omni.worker.gpu_generation_model_runner import ( ExecuteModelState, GPUGenerationModelRunner, @@ -84,3 +90,40 @@ def test_sample_tokens_dict_output(): assert "audio" in output.multimodal_outputs[0] assert "unused" not in output.multimodal_outputs[0] assert output.multimodal_outputs[0]["audio"].shape == (1, 4) + + +class _StubSchedulerOutput: + def __init__(self, total_num_scheduled_tokens): + self.total_num_scheduled_tokens = total_num_scheduled_tokens + self.num_scheduled_tokens = {"req-1": total_num_scheduled_tokens} + self.finished_req_ids = set() + self.kv_connector_metadata = None + + +def _make_guard_runner(): + # Stubbed far enough that a span escaping the guard reaches the real + # `_prepare_inputs`, i.e. fails the way the reported crash does. + runner = object.__new__(GPUGenerationModelRunner) + runner.execute_model_state = None + runner.routed_experts_initialized = False + runner.speculative_config = None + runner.model_config = SimpleNamespace(async_chunk=False) + runner.cache_config = CacheConfig() + runner.input_batch = _DummyInputBatch() + runner.model = object() + runner._update_states = lambda scheduler_output: None + runner.synchronize_input_prep = contextlib.nullcontext + runner.attach_omni_connector_output = lambda result: result + return runner + + +@pytest.mark.parametrize("total", [-1, -512, 0]) +def test_execute_model_skips_non_positive_scheduled_span(monkeypatch, total): + """#5196: a negative span is truthy, so it used to reach `_prepare_inputs`, + whose `assert total_num_scheduled_tokens > 0` killed the stage EngineCore.""" + monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False) + runner = _make_guard_runner() + + output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total)) + + assert output is EMPTY_MODEL_RUNNER_OUTPUT diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index b7b275b12..0817d7bbf 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -463,7 +463,8 @@ class NPUARModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin): output.kv_extracted_req_ids = kv_ids return self.attach_omni_connector_output(output) - if not num_scheduled_tokens: + # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). + if num_scheduled_tokens <= 0: if ( self.parallel_config.distributed_executor_backend == "external_launcher" and self.parallel_config.data_parallel_size > 1 diff --git a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py index d43867c97..f4fc05574 100644 --- a/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_generation_model_runner.py @@ -159,7 +159,8 @@ class NPUGenerationModelRunner(OmniNPUModelRunner, OmniConnectorModelRunnerMixin make_empty_encoder_model_runner_output(scheduler_output) ) - if not num_scheduled_tokens: + # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). + if num_scheduled_tokens <= 0: if ( self.parallel_config.distributed_executor_backend == "external_launcher" and self.parallel_config.data_parallel_size > 1 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 411796d37..67ea34b10 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1057,7 +1057,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): output.kv_extracted_req_ids = kv_ids return self.attach_omni_connector_output(output) - if not num_scheduled_tokens: + # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). + if num_scheduled_tokens <= 0: if ( self.parallel_config.distributed_executor_backend == "external_launcher" and self.parallel_config.data_parallel_size > 1 diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index bcb647605..6208f73e9 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -159,7 +159,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin if scheduler_output.finished_req_ids and hasattr(self.model, "on_requests_finished"): self.model.on_requests_finished(scheduler_output.finished_req_ids) - if not scheduler_output.total_num_scheduled_tokens: + # `<= 0`: upstream can schedule a negative span, which is truthy (#5196). + if scheduler_output.total_num_scheduled_tokens <= 0: return self.attach_omni_connector_output(EMPTY_MODEL_RUNNER_OUTPUT) if has_ec_transfer() and not get_ec_transfer().is_consumer: