912 lines
36 KiB
Python
912 lines
36 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Regression tests for OmniRequestState multimodal DELTA drain and consolidation guard."""
|
|
|
|
from dataclasses import dataclass
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from vllm.outputs import PoolingRequestOutput
|
|
from vllm.sampling_params import RequestOutputKind
|
|
from vllm.v1.engine import EngineCoreEvent, EngineCoreEventType, FinishReason
|
|
from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
|
|
from vllm.v1.metrics.stats import IterationStats, PrefillStats
|
|
|
|
from vllm_omni.engine import OmniEngineCoreOutput
|
|
from vllm_omni.outputs import output_processor
|
|
from vllm_omni.outputs.output_modality import OutputModality, OutputModalityNames
|
|
from vllm_omni.outputs.output_processor import MultimodalOutputProcessor, OmniRequestState
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
# Audio is explicitly listed as a drainable modality
|
|
AUDIO = OutputModalityNames.AUDIO
|
|
|
|
# Latent is explicitly not drainable, but the choice doesn't matter here as
|
|
# long as isn't listed as drainable. I.e., could also be arbitrary keys for
|
|
# the purposes of these tests
|
|
LATENT = OutputModalityNames.LATENT
|
|
|
|
# NOTE: detokenizer and logprobs aren't really used here, but we mock them since
|
|
# some of the utils called in vLLM superclass assert require them to be None.
|
|
_DETOK = MagicMock(
|
|
output_token_ids=[0],
|
|
get_next_output_text=MagicMock(return_value=""),
|
|
num_output_tokens=MagicMock(return_value=1),
|
|
)
|
|
_LOGPROBS = MagicMock(logprobs=None, cumulative_logprob=None, prompt_logprobs=None)
|
|
|
|
|
|
@dataclass
|
|
class _StreamingUpdate:
|
|
final: bool
|
|
prompt: object | None
|
|
prompt_token_ids: list[int]
|
|
arrival_time: float
|
|
|
|
|
|
@dataclass
|
|
class _FinishedRequestMetric:
|
|
request_id: str
|
|
mean_time_per_output_token: float
|
|
|
|
|
|
_DEFAULT_STATE_KWARGS = dict(
|
|
request_id="r",
|
|
external_req_id="r",
|
|
parent_req=None,
|
|
request_index=0,
|
|
lora_request=None,
|
|
prompt=None,
|
|
prompt_token_ids=[0],
|
|
prompt_embeds=None,
|
|
logprobs_processor=_LOGPROBS,
|
|
detokenizer=_DETOK,
|
|
max_tokens_param=None,
|
|
arrival_time=0.0,
|
|
queue=None,
|
|
log_stats=False,
|
|
stream_interval=1,
|
|
)
|
|
|
|
|
|
def _make_state(output_kind: RequestOutputKind):
|
|
return OmniRequestState(**_DEFAULT_STATE_KWARGS, output_kind=output_kind)
|
|
|
|
|
|
def test_output_modality_name_uses_string_value():
|
|
assert str(OutputModalityNames.AUDIO) == "audio"
|
|
|
|
|
|
def test_init_empty_dict():
|
|
"""Ensure mm_accumulated is initially empty."""
|
|
assert _make_state(RequestOutputKind.CUMULATIVE).mm_accumulated == {}
|
|
assert _make_state(RequestOutputKind.DELTA).mm_accumulated == {}
|
|
|
|
|
|
def test_streaming_update_resets_native_text_metrics_for_next_segment():
|
|
state = _make_state(RequestOutputKind.DELTA)
|
|
state.native_text_stats.num_generation_tokens = 12
|
|
state.native_text_stats.first_token_ts = 10.0
|
|
state.native_text_stats.last_token_ts = 10.2
|
|
|
|
state.apply_streaming_update(
|
|
SimpleNamespace(
|
|
final=False,
|
|
prompt=None,
|
|
prompt_token_ids=[],
|
|
arrival_time=20.0,
|
|
)
|
|
)
|
|
|
|
assert state.native_text_stats.arrival_time == 20.0
|
|
assert state.native_text_stats.num_generation_tokens == 0
|
|
assert state.native_text_stats.first_token_ts == 0.0
|
|
assert state.native_text_stats.last_token_ts == 0.0
|
|
|
|
|
|
def test_native_text_tpot_snapshot_matches_vllm_finished_metric_definition():
|
|
calculate = getattr(output_processor, "_mean_time_per_output_token_ms", None)
|
|
assert calculate is not None
|
|
stats = SimpleNamespace(
|
|
num_generation_tokens=4,
|
|
first_token_ts=10.0,
|
|
last_token_ts=10.045,
|
|
)
|
|
|
|
assert calculate(stats) == pytest.approx(15.0)
|
|
|
|
|
|
def test_native_text_metrics_include_segment_generation_token_count(monkeypatch):
|
|
monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
|
|
processor = object.__new__(MultimodalOutputProcessor)
|
|
processor._native_text_metrics_by_request = {}
|
|
processor.lora_states = {}
|
|
state = _make_state(RequestOutputKind.DELTA)
|
|
state.is_prefilling = False
|
|
iteration_stats = MagicMock()
|
|
|
|
def update_native_stats(_output, _timestamp, _was_prefilling, native_stats, *_args):
|
|
native_stats.num_generation_tokens = 27
|
|
native_stats.first_token_ts = 10.0
|
|
native_stats.last_token_ts = 10.26
|
|
|
|
iteration_stats.update_from_output.side_effect = update_native_stats
|
|
|
|
processor._update_stats_from_output(
|
|
state,
|
|
SimpleNamespace(),
|
|
10.26,
|
|
iteration_stats,
|
|
)
|
|
|
|
assert processor.pop_native_text_metrics("r")["num_generation_tokens"] == 27
|
|
|
|
|
|
def test_native_text_metrics_exclude_cross_segment_gap(monkeypatch):
|
|
monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
|
|
processor = object.__new__(MultimodalOutputProcessor)
|
|
processor._native_text_metrics_by_request = {}
|
|
processor.lora_states = {}
|
|
state = _make_state(RequestOutputKind.DELTA)
|
|
iteration_stats = MagicMock()
|
|
|
|
def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args):
|
|
native_stats.num_generation_tokens = 2 if timestamp == 10.03 else 1
|
|
native_stats.first_token_ts = 10.02 if timestamp == 10.03 else timestamp
|
|
native_stats.last_token_ts = timestamp
|
|
native_stats.first_token_latency = 0.01
|
|
|
|
iteration_stats.update_from_output.side_effect = update_segment
|
|
processor._update_stats_from_output(
|
|
state,
|
|
MagicMock(),
|
|
10.0,
|
|
iteration_stats,
|
|
)
|
|
first_segment = processor.pop_native_text_metrics("r")
|
|
assert first_segment["vllm_tpot_ms"] == 0.0
|
|
assert first_segment["vllm_itls_ms"] == []
|
|
|
|
state.apply_streaming_update(
|
|
_StreamingUpdate(
|
|
final=False,
|
|
prompt=None,
|
|
prompt_token_ids=[],
|
|
arrival_time=10.01,
|
|
)
|
|
)
|
|
state.is_prefilling = True
|
|
processor._update_stats_from_output(
|
|
state,
|
|
MagicMock(),
|
|
10.02,
|
|
iteration_stats,
|
|
)
|
|
state.is_prefilling = False
|
|
processor._update_stats_from_output(
|
|
state,
|
|
MagicMock(),
|
|
10.03,
|
|
iteration_stats,
|
|
)
|
|
second_segment = processor.pop_native_text_metrics("r")
|
|
|
|
assert second_segment["num_generation_tokens"] == 2
|
|
assert second_segment["vllm_itls_ms"] == pytest.approx([10.0])
|
|
assert second_segment["vllm_tpot_ms"] == pytest.approx(10.0)
|
|
|
|
|
|
def test_native_text_tpot_weights_multi_token_engine_output(monkeypatch):
|
|
monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_output", lambda *args, **kwargs: None)
|
|
processor = object.__new__(MultimodalOutputProcessor)
|
|
processor._native_text_metrics_by_request = {}
|
|
processor.lora_states = {}
|
|
state = _make_state(RequestOutputKind.DELTA)
|
|
iteration_stats = MagicMock()
|
|
|
|
def update_segment(_output, timestamp, _was_prefilling, native_stats, *_args):
|
|
native_stats.num_generation_tokens = 1 if timestamp == 10.0 else 4
|
|
native_stats.first_token_ts = 10.0
|
|
native_stats.last_token_ts = timestamp
|
|
native_stats.first_token_latency = 0.01
|
|
|
|
iteration_stats.update_from_output.side_effect = update_segment
|
|
state.is_prefilling = True
|
|
processor._update_stats_from_output(state, MagicMock(), 10.0, iteration_stats)
|
|
state.is_prefilling = False
|
|
processor._update_stats_from_output(state, MagicMock(), 10.03, iteration_stats)
|
|
record = processor.pop_native_text_metrics("r")
|
|
|
|
assert record["num_generation_tokens"] == 4
|
|
assert record["vllm_itls_ms"] == pytest.approx([30.0])
|
|
assert record["vllm_itl_ms"] == pytest.approx(30.0)
|
|
assert record["vllm_tpot_ms"] == pytest.approx(10.0)
|
|
assert "_tpot_elapsed_ms" not in record
|
|
assert "_tpot_intervals" not in record
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("finished_tpot_s", "expected_tpot_ms"),
|
|
[(0.0, 19.0), (0.023, 23.0)],
|
|
)
|
|
def test_native_text_tpot_only_accepts_positive_finished_metric(
|
|
monkeypatch,
|
|
finished_tpot_s,
|
|
expected_tpot_ms,
|
|
):
|
|
monkeypatch.setattr(VLLMOutputProcessor, "_update_stats_from_finished", lambda *args, **kwargs: None)
|
|
processor = object.__new__(MultimodalOutputProcessor)
|
|
processor._native_text_metrics_by_request = {"r": {"vllm_tpot_ms": 19.0}}
|
|
state = _make_state(RequestOutputKind.DELTA)
|
|
iteration_stats = MagicMock()
|
|
iteration_stats.finished_requests = [
|
|
_FinishedRequestMetric(
|
|
request_id="r",
|
|
mean_time_per_output_token=finished_tpot_s,
|
|
)
|
|
]
|
|
|
|
processor._update_stats_from_finished(
|
|
state,
|
|
FinishReason.STOP,
|
|
iteration_stats,
|
|
)
|
|
|
|
assert processor.pop_native_text_metrics("r")["vllm_tpot_ms"] == expected_tpot_ms
|
|
|
|
|
|
def test_delta_drains_output_modality_per_step():
|
|
"""DELTA drains the mm_type key (output modality) but preserves hidden-state keys."""
|
|
s = _make_state(RequestOutputKind.DELTA)
|
|
audio1, audio2, hs1, hs2 = [torch.ones(num_elem) for num_elem in range(1, 5)]
|
|
|
|
# Add audio and hidden state tensors
|
|
s.add_multimodal_tensor(audio1, mm_type=AUDIO) # should be drained
|
|
s.add_multimodal_tensor(hs1, mm_type=LATENT) # shouldn't be drained
|
|
|
|
out1 = s._new_completion_output([1], None, None)
|
|
out1_audio = out1.multimodal_output[AUDIO]
|
|
out1_hidden = out1.multimodal_output[LATENT]
|
|
assert isinstance(out1_audio, torch.Tensor)
|
|
assert torch.equal(out1.multimodal_output[AUDIO], audio1)
|
|
assert isinstance(out1_hidden, torch.Tensor)
|
|
assert torch.equal(out1.multimodal_output[LATENT], hs1)
|
|
|
|
# After emission, hidden states should remain, but audio is drained
|
|
assert set(s.mm_accumulated.keys()) == {LATENT}
|
|
|
|
s.add_multimodal_tensor(audio2, AUDIO)
|
|
s.add_multimodal_tensor(hs2, mm_type=LATENT)
|
|
out2 = s._new_completion_output([2], None, None)
|
|
out2_audio = out2.multimodal_output[AUDIO]
|
|
out2_hidden = out2.multimodal_output[LATENT]
|
|
assert isinstance(out2_audio, torch.Tensor)
|
|
assert torch.equal(out2_audio, audio2)
|
|
# Since hidden isn't drained, it's grown to a list
|
|
assert isinstance(out2_hidden, list) and len(out2_hidden) == 2
|
|
assert torch.equal(hs1, out2_hidden[0])
|
|
assert torch.equal(hs2, out2_hidden[1])
|
|
|
|
|
|
def test_cumulative_emits_consolidated_audio_each_step():
|
|
"""Ensure cumulative accumulates and consolidates modality keys every step."""
|
|
s = _make_state(RequestOutputKind.CUMULATIVE)
|
|
# NOTE: audio is usually emitted as (1, size) chunks; we need to be sure
|
|
# to not change the tensor dimension when we consolidate
|
|
audio1 = torch.ones(1, 500)
|
|
s.add_multimodal_tensor(audio1, mm_type=AUDIO)
|
|
req_out = s.make_request_output([1], None, None, None)
|
|
assert req_out is not None
|
|
cons_audio = req_out.outputs[0].multimodal_output[AUDIO]
|
|
# Single chunk keeps original shape [1, 500]
|
|
assert isinstance(cons_audio, torch.Tensor) and cons_audio.shape == audio1.shape
|
|
|
|
audio2 = torch.ones(1, 300)
|
|
s.add_multimodal_tensor(audio2, mm_type=AUDIO)
|
|
req_out = s.make_request_output([2], None, None, None)
|
|
assert req_out is not None
|
|
cons_audio = req_out.outputs[0].multimodal_output[AUDIO]
|
|
# After consolidation, audio chunks are concatenated on last axis,
|
|
# preserving the [1, N] channel dimension
|
|
total_audio_len = audio1.shape[-1] + audio2.shape[-1]
|
|
assert isinstance(audio2, torch.Tensor) and cons_audio.shape == (1, total_audio_len)
|
|
|
|
assert "audio" in s.mm_accumulated
|
|
|
|
|
|
def test_finish_consolidates_hidden_states():
|
|
"""Ensure consolidation merges hidden-state tensor lists on finish."""
|
|
s = _make_state(RequestOutputKind.CUMULATIVE)
|
|
s.add_multimodal_tensor(torch.ones(5, 4), mm_type=LATENT)
|
|
s.add_multimodal_tensor(torch.ones(3, 4), mm_type=LATENT)
|
|
|
|
result = s.make_request_output([1], None, FinishReason.STOP, None)
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
|
|
hs = result.outputs[0].multimodal_output[LATENT]
|
|
assert isinstance(hs, torch.Tensor) and hs.shape[0] == 8
|
|
|
|
|
|
def test_finish_consolidation_for_hs_delta():
|
|
"""Ensure finish doesn't drop the accumulated hidden states."""
|
|
s = _make_state(RequestOutputKind.DELTA)
|
|
# hidden state accumulation (nothing drained)
|
|
s.add_multimodal_tensor({"foo": torch.ones(5, 4)}, mm_type=LATENT)
|
|
result = s.make_request_output([0], None, FinishReason.STOP, None)
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
hs = result.outputs[0].multimodal_output["foo"]
|
|
assert isinstance(hs, torch.Tensor) and hs.shape[0] == 5
|
|
|
|
# Since we don't drain the hidden states, if we add 3 elements, we should get 8
|
|
s.add_multimodal_tensor({"foo": torch.ones(3, 4)}, mm_type=LATENT)
|
|
result = s.make_request_output([0], None, FinishReason.STOP, None)
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
hs = result.outputs[0].multimodal_output["foo"]
|
|
assert isinstance(hs, torch.Tensor) and hs.shape[0] == 8
|
|
assert "foo" in s.mm_accumulated
|
|
|
|
|
|
def test_finish_consolidation_drains_mm_delta():
|
|
"""Ensure making the request output drains modality deltas (e.g., audio)."""
|
|
s = _make_state(RequestOutputKind.DELTA)
|
|
# multimodal data accumulation (drained)
|
|
s.add_multimodal_tensor({AUDIO: torch.ones(5, 4)}, mm_type=AUDIO)
|
|
result = s.make_request_output([0], None, FinishReason.STOP, None)
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
hs = result.outputs[0].multimodal_output[AUDIO]
|
|
assert isinstance(hs, torch.Tensor) and hs.shape[0] == 5
|
|
|
|
# Since we did drain the hidden states, we no longer get the 5 back
|
|
s.add_multimodal_tensor({AUDIO: torch.ones(3, 4)}, mm_type=AUDIO)
|
|
result = s.make_request_output([0], None, FinishReason.STOP, None)
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
hs = result.outputs[0].multimodal_output[AUDIO]
|
|
assert isinstance(hs, torch.Tensor) and hs.shape[0] == 3
|
|
assert AUDIO not in s.mm_accumulated # drained
|
|
|
|
|
|
def test_delta_drains_audio_chunk_metadata_per_step():
|
|
s = _make_state(RequestOutputKind.DELTA)
|
|
segment_text = "你好"
|
|
segment_utf8 = torch.tensor(list(segment_text.encode("utf-8")), dtype=torch.uint8)
|
|
|
|
s.add_multimodal_tensor(
|
|
{
|
|
"model_outputs": torch.ones(5),
|
|
"meta.duplex_epoch": torch.tensor([7], dtype=torch.int32),
|
|
"meta.duplex_turn_id": torch.tensor([1], dtype=torch.int32),
|
|
"meta.llm_output_text_utf8": segment_utf8,
|
|
"meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
|
|
"meta.tts_is_last_chunk": torch.tensor([0], dtype=torch.int32),
|
|
},
|
|
mm_type=AUDIO,
|
|
)
|
|
result1 = s.make_request_output([1], None, FinishReason.STOP, None)
|
|
assert result1 is not None and not isinstance(result1, PoolingRequestOutput)
|
|
meta1 = result1.outputs[0].multimodal_output["meta"]
|
|
assert bytes(meta1["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text
|
|
|
|
s.add_multimodal_tensor(
|
|
{
|
|
"model_outputs": torch.ones(3),
|
|
"meta.duplex_epoch": torch.tensor([8], dtype=torch.int32),
|
|
"meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32),
|
|
"meta.llm_output_text_utf8": segment_utf8,
|
|
"meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
|
|
"meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32),
|
|
},
|
|
mm_type=AUDIO,
|
|
)
|
|
result2 = s.make_request_output([2], None, FinishReason.STOP, None)
|
|
|
|
assert result2 is not None and not isinstance(result2, PoolingRequestOutput)
|
|
meta2 = result2.outputs[0].multimodal_output["meta"]
|
|
assert bytes(meta2["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text
|
|
assert meta2["audio_text_total_chars"].tolist() == [len(segment_text)]
|
|
assert meta2["duplex_epoch"].tolist() == [8]
|
|
assert meta2["duplex_turn_id"].tolist() == [2]
|
|
|
|
|
|
def test_cumulative_audio_replaces_chunk_metadata_per_step():
|
|
s = _make_state(RequestOutputKind.CUMULATIVE)
|
|
segment_text = "你好"
|
|
segment_utf8 = torch.tensor(list(segment_text.encode("utf-8")), dtype=torch.uint8)
|
|
|
|
s.add_multimodal_tensor(
|
|
{
|
|
"model_outputs": torch.ones(5),
|
|
"meta.duplex_epoch": torch.tensor([7], dtype=torch.int32),
|
|
"meta.duplex_turn_id": torch.tensor([1], dtype=torch.int32),
|
|
"meta.llm_output_text_utf8": segment_utf8,
|
|
"meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
|
|
"meta.tts_is_last_chunk": torch.tensor([0], dtype=torch.int32),
|
|
},
|
|
mm_type=AUDIO,
|
|
)
|
|
result1 = s.make_request_output([1], None, None, None)
|
|
assert result1 is not None and not isinstance(result1, PoolingRequestOutput)
|
|
assert result1.outputs[0].multimodal_output["audio"].shape[0] == 5
|
|
|
|
s.add_multimodal_tensor(
|
|
{
|
|
"model_outputs": torch.ones(3),
|
|
"meta.duplex_epoch": torch.tensor([8], dtype=torch.int32),
|
|
"meta.duplex_turn_id": torch.tensor([2], dtype=torch.int32),
|
|
"meta.llm_output_text_utf8": segment_utf8,
|
|
"meta.audio_text_total_chars": torch.tensor([len(segment_text)], dtype=torch.int32),
|
|
"meta.tts_is_last_chunk": torch.tensor([1], dtype=torch.int32),
|
|
},
|
|
mm_type=AUDIO,
|
|
)
|
|
result2 = s.make_request_output([2], None, None, None)
|
|
|
|
assert result2 is not None and not isinstance(result2, PoolingRequestOutput)
|
|
output = result2.outputs[0].multimodal_output
|
|
assert output["audio"].shape[0] == 8
|
|
meta = output["meta"]
|
|
assert bytes(meta["llm_output_text_utf8"].tolist()).decode("utf-8") == segment_text
|
|
assert meta["audio_text_total_chars"].tolist() == [len(segment_text)]
|
|
assert meta["duplex_epoch"].tolist() == [8]
|
|
assert meta["duplex_turn_id"].tolist() == [2]
|
|
assert meta["tts_is_last_chunk"].tolist() == [1]
|
|
|
|
|
|
def test_delta_audio_non_final_tts_chunk_overrides_spurious_finish():
|
|
s = _make_state(RequestOutputKind.DELTA)
|
|
s.add_multimodal_tensor(
|
|
{
|
|
"model_outputs": torch.ones(5),
|
|
"meta.tts_is_last_chunk": torch.tensor([0]),
|
|
},
|
|
mm_type=AUDIO,
|
|
)
|
|
|
|
result = s.make_request_output([1], None, FinishReason.STOP, None)
|
|
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
assert result.finished is False
|
|
assert result.outputs[0].finish_reason is None
|
|
assert result.outputs[0].multimodal_output["meta"]["tts_is_last_chunk"].tolist() == [0]
|
|
assert torch.equal(result.outputs[0].multimodal_output[AUDIO], torch.ones(5))
|
|
|
|
|
|
def test_delta_audio_final_tts_chunk_keeps_finish():
|
|
s = _make_state(RequestOutputKind.DELTA)
|
|
s.add_multimodal_tensor(
|
|
{
|
|
"model_outputs": torch.ones(5),
|
|
"meta.tts_is_last_chunk": torch.tensor([1]),
|
|
},
|
|
mm_type=AUDIO,
|
|
)
|
|
|
|
result = s.make_request_output([151645], None, FinishReason.STOP, None)
|
|
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
assert result.finished is True
|
|
assert result.outputs[0].finish_reason == "stop"
|
|
|
|
|
|
def test_delta_audio_non_final_tts_chunk_accepts_nested_meta():
|
|
s = _make_state(RequestOutputKind.DELTA)
|
|
s.add_multimodal_tensor(
|
|
{
|
|
AUDIO: torch.ones(5),
|
|
"meta": {"tts_is_last_chunk": torch.tensor([0])},
|
|
},
|
|
mm_type=AUDIO,
|
|
)
|
|
|
|
result = s.make_request_output([1], None, FinishReason.STOP, None)
|
|
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
assert result.finished is False
|
|
assert result.outputs[0].multimodal_output["meta"]["tts_is_last_chunk"].tolist() == [0]
|
|
|
|
|
|
@pytest.mark.parametrize("mm_type", [AUDIO, "hidden"])
|
|
def test_final_only_consolidates_drainable_keys(mm_type):
|
|
"""FINAL_ONLY never drains per-step, so modality keys and hidden state
|
|
keys both accumulate and are consolidated on finish."""
|
|
s = _make_state(RequestOutputKind.FINAL_ONLY)
|
|
|
|
# NOTE: Currently there is brittlness in the tensor stacking, so we just
|
|
# test a 1D tensor here. The intention is just to ensure audio /hidden
|
|
# behave the same.
|
|
s.add_multimodal_tensor(torch.ones(500), mm_type=mm_type)
|
|
# Non-finish step returns None without calling _new_completion_output
|
|
assert s.make_request_output([1], None, None, None) is None
|
|
assert mm_type in s.mm_accumulated
|
|
|
|
s.add_multimodal_tensor(torch.ones(300), mm_type=mm_type)
|
|
result = s.make_request_output([2], None, FinishReason.STOP, None)
|
|
assert result is not None and not isinstance(result, PoolingRequestOutput)
|
|
|
|
audio = result.outputs[0].multimodal_output[mm_type]
|
|
assert isinstance(audio, torch.Tensor)
|
|
assert audio.shape == (800,)
|
|
|
|
|
|
def test_cumulative_token_ids_always_set():
|
|
"""cumulative_token_ids is set for all output kinds."""
|
|
for kind in (RequestOutputKind.DELTA, RequestOutputKind.CUMULATIVE, RequestOutputKind.FINAL_ONLY):
|
|
s = _make_state(kind)
|
|
out = s._new_completion_output([42], None, None)
|
|
assert hasattr(out, "cumulative_token_ids")
|
|
# The mock detokenizer has output_token_ids=[0]
|
|
assert list(out.cumulative_token_ids) == [0]
|
|
|
|
|
|
def test_cumulative_token_ids_is_a_copy():
|
|
"""cumulative_token_ids must be a snapshot, not a live reference."""
|
|
s = _make_state(RequestOutputKind.DELTA)
|
|
out = s._new_completion_output([42], None, None)
|
|
assert out.cumulative_token_ids is not _DETOK.output_token_ids
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# No-detokenizer regression tests (Concern 2)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
_NO_DETOK_STATE_KWARGS = {
|
|
**_DEFAULT_STATE_KWARGS,
|
|
"detokenizer": None,
|
|
"logprobs_processor": None,
|
|
}
|
|
|
|
|
|
def _make_no_detok_state(output_kind: RequestOutputKind):
|
|
return OmniRequestState(**_NO_DETOK_STATE_KWARGS, output_kind=output_kind)
|
|
|
|
|
|
def test_no_detokenizer_completion_output():
|
|
"""_new_completion_output works when detokenizer is None (generation stages)."""
|
|
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
|
|
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
|
out = s._new_completion_output([1, 2], None, None)
|
|
# Should still produce output with empty text and the token_ids as-is
|
|
assert out.text == ""
|
|
assert list(out.token_ids) == [1, 2]
|
|
assert hasattr(out, "cumulative_token_ids")
|
|
assert list(out.cumulative_token_ids) == [1, 2]
|
|
# Multimodal data should be attached
|
|
assert AUDIO in out.multimodal_output
|
|
|
|
|
|
def test_no_detokenizer_make_request_output():
|
|
"""make_request_output works without detokenizer when multimodal data is present."""
|
|
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
|
|
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
|
result = s.make_request_output([], None, FinishReason.STOP, None)
|
|
assert result is not None
|
|
assert not isinstance(result, PoolingRequestOutput)
|
|
assert AUDIO in result.outputs[0].multimodal_output
|
|
|
|
|
|
def test_no_detokenizer_process_outputs_returns_nonterminal_audio_chunk(monkeypatch):
|
|
"""Generation-stage MM-only chunks must reach the orchestrator.
|
|
|
|
StagePool registers generation requests without a per-request collector,
|
|
so the processor return value is the only path from Stage2 to
|
|
Orchestrator._handle_processed_outputs().
|
|
"""
|
|
state = _make_no_detok_state(RequestOutputKind.DELTA)
|
|
processor = object.__new__(MultimodalOutputProcessor)
|
|
processor.output_modality = OutputModality.AUDIO
|
|
processor.request_states = {"r": state}
|
|
upstream_result = SimpleNamespace(request_outputs=[], reqs_to_abort=[])
|
|
monkeypatch.setattr(
|
|
VLLMOutputProcessor,
|
|
"process_outputs",
|
|
lambda *_args, **_kwargs: upstream_result,
|
|
)
|
|
engine_output = SimpleNamespace(
|
|
request_id="r",
|
|
multimodal_output={
|
|
"model_outputs": torch.arange(8, dtype=torch.float32),
|
|
"sr": torch.tensor(24000, dtype=torch.int32),
|
|
},
|
|
output_type="audio",
|
|
pooling_output=None,
|
|
new_token_ids=[],
|
|
finish_reason=None,
|
|
stop_reason=None,
|
|
kv_transfer_params=None,
|
|
routed_experts=None,
|
|
num_cached_tokens=0,
|
|
)
|
|
|
|
result = processor.process_outputs([engine_output])
|
|
|
|
assert len(result.request_outputs) == 1
|
|
assert result.request_outputs[0].request_id == "r"
|
|
assert AUDIO in result.request_outputs[0].outputs[0].multimodal_output
|
|
|
|
|
|
def _make_mm_only_output_processor(monkeypatch):
|
|
processor = object.__new__(MultimodalOutputProcessor)
|
|
processor.output_modality = OutputModality.AUDIO
|
|
processor.request_states = {"r": _make_no_detok_state(RequestOutputKind.DELTA)}
|
|
monkeypatch.setattr(
|
|
VLLMOutputProcessor,
|
|
"process_outputs",
|
|
lambda *_args, **_kwargs: SimpleNamespace(
|
|
request_outputs=[],
|
|
reqs_to_abort=[],
|
|
),
|
|
)
|
|
|
|
def finish_request(req_state):
|
|
processor.request_states.pop(req_state.request_id, None)
|
|
|
|
processor._finish_request = finish_request
|
|
return processor
|
|
|
|
|
|
def _audio_engine_output(*, is_segment_finished: bool, is_last_chunk: bool):
|
|
return SimpleNamespace(
|
|
request_id="r",
|
|
multimodal_output={
|
|
"model_outputs": torch.arange(8, dtype=torch.float32),
|
|
"meta.tts_is_last_chunk": torch.tensor([int(is_last_chunk)]),
|
|
},
|
|
output_type="audio",
|
|
pooling_output=None,
|
|
new_token_ids=[],
|
|
finish_reason=FinishReason.STOP,
|
|
stop_reason=None,
|
|
kv_transfer_params=None,
|
|
routed_experts=None,
|
|
num_cached_tokens=0,
|
|
is_segment_finished=is_segment_finished,
|
|
)
|
|
|
|
|
|
def test_mm_only_segment_finish_retains_request_state_for_next_audio_chunk(monkeypatch):
|
|
processor = _make_mm_only_output_processor(monkeypatch)
|
|
|
|
first = processor.process_outputs([_audio_engine_output(is_segment_finished=True, is_last_chunk=False)])
|
|
|
|
assert "r" in processor.request_states
|
|
assert len(first.request_outputs) == 1
|
|
assert first.request_outputs[0].finished is False
|
|
|
|
second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
|
|
|
|
assert len(second.request_outputs) == 1
|
|
assert second.request_outputs[0].finished is True
|
|
assert "r" not in processor.request_states
|
|
|
|
|
|
def test_mm_only_nonfinal_audio_retains_state_without_segment_marker(monkeypatch):
|
|
processor = _make_mm_only_output_processor(monkeypatch)
|
|
|
|
first = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=False)])
|
|
|
|
assert "r" in processor.request_states
|
|
assert len(first.request_outputs) == 1
|
|
assert first.request_outputs[0].finished is False
|
|
|
|
second = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
|
|
|
|
assert len(second.request_outputs) == 1
|
|
assert second.request_outputs[0].finished is True
|
|
assert "r" not in processor.request_states
|
|
|
|
|
|
def test_mm_only_terminal_finish_removes_request_state(monkeypatch):
|
|
processor = _make_mm_only_output_processor(monkeypatch)
|
|
|
|
result = processor.process_outputs([_audio_engine_output(is_segment_finished=False, is_last_chunk=True)])
|
|
|
|
assert len(result.request_outputs) == 1
|
|
assert result.request_outputs[0].finished is True
|
|
assert "r" not in processor.request_states
|
|
|
|
|
|
def test_no_detokenizer_make_request_output_with_routed_experts():
|
|
"""make_request_output accepts the routed_experts arg that the multimodal
|
|
output channel (_process_mm_only_outputs) passes, and attaches it to the
|
|
completion output.
|
|
|
|
routed_experts is omni-specific keyword-only: the 6th positional now
|
|
matches upstream's ec_transfer_params so that super().process_outputs()
|
|
calls do not misroute it.
|
|
"""
|
|
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
|
|
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
|
routed_experts = np.zeros((2, 3), dtype=np.int32)
|
|
# Mirror the exact call shape of _process_mm_only_outputs.
|
|
result = s.make_request_output([], None, FinishReason.STOP, None, None, routed_experts=routed_experts)
|
|
assert result is not None
|
|
assert not isinstance(result, PoolingRequestOutput)
|
|
assert result.outputs[0].routed_experts is routed_experts
|
|
|
|
|
|
def test_no_detokenizer_stream_interval_skipped():
|
|
"""stream_interval > 1 does not crash when detokenizer is None."""
|
|
kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5}
|
|
s = OmniRequestState(**kwargs, output_kind=RequestOutputKind.DELTA)
|
|
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
|
# Should not assert; stream_interval logic is simply skipped
|
|
result = s.make_request_output([], None, FinishReason.STOP, None)
|
|
assert result is not None
|
|
|
|
|
|
def test_no_detokenizer_final_only():
|
|
"""FINAL_ONLY with no detokenizer returns None until finish."""
|
|
s = _make_no_detok_state(RequestOutputKind.FINAL_ONLY)
|
|
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
|
# Non-finish step should return None
|
|
assert s.make_request_output([], None, None, None) is None
|
|
# Finish step should return output
|
|
result = s.make_request_output([], None, FinishReason.STOP, None)
|
|
assert result is not None
|
|
assert AUDIO in result.outputs[0].multimodal_output
|
|
|
|
|
|
def test_mm_only_outputs_update_iteration_stats():
|
|
processor = MultimodalOutputProcessor(
|
|
tokenizer=None,
|
|
log_stats=True,
|
|
engine_core_output_type=AUDIO,
|
|
)
|
|
state = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
|
|
processor.request_states[state.request_id] = state
|
|
processor.external_req_ids[state.external_req_id].append(state.request_id)
|
|
|
|
prefill_stats = PrefillStats()
|
|
prefill_stats.set(
|
|
num_prompt_tokens=4,
|
|
num_local_cached_tokens=0,
|
|
num_external_cached_tokens=0,
|
|
)
|
|
output = OmniEngineCoreOutput(
|
|
request_id=state.request_id,
|
|
new_token_ids=[10, 11],
|
|
finish_reason=FinishReason.STOP,
|
|
events=[
|
|
EngineCoreEvent(EngineCoreEventType.QUEUED, 1.0),
|
|
EngineCoreEvent(EngineCoreEventType.SCHEDULED, 2.0),
|
|
],
|
|
prefill_stats=prefill_stats,
|
|
multimodal_output={AUDIO: torch.ones(1, 8)},
|
|
)
|
|
iteration_stats = IterationStats()
|
|
|
|
processor.process_outputs(
|
|
[output],
|
|
engine_core_timestamp=3.0,
|
|
iteration_stats=iteration_stats,
|
|
)
|
|
|
|
assert iteration_stats.num_prompt_tokens == 4
|
|
assert iteration_stats.num_generation_tokens == 2
|
|
assert len(iteration_stats.time_to_first_tokens_iter) == 1
|
|
assert len(iteration_stats.finished_requests) == 1
|
|
finished = iteration_stats.finished_requests[0]
|
|
assert finished.finish_reason == FinishReason.STOP
|
|
assert finished.num_prompt_tokens == state.prompt_len
|
|
assert finished.num_generation_tokens == 2
|
|
|
|
|
|
def test_ec_transfer_params_survive_both_construction_paths():
|
|
"""v0.28 contract: encoder-cache transfer metadata must reach
|
|
RequestOutput.ec_transfer_params through BOTH omni construction paths —
|
|
the upstream-delegating one (logprobs processor present) and the
|
|
no-detokenizer generation-stage one (direct RequestOutput build).
|
|
Regression: the override accepted the argument but dropped it."""
|
|
ec = {"remote_handle": "enc-cache-1"}
|
|
|
|
state = _make_state(RequestOutputKind.CUMULATIVE)
|
|
out = state.make_request_output([1], None, None, None, {"kv": 1}, ec)
|
|
assert out is not None and out.ec_transfer_params == ec
|
|
|
|
kwargs = dict(_DEFAULT_STATE_KWARGS)
|
|
kwargs.update(logprobs_processor=None, detokenizer=None)
|
|
gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
|
|
completion = gen_state._new_completion_output([1], None, None)
|
|
direct = gen_state._new_request_output("r", [completion], False, {"kv": 1}, ec)
|
|
assert direct.ec_transfer_params == ec
|
|
|
|
|
|
def test_num_cache_creation_tokens_reaches_direct_request_output():
|
|
"""v0.28 contract: prefix-cache creation usage must reach
|
|
RequestOutput.num_cache_creation_tokens through the no-detokenizer
|
|
direct build, which previously passed only num_cached_tokens."""
|
|
kwargs = dict(_DEFAULT_STATE_KWARGS)
|
|
kwargs.update(logprobs_processor=None, detokenizer=None)
|
|
gen_state = OmniRequestState(**kwargs, output_kind=RequestOutputKind.CUMULATIVE)
|
|
gen_state.num_cached_tokens = 3
|
|
gen_state.num_cache_creation_tokens = 2
|
|
|
|
completion = gen_state._new_completion_output([1], None, None)
|
|
direct = gen_state._new_request_output("r", [completion], False, None, None)
|
|
|
|
assert direct.num_cached_tokens == 3
|
|
assert direct.num_cache_creation_tokens == 2
|
|
|
|
|
|
def _abort_processor_with_parent(child_ids: list[str]):
|
|
processor = object.__new__(MultimodalOutputProcessor)
|
|
processor.request_states = {}
|
|
processor.external_req_ids = {}
|
|
processor.parent_requests = {}
|
|
processor._native_text_metrics_by_request = {}
|
|
processor.lora_states = SimpleNamespace(request_finished=lambda *_args, **_kwargs: None)
|
|
parent = SimpleNamespace(request_id="parent", child_requests=set(child_ids))
|
|
processor.parent_requests["parent"] = parent
|
|
for child_id in child_ids:
|
|
req_state = SimpleNamespace(
|
|
parent_req=parent,
|
|
lora_name=None,
|
|
output_kind=RequestOutputKind.CUMULATIVE,
|
|
detokenizer=SimpleNamespace(output_token_ids=[7, 8]),
|
|
queue=None,
|
|
external_req_id="parent",
|
|
make_request_output=MagicMock(return_value=SimpleNamespace(request_id=child_id)),
|
|
)
|
|
processor.request_states[child_id] = req_state
|
|
return processor, parent
|
|
|
|
|
|
def test_abort_last_parallel_child_drops_parent_request():
|
|
processor, parent = _abort_processor_with_parent(["0_parent", "1_parent"])
|
|
|
|
aborted, _outputs = processor.abort_requests_collecting_outputs(["0_parent"], internal=True)
|
|
assert aborted == ["0_parent"]
|
|
assert "parent" in processor.parent_requests
|
|
assert parent.child_requests == {"1_parent"}
|
|
|
|
aborted, _outputs = processor.abort_requests_collecting_outputs(["1_parent"], internal=True)
|
|
assert aborted == ["1_parent"]
|
|
assert "parent" not in processor.parent_requests
|
|
assert not parent.child_requests
|
|
|
|
|
|
def test_abort_snapshot_leaves_state_until_commit():
|
|
processor, parent = _abort_processor_with_parent(["0_parent"])
|
|
processor.external_req_ids["parent"] = ["0_parent"]
|
|
|
|
aborted, outputs = processor.abort_requests_collecting_outputs(
|
|
["parent"],
|
|
internal=False,
|
|
commit_state=False,
|
|
)
|
|
assert aborted == ["0_parent"]
|
|
assert outputs
|
|
assert "0_parent" in processor.request_states
|
|
assert processor.external_req_ids["parent"] == ["0_parent"]
|
|
assert "parent" in processor.parent_requests
|
|
assert parent.child_requests == {"0_parent"}
|
|
|
|
processor.commit_aborted_request_state(["parent"], internal=False)
|
|
assert "0_parent" not in processor.request_states
|
|
assert "parent" not in processor.parent_requests
|
|
assert "parent" not in processor.external_req_ids
|
|
|
|
|
|
def test_cumulative_audio_terminal_output_supplies_full_aligner_waveform():
|
|
from vllm_omni.model_executor.stage_input_processors.forced_aligner import code2wav2aligner
|
|
|
|
state = _make_state(RequestOutputKind.CUMULATIVE)
|
|
for index, count in enumerate((800, 1200, 400)):
|
|
state.add_multimodal_tensor(torch.full((1, count), float(index)), mm_type=AUDIO)
|
|
output = state.make_request_output([index], None, FinishReason.STOP if index == 2 else None, None)
|
|
assert output is not None
|
|
if index < 2:
|
|
assert code2wav2aligner([output], {}) == []
|
|
prompt = {"additional_information": {"text": ["Hello world"], "language": ["English"]}}
|
|
aligned = code2wav2aligner([output], prompt)[0]
|
|
waveform, sr = aligned["multi_modal_data"]["audio"]
|
|
assert sr == 24000
|
|
assert waveform.shape == (2400,)
|
|
assert (waveform[:800] == 0).all()
|
|
assert (waveform[800:2000] == 1).all()
|
|
assert (waveform[2000:] == 2).all()
|
|
assert aligned["additional_information"]["aligner_audio_duration_ms"] == [100.0]
|