vllm-omni/tests/worker/test_gpu_generation_model_runner.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

import contextlib
from types import SimpleNamespace
from unittest.mock import Mock

import pytest
import torch
from vllm.config import CacheConfig, CUDAGraphMode
from vllm.v1.outputs 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,
)

pytestmark = [pytest.mark.core_model, pytest.mark.cpu]


class _StopExecutionError(Exception):
    """Halts execute_model at the model boundary."""


@pytest.mark.parametrize("randomize_inputs", [False, True])
def test_generation_profile_accepts_upstream_randomize_inputs(randomize_inputs):
    """Release workers can pass the profiling flag to codec-stage runners."""
    runner = object.__new__(GPUGenerationModelRunner)
    runner.supports_mm_inputs = False
    runner.max_num_tokens = 128
    runner.encoder_cache = {"temporary": object()}
    runner._dummy_run = Mock(return_value=(torch.zeros(1), None))
    runner._sync_device = Mock()

    runner.profile_run(randomize_inputs=randomize_inputs)

    runner._dummy_run.assert_called_once_with(128, is_profile=True, randomize_inputs=randomize_inputs)
    runner._sync_device.assert_called_once_with()
    assert runner.encoder_cache == {}


class _DummyInputBatch:
    def __init__(self, num_reqs: int = 1):
        self.req_ids = [f"req-{i + 1}" for i in range(num_reqs)]
        self.req_id_to_index = {rid: i for i, rid in enumerate(self.req_ids)}
        self.num_reqs = num_reqs
        self.vocab_size = 10


def _make_runner(multimodal_outputs, num_reqs: int = 1):
    runner = object.__new__(GPUGenerationModelRunner)
    runner.execute_model_state = ExecuteModelState(
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        None,
        multimodal_outputs,
        None,
    )
    runner.kv_connector_output = None
    runner.input_batch = _DummyInputBatch(num_reqs)
    runner.use_async_scheduling = False
    runner.device = torch.device("cpu")
    runner.supports_mm_inputs = False
    runner.speculative_config = None
    runner.routed_experts_initialized = False
    runner._async_chunk = False
    return runner


def test_sample_tokens_tensor_output():
    multimodal_outputs = torch.randn(1, 2, 3)
    runner = _make_runner(multimodal_outputs)

    output = GPUGenerationModelRunner.sample_tokens(runner)

    assert len(output.multimodal_outputs) == 1
    assert output.multimodal_outputs[0]["model_outputs"].shape == (2, 3)


def test_sample_tokens_list_output():
    multimodal_outputs = [torch.randn(2, 1)]
    runner = _make_runner(multimodal_outputs)

    output = GPUGenerationModelRunner.sample_tokens(runner)

    assert len(output.multimodal_outputs) == 1
    assert output.multimodal_outputs[0]["model_outputs"].shape == (2, 1)


def test_sample_tokens_list_allows_none_output():
    multimodal_outputs = [None]
    runner = _make_runner(multimodal_outputs)

    output = GPUGenerationModelRunner.sample_tokens(runner)

    assert len(output.multimodal_outputs) == 1
    assert output.multimodal_outputs[0]["model_outputs"] is None


def test_sample_tokens_dict_output():
    multimodal_outputs = {"audio": torch.randn(1, 4), "unused": None}
    runner = _make_runner(multimodal_outputs)

    output = GPUGenerationModelRunner.sample_tokens(runner)

    assert len(output.multimodal_outputs) == 1
    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
        self.scheduled_encoder_inputs = {}
        self.scheduled_spec_decode_tokens = {}
        self.num_common_prefix_blocks = 0


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.kv_cache_config = SimpleNamespace(kv_cache_groups=[])
    runner.input_batch = _DummyInputBatch()
    runner.model = object()
    runner.parallel_config = SimpleNamespace(
        distributed_executor_backend=None,
        data_parallel_size=1,
    )
    runner._dummy_run_calls = []
    runner._dummy_run = lambda n: runner._dummy_run_calls.append(n)
    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_ec_transfer", lambda: False)
    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
    assert runner._dummy_run_calls == []


@pytest.mark.parametrize("total", [-1, 0])
def test_execute_model_zero_tokens_runs_dp_dummy_run(monkeypatch, total):
    """A 0-token step under external_launcher + DP>1 must `_dummy_run(1)` so
    `coordinate_batch_across_dp` stays in sync (it used to be shadowed by an
    unconditional early return and never ran -> DP hang risk)."""
    monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False)
    monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False)
    runner = _make_guard_runner()
    runner.parallel_config = SimpleNamespace(
        distributed_executor_backend="external_launcher",
        data_parallel_size=2,
    )

    output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(total))

    assert output is EMPTY_MODEL_RUNNER_OUTPUT
    assert runner._dummy_run_calls == [1]


def test_execute_model_zero_tokens_kv_connector_no_forward(monkeypatch):
    """With a KV-transfer group, a 0-token step must go through
    `kv_connector_no_forward` (also shadowed by the old early return)."""
    monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False)
    monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: True)
    runner = _make_guard_runner()
    sentinel = object()
    calls = []
    runner.kv_connector_no_forward = lambda scheduler_output, vllm_config: (
        calls.append(scheduler_output),
        sentinel,
    )[1]
    runner.vllm_config = object()

    output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0))

    assert output is sentinel
    assert len(calls) == 1


