138 lines
4.9 KiB
Python
138 lines
4.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Regression tests for the dots.tts talker's stop-signal pairing and
|
|
per-request state lifecycle.
|
|
|
|
Pure-CPU tests on a bare talker (``__new__`` + the handful of attributes
|
|
the methods under test read) — no weights, no GPU, no engine. Mirrors
|
|
tests/model_executor/models/voxcpm2/test_talker_state_eviction.py.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
from types import SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
|
|
|
|
|
torch = pytest.importorskip("torch")
|
|
|
|
_VOCAB_SIZE = 151672
|
|
|
|
|
|
@functools.lru_cache(maxsize=1)
|
|
def _dots_tts_talker_mod():
|
|
"""Defer talker import (pulls vLLM model_executor) until first use."""
|
|
from vllm_omni.model_executor.models.dots_tts.dots_tts_talker import (
|
|
DotsTTSForConditionalGeneration,
|
|
_RequestState,
|
|
)
|
|
|
|
return DotsTTSForConditionalGeneration, _RequestState
|
|
|
|
|
|
def _make_bare_talker():
|
|
DotsTTSForConditionalGeneration, _ = _dots_tts_talker_mod()
|
|
talker = DotsTTSForConditionalGeneration.__new__(DotsTTSForConditionalGeneration)
|
|
talker.config = SimpleNamespace(vocab_size=_VOCAB_SIZE)
|
|
talker._active_states = {}
|
|
talker._pending_requests = []
|
|
talker._results_queue = []
|
|
talker._audio_queue = []
|
|
talker._deferred_cleanup_ids = set()
|
|
talker._beta_trace = False
|
|
return talker
|
|
|
|
|
|
def _add_state(talker, req_id: str, *, is_stopping: bool = False):
|
|
_, _RequestState = _dots_tts_talker_mod()
|
|
state = _RequestState(request_id=req_id)
|
|
state.prefill_completed = True
|
|
state.is_stopping = is_stopping
|
|
talker._active_states[req_id] = state
|
|
return state
|
|
|
|
|
|
def _hidden(bsz: int) -> torch.Tensor:
|
|
return torch.zeros(bsz, 1536, dtype=torch.float32)
|
|
|
|
|
|
class TestComputeLogitsPairing:
|
|
"""compute_logits pairs _results_queue entries with batch rows
|
|
positionally, so every request must push exactly one entry per step —
|
|
including prefills, which push a (req_id, None) placeholder."""
|
|
|
|
def test_mixed_batch_prefill_placeholder_keeps_rows_aligned(self) -> None:
|
|
talker = _make_bare_talker()
|
|
_add_state(talker, "prefill-req")
|
|
decode_state = _add_state(talker, "decode-req", is_stopping=True)
|
|
decode_state.precomputed_stop_logits = torch.tensor([[0.1, 0.9]])
|
|
|
|
# Row 0 is the prefill (None placeholder), row 1 the stopping decode.
|
|
talker._results_queue = [
|
|
("prefill-req", None),
|
|
("decode-req", torch.tensor([[0.1, 0.9]])),
|
|
]
|
|
logits = talker.compute_logits(_hidden(2))
|
|
|
|
assert logits.shape == (2, _VOCAB_SIZE)
|
|
# Prefill row: forced continue.
|
|
assert logits[0, 0].item() == 1.0
|
|
assert logits[0, 1].item() == float("-inf")
|
|
# Stopping decode row: stop slot wins.
|
|
assert logits[1, 1].item() == 1.0
|
|
assert logits[1, 0].item() == 0.0
|
|
# Everything else stays masked.
|
|
assert logits[0, 2].item() == float("-inf")
|
|
assert logits[1, 2].item() == float("-inf")
|
|
# Queue drained.
|
|
assert talker._results_queue == []
|
|
|
|
def test_below_threshold_stop_logits_force_continue(self) -> None:
|
|
"""Raw (continue, stop) softmax must not reach the sampler: a
|
|
prob_stop in (0.5, 0.8] would flip greedy argmax even though
|
|
is_stopping stayed False (the 0.8-threshold contract)."""
|
|
talker = _make_bare_talker()
|
|
_add_state(talker, "req", is_stopping=False)
|
|
talker._results_queue = [("req", torch.tensor([[0.3, 0.7]]))]
|
|
|
|
logits = talker.compute_logits(_hidden(1))
|
|
|
|
assert logits[0, 0].item() == 1.0
|
|
assert logits[0, 1].item() == float("-inf")
|
|
|
|
def test_empty_queue_defaults_to_all_continue(self) -> None:
|
|
talker = _make_bare_talker()
|
|
logits = talker.compute_logits(_hidden(3))
|
|
|
|
assert torch.all(logits[:, 0] == 1.0)
|
|
assert torch.all(logits[:, 1] == float("-inf"))
|
|
|
|
|
|
class TestStateEviction:
|
|
"""on_requests_finished fires BEFORE forward(), so eviction must be
|
|
deferred to _flush_deferred_cleanup at the end of the step — the
|
|
finishing request's audio still drains through the current forward."""
|
|
|
|
def test_finished_request_survives_until_flush(self) -> None:
|
|
talker = _make_bare_talker()
|
|
_add_state(talker, "req-a")
|
|
_add_state(talker, "req-b")
|
|
|
|
talker.on_requests_finished(["req-a"])
|
|
assert "req-a" in talker._active_states # not evicted yet
|
|
|
|
talker._flush_deferred_cleanup()
|
|
assert "req-a" not in talker._active_states
|
|
assert "req-b" in talker._active_states
|
|
assert talker._deferred_cleanup_ids == set()
|
|
|
|
def test_unknown_request_id_is_ignored(self) -> None:
|
|
talker = _make_bare_talker()
|
|
talker.on_requests_finished(["never-seen"])
|
|
talker._flush_deferred_cleanup()
|
|
assert talker._active_states == {}
|