815 lines
36 KiB
Python
815 lines
36 KiB
Python
# 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 dataclasses import replace
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
|
|
from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner
|
|
from vllm.v1.worker.gpu.model_states.default import DefaultModelState
|
|
from vllm.v1.worker.gpu.states import RequestState
|
|
|
|
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]
|
|
|
|
|
|
def test_prepare_attn_forwards_release_model_specific_metadata():
|
|
state = object.__new__(OmniModelState)
|
|
metadata = object()
|
|
output = object()
|
|
batch = SimpleNamespace()
|
|
with patch.object(DefaultModelState, "prepare_attn", return_value=output) as prepare:
|
|
assert state.prepare_attn(batch, None, (), None, [], None, model_specific_attn_metadata=metadata) is output
|
|
assert prepare.call_args.kwargs["model_specific_attn_metadata"] is metadata
|
|
|
|
|
|
class _DummyInputBatch:
|
|
is_prefilling_np: np.ndarray
|
|
req_ids: list[str]
|
|
|
|
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.stream_decoder = None
|
|
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
|
|
from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
|
|
|
|
state._eager_state = EagerMTPState(state)
|
|
state._stream_pos = {}
|
|
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
|
|
|
|
|
|
def test_moss_local_decode_runs_depth_predictor_and_routes_eos(mocker):
|
|
"""Exercise MRV2 dispatch through the real Local hook, output and logits."""
|
|
from vllm.sampling_params import SamplingParams
|
|
from vllm.v1.worker.gpu.states import RequestState
|
|
|
|
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_local_depth import MossTTSLocalDepthTransformer
|
|
from vllm_omni.model_executor.models.moss_tts.modeling_moss_tts_talker import (
|
|
MossTTSLocalTalkerForGeneration,
|
|
)
|
|
|
|
model = MossTTSLocalTalkerForGeneration.__new__(MossTTSLocalTalkerForGeneration)
|
|
torch.nn.Module.__init__(model)
|
|
model.model = torch.nn.Module()
|
|
model.model.embed_tokens = torch.nn.Embedding(8, 4, _weight=torch.zeros(8, 4))
|
|
model.hidden_size = 4
|
|
model.n_vq = 2
|
|
model.audio_pad_token_id = 8
|
|
model.audio_assistant_slot_token_id = 2
|
|
model.im_end_token_id = 3
|
|
model.text_vocab_size = 8
|
|
model.talker_mtp_output_key = ("audio_codes", "current")
|
|
model.talker_mtp_graph_safe = False
|
|
model.gpu_resident_buffer_keys = set()
|
|
model.audio_lm_heads = model.audio_embeddings = model.local_text_lm_head = None
|
|
model._audio_embed = lambda codes: codes[:, :1].expand(-1, 4).float()
|
|
frame = mocker.Mock(return_value=(torch.tensor([True, False]), torch.tensor([[1, 2], [3, 4]])))
|
|
model.local_transformer = mocker.Mock(spec=MossTTSLocalDepthTransformer, generate_frame=frame)
|
|
|
|
state = _make_state(max_num_reqs=2, has_preprocess=True, have_multimodal_outputs=True)
|
|
state.model = model
|
|
for idx in range(2):
|
|
state.intermediate_buffer.buffers[idx] = {
|
|
"req_id": f"r{idx}",
|
|
"audio_state": {"is_stopping": False},
|
|
"hidden_states": {"last": torch.full((4,), float(idx + 1))},
|
|
}
|
|
batch = _DummyInputBatch([1, 0], num_computed_tokens_cpu=[1, 1])
|
|
inputs = {"input_ids": torch.tensor([2, 2]), "inputs_embeds": torch.zeros(2, 4)}
|
|
request_state = mocker.Mock(spec=RequestState, prompt_len=np.array([1, 1]), num_computed_tokens=None)
|
|
state.run_preprocess(batch, inputs, request_state)
|
|
|
|
frame.assert_called_once()
|
|
torch.testing.assert_close(frame.call_args.args[0], torch.tensor([[2.0] * 4, [1.0] * 4]))
|
|
assert frame.call_args.kwargs["temperature"] == 1.7
|
|
assert frame.call_args.kwargs["top_k"] == 25
|
|
assert frame.call_args.kwargs["top_p"] == 0.8
|
|
torch.testing.assert_close(inputs["inputs_embeds"], torch.tensor([[1.0] * 4, [0.0] * 4]))
|
|
_, payload = state.postprocess_model_output(torch.zeros(2, 4), batch, request_state)
|
|
assert [codes.tolist() for codes in payload["codes"]["audio"]] == [[[1, 2]], [[8, 8]]]
|
|
assert model.compute_logits(torch.zeros(2, 4)).argmax(-1).tolist() == [2, 3]
|
|
# An explicit local seed must reach the request-owned MRV2 generator.
|
|
params = SamplingParams(extra_args={"tts_local_seed": 17}, seed=99)
|
|
assert state._get_mtp_generator("seeded", params, torch.device("cpu")).initial_seed() == 17
|
|
|
|
|
|
@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
|
|
state._mtp_generators.clear() # The following rows are new requests.
|
|
# 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)
|
|
|
|
|
|
_EAGER_DIM = 3
|
|
_CODEBOOK = 2048
|
|
_EOS = 2150
|
|
|
|
|
|
class _EagerBatch(_DummyInputBatch):
|
|
def __init__(self, spans, indices=None):
|
|
indices = list(range(len(spans))) if indices is None else indices
|
|
super().__init__(indices)
|
|
starts = [0]
|
|
for n_tok in spans:
|
|
starts.append(starts[-1] + n_tok)
|
|
self.num_scheduled_tokens = list(spans)
|
|
self.query_start_loc_np = starts[:-1]
|
|
self.query_start_loc = torch.tensor(starts, dtype=torch.int32)
|
|
|
|
|
|
def _make_eager_state(max_num_reqs=4):
|
|
state = _make_state(max_num_reqs=max_num_reqs, has_preprocess=True)
|
|
_init_static(state, max_num_reqs, dim=_EAGER_DIM)
|
|
from vllm_omni.worker_v2.model_states.eager_mtp import EagerMTPState
|
|
|
|
state._eager_state = EagerMTPState(state)
|
|
state._first_audio_requests = set()
|
|
state._eager_mtp = True
|
|
state._eager_fastpath = False
|
|
state._eager_rows = None
|
|
state._eager_ready = {}
|
|
state._eager_settled = {}
|
|
state._eager_embeds = torch.zeros((max_num_reqs, _EAGER_DIM))
|
|
state._decode_preprocess = None
|
|
state._mtp_sample_uniforms = None
|
|
state.vllm_config.cache_config = SimpleNamespace(enable_prefix_caching=False)
|
|
model = state.model
|
|
model.first_frame_decoder = None
|
|
model.mtp_frame_valid = lambda layer0: (layer0 >= 0) & (layer0 < _CODEBOOK)
|
|
model.embed_input_ids = lambda ids: ids.float().reshape(-1, 1, 1).expand(-1, 1, _EAGER_DIM)
|
|
model.mtp_calls = []
|
|
|
|
def mtp(input_ids, input_embeds, last_hidden, text_step, **kwargs):
|
|
model.mtp_calls.append((input_ids.clone(), last_hidden.clone(), text_step.clone()))
|
|
codes = torch.stack([input_ids, input_ids + 1, input_ids + 2], dim=1)
|
|
# Frame embedding sum plus the text step, like the Qwen3 talkers.
|
|
return input_embeds + 100 + text_step, codes
|
|
|
|
model.mtp = mtp
|
|
return state
|
|
|
|
|
|
def _eager_outputs(num_tokens):
|
|
return {
|
|
"codes": {"audio": torch.zeros((num_tokens, 3), dtype=torch.long)},
|
|
"meta": {"codec_frame_valid": torch.zeros((num_tokens,), dtype=torch.int8)},
|
|
}
|
|
|
|
|
|
def test_eager_mtp_publishes_frames_in_the_step_that_sampled_cb0():
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "prefill", "decode", "ended")
|
|
# prefill: final 3-token chunk; decode: input CB0 42; ended: input CB0 was
|
|
# EOS (async scheduling ran the finished request once more).
|
|
batch = _EagerBatch([3, 1, 1])
|
|
input_ids = torch.tensor([0, 0, 0, 42, _EOS])
|
|
state._eager_rows = (batch, [(0, 0, "prefill", True), (1, 1, "decode", False), (2, 2, "ended", False)], input_ids)
|
|
text_hidden = torch.arange(5 * _EAGER_DIM, dtype=torch.float32).reshape(5, _EAGER_DIM)
|
|
sampled = torch.tensor([[7], [_EOS], [9]])
|
|
outputs = _eager_outputs(5)
|
|
|
|
state.run_eager_mtp(batch, text_hidden, sampled, outputs)
|
|
|
|
ids, hidden, text_step = state.model.mtp_calls[0]
|
|
assert ids.tolist() == [7, _EOS, 9]
|
|
# The frame's hidden is this step's last-token hidden of each span.
|
|
assert torch.equal(hidden, text_hidden[[2, 3, 4]])
|
|
# The next step adds its own text step.
|
|
assert not text_step.any()
|
|
codes = outputs["codes"]["audio"]
|
|
assert codes[2].tolist() == [7, 8, 9]
|
|
assert codes[3].tolist() == [_EOS, _EOS + 1, _EOS + 2]
|
|
assert codes[4].tolist() == [9, 10, 11]
|
|
assert not codes[:2].any()
|
|
# Valid only when both the sampled and (for decode) the input CB0 are codec ids.
|
|
assert outputs["meta"]["codec_frame_valid"].tolist() == [0, 0, 1, 0, 0]
|
|
assert state._eager_ready == {0: "prefill", 1: "decode", 2: "ended"}
|
|
assert torch.equal(state._eager_embeds[0], torch.full((_EAGER_DIM,), 107.0))
|
|
assert state._eager_rows is None
|
|
|
|
|
|
def test_eager_mtp_ignores_rows_recorded_for_another_batch():
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "r0")
|
|
state._eager_rows = (_EagerBatch([1]), [(0, 0, "r0", False)], torch.tensor([5]))
|
|
outputs = _eager_outputs(1)
|
|
state.run_eager_mtp(_EagerBatch([1]), torch.zeros(1, _EAGER_DIM), torch.tensor([[3]]), outputs)
|
|
assert state.model.mtp_calls == []
|
|
assert not outputs["codes"]["audio"].any()
|
|
|
|
|
|
def test_eager_mtp_requires_retained_frame_outputs():
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "r0")
|
|
batch = _EagerBatch([1])
|
|
state._eager_rows = (batch, [(0, 0, "r0", False)], torch.tensor([5]))
|
|
with pytest.raises(RuntimeError, match="requires retained codes.audio"):
|
|
state.run_eager_mtp(batch, torch.zeros(1, _EAGER_DIM), torch.tensor([[3]]), {"codes": {"audio": None}})
|
|
|
|
|
|
def test_eager_decode_input_is_previous_frame_plus_text_step():
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "r0", "r1")
|
|
state._eager_ready = {0: "r0", 1: "r1"}
|
|
state._eager_embeds[0] = 5.0
|
|
state._eager_embeds[1] = 7.0
|
|
embeds = torch.zeros((2, _EAGER_DIM))
|
|
batches = [
|
|
(0, 0, (torch.zeros(_EAGER_DIM), torch.full((_EAGER_DIM,), 2.0))),
|
|
(1, 1, (None, torch.ones(_EAGER_DIM))),
|
|
]
|
|
state._eager_state._apply_eager_frames(batches, embeds, _EagerBatch([1, 1]), None)
|
|
assert embeds.tolist() == [[7.0] * _EAGER_DIM, [8.0] * _EAGER_DIM]
|
|
# The deferred MTP must not run for an eager decode row.
|
|
assert state.model.mtp_calls == []
|
|
|
|
|
|
def test_eager_decode_without_a_frame_fails_loudly():
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "r0")
|
|
batches = [(0, 0, (torch.zeros(_EAGER_DIM), torch.zeros(_EAGER_DIM)))]
|
|
with pytest.raises(RuntimeError, match="frame missing"):
|
|
state._eager_state._apply_eager_frames(batches, torch.zeros((1, _EAGER_DIM)), _EagerBatch([1]), None)
|
|
|
|
|
|
def test_run_preprocess_records_rows_that_keep_a_sample():
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "chunk", "final", "decode")
|
|
state._eager_ready = {2: "decode"}
|
|
state._eager_embeds[2] = 4.0
|
|
text_step = torch.ones(_EAGER_DIM)
|
|
|
|
def preprocess(input_ids, input_embeds, **info):
|
|
updates = {"mtp_inputs": (torch.zeros(_EAGER_DIM), text_step)} if input_ids.shape[0] == 1 else {}
|
|
return input_ids, input_embeds, updates
|
|
|
|
state.model.preprocess = preprocess
|
|
batch = _EagerBatch([3, 3, 1])
|
|
model_inputs = {"input_ids": torch.zeros(7, dtype=torch.long), "inputs_embeds": torch.zeros((7, _EAGER_DIM))}
|
|
req_states = SimpleNamespace(
|
|
prompt_len=np.array([10, 10, 10], dtype=np.int32),
|
|
num_computed_tokens=np.array([0, 7, 10], dtype=np.int32),
|
|
)
|
|
state.run_preprocess(batch, model_inputs, req_states)
|
|
|
|
recorded_batch, entries, _ids = state._eager_rows
|
|
assert recorded_batch is batch
|
|
# A non-final prefill chunk samples nothing that is kept.
|
|
assert entries == [(1, 1, "final", True), (2, 2, "decode", False)]
|
|
assert model_inputs["inputs_embeds"][6].tolist() == [5.0] * _EAGER_DIM
|
|
assert state.model.mtp_calls == []
|
|
|
|
|
|
@pytest.mark.parametrize("first_was_valid", [False, True])
|
|
def test_first_audio_requirement_survives_eos_only_if_audio_was_queued(first_was_valid):
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "r0")
|
|
state._first_audio_requests.add("r0")
|
|
state._eager_state._first_audio_valid = torch.tensor([first_was_valid, False, False, False])
|
|
batch = _EagerBatch([1])
|
|
state._eager_rows = (batch, [(0, 0, "r0", False)], torch.tensor([5]))
|
|
outputs = _eager_outputs(1)
|
|
outputs["meta"]["first_audio"] = torch.zeros(1, dtype=torch.bool)
|
|
state.run_eager_mtp(batch, torch.zeros(1, _EAGER_DIM), torch.tensor([[_EOS]]), outputs)
|
|
assert outputs["meta"]["first_audio"].tolist() == [first_was_valid]
|
|
|
|
|
|
@pytest.mark.parametrize("accepted", [[], ["r1"], ["r0", "r1"]])
|
|
def test_first_audio_marker_requires_accepted_delivery(monkeypatch, accepted):
|
|
from contextlib import nullcontext
|
|
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "r0", "r1")
|
|
state.model.first_frame_decoder = SimpleNamespace(
|
|
sample_rate=24000, decode=lambda codes: torch.ones(codes.shape[0], 2)
|
|
)
|
|
state._first_audio_stream = SimpleNamespace(wait_stream=lambda stream: None)
|
|
state._first_audio_sender = SimpleNamespace(submit=lambda ids, pcm, sr, valid: accepted)
|
|
monkeypatch.setattr(torch.cuda, "current_stream", lambda device: None)
|
|
monkeypatch.setattr(torch.cuda, "stream", lambda stream: nullcontext())
|
|
monkeypatch.setattr(torch.Tensor, "record_stream", lambda tensor, stream: None)
|
|
batch = _EagerBatch([1, 1])
|
|
state._eager_rows = (batch, [(0, 0, "r0", True), (1, 1, "r1", True)], torch.zeros(2, dtype=torch.long))
|
|
outputs = _eager_outputs(2)
|
|
outputs["meta"]["first_audio"] = torch.zeros(2, dtype=torch.bool)
|
|
|
|
state.run_eager_mtp(batch, torch.zeros(2, _EAGER_DIM), torch.tensor([[7], [8]]), outputs)
|
|
|
|
assert state._first_audio_requests == set(accepted)
|
|
# A missing route must leave the normal codec path responsible for frame 0;
|
|
# otherwise it skips the frame and the orchestrator waits forever for it.
|
|
assert outputs["meta"]["first_audio"].tolist() == ["r0" in accepted, "r1" in accepted]
|
|
|
|
|
|
class _FakeEvent:
|
|
def record(self, stream=None) -> None:
|
|
self.stream = stream
|
|
|
|
|
|
def test_publish_sampled_embeddings_for_rows_whose_sample_is_kept(monkeypatch) -> None:
|
|
monkeypatch.setattr(torch.cuda, "Event", _FakeEvent)
|
|
state = _make_state()
|
|
state.model.publishes_sampled_embeddings = True
|
|
state.model.embed_input_ids = lambda ids: ids.to(torch.float32).reshape(-1, 1).repeat(1, 3)
|
|
# Rows: final prefill chunk, decode, non-final prefill chunk, one-chunk prefill.
|
|
batch = SimpleNamespace(
|
|
num_reqs=4,
|
|
num_scheduled_tokens=[4, 1, 2, 6],
|
|
num_computed_prefill_tokens_np=np.array([2, 9, 0, 0], dtype=np.int32),
|
|
prefill_len_np=np.array([6, 9, 5, 6], dtype=np.int32),
|
|
is_prefilling_np=np.array([True, False, True, True]),
|
|
)
|
|
extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]]))
|
|
sampled = extra["embed"]["sampled"]
|
|
# A non-final prefill chunk's sample is discarded, so it publishes nothing.
|
|
assert [tuple(row.shape) for row in sampled] == [(1, 3), (1, 3), (0,), (1, 3)]
|
|
assert [row.tolist() for row in sampled if row.numel()] == [[[v] * 3] for v in (11.0, 12.0, 14.0)]
|
|
|
|
batch.is_prefilling_np[:] = True
|
|
batch.num_computed_prefill_tokens_np[:] = 0
|
|
batch.num_scheduled_tokens = [1] * 4
|
|
extra, _done = state.publish_sampled_embeddings(batch, torch.tensor([[11], [12], [13], [14]]))
|
|
assert all(row.numel() == 0 for row in extra["embed"]["sampled"])
|
|
|
|
# Speculative steps sample several tokens per row: not published.
|
|
assert state.publish_sampled_embeddings(batch, torch.tensor([[11, 1], [12, 1], [13, 1], [14, 1]])) is None
|
|
|
|
|
|
def test_publish_sampled_embeddings_is_opt_in() -> None:
|
|
state = _make_state()
|
|
state.model.publishes_sampled_embeddings = False
|
|
assert state.publish_sampled_embeddings(SimpleNamespace(num_reqs=1), torch.tensor([[1]])) is None
|
|
|
|
|
|
def test_split_settled_rows_matches_per_row_check():
|
|
batch = _DummyInputBatch(np.array([3, 0, 2, 1, 4]))
|
|
batch.req_ids = ["c", "a", "x", "b", "e"]
|
|
batch.num_scheduled_tokens = [1, 1, 1, 4, 1]
|
|
batch.query_start_loc_np = [0, 1, 2, 3, 7]
|
|
req_indices = batch.idx_mapping_np.tolist()
|
|
# Slot 2 still names a finished request, slot 1 schedules a prefill
|
|
# chunk, slot 4 is not settled and slot 7 is not scheduled.
|
|
settled = {3: "c", 0: "a", 2: "old", 1: "b", 7: "z"}
|
|
split = OmniModelState._split_settled_rows(batch, req_indices, settled)
|
|
assert split is not None
|
|
settled_rows, remaining = split
|
|
assert settled_rows == [(0, 3, 0, "c"), (1, 0, 1, "a")]
|
|
assert remaining == [(2, 2), (3, 1), (4, 4)]
|
|
assert OmniModelState._split_settled_rows(batch, req_indices, {2: "old"}) is None
|
|
|
|
|
|
def test_settled_row_split_preserves_mixed_batch_and_pending_replay(monkeypatch):
|
|
"""Exercise run_preprocess, including the scalar fallback during replay."""
|
|
|
|
def prepare():
|
|
state = _make_eager_state()
|
|
_fill_buffers(state, "prefill", "decode")
|
|
state._eager_fastpath = True
|
|
state._eager_settled = {1: "decode"}
|
|
state._eager_ready = {1: "decode"}
|
|
state._eager_embeds[1].fill_(7)
|
|
state.model.eager_settled_text_step = lambda: torch.ones(_EAGER_DIM)
|
|
state.model.eager_decode_settled = lambda info: False
|
|
state.model.preprocess = lambda input_ids, input_embeds, **info: (input_ids, input_embeds, {})
|
|
batch = _EagerBatch([3, 1])
|
|
batch.req_ids = ["prefill", "decode"]
|
|
inputs = {"input_ids": torch.zeros(4, dtype=torch.long), "inputs_embeds": torch.zeros(4, _EAGER_DIM)}
|
|
return state, batch, inputs
|
|
|
|
vector, batch, inputs = prepare()
|
|
vector.run_preprocess(batch, inputs)
|
|
reference, ref_batch, ref_inputs = prepare()
|
|
monkeypatch.setattr(reference, "_split_settled_rows", lambda *args: None)
|
|
reference.run_preprocess(ref_batch, ref_inputs)
|
|
torch.testing.assert_close(inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
|
|
assert vector._eager_rows[1] == reference._eager_rows[1]
|
|
assert inputs["inputs_embeds"][3].tolist() == [8.0] * _EAGER_DIM
|
|
|
|
replay, replay_batch, replay_inputs = prepare()
|
|
replay._eager_state._restore_audio["prefill"] = [torch.ones(1)]
|
|
assert replay._eager_state.has_pending_replay()
|
|
monkeypatch.setattr(replay, "_split_settled_rows", MagicMock(side_effect=AssertionError("replay must skip split")))
|
|
replay.run_preprocess(replay_batch, replay_inputs)
|
|
torch.testing.assert_close(replay_inputs["inputs_embeds"], ref_inputs["inputs_embeds"], rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("prefilling", [False, True])
|
|
def test_identity_preprocess_skips_only_decode_rows(prefilling):
|
|
state = _make_state(has_preprocess=True)
|
|
state._decode_preprocess_is_identity = True
|
|
state._decode_preprocess = None
|
|
state.intermediate_buffer.buffers[0] = {"req_id": "r1"}
|
|
seen = []
|
|
|
|
def preprocess(input_ids, input_embeds, **info):
|
|
seen.append(info["req_id"])
|
|
return input_ids, input_embeds, {}
|
|
|
|
state.model.preprocess = preprocess
|
|
batch = _DummyInputBatch([0])
|
|
batch.is_prefilling_np = np.array([prefilling])
|
|
embeds = torch.ones(1, 4)
|
|
state.run_preprocess(batch, {"input_ids": torch.tensor([1]), "inputs_embeds": embeds})
|
|
assert seen == (["r1"] if prefilling else [])
|
|
assert torch.equal(embeds, torch.ones(1, 4))
|
|
|
|
|
|
def test_mm_embeddings_exclude_zero_length_graph_padding_rows():
|
|
state = object.__new__(OmniModelState)
|
|
state.supports_mm_inputs = True
|
|
state.mm_pruner = None
|
|
state.prompt_embeds_state = None
|
|
state.encoder_runner = MagicMock(spec=EncoderRunner, inputs_embeds=torch.zeros(8, 2))
|
|
state.execute_mm_encoder = lambda _: None
|
|
state.gather_mm_embeddings = lambda _: ([torch.ones(1, 2)], torch.ones(5, dtype=torch.bool))
|
|
|
|
def embed(input_ids, *, query_start_loc, multimodal_embeddings, is_multimodal):
|
|
assert query_start_loc == [0, 2, 5]
|
|
return input_ids[:, None].expand(-1, 2).float()
|
|
|
|
state.model = MagicMock(supports_embed_input_ids_query_start_loc=True, embed_input_ids=embed)
|
|
buffers = InputBuffers(4, 8, torch.device("cpu"))
|
|
batch = replace(
|
|
InputBatch.make_dummy(2, 5, buffers),
|
|
num_reqs=2,
|
|
num_reqs_after_padding=4,
|
|
num_tokens=5,
|
|
num_tokens_after_padding=8,
|
|
input_ids=torch.arange(8),
|
|
query_start_loc_np=np.array([0, 2, 5, 5, 5], dtype=np.int32),
|
|
)
|
|
result = state.prepare_inputs_embeds({}, batch, MagicMock(spec=RequestState))
|
|
torch.testing.assert_close(result[:5], torch.arange(5).float()[:, None].expand(-1, 2))
|
|
assert result.shape == (8, 2)
|