def test_execute_model_zero_tokens_ec_producer_takes_encoder_path(monkeypatch):
    """AR lock-step: the EC-producer block sits ABOVE the zero-token block
    (matching GPUARModelRunner), so a 0-token step on an EC producer returns
    the empty ENCODER output. Before the C1 reorder, the unconditional early
    return won and produced EMPTY_MODEL_RUNNER_OUTPUT instead — this is the
    one cell whose behaviour the reorder changed."""
    monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: True)
    monkeypatch.setattr(gen_runner_module, "get_ec_transfer", lambda: SimpleNamespace(is_consumer=False))
    monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False)
    sentinel = object()
    monkeypatch.setattr(gen_runner_module, "make_empty_encoder_model_runner_output", lambda so: sentinel)
    runner = _make_guard_runner()
    runner.encoder_cache = {}
    runner.maybe_get_ec_connector_output = lambda scheduler_output, encoder_cache: contextlib.nullcontext()
    encoder_calls = []
    runner._execute_mm_encoder = lambda so: encoder_calls.append(so)

    output = GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(0))

    assert output is sentinel
    assert len(encoder_calls) == 1
    assert runner._dummy_run_calls == []


class TestSampleTokensBatched:
    """RFC #5450 C5: `sample_tokens` must slice per request for num_reqs > 1.
    The tensor branch's asserts used to force num_reqs == 1 and the list
    branch built a length-1 payload misaligned with `req_ids`."""

    def test_tensor_output_sliced_per_request(self):
        multimodal_outputs = torch.arange(12, dtype=torch.float32).reshape(2, 2, 3)
        runner = _make_runner(multimodal_outputs, num_reqs=2)

        output = GPUGenerationModelRunner.sample_tokens(runner)

        assert len(output.multimodal_outputs) == 2
        assert torch.equal(output.multimodal_outputs[0]["model_outputs"], multimodal_outputs[0])
        assert torch.equal(output.multimodal_outputs[1]["model_outputs"], multimodal_outputs[1])

    def test_list_output_one_entry_per_request(self):
        entries = [torch.zeros(2, 1), torch.ones(2, 1)]
        runner = _make_runner(list(entries), num_reqs=2)

        output = GPUGenerationModelRunner.sample_tokens(runner)

        assert len(output.multimodal_outputs) == 2
        assert torch.equal(output.multimodal_outputs[0]["model_outputs"], entries[0])
        assert torch.equal(output.multimodal_outputs[1]["model_outputs"], entries[1])

    def test_tensor_output_batch_mismatch_raises(self):
        runner = _make_runner(torch.randn(1, 2, 3), num_reqs=2)

        with pytest.raises(ValueError, match="leading dim 1 .* 2 requests"):
            GPUGenerationModelRunner.sample_tokens(runner)

    def test_list_output_batch_mismatch_raises(self):
        runner = _make_runner([torch.randn(2, 1)], num_reqs=2)

        with pytest.raises(ValueError, match="length 1 .* 2 requests"):
            GPUGenerationModelRunner.sample_tokens(runner)


@pytest.mark.parametrize("exact_shape, expected_len", [(True, 3), (False, 4)])
def test_exact_shape_models_receive_unpadded_input_ids(monkeypatch, exact_shape, expected_len):
    """#6712: the cudagraph dispatcher rounds a 3-token batch up to a capture
    size of 4, so ``_preprocess`` returns a 4-element ``input_ids`` while
    ``seq_token_counts`` stays ``[3]``. Models opting into
    ``requires_exact_input_shape`` must receive the trimmed view; everything
    else must keep the padded buffer.
    """
    monkeypatch.setattr(gen_runner_module, "has_ec_transfer", lambda: False)
    monkeypatch.setattr(gen_runner_module, "has_kv_transfer_group", lambda: False)

    runner = _make_guard_runner()
    runner.model = SimpleNamespace(
        has_preprocess=True,
        requires_request_ids=False,
        requires_exact_input_shape=exact_shape,
    )
    runner.cascade_attn_enabled = False
    runner.num_prompt_logprobs = None
    runner.vllm_config = None
    runner.parallel_config = SimpleNamespace(
        distributed_executor_backend=None,
        data_parallel_size=1,
        use_ubatching=False,
        num_ubatches=1,
    )

    monkeypatch.setattr(GPUGenerationModelRunner, "_prepare_inputs", lambda self, so, nst: (None, None, 1))
    monkeypatch.setattr(
        GPUGenerationModelRunner,
        "_determine_batch_execution_and_padding",
        lambda self, **kw: (
            CUDAGraphMode.PIECEWISE,
            SimpleNamespace(num_tokens=4, num_reqs=1),
            False,
            None,
            None,
        ),
    )
    monkeypatch.setattr(GPUGenerationModelRunner, "_get_slot_mappings", lambda self, **kw: (None, None))
    monkeypatch.setattr(GPUGenerationModelRunner, "_build_attention_metadata", lambda self, **kw: (None, None))
    monkeypatch.setattr(
        GPUGenerationModelRunner,
        "_maybe_attach_attention_metadata_extensions",
        lambda self, **kw: None,
    )
    monkeypatch.setattr(gen_runner_module, "set_forward_context", lambda *a, **kw: contextlib.nullcontext())

    def _fake_preprocess(self, so, num_input_tokens, itensors=None):
        return (torch.zeros(num_input_tokens, dtype=torch.int32), None, None, None, {}, None)

    monkeypatch.setattr(GPUGenerationModelRunner, "_preprocess", _fake_preprocess)

    seen = {}

    def _capture(self, **kwargs):
        seen["input_ids"] = kwargs["input_ids"]
        seen["counts"] = kwargs["model_kwargs"]["seq_token_counts"]
        raise _StopExecutionError

    monkeypatch.setattr(GPUGenerationModelRunner, "_run_generation_model", _capture)

    with pytest.raises(_StopExecutionError):
        GPUGenerationModelRunner.execute_model(runner, _StubSchedulerOutput(3))

    assert seen["input_ids"].numel() == expected_len
    assert seen["counts"] == [3]
    if exact_shape:
        assert sum(seen["counts"]) == seen["input_ids"].numel()