[Bugfix] Skip non-positive scheduled token spans in omni stage runners (#5269)

Signed-off-by: cr-gao <gaochenrui@sjtu.edu.cn>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

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

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

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

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

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: