vllm-omni/tests/worker_v2/test_omni_model_state.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""OmniModelState: mixed-batch reorder, state isolation/slot reuse, MTP
graph/eager paths with per-request seed independence, async snapshot ownership."""

from contextlib import nullcontext
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import numpy as np
import pytest
import torch
from vllm.v1.worker.gpu.model_states.default import DefaultModelState

from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor
from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _make_safe_get_rope

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


class _DummyInputBatch:
    input_ids: SimpleNamespace
    query_start_loc: torch.Tensor

    def __init__(self, indices, *, num_computed_tokens_cpu=None):
        self.idx_mapping_np = indices
        self.num_reqs = len(indices)
        self.num_scheduled_tokens = [1] * len(indices)
        self.query_start_loc_np = list(range(len(indices)))
        if num_computed_tokens_cpu is not None:
            self.num_computed_tokens_cpu = np.array(num_computed_tokens_cpu, dtype=np.int32)


def _make_state(max_num_reqs=4, has_preprocess=False, has_postprocess=False, have_multimodal_outputs=False):
    state = object.__new__(OmniModelState)
    model = MagicMock()
    model.has_preprocess = has_preprocess
    model.has_postprocess = has_postprocess
    model.have_multimodal_outputs = have_multimodal_outputs
    model.gpu_resident_buffer_keys = set()
    model.batched_gpu_staging_keys = set()
    model.mtp_accepts_per_row_generators = False
    model.mtp_accepts_req_infos = False
    model.mtp_output_key = ("codes", "audio")
    model.mtp_validity_key = None
    model.mtp_sampling_params = {}
    model.get_mtp_seed = lambda params: (getattr(params, "extra_args", None) or {}).get("test_seed")
    model.preprocess_batch_mrv2 = None
    model.preprocess_decode_batch_mrv2 = None
    model.preprocess_decode_batch = None
    model.postprocess_batch_mrv2 = None
    state.model = model
    state.scheduler_config = SimpleNamespace(max_num_seqs=max_num_reqs)
    state.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
    state.has_preprocess = has_preprocess
    state.has_postprocess = has_postprocess
    state.have_multimodal_outputs = have_multimodal_outputs

    from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer

    state.intermediate_buffer = OmniIntermediateBuffer(max_num_reqs)
    state._static_inputs_embeds = None
    state._mtp_generators = {}
    state._mtp_runner = None
    for name in ("_mtp_input_ids", "_mtp_input_embeds", "_mtp_hidden", "_mtp_text_step", "_mtp_offsets"):
        setattr(state, name, None)
    return state


def _add(state, req_id="r1", idx=0, **req_kwargs):
    with patch.object(DefaultModelState, "add_request", return_value=None):
        state.add_request(idx, SimpleNamespace(req_id=req_id, mm_features=[], **req_kwargs))


def _fill_buffers(state, *req_ids):
    for idx, req_id in enumerate(req_ids):
        state.intermediate_buffer.buffers[idx] = {"req_id": req_id}


def _seeded(seed):
    return SimpleNamespace(extra_args={"test_seed": seed}, seed=None)


def _init_static(state, bsz, dim=3):
    state._mtp_input_ids = torch.zeros(bsz, dtype=torch.long)
    state._mtp_input_embeds = torch.zeros((bsz, dim))
    state._mtp_hidden = torch.zeros((bsz, dim))
    state._mtp_text_step = torch.zeros((bsz, dim))
    state._mtp_offsets = torch.zeros(bsz, dtype=torch.long)


def _mtp_batches(dim=3):
    return [(0, 0, (torch.ones(dim), torch.ones(dim) * 2)), (1, 1, (torch.ones(dim) * 3, torch.ones(dim) * 4))]


@pytest.fixture(autouse=True)
def _fwd_ctx():
    with patch("vllm.forward_context.set_forward_context", return_value=nullcontext()):
        yield


@pytest.mark.parametrize("req_id,expect_validity", [("_warmup_0_", True), ("r1", False)])
def test_add_request_declared_validity_only_for_warmup(req_id, expect_validity):
    # Real requests must not be masked with a fabricated validity key.
    state = _make_state()
    state.model.mtp_validity_key = ("meta", "codec_frame_valid")
    _add(state, req_id, 0)
    buf = state.intermediate_buffer.buffers[0]
    assert ("meta" in buf) is expect_validity
    if expect_validity:
        assert buf["meta"]["codec_frame_valid"].item() is False


@pytest.mark.parametrize("mode", ["index", "req_id", "unknown"])
def test_remove_request_state_isolation(mode):
    state = _make_state()
    _add(state, "r1", 0)
    state._mtp_generators["r1"] = torch.Generator(device="cpu")
    if mode == "unknown":
        state.remove_request("missing")
        assert state.intermediate_buffer.buffers[0]["req_id"] == "r1"
        assert "r1" in state._mtp_generators
        return
    state.remove_request(0 if mode == "index" else "r1")
    assert state.intermediate_buffer.buffers[0] == {}
    assert "r1" not in state._mtp_generators  # recycled slot inherits no seed stream


def test_output_spans_follow_reordered_mixed_batch():
    # batch=[2, 0]: prefill (3 tokens) reordered ahead of decode (1 token).
    state = _make_state(have_multimodal_outputs=True)
    batch = _DummyInputBatch([2, 0])
    batch.num_scheduled_tokens = [3, 1]
    batch.query_start_loc_np = [0, 3]
    state.intermediate_buffer.buffers[2] = {"req_id": "prefill"}
    state.intermediate_buffer.buffers[0] = {"req_id": "decode"}
    seen = {}
    state.model.make_omni_output = lambda hidden, **kwargs: (
        seen.update(kwargs) or OmniOutput(text_hidden_states=hidden, multimodal_outputs={})
    )
    state.postprocess_model_output(torch.zeros(4, 2), batch, SimpleNamespace())
    assert seen["request_token_spans"] == [(0, 3), (3, 4)]
    assert [info["req_id"] for info in seen["model_intermediate_buffer"]] == ["prefill", "decode"]


@pytest.mark.parametrize("hook_name", ["preprocess_decode_batch_mrv2", "preprocess_decode_batch", "legacy"])
def test_decode_cohort_uses_batch_hook(hook_name):
    state = _make_state(max_num_reqs=2, has_preprocess=True)
    state.intermediate_buffer.buffers[0] = {"req_id": "r1", "meta": {"step": 1}}
    state.intermediate_buffer.buffers[1] = {"req_id": "r2", "meta": {"step": 2}}
    batch_calls = []

    def preprocess_decode_batch(*, input_ids, input_embeds, req_infos):
        batch_calls.append(list(req_infos))
        return (
            input_ids + 10,
            input_embeds + 20,
            torch.tensor([[31.0, 32.0], [41.0, 42.0]]),
            torch.tensor([[51.0, 52.0], [61.0, 62.0]]),
            [{"meta": {"step": 11}}, {"meta": {"step": 12}}],
        )

    if hook_name == "legacy":
        # Legacy V1 hook without input_embeds is wrapped by the adapter.
        state.model.preprocess_decode_batch = lambda *, input_ids, req_infos: preprocess_decode_batch(
            input_ids=input_ids, input_embeds=model_inputs["inputs_embeds"], req_infos=req_infos
        )
    else:
        setattr(state.model, hook_name, preprocess_decode_batch)
        if hook_name == "preprocess_decode_batch_mrv2":
            state.model.preprocess_decode_batch = MagicMock(
                side_effect=AssertionError("MRV2 hook must take precedence")
            )
    state.model.preprocess = MagicMock(side_effect=AssertionError("decode requests must use the batch hook"))
    seen_prepacked = []

    def capture(batches, *_a, prepacked_mtp_inputs=None):
        seen_prepacked.append(prepacked_mtp_inputs)

    state._run_batched_mtp = capture
    model_inputs = {
        "input_ids": torch.tensor([101, 202], dtype=torch.long),
        "inputs_embeds": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),
    }
    state.run_preprocess(
        _DummyInputBatch([0, 1], num_computed_tokens_cpu=[1, 1]),
        model_inputs,
        SimpleNamespace(prompt_len=np.array([1, 1], dtype=np.int32)),
    )
    state.model.preprocess.assert_not_called()
    assert [info["req_id"] for info in batch_calls[0]] == ["r1", "r2"]
    assert torch.equal(model_inputs["input_ids"], torch.tensor([111, 212]))
    assert state.intermediate_buffer.buffers[0]["meta"]["step"] == 11
    assert torch.equal(seen_prepacked[0][0], torch.tensor([[31.0, 32.0], [41.0, 42.0]]))


def test_static_decode_embeddings_refresh_from_input_ids():
    # FULL-graph replay reads the static buffer: it must hold fresh embeddings.
    state = _make_state(has_preprocess=True)
    static_embeds = torch.zeros(1, 4)
    state._static_inputs_embeds = static_embeds
    expected = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
    state.model.embed_input_ids = lambda input_ids: expected.to(input_ids.device)
    state.model.preprocess = lambda input_ids, input_embeds, **_info: (input_ids, input_embeds, {})
    state.intermediate_buffer.buffers[0] = {"req_id": "r1"}
    state.run_preprocess(
        _DummyInputBatch([0]),
        {"input_ids": torch.tensor([1049], dtype=torch.long), "inputs_embeds": static_embeds[:1]},
    )
    assert torch.equal(static_embeds, expected)
    original = torch.tensor([[1.0, 2.0]])
    assert OmniModelState._preprocess_result_needs_writeback(original, original) is False
    assert OmniModelState._preprocess_result_needs_writeback(original, original.view_as(original)) is True


@pytest.mark.parametrize("owned", [False, True])
def test_batched_postprocess_gpu_snapshot_writeback(owned):
    # batch=[1, 0]: last-token indices follow the reordered batch, the scalar
    # hook is not called, and both rows share one snapshot storage so a later
    # graph replay cannot overwrite published rows. Owned producer batches
    # skip the extra snapshot while borrowed graph output still gets one.
    state = _make_state(max_num_reqs=2, has_postprocess=True)
    state.model.gpu_resident_buffer_keys = {("hidden_states", "last")}
    _fill_buffers(state, "r0", "r1")
    state.model.postprocess = MagicMock(side_effect=AssertionError("batch hook must replace scalar postprocess"))
    calls = []
    produced = []

    def postprocess_batch(*, hidden_states, last_token_indices):
        calls.append(last_token_indices.clone())
        gathered = hidden_states.index_select(0, last_token_indices)
        produced.append(gathered)
        if owned:
            return (("hidden_states", "last"), OwnedBatchTensor(gathered))
        return (("hidden_states", "last"), gathered)

    state.model.postprocess_batch_mrv2 = postprocess_batch
    batch = _DummyInputBatch([1, 0])
    batch.num_scheduled_tokens = [2, 2]
    batch.query_start_loc_np = [0, 2]
    batch.query_start_loc = torch.tensor([0, 2, 4], dtype=torch.int32)
    hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])
    state.run_postprocess(hidden, batch)

    state.model.postprocess.assert_not_called()
    assert torch.equal(calls[0], torch.tensor([1, 3], dtype=torch.int32))
    first = state.intermediate_buffer.buffers[0]["hidden_states"]["last"]
    second = state.intermediate_buffer.buffers[1]["hidden_states"]["last"]
    assert torch.equal(first, hidden[3])
    assert torch.equal(second, hidden[1])
    assert first.untyped_storage().data_ptr() == second.untyped_storage().data_ptr()
    if owned:
        # Owned output: rows view into the producer batch, no extra snapshot.
        assert first.untyped_storage().data_ptr() == produced[0].untyped_storage().data_ptr()
    else:
        # Borrowed/graph output: the snapshot owns its storage, not the producer.
        assert first.untyped_storage().data_ptr() != produced[0].untyped_storage().data_ptr()


def test_seed_independence_resolve_once_and_sampling_kwargs():
    state = _make_state(max_num_reqs=2)
    # vLLM sampling seed must not produce a talker generator.
    cpu = torch.device("cpu")
    assert state._get_mtp_generator("r1", SimpleNamespace(extra_args={}, seed=42), cpu) is None
    # Same model-local seed reproduces identical uniforms regardless of batch makeup.
    state._mtp_sample_uniforms = torch.empty((2, 2, 4))
    assert torch.equal(
        state._prepare_mtp_sample_uniforms([torch.Generator().manual_seed(11)], 1).clone(),
        state._prepare_mtp_sample_uniforms([torch.Generator().manual_seed(11)], 1),
    )
    # Per-request generators are resolved once per step and ride sampling kwargs.
    state.model.mtp_sampling_params = {"do_sample": True, "temperature": 0.7}
    state.model.mtp_accepts_per_row_generators = True
    _init_static(state, 2)
    state.intermediate_buffer.buffers[0] = {"req_id": "r0", "sampling_params": _seeded(11)}
    state.intermediate_buffer.buffers[1] = {"req_id": "r1", "sampling_params": _seeded(22)}
    resolver = MagicMock(wraps=state._get_mtp_generator)
    state._get_mtp_generator = resolver

    def mtp(input_ids, input_embeds, last_hidden, text_step, **kwargs):
        assert (kwargs["do_sample"], kwargs["temperature"]) == (True, 0.7)
        assert all(isinstance(g, torch.Generator) for g in kwargs["generators"])
        return input_embeds + 10, torch.tensor([[1, 2, 3], [4, 5, 6]])

    state.model.mtp = mtp
    state._run_batched_mtp(
        _mtp_batches(), torch.tensor([101, 202]), torch.zeros(2, 3), _DummyInputBatch([0, 1]), {("codes", "audio")}
    )
    assert resolver.call_count == 2
    assert torch.equal(state.intermediate_buffer.buffers[0]["codes"]["audio"], torch.tensor([[1, 2, 3]]))
    assert torch.equal(state.intermediate_buffer.buffers[1]["codes"]["audio"], torch.tensor([[4, 5, 6]]))


def test_seeded_mtp_bypasses_outer_graph_runner():
    state = _make_state(max_num_reqs=2)
    raw_calls = []

    def mtp(*args, **kwargs):
        raw_calls.append(kwargs)
        return args[1], args[0].reshape(-1, 1)

    state.model.mtp = mtp
    state._mtp_runner = MagicMock(side_effect=AssertionError("seeded sampling must not replay the graph"))
    generators = [torch.Generator().manual_seed(11), torch.Generator().manual_seed(22)]
    state._call_mtp_runner(
        torch.tensor([101, 202]), torch.zeros(2, 3), torch.zeros(2, 3), torch.zeros(2, 3), generators=generators
    )
    assert raw_calls == [{"generators": generators}]
    state._mtp_runner.assert_not_called()


def test_run_batched_mtp_uses_dispatched_graph_descriptor():
    state = _make_state(max_num_reqs=4)

    class _FakeGraphRunner:
        def __call__(self, input_ids, input_embeds, last_hidden, text_step, **kwargs):
            assert input_ids.shape[0] == 4  # graph runs the padded descriptor size
            return input_embeds + 2, torch.arange(12, dtype=torch.long).reshape(4, 3)

    state._mtp_runner = _FakeGraphRunner()
    _init_static(state, 4)
    _fill_buffers(state, "r0", "r1")
    graph_desc = SimpleNamespace(cg_mode="FULL", num_tokens=4)
    set_ctx = MagicMock(return_value=nullcontext())
    with (
        patch(
            "vllm_omni.worker_v2.model_states.omni_model_state.current_omni_platform.get_graph_wrapper_cls",
            return_value=_FakeGraphRunner,
        ),
        patch("vllm.forward_context.set_forward_context", set_ctx),
    ):
        state._run_batched_mtp(
            _mtp_batches(),
            torch.tensor([101, 202]),
            torch.zeros(2, 3),
            _DummyInputBatch([0, 1]),
            {("codes", "audio")},
            lambda bsz: graph_desc,
        )
    _, kwargs = set_ctx.call_args
    assert kwargs["batch_descriptor"] is graph_desc
    assert kwargs["num_tokens"] == 4
    assert torch.equal(state.intermediate_buffer.buffers[0]["codes"]["audio"], torch.tensor([[0, 1, 2]]))
    assert torch.equal(state.intermediate_buffer.buffers[1]["codes"]["audio"], torch.tensor([[3, 4, 5]]))


def test_rope_shim_propagates_type_error():
    # vLLM rope API drift (TypeError) must not be swallowed by the shim.
    def broken_get_rope(*_args, **_kwargs):
        raise TypeError("vLLM rope API drift")

    with pytest.raises(TypeError, match="rope API drift"):
        _make_safe_get_rope(broken_get_rope)(SimpleNamespace(uses_mrope=False), object())


def test_rope_shim_constructs_sequential_mrope_state():
    from vllm.platforms.cpu import CpuPlatform
    from vllm.v1.worker.gpu.mm.rope import RopeState, get_rope_state

    model = torch.nn.Module()
    config = SimpleNamespace(uses_mrope=True, mrope_num_dims=3)
    # Exercise the upstream constructor, including its argument contract, while
    # keeping the backing buffers on CPU for this unit test.
    with (
        patch("vllm.platforms.current_platform", CpuPlatform()),
        patch("vllm.v1.worker.gpu.buffer_utils.is_uva_available", return_value=False),
    ):
        rope = _make_safe_get_rope(get_rope_state)(
            config,
            model,
            max_num_reqs=2,
            max_num_tokens=8,
            max_model_len=16,
            device=torch.device("cpu"),
        )

    assert isinstance(rope, RopeState)
    assert rope.get_positions(4).shape == (3, 4)
    positions, delta = model.get_mrope_input_positions([11, 22, 33, 44], [])
    torch.testing.assert_close(positions, torch.arange(4).expand(3, -1))
    assert delta == 0
    rope.init_prefill_positions(1, model, [11, 22, 33, 44], [])
    assert rope.prefill_delta.np[1] == 0
    assert rope.prefill_positions._staged_write_indices == [3, 4, 5]
    assert rope.prefill_positions._staged_write_contents == [0, 1, 2, 3] * 3


@pytest.mark.parametrize("output_key", [None, ("only_one",)])
def test_mtp_requires_model_declared_output_key(output_key):
    state = _make_state()
    model = SimpleNamespace(mtp=lambda: None, mtp_output_key=output_key)
    with pytest.raises(TypeError, match="must declare mtp_output_key"):
        state._init_mtp_runner(model)