[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