NATUREandamy-why-3459
4f7c08c0eb
Signed-off-by: natureofnature <wzliu@connect.hku.hk> Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
860 lines
32 KiB
Python
860 lines
32 KiB
Python
from contextlib import contextmanager
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
import torch
|
|
from vllm.v1.cudagraph_dispatcher import CUDAGraphMode
|
|
|
|
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
|
|
from vllm_omni.worker.gpu_model_runner import (
|
|
OmniGPUModelRunner,
|
|
_filter_mrope_kwargs_for_model,
|
|
)
|
|
from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
def _runner_for_talker_graph_init(
|
|
*,
|
|
talker_mtp_graph_safe: bool | None,
|
|
has_separate_talker: bool = True,
|
|
model_stage: str = "talker",
|
|
) -> OmniGPUModelRunner:
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
runner.model = SimpleNamespace(
|
|
talker=object() if has_separate_talker else None,
|
|
talker_mtp=object(),
|
|
model_stage=model_stage,
|
|
)
|
|
if talker_mtp_graph_safe is not None:
|
|
runner.model.talker_mtp_graph_safe = talker_mtp_graph_safe
|
|
runner.model_config = SimpleNamespace(hf_text_config=SimpleNamespace(hidden_size=4))
|
|
runner.compilation_config = SimpleNamespace(
|
|
cudagraph_mode=CUDAGraphMode.FULL,
|
|
max_cudagraph_capture_size=1,
|
|
)
|
|
runner.vllm_config = object()
|
|
runner.max_num_reqs = 1
|
|
runner.dtype = torch.float32
|
|
runner._make_buffer = lambda *args, **kwargs: SimpleNamespace(args=args, kwargs=kwargs)
|
|
return runner
|
|
|
|
|
|
def test_talker_mtp_skips_graph_when_model_declares_unsafe(monkeypatch):
|
|
runner = _runner_for_talker_graph_init(talker_mtp_graph_safe=False)
|
|
talker_mtp = runner.model.talker_mtp
|
|
monkeypatch.setattr(
|
|
"vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
|
|
lambda: pytest.fail("graph wrapper must not be selected"),
|
|
)
|
|
|
|
OmniGPUModelRunner._init_talker_mtp(runner)
|
|
|
|
assert runner.talker_mtp is talker_mtp
|
|
|
|
|
|
@pytest.mark.parametrize("model_stage", ["thinker", "code2wav"])
|
|
def test_non_talker_stage_does_not_use_talker_mtp_graph(monkeypatch, model_stage: str):
|
|
runner = _runner_for_talker_graph_init(
|
|
talker_mtp_graph_safe=None,
|
|
has_separate_talker=False,
|
|
model_stage=model_stage,
|
|
)
|
|
talker_mtp = runner.model.talker_mtp
|
|
monkeypatch.setattr(
|
|
"vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
|
|
lambda: pytest.fail("graph wrapper must not be selected"),
|
|
)
|
|
|
|
OmniGPUModelRunner._init_talker_mtp(runner)
|
|
|
|
assert runner.talker_mtp is talker_mtp
|
|
|
|
|
|
@pytest.mark.parametrize("talker_mtp_graph_safe", [None, True])
|
|
def test_talker_mtp_uses_graph_for_legacy_or_explicit_safe_model(monkeypatch, talker_mtp_graph_safe):
|
|
runner = _runner_for_talker_graph_init(talker_mtp_graph_safe=talker_mtp_graph_safe)
|
|
wrapped = object()
|
|
monkeypatch.setattr(
|
|
"vllm_omni.worker.gpu_model_runner.current_omni_platform.get_graph_wrapper_cls",
|
|
lambda: lambda *args, **kwargs: wrapped,
|
|
)
|
|
|
|
OmniGPUModelRunner._init_talker_mtp(runner)
|
|
|
|
assert runner.talker_mtp is wrapped
|
|
|
|
|
|
class DummyBuffer:
|
|
"""A minimal buffer wrapper that exposes the `.gpu` attribute."""
|
|
|
|
def __init__(self, t: torch.Tensor):
|
|
self.gpu = t
|
|
|
|
|
|
class DummyInputBatch:
|
|
"""A minimal input batch that only provides `req_ids`."""
|
|
|
|
def __init__(self, req_ids):
|
|
self.req_ids = req_ids
|
|
self.req_id_to_index = {r: i for i, r in enumerate(req_ids)}
|
|
|
|
|
|
class DummyReqState:
|
|
"""A minimal request state container."""
|
|
|
|
pass
|
|
|
|
|
|
class MiMoAudioForConditionalGeneration(torch.nn.Module):
|
|
"""Dummy model whose class name must exactly match the production check."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
# No real forward needed for these tests.
|
|
|
|
|
|
class DummyTalkerMTP(torch.nn.Module):
|
|
"""A fake talker_mtp module for deterministic CPU testing."""
|
|
|
|
def forward(
|
|
self,
|
|
req_input_ids,
|
|
req_embeds,
|
|
last_talker_hidden,
|
|
text_step,
|
|
do_sample=None,
|
|
temperature=None,
|
|
top_k=None,
|
|
top_p=None,
|
|
):
|
|
# Deterministic behavior:
|
|
# - output embeds = input embeds + 1
|
|
# - output codes = [[0], [1], ...]
|
|
bsz = req_embeds.shape[0]
|
|
new_embeds = req_embeds + 1.0
|
|
codes = torch.arange(bsz, dtype=torch.int64).view(bsz, 1)
|
|
return new_embeds, codes
|
|
|
|
|
|
class CaptureTalkerMTP(torch.nn.Module):
|
|
"""A fake talker_mtp module that records sampling kwargs."""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
self.calls = []
|
|
|
|
def forward(
|
|
self,
|
|
req_input_ids,
|
|
req_embeds,
|
|
last_talker_hidden,
|
|
text_step,
|
|
do_sample=None,
|
|
temperature=None,
|
|
top_k=None,
|
|
top_p=None,
|
|
generator=None,
|
|
generators=None,
|
|
):
|
|
self.calls.append(
|
|
{
|
|
"batch_size": int(req_embeds.shape[0]),
|
|
"do_sample": do_sample,
|
|
"temperature": temperature,
|
|
"top_k": top_k,
|
|
"top_p": top_p,
|
|
"generator": generator,
|
|
"generators": generators,
|
|
}
|
|
)
|
|
codes = torch.zeros((req_embeds.shape[0], 1), dtype=torch.int64)
|
|
return req_embeds, codes
|
|
|
|
|
|
class StrictMRoPEModel:
|
|
def get_mrope_input_positions(self, input_tokens, mm_features):
|
|
raise NotImplementedError
|
|
|
|
|
|
class FlexibleMRoPEModel:
|
|
def get_mrope_input_positions(self, input_tokens, mm_features=None, **kwargs):
|
|
raise NotImplementedError
|
|
|
|
|
|
@contextmanager
|
|
def _noop_forward_context(*args, **kwargs):
|
|
"""A no-op context manager to replace vLLM forward context in CPU tests."""
|
|
yield
|
|
|
|
|
|
def test_filter_mrope_kwargs_for_strict_model_signature():
|
|
kwargs = {
|
|
"mm_features": ["audio"],
|
|
"hf_config": object(),
|
|
"image_grid_thw": [],
|
|
}
|
|
|
|
assert _filter_mrope_kwargs_for_model(StrictMRoPEModel(), kwargs) == {
|
|
"mm_features": ["audio"],
|
|
}
|
|
|
|
|
|
def test_filter_mrope_kwargs_preserves_flexible_model_kwargs():
|
|
kwargs = {
|
|
"mm_features": ["video"],
|
|
"hf_config": object(),
|
|
"video_grid_thw": [[1, 2, 3]],
|
|
}
|
|
|
|
assert _filter_mrope_kwargs_for_model(FlexibleMRoPEModel(), kwargs) is kwargs
|
|
|
|
|
|
def _make_runner(req_ids=("r1", "r2"), hidden_size=4):
|
|
# Create an instance without calling OmniGPUModelRunner.__init__
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
|
|
# Minimal attributes used by OmniGPUModelRunner._talker_mtp_forward
|
|
runner.input_batch = DummyInputBatch(list(req_ids))
|
|
runner.requests = {rid: DummyReqState() for rid in req_ids}
|
|
runner.model_intermediate_buffer = {}
|
|
|
|
# query_start_loc.cpu[req_index] is used to locate the token position
|
|
# in the flattened `inputs_embeds`.
|
|
runner.query_start_loc = type("QSL", (), {})()
|
|
# Map: r1 -> offset 0, r2 -> offset 3
|
|
runner.query_start_loc.cpu = torch.tensor([0, 3], dtype=torch.int32)
|
|
|
|
bsz = len(req_ids)
|
|
runner.talker_mtp_input_ids = DummyBuffer(torch.zeros((bsz,), dtype=torch.int64))
|
|
runner.talker_mtp_inputs_embeds = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
|
|
runner.last_talker_hidden = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
|
|
runner.text_step = DummyBuffer(torch.zeros((bsz, hidden_size), dtype=torch.float32))
|
|
|
|
runner.talker_mtp = DummyTalkerMTP()
|
|
runner.model = SimpleNamespace(talker_mtp_output_key=("codes", "audio"))
|
|
runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace())
|
|
|
|
# Provide a minimal implementation that returns the expected 4-tuple.
|
|
def _determine_batch_execution_and_padding(**kwargs):
|
|
return None, object(), None, None, None
|
|
|
|
runner._determine_batch_execution_and_padding = _determine_batch_execution_and_padding
|
|
|
|
# Use the real merge method from OmniGPUModelRunner.
|
|
return runner
|
|
|
|
|
|
def _make_runner_for_mimo(req_id="r_mimo"):
|
|
"""Create a minimal runner with MiMoAudio-like model and request state."""
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
runner.model = MiMoAudioForConditionalGeneration()
|
|
|
|
# Minimal vllm_config / model_config used by helper.
|
|
class _DummyModelConfig:
|
|
async_chunk = False
|
|
|
|
class _DummyVllmConfig:
|
|
model_config = _DummyModelConfig()
|
|
|
|
runner.vllm_config = _DummyVllmConfig()
|
|
|
|
# Attach a single request state with mm_features and additional_information_cpu.
|
|
req_state = DummyReqState()
|
|
req_state.mm_features = ["mm_feature_obj"]
|
|
req_state.additional_information_cpu = {"some_key": "some_value"}
|
|
|
|
runner.requests = {req_id: req_state}
|
|
|
|
return runner
|
|
|
|
|
|
def test_talker_mtp_forward_cpu_updates_inputs_and_info(monkeypatch):
|
|
# `_talker_mtp_forward` dispatches through the active platform.
|
|
import vllm_omni.worker.gpu_model_runner as mod # Must be the same module that defines OmniGPUModelRunner
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
# Initialize per-request embeds (batch-major inside talker_mtp_inputs_embeds)
|
|
runner.talker_mtp_inputs_embeds.gpu[0] = torch.tensor([1.0, 2.0, 3.0, 4.0])
|
|
runner.talker_mtp_inputs_embeds.gpu[1] = torch.tensor([10.0, 20.0, 30.0, 40.0])
|
|
|
|
# Flattened `inputs_embeds`: offsets 0 and 3 will be overwritten
|
|
inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
|
|
|
|
# Call the original implementation from OmniGPUModelRunner (no re-implementation)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
|
|
|
|
# Validate embeds were written back (+1)
|
|
assert torch.allclose(inputs_embeds[0], torch.tensor([2.0, 3.0, 4.0, 5.0]))
|
|
assert torch.allclose(inputs_embeds[3], torch.tensor([11.0, 21.0, 31.0, 41.0]))
|
|
|
|
# Validate per-request additional_information_cpu was updated
|
|
info_r1 = runner.requests["r1"].additional_information_cpu
|
|
info_r2 = runner.requests["r2"].additional_information_cpu
|
|
assert int(info_r1["codes"]["audio"][0, 0]) == 0
|
|
assert int(info_r2["codes"]["audio"][0, 0]) == 1
|
|
|
|
|
|
def test_talker_mtp_forward_cpu_empty_batch_noop(monkeypatch):
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
inputs_embeds = torch.randn((2, 4))
|
|
before = inputs_embeds.clone()
|
|
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, [], inputs_embeds)
|
|
|
|
# Ensure no changes were made
|
|
assert torch.allclose(inputs_embeds, before)
|
|
|
|
|
|
def test_talker_mtp_forward_ignores_default_sampling_seed_without_request_marker(monkeypatch):
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
runner.requests["r1"].sampling_params = SimpleNamespace(seed=42)
|
|
runner.talker_mtp = CaptureTalkerMTP()
|
|
runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
inputs_embeds = torch.zeros((2, 4), dtype=torch.float32)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
|
|
|
|
assert runner.talker_mtp.calls[0]["generator"] is None
|
|
|
|
|
|
def test_talker_mtp_forward_passes_qwen3_tts_subtalker_sampling_params_to_talker(monkeypatch):
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
runner.requests["r1"].sampling_params = SimpleNamespace(
|
|
seed=42,
|
|
extra_args={"tts_local_seed": 42},
|
|
)
|
|
runner.talker_mtp = CaptureTalkerMTP()
|
|
runner.vllm_config = SimpleNamespace(
|
|
model_config=SimpleNamespace(
|
|
subtalker_sampling_params={
|
|
"do_sample": False,
|
|
"temperature": 0.2,
|
|
"top_k": 9,
|
|
"top_p": 0.55,
|
|
}
|
|
)
|
|
)
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
inputs_embeds = torch.zeros((2, 4), dtype=torch.float32)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
|
|
|
|
assert runner.talker_mtp.calls == [
|
|
{
|
|
"batch_size": 1,
|
|
"do_sample": False,
|
|
"temperature": 0.2,
|
|
"top_k": 9,
|
|
"top_p": 0.55,
|
|
"generator": runner.talker_mtp.calls[0]["generator"],
|
|
"generators": None,
|
|
}
|
|
]
|
|
assert runner.talker_mtp.calls[0]["generator"] is not None
|
|
|
|
|
|
def test_talker_mtp_forward_keeps_explicit_seeded_requests_scalar(monkeypatch):
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
|
|
runner.requests["r1"].sampling_params = SimpleNamespace(
|
|
seed=11,
|
|
extra_args={"tts_local_seed": 11},
|
|
)
|
|
runner.requests["r2"].sampling_params = SimpleNamespace(
|
|
seed=22,
|
|
extra_args={"tts_local_seed": 22},
|
|
)
|
|
runner.talker_mtp = CaptureTalkerMTP()
|
|
runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
runner.talker_mtp_input_ids.gpu[:] = torch.tensor([101, 202], dtype=torch.int64)
|
|
runner.talker_mtp_inputs_embeds.gpu[0] = torch.tensor([1.0, 2.0, 3.0, 4.0])
|
|
runner.talker_mtp_inputs_embeds.gpu[1] = torch.tensor([10.0, 20.0, 30.0, 40.0])
|
|
saved_input_ids = runner.talker_mtp_input_ids.gpu.clone()
|
|
saved_embeds = runner.talker_mtp_inputs_embeds.gpu.clone()
|
|
|
|
inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
|
|
|
|
assert [call["batch_size"] for call in runner.talker_mtp.calls] == [1, 1]
|
|
assert all(call["generator"] is not None for call in runner.talker_mtp.calls)
|
|
assert runner.talker_mtp.calls[0]["generator"] is not runner.talker_mtp.calls[1]["generator"]
|
|
assert torch.equal(runner.talker_mtp_input_ids.gpu, saved_input_ids)
|
|
assert torch.equal(runner.talker_mtp_inputs_embeds.gpu, saved_embeds)
|
|
|
|
|
|
def test_talker_mtp_forward_batches_seeded_requests_for_opted_in_models(monkeypatch):
|
|
"""Models with talker_mtp_accepts_per_row_generators get one batched call (#4883)."""
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
|
|
runner.requests["r1"].sampling_params = SimpleNamespace(
|
|
seed=11,
|
|
extra_args={"tts_local_seed": 11},
|
|
)
|
|
runner.requests["r2"].sampling_params = SimpleNamespace(
|
|
seed=22,
|
|
extra_args={"tts_local_seed": 22},
|
|
)
|
|
runner.talker_mtp = CaptureTalkerMTP()
|
|
runner.model = SimpleNamespace(
|
|
talker_mtp_output_key=("codes", "audio"),
|
|
talker_mtp_accepts_per_row_generators=True,
|
|
)
|
|
runner.vllm_config = SimpleNamespace(model_config=SimpleNamespace(subtalker_sampling_params={}))
|
|
|
|
def fake_determine(self, num_tokens, num_reqs, num_scheduled_tokens_np, max_num_scheduled_tokens, use_cascade_attn):
|
|
batch_desc = SimpleNamespace(num_tokens=int(num_tokens))
|
|
return (False, batch_desc, None, None, None)
|
|
|
|
monkeypatch.setattr(runner, "_determine_batch_execution_and_padding", fake_determine.__get__(runner, type(runner)))
|
|
|
|
inputs_embeds = torch.zeros((6, 4), dtype=torch.float32)
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
|
|
|
|
# One batched call with distinct per-row generators, not two scalar calls.
|
|
assert [call["batch_size"] for call in runner.talker_mtp.calls] == [2]
|
|
row_generators = runner.talker_mtp.calls[0]["generators"]
|
|
assert runner.talker_mtp.calls[0]["generator"] is None
|
|
assert len(row_generators) == 2
|
|
assert all(generator is not None for generator in row_generators)
|
|
assert row_generators[0] is not row_generators[1]
|
|
|
|
# The per-request generator stream persists across steps...
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1", "r2"], inputs_embeds)
|
|
assert runner.talker_mtp.calls[1]["generators"][0] is row_generators[0]
|
|
assert runner.talker_mtp.calls[1]["generators"][1] is row_generators[1]
|
|
|
|
# ...and is evicted once its request finishes.
|
|
del runner.requests["r2"]
|
|
OmniGPUModelRunner._talker_mtp_forward(runner, ["r1"], inputs_embeds)
|
|
assert set(runner._talker_mtp_generators) == {"r1"}
|
|
assert runner.talker_mtp.calls[2]["generator"] is row_generators[0]
|
|
|
|
|
|
def test_update_intermediate_buffer_writes_to_buffer_and_setattr(monkeypatch):
|
|
"""Validate that _update_intermediate_buffer writes to model_intermediate_buffer
|
|
(forward path) and mirrors to additional_information_cpu setattr (backward compat)."""
|
|
import vllm_omni.worker.gpu_model_runner as mod
|
|
|
|
monkeypatch.setattr(mod.current_omni_platform, "set_forward_context", _noop_forward_context)
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
update = {"my_tensor": torch.tensor([1.0, 2.0]), "my_list": [3, 4]}
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", update)
|
|
|
|
# Forward: buffer is populated
|
|
assert "r1" in runner.model_intermediate_buffer
|
|
buf = runner.model_intermediate_buffer["r1"]
|
|
assert torch.allclose(buf["my_tensor"], torch.tensor([1.0, 2.0]))
|
|
assert buf["my_list"] == [3, 4]
|
|
|
|
# Backward compat: setattr is also populated
|
|
info_cpu = runner.requests["r1"].additional_information_cpu
|
|
assert torch.allclose(info_cpu["my_tensor"], torch.tensor([1.0, 2.0]))
|
|
assert info_cpu["my_list"] == [3, 4]
|
|
|
|
|
|
def test_update_intermediate_buffer_accumulates():
|
|
"""Validate that successive merges accumulate keys in the buffer."""
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {"a": torch.tensor([1.0])})
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {"b": torch.tensor([2.0])})
|
|
|
|
buf = runner.model_intermediate_buffer["r1"]
|
|
assert "a" in buf and "b" in buf
|
|
assert torch.allclose(buf["a"], torch.tensor([1.0]))
|
|
assert torch.allclose(buf["b"], torch.tensor([2.0]))
|
|
|
|
|
|
def test_update_additional_information_deserializes_new_request_payload():
|
|
from vllm_omni.engine.serialization import serialize_additional_information
|
|
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
conditioning = {
|
|
"tts_token_ids": torch.tensor([1, 2]),
|
|
"tts_hidden_states": torch.ones(2, 4),
|
|
}
|
|
scheduler_output = SimpleNamespace(
|
|
scheduled_new_reqs=[
|
|
SimpleNamespace(
|
|
req_id="r1",
|
|
additional_information=serialize_additional_information(conditioning),
|
|
)
|
|
],
|
|
scheduled_cached_reqs=SimpleNamespace(),
|
|
)
|
|
|
|
OmniGPUModelRunner._update_additional_information(runner, scheduler_output)
|
|
|
|
assert torch.equal(runner.model_intermediate_buffer["r1"]["tts_token_ids"], conditioning["tts_token_ids"])
|
|
assert torch.equal(
|
|
runner.model_intermediate_buffer["r1"]["tts_hidden_states"],
|
|
conditioning["tts_hidden_states"],
|
|
)
|
|
|
|
|
|
def test_streaming_new_request_marker_replaces_terminal_chunk_snapshot():
|
|
from vllm_omni.engine.serialization import serialize_additional_information
|
|
|
|
runner = _make_runner(req_ids=("r1", "r2"), hidden_size=4)
|
|
runner.model.replace_runtime_additional_information = True
|
|
terminal = {
|
|
"codes": {"audio": torch.tensor([1, 2])},
|
|
"meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
|
|
}
|
|
peer = {
|
|
"codes": {"audio": torch.tensor([9])},
|
|
"meta": {"cache_epoch": 3, "chunk_seq": 1, "last_chunk": False},
|
|
}
|
|
runner.model_intermediate_buffer.update(r1=terminal, r2=peer)
|
|
marker = {
|
|
"meta": {
|
|
"finished": False,
|
|
"is_segment_finished": True,
|
|
"request_finished": False,
|
|
"replace_runtime_additional_information": True,
|
|
}
|
|
}
|
|
new_req = SimpleNamespace(
|
|
req_id="r1",
|
|
model_intermediate_buffer=marker,
|
|
additional_information=serialize_additional_information(terminal),
|
|
)
|
|
|
|
OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req, "r1")
|
|
OmniGPUModelRunner._update_additional_information(
|
|
runner,
|
|
SimpleNamespace(
|
|
scheduled_new_reqs=[new_req],
|
|
scheduled_cached_reqs=SimpleNamespace(),
|
|
),
|
|
)
|
|
|
|
info = runner.model_intermediate_buffer["r1"]
|
|
assert "codes" not in info
|
|
assert info["meta"] == {
|
|
**marker["meta"],
|
|
"num_processed_tokens": 0,
|
|
"resumable": True,
|
|
}
|
|
assert runner.requests["r1"].additional_information_cpu == info
|
|
assert runner.model_intermediate_buffer["r2"] == peer
|
|
|
|
|
|
def test_cached_empty_marker_replaces_terminal_chunk_snapshot():
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
runner.model.replace_runtime_additional_information = True
|
|
runner.model_intermediate_buffer["r1"] = {
|
|
"codes": {"audio": torch.tensor([1, 2])},
|
|
"meta": {"cache_epoch": 0, "chunk_seq": 2, "last_chunk": True},
|
|
}
|
|
marker = {
|
|
"meta": {
|
|
"is_segment_finished": torch.tensor(True, dtype=torch.bool),
|
|
"replace_runtime_additional_information": True,
|
|
}
|
|
}
|
|
|
|
OmniGPUModelRunner._update_additional_information(
|
|
runner,
|
|
SimpleNamespace(
|
|
scheduled_new_reqs=[],
|
|
scheduled_cached_reqs=SimpleNamespace(additional_information={"r1": marker}),
|
|
),
|
|
)
|
|
|
|
assert runner.model_intermediate_buffer["r1"] == marker
|
|
assert runner.requests["r1"].additional_information_cpu == marker
|
|
|
|
|
|
def test_update_intermediate_buffer_skips_empty_update():
|
|
"""Validate that an empty update dict is a no-op."""
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "r1", {})
|
|
|
|
assert "r1" not in runner.model_intermediate_buffer
|
|
|
|
|
|
def test_update_intermediate_buffer_skips_unknown_req_id():
|
|
"""Validate that merge is a no-op when req_id is not in self.requests."""
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
|
|
OmniGPUModelRunner._update_intermediate_buffer(runner, "unknown_req", {"key": torch.tensor([1.0])})
|
|
|
|
assert "unknown_req" not in runner.model_intermediate_buffer
|
|
|
|
|
|
def test_streaming_input_update_merges_model_intermediate_buffer():
|
|
runner = _make_runner(req_ids=("r1",), hidden_size=4)
|
|
runner.model_intermediate_buffer["r1"] = {
|
|
"duplex": {
|
|
"session_id": "sid",
|
|
"seq": 1,
|
|
}
|
|
}
|
|
runner.requests["r1"].additional_information_cpu = runner.model_intermediate_buffer["r1"]
|
|
new_req_data = SimpleNamespace(
|
|
model_intermediate_buffer={
|
|
"duplex": {
|
|
"session_id": "sid",
|
|
"seq": 2,
|
|
"payload": {"type": "audio"},
|
|
}
|
|
},
|
|
additional_information=None,
|
|
)
|
|
|
|
OmniGPUModelRunner._update_streaming_input_additional_info(runner, new_req_data, "r1")
|
|
|
|
info = runner.model_intermediate_buffer["r1"]
|
|
assert info["duplex"]["session_id"] == "sid"
|
|
assert info["duplex"]["seq"] == 2
|
|
assert info["duplex"]["payload"] == {"type": "audio"}
|
|
assert runner.requests["r1"].additional_information_cpu is info
|
|
|
|
|
|
def _make_full_payload_accumulation_runner(
|
|
model_arch="Qwen3OmniMoeForConditionalGeneration",
|
|
model_stage="talker",
|
|
async_chunk=False,
|
|
final_output=False,
|
|
custom_process_next_stage_input_func="module.full_payload",
|
|
):
|
|
runner = object.__new__(OmniConnectorModelRunnerMixin)
|
|
runner.model_config = SimpleNamespace(
|
|
model_arch=model_arch,
|
|
model_stage=model_stage,
|
|
async_chunk=async_chunk,
|
|
final_output=final_output,
|
|
custom_process_next_stage_input_func=custom_process_next_stage_input_func,
|
|
)
|
|
runner._custom_process_func = object()
|
|
runner._pending_full_payload_send = {}
|
|
runner._stage_id = 1
|
|
runner._omni_connector = object()
|
|
return runner
|
|
|
|
|
|
def test_accumulate_full_payload_output_preserves_aligned_all_zero_qwen3_omni_codec_rows():
|
|
runner = _make_full_payload_accumulation_runner()
|
|
request = SimpleNamespace(output_token_ids=[0, 1])
|
|
codes = torch.zeros((2, 3), dtype=torch.long)
|
|
|
|
OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
|
|
|
|
stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
|
|
assert torch.equal(stored["codes.audio"], codes)
|
|
|
|
|
|
def test_accumulate_full_payload_output_keeps_misaligned_all_zero_qwen3_omni_codec_rows():
|
|
# After removing the sender-side zero filter, the accumulator keeps every
|
|
# codec row including misaligned all-zero rows. The downstream consumer
|
|
# (_extract_qwen3_full_payload_codec_rows) is the authoritative crop and
|
|
# filters by output_token_ids.
|
|
runner = _make_full_payload_accumulation_runner()
|
|
request = SimpleNamespace(output_token_ids=[0, 1])
|
|
codes = torch.zeros((1, 3), dtype=torch.long)
|
|
|
|
OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
|
|
|
|
stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
|
|
assert "codes.audio" in stored
|
|
assert torch.equal(stored["codes.audio"], codes)
|
|
|
|
|
|
def test_accumulate_full_payload_output_preserves_incremental_aligned_all_zero_qwen3_omni_codec_rows():
|
|
runner = _make_full_payload_accumulation_runner()
|
|
request = SimpleNamespace(output_token_ids=[0, 1])
|
|
runner._pending_full_payload_send["r1"] = (
|
|
{"codes.audio": torch.ones((1, 3), dtype=torch.long)},
|
|
request,
|
|
)
|
|
codes = torch.zeros((1, 3), dtype=torch.long)
|
|
|
|
OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
|
|
|
|
stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
|
|
assert stored["codes.audio"].shape == (2, 3)
|
|
assert torch.equal(stored["codes.audio"][1], torch.zeros(3, dtype=torch.long))
|
|
|
|
|
|
def test_accumulate_full_payload_output_keeps_all_zero_qwen3_omni_prefill_placeholder():
|
|
# Prefill placeholder rows (output_token_ids empty) are no longer dropped
|
|
# at the sender. The consumer-side crop trims them off using
|
|
# output_token_ids, so the end-to-end semantics are unchanged.
|
|
runner = _make_full_payload_accumulation_runner()
|
|
request = SimpleNamespace(output_token_ids=[])
|
|
codes = torch.zeros((2, 3), dtype=torch.long)
|
|
|
|
OmniConnectorModelRunnerMixin.accumulate_full_payload_output(runner, "r1", {"codes.audio": codes}, request)
|
|
|
|
stored, _ = OmniConnectorModelRunnerMixin._materialize_full_payload_entry(runner._pending_full_payload_send["r1"])
|
|
assert "codes.audio" in stored
|
|
assert torch.equal(stored["codes.audio"], codes)
|
|
|
|
|
|
def test_full_payload_output_accumulation_hook_matrix():
|
|
assert _make_full_payload_accumulation_runner(model_stage="thinker")._should_accumulate_full_payload_output()
|
|
assert _make_full_payload_accumulation_runner(model_stage="talker")._should_accumulate_full_payload_output()
|
|
assert not _make_full_payload_accumulation_runner(
|
|
model_stage="code2wav", final_output=True
|
|
)._should_accumulate_full_payload_output()
|
|
assert not _make_full_payload_accumulation_runner(
|
|
model_stage="token2audio",
|
|
custom_process_next_stage_input_func=None,
|
|
)._should_accumulate_full_payload_output()
|
|
assert not _make_full_payload_accumulation_runner(
|
|
model_stage="talker", async_chunk=True
|
|
)._should_accumulate_full_payload_output()
|
|
for model_arch in (
|
|
"Qwen3TTSForConditionalGeneration",
|
|
"Qwen2_5OmniForConditionalGeneration",
|
|
):
|
|
runner = _make_full_payload_accumulation_runner(model_arch=model_arch)
|
|
runner._custom_process_func = None
|
|
assert not runner._should_accumulate_full_payload_output()
|
|
|
|
|
|
def _make_request_end_payload_runner(*, enabled=True, prefix_cache=None):
|
|
runner = object.__new__(GPUARModelRunner)
|
|
runner.model = SimpleNamespace(omni_payload_at_request_end=enabled)
|
|
runner.omni_prefix_cache = prefix_cache
|
|
runner.model_config = SimpleNamespace(
|
|
model_arch="IndexTTS25TalkerForConditionalGeneration",
|
|
model_stage="indextts2_5_talker",
|
|
async_chunk=False,
|
|
final_output=False,
|
|
custom_process_next_stage_input_func="module.full_payload",
|
|
)
|
|
runner._custom_process_func = object()
|
|
runner._pending_full_payload_send = {}
|
|
runner._stage_id = 0
|
|
runner._omni_connector = object()
|
|
return runner
|
|
|
|
|
|
def test_request_end_payload_d2h_gate_requires_opt_in_and_no_prefix_cache():
|
|
assert _make_request_end_payload_runner()._should_defer_full_payload_d2h()
|
|
assert not _make_request_end_payload_runner(enabled=False)._should_defer_full_payload_d2h()
|
|
assert not _make_request_end_payload_runner(prefix_cache=object())._should_defer_full_payload_d2h()
|
|
|
|
|
|
def test_request_end_payload_suppresses_per_step_multimodal_outputs():
|
|
runner = _make_request_end_payload_runner()
|
|
|
|
def unexpected_build(_payload):
|
|
raise AssertionError("request-end payloads must stay inside the GPU accumulator")
|
|
|
|
runner._build_multimodal_outputs = unexpected_build
|
|
pooler_inter = [{"codes.mel": torch.tensor([[7]])}]
|
|
|
|
inter_stage, client = runner._build_omni_step_outputs(
|
|
pooler_inter,
|
|
pooler_inter,
|
|
defer_full_payload_d2h=True,
|
|
)
|
|
|
|
assert inter_stage is None
|
|
assert client is None
|
|
|
|
|
|
def test_sync_local_stage_payloads_retains_payload_until_request_is_active():
|
|
runner = object.__new__(OmniGPUModelRunner)
|
|
payload = {"codes": {"audio": [1, 2, 3]}}
|
|
runner._local_stage_payload_cache = {"late": payload}
|
|
runner._full_payload_pending_broadcast_req_ids = set()
|
|
runner.requests = {}
|
|
runner.model_intermediate_buffer = {}
|
|
|
|
OmniGPUModelRunner._sync_local_stage_payloads(runner)
|
|
|
|
assert runner._local_stage_payload_cache == {"late": payload}
|
|
assert runner.model_intermediate_buffer == {}
|
|
|
|
runner.requests = {"late": DummyReqState()}
|
|
OmniGPUModelRunner._sync_local_stage_payloads(runner)
|
|
|
|
assert runner._local_stage_payload_cache == {}
|
|
assert runner.model_intermediate_buffer["late"] == payload
|
|
assert runner.requests["late"].additional_information_cpu == payload
|
|
|
|
|
|
def test_maybe_attach_mimo_audio_req_infos_enriches_dict():
|
|
runner = _make_runner_for_mimo()
|
|
req_id = "r_mimo"
|
|
req_state = runner.requests[req_id]
|
|
|
|
# Existing req_infos should be copied and enriched, not mutated in place.
|
|
original_req_infos = {"existing": 1}
|
|
enriched = OmniGPUModelRunner._maybe_attach_mimo_audio_req_infos(runner, req_state, original_req_infos, req_id)
|
|
|
|
assert enriched is not original_req_infos
|
|
assert enriched["existing"] == 1
|
|
# mm_features should be filled from req_state when missing
|
|
assert enriched["mm_features"] == req_state.mm_features
|
|
# req_id should always be attached
|
|
assert enriched["req_id"] == req_id
|
|
|
|
|
|
def test_maybe_attach_mimo_audio_req_infos_no_req_state_returns_input():
|
|
runner = _make_runner_for_mimo()
|
|
req_id = "missing"
|
|
req_state = None
|
|
req_infos = {"k": "v"}
|
|
|
|
result = OmniGPUModelRunner._maybe_attach_mimo_audio_req_infos(runner, req_state, req_infos, req_id)
|
|
|
|
# When no req_state, helper should be a no-op.
|
|
assert result is req_infos
|