[feat][Qwen3TTS] Simple dynamic TTFA based on Code2Wav load (#1714)
Signed-off-by: pablo <pablo@agigo.ai> Co-authored-by: pablo <pablo@agigo.ai>
This commit is contained in:
Juan Pablo Zuluagaandpablo authored and GitHub committed
1 parent
2fdc647ed7
commit
1d0a97f0d1
14 files changed
+221
-121
No files matched your search
@@ -7,16 +7,20 @@ from types import SimpleNamespace
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import (
|
||||
compute_dynamic_initial_chunk_size,
|
||||
max_ic_for_chunk_size,
|
||||
)
|
||||
from vllm_omni.model_executor.stage_input_processors.qwen3_tts import (
|
||||
talker2code2wav,
|
||||
talker2code2wav_async_chunk,
|
||||
)
|
||||
|
||||
_FRAME = [1, 2, 3, 4] # 4-codebook frame
|
||||
_Q = len(_FRAME) # num quantizers
|
||||
_FRAME = [1, 2, 3, 4]
|
||||
_Q = len(_FRAME)
|
||||
|
||||
|
||||
def _req(rid: str, *, finished: bool, initial_codec_chunk_frames: int | None = None):
|
||||
def _req(rid, *, finished, initial_codec_chunk_frames=None):
|
||||
ai = None
|
||||
if initial_codec_chunk_frames is not None:
|
||||
entry = SimpleNamespace(list_data=[initial_codec_chunk_frames])
|
||||
@@ -28,9 +32,10 @@ def _req(rid: str, *, finished: bool, initial_codec_chunk_frames: int | None = N
|
||||
)
|
||||
|
||||
|
||||
def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0):
|
||||
def _tm(*, chunk_frames=25, left_context=25, max_num_seqs=1):
|
||||
return SimpleNamespace(
|
||||
code_prompt_token_ids=defaultdict(list),
|
||||
scheduler_max_num_seqs=max_num_seqs,
|
||||
put_req_chunk=defaultdict(int),
|
||||
request_payload={},
|
||||
connector=SimpleNamespace(
|
||||
@@ -38,16 +43,14 @@ def _tm(*, chunk_frames=25, left_context=25, initial_chunk=0):
|
||||
"extra": {
|
||||
"codec_chunk_frames": chunk_frames,
|
||||
"codec_left_context_frames": left_context,
|
||||
"initial_codec_chunk_frames": initial_chunk,
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _call(tm, rid, *, n_frames, put_req=0, finished=False, req_ic=None):
|
||||
def _call(tm, rid, *, n_frames, finished=False, req_ic=None):
|
||||
tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(n_frames)]
|
||||
tm.put_req_chunk[rid] = put_req
|
||||
return talker2code2wav_async_chunk(
|
||||
transfer_manager=tm,
|
||||
pooling_output={"audio_codes": torch.zeros((0,))},
|
||||
@@ -56,84 +59,68 @@ def _call(tm, rid, *, n_frames, put_req=0, finished=False, req_ic=None):
|
||||
)
|
||||
|
||||
|
||||
def test_does_not_emit_empty_chunk_when_not_finished():
|
||||
def test_empty_returns_none():
|
||||
tm = _tm()
|
||||
request = _req("rid-empty", finished=False)
|
||||
|
||||
payload = talker2code2wav_async_chunk(
|
||||
p = talker2code2wav_async_chunk(
|
||||
transfer_manager=tm,
|
||||
pooling_output={"audio_codes": torch.zeros((0,))},
|
||||
request=request,
|
||||
request=_req("r", finished=False),
|
||||
)
|
||||
|
||||
assert payload is None
|
||||
assert p is None
|
||||
|
||||
|
||||
def test_flushes_tail_when_finished_without_pooler_output():
|
||||
def test_eof_marker_when_finished_empty():
|
||||
tm = _tm()
|
||||
rid = "rid-tail"
|
||||
tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(24)]
|
||||
request = _req(rid, finished=True)
|
||||
|
||||
payload = talker2code2wav_async_chunk(
|
||||
p = talker2code2wav_async_chunk(
|
||||
transfer_manager=tm,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
request=_req("r", finished=True),
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert payload is not None
|
||||
assert payload["finished"].item() is True
|
||||
assert len(payload["code_predictor_codes"]) == _Q * 24
|
||||
assert p == {"code_predictor_codes": [], "finished": torch.tensor(True, dtype=torch.bool)}
|
||||
|
||||
|
||||
def test_emits_eof_marker_when_finished_with_no_frames():
|
||||
def test_flush_on_finish():
|
||||
tm = _tm()
|
||||
request = _req("rid-eof", finished=True)
|
||||
|
||||
payload = talker2code2wav_async_chunk(
|
||||
tm.code_prompt_token_ids["r"] = [_FRAME[:] for _ in range(24)]
|
||||
p = talker2code2wav_async_chunk(
|
||||
transfer_manager=tm,
|
||||
pooling_output=None,
|
||||
request=request,
|
||||
request=_req("r", finished=True),
|
||||
is_finished=True,
|
||||
)
|
||||
|
||||
assert payload == {
|
||||
"code_predictor_codes": [],
|
||||
"finished": torch.tensor(True, dtype=torch.bool),
|
||||
}
|
||||
assert p is not None
|
||||
assert p["finished"].item() is True
|
||||
assert len(p["code_predictor_codes"]) == _Q * 24
|
||||
|
||||
|
||||
_CASES = [
|
||||
# Normal path (initial=0): emit at chunk_size boundaries
|
||||
((25, 25, 0), (24, 0, False), None),
|
||||
((25, 25, 0), (25, 0, False), (0, 25)),
|
||||
# Initial-chunk phase: hold, first emit, second emit
|
||||
((25, 25, 10), (9, 0, False), None),
|
||||
((25, 25, 10), (10, 0, False), (0, 10)),
|
||||
((25, 25, 10), (20, 1, False), (10, 20)),
|
||||
# Non-divisible: holds at chunk boundary
|
||||
((25, 25, 12), (25, 2, False), None),
|
||||
# Normal phase: offset by initial_coverage (chunk//initial * initial)
|
||||
((25, 25, 10), (45, 2, False), (20, 45)),
|
||||
# Second normal emit (offset must stay stable)
|
||||
((25, 25, 10), (70, 3, False), (25, 50)),
|
||||
# initial >= chunk clamps to chunk_size (behaves as normal)
|
||||
((25, 25, 30), (25, 0, False), (0, 25)),
|
||||
# finished=True flushes IC tail
|
||||
((25, 25, 10), (5, 0, True), (0, 5)),
|
||||
# finished=True flushes non-divisible IC residual
|
||||
((25, 25, 12), (25, 2, True), (24, 25)),
|
||||
# finished=True flushes normal phase tail
|
||||
((25, 25, 10), (30, 2, True), (20, 30)),
|
||||
# Dynamic IC=16, cs=25, initial_coverage=16
|
||||
((25, 25, 0), 24, False, None), # IC phase: 24%16!=0 -> hold
|
||||
((25, 25, 0), 25, False, None), # transition: adjusted=9, hold (no replay)
|
||||
((25, 25, 0), 41, False, (16, 41)), # first normal emit, lc=16
|
||||
# Per-request IC=10, cs=25, initial_coverage=20
|
||||
((25, 25, 10), 9, False, None), # IC: hold
|
||||
((25, 25, 10), 10, False, (0, 10)), # IC: emit at boundary
|
||||
((25, 25, 10), 25, False, None), # transition: hold (no replay)
|
||||
((25, 25, 10), 45, False, (20, 45)), # first normal emit, lc=20
|
||||
((25, 25, 10), 5, True, (0, 5)), # finished flushes IC tail
|
||||
((25, 25, 10), 33, True, (20, 33)), # finished flushes normal tail
|
||||
# IC=8, cs=16: IC divides chunk_size evenly (edge case)
|
||||
((16, 25, 8), 8, False, (0, 8)), # IC: emit
|
||||
((16, 25, 8), 16, False, None), # transition: hold (no replay)
|
||||
((16, 25, 8), 24, False, (8, 24)), # first normal emit, lc=8
|
||||
# Per-request override: IC=15 at n_frames=10 -> 10%15!=0 -> hold
|
||||
((25, 25, 15), 10, False, None),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("config, state, expected", _CASES)
|
||||
def test_streaming_decoding_with_variable_initial(config, state, expected):
|
||||
chunk_frames, left_context, initial_chunk = config
|
||||
n_frames, put_req, finished = state
|
||||
|
||||
tm = _tm(chunk_frames=chunk_frames, left_context=left_context, initial_chunk=initial_chunk)
|
||||
payload = _call(tm, "r", n_frames=n_frames, put_req=put_req, finished=finished)
|
||||
@pytest.mark.parametrize("config, n_frames, finished, expected", _CASES)
|
||||
def test_streaming_phases(config, n_frames, finished, expected):
|
||||
chunk_frames, left_context, req_ic_val = config
|
||||
tm = _tm(chunk_frames=chunk_frames, left_context=left_context)
|
||||
req_ic = req_ic_val if req_ic_val > 0 else None
|
||||
payload = _call(tm, "r", n_frames=n_frames, finished=finished, req_ic=req_ic)
|
||||
|
||||
if expected is None:
|
||||
assert payload is None
|
||||
@@ -144,22 +131,86 @@ def test_streaming_decoding_with_variable_initial(config, state, expected):
|
||||
assert len(payload["code_predictor_codes"]) == _Q * exp_window
|
||||
|
||||
|
||||
def test_per_request_override_activates_initial_phase():
|
||||
tm = _tm(initial_chunk=0)
|
||||
payload = _call(tm, "r-override", n_frames=10, req_ic=10)
|
||||
assert payload is not None
|
||||
assert payload["left_context_size"] == 0
|
||||
assert len(payload["code_predictor_codes"]) == _Q * 10
|
||||
def test_dynamic_ic_adapts_to_load():
|
||||
# chunk_size=25 -> max_ic=16, steps=[2,4,8,16]
|
||||
tm = _tm(max_num_seqs=8)
|
||||
|
||||
# Low load (1/8) -> IC=2 -> emit at 2
|
||||
p1 = _call(tm, "r", n_frames=2)
|
||||
assert p1 is not None
|
||||
assert len(p1["code_predictor_codes"]) == _Q * 2
|
||||
|
||||
# High load: add 4 others -> active=5/8 -> IC=8 -> emit at 8
|
||||
for i in range(4):
|
||||
tm.code_prompt_token_ids[f"other-{i}"] = [[0]]
|
||||
p2 = _call(tm, "r", n_frames=8)
|
||||
assert p2 is not None
|
||||
assert len(p2["code_predictor_codes"]) == _Q * 8
|
||||
|
||||
# Requests past initial phase still count in load factor
|
||||
tm2 = _tm(max_num_seqs=4)
|
||||
for i in range(3):
|
||||
tm2.code_prompt_token_ids[f"long-{i}"] = [[0]] * 50 # well past cs=25
|
||||
# active=4/4=1.0 -> IC=16
|
||||
p3 = _call(tm2, "new", n_frames=16)
|
||||
assert p3 is not None
|
||||
assert len(p3["code_predictor_codes"]) == _Q * 16
|
||||
|
||||
|
||||
def test_per_request_override_wins_over_stage_config():
|
||||
tm = _tm(initial_chunk=5)
|
||||
payload = _call(tm, "r-override2", n_frames=10, put_req=0, req_ic=15)
|
||||
assert payload is None
|
||||
def test_ic_load_change_mid_request():
|
||||
"""IC stateless: load spike mid-request shifts initial_coverage."""
|
||||
tm = _tm(chunk_frames=25, left_context=25, max_num_seqs=8)
|
||||
|
||||
# Low load -> IC=2 -> emit at frame 2
|
||||
p1 = _call(tm, "r", n_frames=2)
|
||||
assert p1 is not None
|
||||
|
||||
# Spike load: 6 others -> IC=16 -> initial_coverage=16
|
||||
for i in range(6):
|
||||
tm.code_prompt_token_ids[f"other-{i}"] = [[0]] * 10
|
||||
|
||||
# adjusted=25-16=9, 9%25!=0 -> hold
|
||||
assert _call(tm, "r", n_frames=25) is None
|
||||
|
||||
# First normal emit at 16+25=41
|
||||
p3 = _call(tm, "r", n_frames=41)
|
||||
assert p3 is not None
|
||||
assert p3["left_context_size"] == 16
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"active,max_bs,max_ic,expected",
|
||||
[
|
||||
(0, 4, 32, 2), # zero load -> min step
|
||||
(2, 4, 32, 8), # mid load
|
||||
(4, 4, 32, 32), # full load
|
||||
(10, 4, 16, 16), # over capacity, capped
|
||||
(0, 4, 1, 1), # max_ic below min step
|
||||
(0, 0, 16, 2), # zero capacity edge case
|
||||
],
|
||||
)
|
||||
def test_compute_dynamic_initial_chunk_size(active, max_bs, max_ic, expected):
|
||||
assert compute_dynamic_initial_chunk_size(active, max_bs, max_ic) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"chunk_size,expected",
|
||||
[
|
||||
(25, 16),
|
||||
(50, 32),
|
||||
(70, 64),
|
||||
(8, 4),
|
||||
(4, 2),
|
||||
(2, 1),
|
||||
(1, 1),
|
||||
],
|
||||
)
|
||||
def test_max_ic_for_chunk_size(chunk_size, expected):
|
||||
assert max_ic_for_chunk_size(chunk_size) == expected
|
||||
|
||||
|
||||
def test_first_streaming_chunk_prepends_ref_code_context():
|
||||
tm = _tm(initial_chunk=10)
|
||||
tm = _tm()
|
||||
rid = "r-ref"
|
||||
tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(10)]
|
||||
ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
|
||||
@@ -177,7 +228,7 @@ def test_first_streaming_chunk_prepends_ref_code_context():
|
||||
|
||||
|
||||
def test_ref_code_context_only_applies_to_first_streaming_chunk():
|
||||
tm = _tm(initial_chunk=10)
|
||||
tm = _tm()
|
||||
rid = "r-ref2"
|
||||
tm.code_prompt_token_ids[rid] = [_FRAME[:] for _ in range(20)]
|
||||
tm.put_req_chunk[rid] = 1
|
||||
@@ -196,7 +247,7 @@ def test_ref_code_context_only_applies_to_first_streaming_chunk():
|
||||
|
||||
|
||||
def test_ref_code_context_can_be_buffered_before_first_emit():
|
||||
tm = _tm(initial_chunk=10)
|
||||
tm = _tm()
|
||||
rid = "r-ref-buffered"
|
||||
ref_code = torch.tensor([[9, 9, 9, 9], [8, 8, 8, 8]], dtype=torch.long)
|
||||
|
||||
|
||||
Reference in new issue
Block a user