[Model] Optimize Qwen3-Omni MRv2 handoff, code prediction and audio decoding (#8223)
Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
35 files changed
+2363
-161
No files matched your search
@@ -0,0 +1,90 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.cuda,
|
||||
pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
|
||||
]
|
||||
|
||||
|
||||
def _reference(logits, uniforms, top_k, top_p):
|
||||
if top_k > 0:
|
||||
threshold = logits.topk(top_k, dim=-1).values[:, -1:]
|
||||
logits = logits.masked_fill(logits < threshold, -torch.inf)
|
||||
if top_p < 1:
|
||||
values, indices = logits.sort(dim=-1, descending=True)
|
||||
probs = values.softmax(dim=-1, dtype=torch.float32)
|
||||
values[(probs.cumsum(-1) - probs) >= top_p] = -torch.inf
|
||||
logits = values.scatter(1, indices, values)
|
||||
return (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch,vocab", [(1, 128), (7, 2048), (64, 2048)])
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("top_k,top_p", [(-1, 0.8), (0, 1.0), (50, 0.8), (50, 0.01)])
|
||||
def test_fixed_uniform_sampling_and_graph(batch, vocab, dtype, top_k, top_p):
|
||||
from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
|
||||
|
||||
torch.manual_seed(173)
|
||||
logits = torch.randn(batch, vocab, device="cuda", dtype=dtype)
|
||||
# The MTP caller supplies a noncontiguous row slice of [B, 15, vocab].
|
||||
uniforms = torch.rand(batch, 15, vocab, device="cuda")[:, 3].clamp_min_(1e-7)
|
||||
expected = _reference(logits, uniforms, top_k, top_p)
|
||||
actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
|
||||
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
captured = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
|
||||
logits.add_(torch.randn_like(logits) * 0.1)
|
||||
uniforms.uniform_(1e-7, 1 - 1e-7)
|
||||
graph.replay()
|
||||
torch.testing.assert_close(captured, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("top_k,top_p", [(3, 0.5), (3, 1.0), (0, 0.8)])
|
||||
def test_threshold_ties_and_degenerate_rows(top_k, top_p):
|
||||
from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
|
||||
|
||||
logits = torch.tensor(
|
||||
[[2, 2, 2, 2, 2, 1, 0, -1], [1] * 8, [-torch.inf] * 8, [torch.nan, 1, 2, 3, 4, 5, 6, 7]],
|
||||
device="cuda",
|
||||
dtype=torch.float32,
|
||||
)
|
||||
uniforms = torch.tensor([[0.4, 0.7, 0.8, 0.9, 0.2, 0.3, 0.6, 0.5]], device="cuda").expand(4, -1)
|
||||
actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=top_k, top_p=top_p)
|
||||
torch.testing.assert_close(actual, _reference(logits, uniforms, top_k, top_p), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
|
||||
def test_native_nucleus_boundaries(dtype):
|
||||
from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
|
||||
|
||||
torch.manual_seed(1923)
|
||||
logits = torch.randn(1, 2048, device="cuda", dtype=dtype)
|
||||
values, indices = logits.sort(descending=True)
|
||||
probabilities = values.softmax(-1, dtype=torch.float32)
|
||||
before = probabilities.cumsum(-1) - probabilities
|
||||
for rank in [1, 19, 49, 1000, 2000]:
|
||||
boundary = before[0, rank]
|
||||
uniforms = torch.full_like(logits, 0.001, dtype=torch.float32)
|
||||
uniforms[0, indices[0, rank]] = 1 - 1e-7
|
||||
for direction in [-torch.inf, torch.inf]:
|
||||
top_p = torch.nextafter(boundary, boundary.new_tensor(direction)).item()
|
||||
actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=0, top_p=top_p)
|
||||
torch.testing.assert_close(actual, _reference(logits, uniforms, 0, top_p), rtol=0, atol=0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("top_p", [0.8, 1.0])
|
||||
def test_strided_logits_and_non_power_of_two_vocabulary(top_p):
|
||||
from vllm_omni.model_executor.models.common.fused_sampling import sample_top_k_top_p_gumbel
|
||||
|
||||
for seed in range(16):
|
||||
torch.manual_seed(seed)
|
||||
logits = torch.randn(128, 2050, device="cuda")[:, ::2]
|
||||
uniforms = torch.rand(128, 3, 1025, device="cuda")[:, 1].clamp_min_(1e-7)
|
||||
actual = sample_top_k_top_p_gumbel(logits, uniforms, top_k=-1, top_p=top_p)
|
||||
torch.testing.assert_close(actual, _reference(logits, uniforms, -1, top_p), rtol=0, atol=0)
|
||||
+1
-1
@@ -16,7 +16,7 @@ pytestmark = [
|
||||
@pytest.mark.parametrize("heads,kv_heads,dim", [(16, 8, 128), (8, 8, 64)])
|
||||
@pytest.mark.parametrize("positions", [[0, 1], [2], [8], [15]])
|
||||
def test_short_attention_causal_gqa_and_graph(batch, heads, kv_heads, dim, positions):
|
||||
from vllm_omni.model_executor.models.qwen3_tts.short_kv_attention import short_kv_attention
|
||||
from vllm_omni.model_executor.models.common.short_kv_attention import short_kv_attention
|
||||
|
||||
with torch.inference_mode():
|
||||
torch.manual_seed(42)
|
||||
@@ -0,0 +1,32 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cuda]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
@pytest.mark.parametrize("rows", [16, 19])
|
||||
@torch.inference_mode()
|
||||
def test_first_frame_graph_outputs_own_storage(dtype, rows):
|
||||
class Decoder(torch.nn.Linear):
|
||||
config = SimpleNamespace(num_quantizers=1)
|
||||
|
||||
def forward(self, codes):
|
||||
return codes.to(dtype) * self.weight
|
||||
|
||||
model = Decoder(1, 1, bias=False, device="cuda", dtype=dtype)
|
||||
model.weight.fill_(1)
|
||||
decoder = Qwen3OmniFirstFrameDecoder(model, sample_rate=24000)
|
||||
decoder.capture()
|
||||
codes = torch.arange(rows, device="cuda").reshape(-1, 1)
|
||||
first = decoder.decode(codes)
|
||||
second = decoder.decode(codes + 32)
|
||||
assert first.dtype == second.dtype == torch.float32
|
||||
torch.testing.assert_close(first, codes.float(), rtol=0, atol=0)
|
||||
torch.testing.assert_close(second, (codes + 32).float(), rtol=0, atol=0)
|
||||
@@ -0,0 +1,87 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
|
||||
|
||||
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.core_model,
|
||||
pytest.mark.cuda,
|
||||
pytest.mark.skipif(not torch.cuda.is_available() or torch.version.hip is not None, reason="requires CUDA"),
|
||||
]
|
||||
|
||||
|
||||
def _codec(enabled):
|
||||
config = Qwen3OmniMoeCode2WavConfig(
|
||||
hidden_size=32,
|
||||
intermediate_size=64,
|
||||
num_hidden_layers=1,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=4,
|
||||
head_dim=8,
|
||||
decoder_dim=32,
|
||||
upsample_rates=[2, 2],
|
||||
upsampling_ratios=[2],
|
||||
codebook_size=16,
|
||||
num_quantizers=2,
|
||||
)
|
||||
config._attn_implementation = "sdpa"
|
||||
vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
hf_config=config, stage_connector_config={"extra": {"codec_cudnn_benchmark": enabled}}
|
||||
)
|
||||
)
|
||||
model = Qwen3OmniMoeCode2Wav(vllm_config=vllm_config).cuda().eval()
|
||||
model.precompute_snake_caches()
|
||||
return model
|
||||
|
||||
|
||||
def _flags():
|
||||
c = torch.backends.cudnn
|
||||
return c.enabled, c.benchmark, c.benchmark_limit, c.deterministic, c.allow_tf32
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_autotuned_codec_graph_replay_restores_flags():
|
||||
torch.manual_seed(21)
|
||||
model = _codec(True)
|
||||
codes = torch.randint(0, 16, (2, 2, 8), device="cuda")
|
||||
original = _flags()
|
||||
for _ in range(3):
|
||||
model(codes)
|
||||
assert _flags() == original
|
||||
graph = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(graph):
|
||||
captured = model(codes)
|
||||
assert _flags() == original
|
||||
codes.random_(0, 16)
|
||||
graph.replay()
|
||||
expected = model(codes)
|
||||
torch.testing.assert_close(captured, expected)
|
||||
assert _flags() == original
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enabled", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_codec_failure_restores_cudnn_settings(mocker, enabled):
|
||||
model = _codec(enabled)
|
||||
codes = torch.randint(0, 16, (1, 2, 8), device="cuda")
|
||||
original = _flags()
|
||||
|
||||
def fail(hidden):
|
||||
if enabled:
|
||||
assert torch.backends.cudnn.benchmark
|
||||
assert torch.backends.cudnn.benchmark_limit == 10
|
||||
else:
|
||||
assert _flags() == original
|
||||
raise RuntimeError("decode failed")
|
||||
|
||||
mocker.patch.object(model, "_decode_waveform", side_effect=fail)
|
||||
with pytest.raises(RuntimeError, match="decode failed"):
|
||||
model(codes)
|
||||
assert _flags() == original
|
||||
@@ -0,0 +1,88 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
import copy
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
|
||||
from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (
|
||||
Qwen3OmniMoeCode2WavDecoderBlock,
|
||||
Qwen3OmniMoeSnakeBeta,
|
||||
)
|
||||
|
||||
from vllm_omni.model_executor.models.common.snake_activation import SnakeBeta
|
||||
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import Qwen3OmniMoeCode2Wav, use_fused_snake
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _blocks():
|
||||
config = Qwen3OmniMoeCode2WavConfig(decoder_dim=32, upsample_rates=[2, 2])
|
||||
torch.manual_seed(0)
|
||||
blocks = torch.nn.ModuleList(Qwen3OmniMoeCode2WavDecoderBlock(config, i) for i in range(2)).eval()
|
||||
with torch.no_grad():
|
||||
for module in blocks.modules():
|
||||
if isinstance(module, Qwen3OmniMoeSnakeBeta):
|
||||
module.alpha.uniform_(-0.5, 0.5)
|
||||
module.beta.uniform_(-0.5, 0.5)
|
||||
return blocks
|
||||
|
||||
|
||||
def test_decoder_blocks_swap_every_hf_snake_and_keep_weight_names():
|
||||
blocks = _blocks()
|
||||
hf_count = sum(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
|
||||
keys = set(blocks.state_dict())
|
||||
|
||||
assert use_fused_snake(blocks) == hf_count > 0
|
||||
assert not any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in blocks.modules())
|
||||
assert sum(isinstance(m, SnakeBeta) for m in blocks.modules()) == hf_count
|
||||
assert set(blocks.state_dict()) == keys # checkpoint weights load unchanged
|
||||
|
||||
|
||||
def test_fused_snake_blocks_match_hf_blocks():
|
||||
reference = _blocks()
|
||||
fused = copy.deepcopy(reference)
|
||||
use_fused_snake(fused)
|
||||
fused.load_state_dict(reference.state_dict())
|
||||
x = torch.randn(2, 32, 20)
|
||||
with torch.no_grad():
|
||||
expected = x
|
||||
actual = x
|
||||
for ref_block, fused_block in zip(reference, fused):
|
||||
expected = ref_block(expected)
|
||||
actual = fused_block(actual)
|
||||
torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("option", "device", "cuda", "fused"),
|
||||
[
|
||||
(False, "cuda", True, False),
|
||||
(True, "cuda", True, True),
|
||||
(True, "cpu", True, False),
|
||||
(True, "cuda", False, False),
|
||||
],
|
||||
)
|
||||
def test_decoder_block_fusion_requires_cuda_opt_in(monkeypatch, option, device, cuda, fused):
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: cuda)
|
||||
config = Qwen3OmniMoeCode2WavConfig(
|
||||
hidden_size=32,
|
||||
intermediate_size=64,
|
||||
num_hidden_layers=1,
|
||||
num_attention_heads=2,
|
||||
num_key_value_heads=2,
|
||||
decoder_dim=32,
|
||||
)
|
||||
model = Qwen3OmniMoeCode2Wav(
|
||||
vllm_config=SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
hf_config=config, stage_connector_config={"extra": {"codec_fused_snake": option}}
|
||||
),
|
||||
device_config=SimpleNamespace(device=device),
|
||||
)
|
||||
)
|
||||
assert any(isinstance(m, Qwen3OmniMoeSnakeBeta) for m in model.decoder.modules()) is (not fused)
|
||||
@@ -0,0 +1,208 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeCode2WavConfig
|
||||
|
||||
from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import Qwen3OmniFirstFrameDecoder
|
||||
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
|
||||
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav import (
|
||||
Qwen3OmniMoeCode2Wav,
|
||||
plan_decode_groups,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_Q = 4
|
||||
_CODEBOOK = 16
|
||||
|
||||
|
||||
def _code2wav(quantizers=_Q) -> Qwen3OmniMoeCode2Wav:
|
||||
torch.manual_seed(0)
|
||||
config = Qwen3OmniMoeCode2WavConfig(
|
||||
codebook_size=_CODEBOOK,
|
||||
num_quantizers=quantizers,
|
||||
hidden_size=32,
|
||||
intermediate_size=64,
|
||||
num_hidden_layers=2,
|
||||
num_attention_heads=2,
|
||||
num_key_value_heads=2,
|
||||
decoder_dim=32,
|
||||
sliding_window=8,
|
||||
)
|
||||
model = Qwen3OmniMoeCode2Wav(vllm_config=SimpleNamespace(model_config=SimpleNamespace(hf_config=config)))
|
||||
return model.eval()
|
||||
|
||||
|
||||
def test_first_frame_pcm_equals_code2wav_streaming_chunk0():
|
||||
code2wav = _code2wav()
|
||||
decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
|
||||
codes = torch.randint(0, _CODEBOOK, (3, _Q))
|
||||
with torch.inference_mode():
|
||||
for row in range(codes.shape[0]):
|
||||
pcm = decoder.decode(codes[row : row + 1])
|
||||
chunk0 = code2wav.chunked_decode_streaming(
|
||||
codes[row].reshape(1, _Q, 1), left_context_size=[0], seq_token_counts=[_Q]
|
||||
)[0]
|
||||
assert pcm.dtype == torch.float32
|
||||
assert torch.equal(pcm[0], chunk0.reshape(-1).float())
|
||||
assert 0 < pcm.shape[-1] < int(code2wav.total_upsample)
|
||||
batched = decoder.decode(codes)
|
||||
assert batched.shape == (3, pcm.shape[-1])
|
||||
|
||||
|
||||
def test_first_frame_is_a_prefix_of_a_longer_first_chunk():
|
||||
code2wav = _code2wav()
|
||||
decoder = Qwen3OmniFirstFrameDecoder(code2wav, sample_rate=24000)
|
||||
codes = torch.randint(0, _CODEBOOK, (_Q, 4))
|
||||
with torch.inference_mode():
|
||||
first = decoder.decode(codes[:, :1].T)[0]
|
||||
chunk0 = code2wav.chunked_decode_streaming(
|
||||
codes.reshape(1, _Q, 4), left_context_size=[0], seq_token_counts=[4 * _Q]
|
||||
)[0]
|
||||
torch.testing.assert_close(chunk0.reshape(-1)[: first.numel()].float(), first)
|
||||
|
||||
|
||||
def test_chunk_ramp_adds_exact_code2wav_graph_sizes(monkeypatch):
|
||||
import vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper as wrapper_module
|
||||
|
||||
real = wrapper_module.CUDAGraphDecoderWrapper
|
||||
wrapper = MagicMock(compute_capture_sizes=real.compute_capture_sizes)
|
||||
monkeypatch.setattr(wrapper_module, "CUDAGraphDecoderWrapper", wrapper)
|
||||
extra = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "codec_chunk_ramp": [1, 2, 4, 8, 16, 25]}
|
||||
code2wav = SimpleNamespace(config=SimpleNamespace(num_quantizers=_Q))
|
||||
code2wav.enable_cudagraph = lambda **kw: Qwen3OmniMoeCode2Wav.enable_cudagraph(
|
||||
code2wav, device=torch.device("cuda"), **kw
|
||||
)
|
||||
model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.code2wav = code2wav
|
||||
model.vllm_config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(async_chunk=True, enforce_eager=False, stage_connector_config={"extra": extra})
|
||||
)
|
||||
for batches in ([], [2, 4, 8]):
|
||||
extra["codec_graph_batch_sizes"] = batches
|
||||
Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
|
||||
args = wrapper.call_args.kwargs
|
||||
sizes = args["capture_sizes"]
|
||||
assert {1, 3, 7, 15, 31, 50} <= set(sizes)
|
||||
assert set(real.compute_capture_sizes(codec_chunk_frames=25, codec_left_context_frames=25)) <= set(sizes)
|
||||
assert set(args["extra_capture_shapes"]) == {
|
||||
(b, size) for b in batches for size in sizes if size <= 50 and b * size <= 325
|
||||
}
|
||||
assert code2wav._streaming_batch_sizes[50] == [1, 2, 4]
|
||||
extra.pop("codec_chunk_ramp")
|
||||
extra.pop("codec_graph_batch_sizes")
|
||||
Qwen3OmniMoeForConditionalGeneration._maybe_enable_code2wav_cudagraph(model)
|
||||
assert wrapper.call_args.kwargs["capture_sizes"] is None
|
||||
|
||||
|
||||
_SIZES = [1, 2, 3, 4, 7, 8, 15, 16, 25, 31, 32, 50, 64]
|
||||
|
||||
|
||||
def _bucket(length):
|
||||
return next((size for size in _SIZES if size >= length), None)
|
||||
|
||||
|
||||
def test_decode_groups_keep_a_uniform_batch_whole_and_split_by_the_largest_graph():
|
||||
assert plan_decode_groups([50] * 6, _bucket, lambda size: [1, 2, 4, 8]) == [(list(range(6)), 50)]
|
||||
calls = plan_decode_groups([25] * 20, _bucket, lambda size: [1, 2, 4, 8, 16])
|
||||
assert [len(rows) for rows, _length in calls] == [16, 4]
|
||||
calls = plan_decode_groups([1] * 8 + [50] * 8, _bucket, lambda size: [1, 2, 4, 8] if size < 50 else [1, 2, 4])
|
||||
assert [(len(rows), length) for rows, length in calls] == [(8, 1), (4, 50), (4, 50)]
|
||||
assert plan_decode_groups([], _bucket, lambda size: [1, 2]) == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("frames", [65, 129])
|
||||
def test_decode_groups_beyond_largest_graph_use_eager(monkeypatch, frames):
|
||||
from vllm_omni.model_executor.models.qwen3_tts.cuda_graph_decoder_wrapper import CUDAGraphDecoderWrapper
|
||||
|
||||
wrapper = CUDAGraphDecoderWrapper.__new__(CUDAGraphDecoderWrapper)
|
||||
wrapper.enabled = wrapper._warmed_up = True
|
||||
wrapper._bucket_sizes = [8, 32, 64]
|
||||
wrapper._compiled_shapes = {(1, 64), (8, 64)}
|
||||
wrapper._compiled_graphs = wrapper.graphs = dict.fromkeys(wrapper._compiled_shapes, MagicMock())
|
||||
wrapper.decoder = MagicMock(side_effect=lambda codes: codes.float() * 2)
|
||||
monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False)
|
||||
groups = plan_decode_groups([frames, frames], wrapper._get_padded_size, lambda _: (1,))
|
||||
assert groups == [([0], frames), ([1], frames)]
|
||||
codes = torch.arange(2 * frames).reshape(2, 1, frames)
|
||||
for rows, width in groups:
|
||||
actual = wrapper.decode(codes[rows, :, :width])
|
||||
torch.testing.assert_close(actual, codes[rows].float() * 2, rtol=0, atol=0)
|
||||
assert wrapper.decoder.call_count == 2
|
||||
for graph in (*wrapper.graphs.values(), *wrapper._compiled_graphs.values()):
|
||||
graph.replay.assert_not_called()
|
||||
|
||||
|
||||
class _PaddingGraphs:
|
||||
def __init__(self, model, batch_sizes):
|
||||
self.model = model
|
||||
self.batch_sizes = sorted(batch_sizes)
|
||||
self.calls = []
|
||||
|
||||
def _get_padded_size(self, length):
|
||||
return _bucket(length)
|
||||
|
||||
def decode(self, codes):
|
||||
rows, _q, frames = codes.shape
|
||||
size = _bucket(frames)
|
||||
batch = next(b for b in self.batch_sizes if b >= rows)
|
||||
self.calls.append((rows, frames))
|
||||
padded = torch.zeros(batch, codes.shape[1], size, dtype=codes.dtype)
|
||||
padded[:rows, :, :frames] = codes
|
||||
out = self.model(padded)
|
||||
return out[:rows, :, : out.shape[-1] - (size - frames) * int(self.model.total_upsample)].clone()
|
||||
|
||||
|
||||
def test_grouped_streaming_decode_matches_one_padded_batch():
|
||||
code2wav = _code2wav()
|
||||
lengths = [1, 3, 50, 7, 31]
|
||||
lefts = [0, 1, 25, 3, 15]
|
||||
codes = torch.zeros(len(lengths), _Q, max(lengths), dtype=torch.long)
|
||||
for row, length in enumerate(lengths):
|
||||
codes[row, :, :length] = torch.randint(0, _CODEBOOK, (_Q, length))
|
||||
counts = [length * _Q for length in lengths]
|
||||
with torch.inference_mode():
|
||||
reference = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
|
||||
graphs = _PaddingGraphs(code2wav, [1, 2, 4])
|
||||
code2wav._cudagraph_wrapper, code2wav._cudagraph_enabled = graphs, True
|
||||
code2wav._streaming_batch_sizes = {size: [1, 2, 4] for size in _SIZES}
|
||||
grouped = code2wav.chunked_decode_streaming(codes, left_context_size=lefts, seq_token_counts=counts)
|
||||
|
||||
for want, got in zip(reference, grouped, strict=True):
|
||||
assert want.shape == got.shape
|
||||
torch.testing.assert_close(got, want, rtol=1e-4, atol=1e-5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("frames", "quantizers"), [(1, 16), (4, 16), (4, 4)])
|
||||
def test_delivered_first_frame_is_removed_from_code2wav_output(frames, quantizers):
|
||||
from vllm_omni.data_entry_keys import FIRST_AUDIO_REQUIRED_KEY
|
||||
|
||||
codec = _code2wav(quantizers)
|
||||
model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
|
||||
torch.nn.Module.__init__(model)
|
||||
model.model_stage, model.code2wav, model.code2wav_config = "code2wav", codec, codec.config
|
||||
model.generate_audio = lambda codes, left, counts: codec.chunked_decode_streaming(
|
||||
codes, left_context_size=left, seq_token_counts=counts
|
||||
)
|
||||
codes = torch.randint(0, _CODEBOOK, (2, quantizers, frames))
|
||||
with torch.inference_mode():
|
||||
reference = model.generate_audio(codes, [0, 0], [quantizers * frames] * 2)
|
||||
first = Qwen3OmniFirstFrameDecoder(codec, sample_rate=24000).decode(codes[0, :, :1].T)[0]
|
||||
out = model.forward(
|
||||
input_ids=codes.flatten(),
|
||||
positions=None,
|
||||
seq_token_counts=[quantizers * frames] * 2,
|
||||
runtime_additional_information=[
|
||||
{"meta": {"left_context_size": 0, "first_audio": True}},
|
||||
{"meta": {"left_context_size": 0, "first_audio": False}},
|
||||
],
|
||||
).multimodal_outputs
|
||||
torch.testing.assert_close(torch.cat([first, out["model_outputs"][0].flatten()]), reference[0].flatten())
|
||||
torch.testing.assert_close(out["model_outputs"][1], reference[1].reshape(1, -1))
|
||||
assert [bool(flag) for flag in out[FIRST_AUDIO_REQUIRED_KEY]] == [True, False]
|
||||
@@ -0,0 +1,116 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni import Qwen3OmniMoeForConditionalGeneration
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_H = 4
|
||||
_Q = 3
|
||||
|
||||
|
||||
class _Talker(nn.Module):
|
||||
num_code_groups = _Q
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.text_projection = nn.Identity()
|
||||
self.code_predictor = SimpleNamespace(
|
||||
_top_k=50,
|
||||
_top_p=0.8,
|
||||
set_sampling_params=lambda **kwargs: setattr(self, "sampling", kwargs),
|
||||
)
|
||||
|
||||
|
||||
def _talker(*, async_chunk: bool = True) -> Qwen3OmniMoeForConditionalGeneration:
|
||||
model = object.__new__(Qwen3OmniMoeForConditionalGeneration)
|
||||
nn.Module.__init__(model)
|
||||
model.model_stage = "talker"
|
||||
model.talker = _Talker()
|
||||
model.vllm_config = SimpleNamespace(model_config=SimpleNamespace(async_chunk=async_chunk, use_v2_model_runner=True))
|
||||
model.tts_eos_embed = torch.full((_H,), -1.0)
|
||||
model.tts_pad_embed = torch.full((_H,), -2.0)
|
||||
model._codec_codebook_size = 2048
|
||||
return model
|
||||
|
||||
|
||||
def _rows(*values: float) -> torch.Tensor:
|
||||
return torch.tensor([[value] * _H for value in values])
|
||||
|
||||
|
||||
def _step(model, payload: dict) -> torch.Tensor:
|
||||
update: dict = {}
|
||||
text_step = model._thinker_decode_to_talker_decode(payload, torch.device("cpu"), update)
|
||||
for key, value in update.items():
|
||||
payload.setdefault(key, {}).update(value)
|
||||
return text_step.reshape(-1)
|
||||
|
||||
|
||||
def test_mtp_frame_valid_rejects_codec_eos_and_special_ids():
|
||||
layer0 = torch.tensor([0, 2047, 2048, 2150, -1])
|
||||
assert _talker().mtp_frame_valid(layer0).tolist() == [True, True, False, False, False]
|
||||
|
||||
|
||||
def test_decode_rows_are_consumed_one_per_step_in_arrival_order():
|
||||
model = _talker()
|
||||
# Two rows were already pending (cached at prefill); two more arrive at once.
|
||||
payload = {"embed": {"cached_decode": _rows(1, 2), "decode": _rows(3, 4)}, "meta": {}}
|
||||
steps = [_step(model, payload)[0].item()]
|
||||
for arrival in (None, _rows(5), None, None, None, None):
|
||||
payload["embed"]["decode"] = arrival
|
||||
steps.append(_step(model, payload)[0].item())
|
||||
assert steps == [1, 2, 3, 4, 5, -1, -2]
|
||||
assert payload["embed"]["cached_decode"] is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("v2", "resumable"), [(True, True), (False, False), (False, True)])
|
||||
def test_v1_and_resumable_requests_keep_the_indexed_cache_path(v2, resumable):
|
||||
model = _talker()
|
||||
model.vllm_config.model_config.use_v2_model_runner = v2
|
||||
payload = {"embed": {"cached_decode": _rows(1, 2, 3), "decode": None}, "meta": {"resumable": resumable}}
|
||||
payload["meta"]["num_processed_tokens"] = 2
|
||||
assert _step(model, payload)[0].item() == 3
|
||||
|
||||
|
||||
def test_make_omni_output_aligns_codes_to_request_spans():
|
||||
model = _talker()
|
||||
hidden = torch.zeros((8, _H)) # 5 real tokens + graph padding
|
||||
decode_codes = torch.tensor([[5, 6, 7]])
|
||||
buffers = [
|
||||
{"req_id": "_warmup_0_"}, # no preprocess, no codes
|
||||
{"req_id": "a", "codes": {"audio": torch.zeros((3, _Q), dtype=torch.long)}},
|
||||
{"req_id": "b", "codes": {"audio": decode_codes}},
|
||||
]
|
||||
out = model.make_omni_output(
|
||||
hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 1), (1, 4), (4, 5)]
|
||||
)
|
||||
codes = out.multimodal_outputs["codes"]["audio"]
|
||||
assert codes.shape == (5, _Q)
|
||||
assert codes[4].tolist() == [5, 6, 7]
|
||||
assert out.text_hidden_states.shape[0] == 5
|
||||
|
||||
|
||||
def test_make_omni_output_without_spans_keeps_v1_concat():
|
||||
model = _talker()
|
||||
buffers = [{"codes": {"audio": torch.tensor([[1, 2, 3]])}}, {"codes": {"audio": torch.tensor([[4, 5, 6]])}}]
|
||||
out = model.make_omni_output(torch.zeros((4, _H)), model_intermediate_buffer=buffers)
|
||||
assert out.multimodal_outputs["codes"]["audio"].tolist() == [[1, 2, 3], [4, 5, 6]]
|
||||
assert out.text_hidden_states.shape[0] == 2
|
||||
|
||||
|
||||
def test_eager_output_is_token_major_with_invalid_rows():
|
||||
model = _talker()
|
||||
model.eager_frames_active = True
|
||||
hidden = torch.zeros((6, _H))
|
||||
buffers = [{"req_id": "a"}, {"req_id": "b"}]
|
||||
out = model.make_omni_output(hidden, model_intermediate_buffer=buffers, request_token_spans=[(0, 3), (3, 4)])
|
||||
mm = out.multimodal_outputs
|
||||
assert mm["codes"]["audio"].shape == (4, _Q)
|
||||
assert mm["meta"]["codec_frame_valid"].tolist() == [0, 0, 0, 0]
|
||||
assert out.text_hidden_states is hidden
|
||||
@@ -104,6 +104,7 @@ def test_factory_declares_only_captures_from_earlier_ranks(monkeypatch, start, e
|
||||
if staged:
|
||||
model_config.model_stage = "thinker"
|
||||
vllm_config = SimpleNamespace(model_config=model_config, with_hf_config=lambda *args, **kwargs: None)
|
||||
vllm_config.speculative_config = None
|
||||
thinker = nn.Module()
|
||||
thinker.weight = nn.Parameter(torch.ones(1))
|
||||
thinker.language_model = SimpleNamespace(model=SimpleNamespace(start_layer=start))
|
||||
|
||||
@@ -1432,3 +1432,24 @@ def test_shared_predictor_ignores_tts_fast_path_options(mocker, loaded_target_cl
|
||||
expected = (logits.float() - torch.log(-torch.log(uniforms))).argmax(-1, keepdim=True)
|
||||
actual = predictor._sample_per_call(logits, 1.0, 0, None, uniforms)
|
||||
torch.testing.assert_close(actual, expected)
|
||||
|
||||
|
||||
def test_tts_cached_predictor_keeps_its_own_step_dispatch(mocker: MockerFixture, loaded_target_classes) -> None:
|
||||
"""The TTS cache must not enter Omni's uninitialized incremental loop."""
|
||||
_, _, wrapper, _, _ = loaded_target_classes
|
||||
common_mod = sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"]
|
||||
mocker.patch.object(common_mod.current_omni_platform, "is_cuda", return_value=True)
|
||||
mocker.patch.object(common_mod.current_omni_platform, "is_npu", return_value=False)
|
||||
mocker.patch.object(wrapper, "_stage_connector_extra_config", return_value={"code_predictor_kv_cache": True})
|
||||
cp_config, talker_config = _make_tiny_config(loaded_target_classes)
|
||||
predictor = wrapper(vllm_config=_make_vllm_config(mocker), config=cp_config, talker_config=talker_config)
|
||||
assert predictor._kv_requested
|
||||
predictor._model_dtype = torch.float32
|
||||
predictor._codec_embeds_list = list(predictor.model.codec_embedding)
|
||||
predictor._bucket_sizes = [1]
|
||||
mocker.patch.object(predictor, "_setup_compile")
|
||||
step = mocker.patch.object(predictor, "_predict_step_logits", return_value=torch.ones(1, cp_config.vocab_size))
|
||||
codes = predictor(torch.zeros(1, dtype=torch.long), torch.zeros(1, 32), torch.zeros(1, 32), do_sample=False)
|
||||
assert step.call_count == cp_config.num_code_groups - 1
|
||||
assert codes.shape == (1, cp_config.num_code_groups)
|
||||
assert predictor._kv_buf is None
|
||||
@@ -0,0 +1,70 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.model_executor.models.qwen3_omni.first_frame_decoder import talker_first_audio_enabled
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("env", "extra", "async_chunk", "v2", "expected"),
|
||||
[
|
||||
(None, {}, True, True, False),
|
||||
("1", {}, True, True, False),
|
||||
(None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, True),
|
||||
("0", {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, True, False),
|
||||
(None, {"talker_first_audio": True, "codec_chunk_ramp": [4, 8]}, True, True, False),
|
||||
(None, {"talker_first_audio": True, "initial_codec_chunk_frames": 1}, True, True, True),
|
||||
(None, {"talker_first_audio": True, "initial_codec_chunk_frames": 4}, True, True, False),
|
||||
(None, {"talker_first_audio": True, "codec_chunk_ramp": [1]}, True, True, False),
|
||||
(None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, False, True, False),
|
||||
(None, {"talker_first_audio": True, "codec_chunk_ramp": [1, 2]}, True, False, False),
|
||||
],
|
||||
)
|
||||
def test_talker_first_audio_default_and_conditions(monkeypatch, env, extra, async_chunk, v2, expected):
|
||||
if env is None:
|
||||
monkeypatch.delenv("VLLM_OMNI_TALKER_FIRST_AUDIO", raising=False)
|
||||
else:
|
||||
monkeypatch.setenv("VLLM_OMNI_TALKER_FIRST_AUDIO", env)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: True)
|
||||
config = SimpleNamespace(
|
||||
model_config=SimpleNamespace(
|
||||
async_chunk=async_chunk, use_v2_model_runner=v2, stage_connector_config={"extra": extra}
|
||||
),
|
||||
device_config=SimpleNamespace(device="cuda"),
|
||||
parallel_config=SimpleNamespace(
|
||||
tensor_parallel_size=1, pipeline_parallel_size=1, distributed_executor_backend=None
|
||||
),
|
||||
cache_config=SimpleNamespace(enable_prefix_caching=False),
|
||||
)
|
||||
assert talker_first_audio_enabled(config) is expected
|
||||
for section, field, value in [
|
||||
("device_config", "device", "cpu"),
|
||||
("parallel_config", "tensor_parallel_size", 2),
|
||||
("parallel_config", "pipeline_parallel_size", 2),
|
||||
("parallel_config", "distributed_executor_backend", "mp"),
|
||||
("cache_config", "enable_prefix_caching", True),
|
||||
]:
|
||||
with monkeypatch.context() as patch:
|
||||
patch.setattr(getattr(config, section), field, value)
|
||||
assert not talker_first_audio_enabled(config)
|
||||
monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False)
|
||||
assert not talker_first_audio_enabled(config)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("option", [False, True])
|
||||
def test_tts_requires_connector_opt_in(monkeypatch, option):
|
||||
from vllm_omni.model_executor.models.qwen3_tts import first_audio
|
||||
from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm import (
|
||||
Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM as Predictor,
|
||||
)
|
||||
|
||||
monkeypatch.setattr(Predictor, "_stage_connector_extra_config", lambda _: {"talker_first_audio": option})
|
||||
monkeypatch.setattr(first_audio, "supports_talker_first_audio", lambda _: True)
|
||||
assert first_audio.talker_first_audio_enabled(None) is option
|
||||
@@ -0,0 +1,116 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
|
||||
from vllm_omni.model_executor.stage_input_processors.chunk_size_utils import ramp_decode_windows
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
_Q = 2
|
||||
_RAMP = [1, 2, 4, 8, 16, 25]
|
||||
|
||||
|
||||
def _manager(**extra):
|
||||
config = {"codec_chunk_frames": 25, "codec_left_context_frames": 25, "initial_codec_chunk_frames": 4, **extra}
|
||||
return SimpleNamespace(
|
||||
code_prompt_token_ids=defaultdict(list),
|
||||
put_req_chunk=defaultdict(int),
|
||||
ramp_chunk_count=defaultdict(int),
|
||||
connector=SimpleNamespace(config={"extra": config}),
|
||||
)
|
||||
|
||||
|
||||
def _request(request_id="req"):
|
||||
return SimpleNamespace(external_req_id=request_id, sampling_params=SimpleNamespace(stop_token_ids=[]))
|
||||
|
||||
|
||||
def _frame(index: int) -> dict:
|
||||
return {
|
||||
"codes": {"audio": torch.tensor([[index, 1000 + index]], dtype=torch.long)},
|
||||
"meta": {"codec_frame_valid": torch.ones(1, dtype=torch.int8)},
|
||||
}
|
||||
|
||||
|
||||
def _stream(manager, num_frames, finish_after_last=True):
|
||||
request = _request()
|
||||
chunks = []
|
||||
|
||||
def emit(payload):
|
||||
if payload is None:
|
||||
return
|
||||
codes = payload.codes.audio.reshape(_Q, -1)
|
||||
chunks.append((codes[0].tolist(), int(payload.meta.left_context_size), bool(payload.meta.finished)))
|
||||
# The transport counts every enqueued chunk.
|
||||
manager.put_req_chunk[request.external_req_id] += 1
|
||||
manager.ramp_chunk_count[request.external_req_id] += 1
|
||||
|
||||
for index in range(num_frames):
|
||||
last = finish_after_last and index == num_frames - 1
|
||||
emit(q3.talker2code2wav_async_chunk(manager, _frame(index), request, is_finished=last))
|
||||
return chunks
|
||||
|
||||
|
||||
def test_ramp_decode_windows_include_left_context():
|
||||
assert ramp_decode_windows(_RAMP, 25) == [1, 3, 7, 15, 31, 50]
|
||||
assert ramp_decode_windows([1, 2], 0) == [1, 2]
|
||||
|
||||
|
||||
def test_ramp_chunks_grow_and_resend_left_context():
|
||||
chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 56 + 25, finish_after_last=False)
|
||||
|
||||
new_frames = [len(window) - left for window, left, _finished in chunks]
|
||||
assert new_frames == [1, 2, 4, 8, 16, 25, 25]
|
||||
assert [left for _window, left, _finished in chunks] == [0, 1, 3, 7, 15, 25, 25]
|
||||
assert [len(window) for window, _left, _finished in chunks] == [1, 3, 7, 15, 31, 50, 50]
|
||||
emitted = 0
|
||||
for window, left, _finished in chunks:
|
||||
# Each window ends at the newest frame and re-sends `left` earlier ones.
|
||||
assert window == list(range(emitted - left, emitted + len(window) - left))
|
||||
emitted += len(window) - left
|
||||
# The first chunk is the stream's first frame alone (no context).
|
||||
assert chunks[0][0] == [0]
|
||||
|
||||
|
||||
def test_ramp_flushes_partial_tail_with_left_context():
|
||||
chunks = _stream(_manager(codec_chunk_ramp=_RAMP), 10)
|
||||
|
||||
assert [len(window) - left for window, left, _finished in chunks] == [1, 2, 4, 3]
|
||||
window, left, finished = chunks[-1]
|
||||
assert finished and left == 7 and window == list(range(0, 10))
|
||||
|
||||
|
||||
def test_ramp_finish_on_chunk_boundary_leaves_the_marker_to_the_transport():
|
||||
manager = _manager(codec_chunk_ramp=_RAMP)
|
||||
request = _request()
|
||||
for index in range(3):
|
||||
payload = q3.talker2code2wav_async_chunk(manager, _frame(index), request)
|
||||
if payload is not None:
|
||||
manager.ramp_chunk_count[request.external_req_id] += 1
|
||||
# Frames 0 and 1-2 went out as chunks 0 and 1; the finish step has no frame.
|
||||
finish = {"codes": {"audio": torch.zeros((1, _Q), dtype=torch.long)}, "meta": {"codec_frame_valid": torch.zeros(1)}}
|
||||
assert q3.talker2code2wav_async_chunk(manager, finish, request, is_finished=True) is None
|
||||
|
||||
|
||||
def test_without_ramp_the_initial_chunk_setting_is_unchanged():
|
||||
chunks = _stream(_manager(), 29 + 3)
|
||||
|
||||
assert [(len(window), left) for window, left, _finished in chunks] == [(4, 0), (29, 4), (28, 25)]
|
||||
|
||||
|
||||
def test_first_audio_marker_only_reaches_initial_codec_chunk():
|
||||
manager = _manager(codec_chunk_ramp=[1, 1])
|
||||
for index in range(2):
|
||||
frame = _frame(index)
|
||||
frame["meta"]["first_audio"] = torch.tensor([True])
|
||||
payload = q3.talker2code2wav_async_chunk(manager, frame, _request())
|
||||
assert bool(payload.meta.first_audio) is (index == 0)
|
||||
manager.put_req_chunk["req"] += 1
|
||||
manager.ramp_chunk_count["req"] += 1
|
||||
@@ -0,0 +1,146 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import defaultdict
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm_omni.model_executor.stage_input_processors.qwen3_omni as q3
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
ACCEPT = 24
|
||||
HIDDEN = 8
|
||||
PROMPT = [151644, 872, 5, 6, 151645, 151644, 77091, 198]
|
||||
EOS = 151645
|
||||
|
||||
|
||||
class _Manager:
|
||||
def __init__(self) -> None:
|
||||
self.put_req_chunk: defaultdict[str, int] = defaultdict(int)
|
||||
self.request_payload: dict[str, dict] = {}
|
||||
self.config = SimpleNamespace(
|
||||
hf_config=SimpleNamespace(talker_config=SimpleNamespace(accept_hidden_layer=ACCEPT))
|
||||
)
|
||||
|
||||
def process(self, output, request, is_finished=False):
|
||||
payload = q3.thinker2talker_async_chunk(self, output, request, is_finished=is_finished)
|
||||
if payload is not None:
|
||||
self.put_req_chunk[request.external_req_id] += 1
|
||||
return payload
|
||||
|
||||
|
||||
def _request(output_token_ids, *, history=True, max_tokens=64):
|
||||
# The native data plane keeps token history only while chunk 0 is pending.
|
||||
return SimpleNamespace(
|
||||
external_req_id="r",
|
||||
prompt_token_ids=list(PROMPT) if history else [],
|
||||
output_token_ids=list(output_token_ids) if history else [],
|
||||
all_token_ids=list(PROMPT) + list(output_token_ids) if history else [],
|
||||
output_token_count=len(output_token_ids),
|
||||
last_output_token_id=output_token_ids[-1] if output_token_ids else None,
|
||||
resumable=False,
|
||||
additional_information=None,
|
||||
sampling_params=SimpleNamespace(all_stop_token_ids={EOS}, max_tokens=max_tokens),
|
||||
)
|
||||
|
||||
|
||||
def _embed(token_id: int) -> torch.Tensor:
|
||||
return torch.full((1, HIDDEN), float(token_id), dtype=torch.bfloat16)
|
||||
|
||||
|
||||
def _output(token_ids, sampled=None):
|
||||
layers = {
|
||||
0: torch.cat([_embed(t) for t in token_ids]),
|
||||
ACCEPT: torch.randn(len(token_ids), HIDDEN).to(torch.bfloat16),
|
||||
}
|
||||
embed = {"tts_bos": _embed(1), "tts_eos": _embed(2), "tts_pad": _embed(3)}
|
||||
if sampled is not None:
|
||||
embed["sampled"] = sampled
|
||||
return {"hidden_states": {"layers": layers}, "embed": embed}
|
||||
|
||||
|
||||
def _run_turn(tokens, *, publish, max_tokens=64, prefill_chunks=(len(PROMPT),)):
|
||||
manager = _Manager()
|
||||
rows, chunks, start = [], 0, 0
|
||||
for i, size in enumerate(prefill_chunks):
|
||||
last = i == len(prefill_chunks) - 1
|
||||
sampled = _embed(tokens[0]) if last else torch.empty(0, dtype=torch.bfloat16)
|
||||
request = _request(tokens[:1] if last else [], max_tokens=max_tokens)
|
||||
payload = manager.process(_output(PROMPT[start : start + size], sampled if publish else None), request)
|
||||
start += size
|
||||
if payload is not None:
|
||||
chunks += 1
|
||||
rows.extend(payload.embed.prefill[len(PROMPT) :])
|
||||
for step in range(1, len(tokens)):
|
||||
history = manager.put_req_chunk["r"] == 0
|
||||
request = _request(tokens[: step + 1], history=history, max_tokens=max_tokens)
|
||||
output = _output([tokens[step - 1]], _embed(tokens[step]) if publish else None)
|
||||
payload = manager.process(output, request)
|
||||
if payload is None:
|
||||
continue
|
||||
chunks += 1
|
||||
rows.extend(payload.embed.prefill[len(PROMPT) :] if payload.embed.prefill is not None else payload.embed.decode)
|
||||
return [int(row[0]) for row in rows], chunks, manager
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("tokens", "max_tokens"),
|
||||
[
|
||||
([11, 12, 13, EOS], 64), # ends on a stop token
|
||||
([11, 12, 13], 3), # ends on the max_tokens-th token
|
||||
([11, EOS], 64),
|
||||
],
|
||||
)
|
||||
def test_sampled_stream_matches_captured_stream(tokens, max_tokens):
|
||||
late, _chunks, _manager = _run_turn(tokens, publish=False, max_tokens=max_tokens)
|
||||
early, _chunks, manager = _run_turn(tokens, publish=True, max_tokens=max_tokens)
|
||||
# The final token is never processed, so neither path feeds it to the Talker.
|
||||
assert early == late == tokens[:-1]
|
||||
assert q3._SAMPLED_TEXT_STREAM in manager.request_payload["r"]
|
||||
|
||||
|
||||
def test_sampled_stream_sends_each_token_one_step_earlier():
|
||||
manager = _Manager()
|
||||
chunk0 = manager.process(_output(PROMPT, sampled=_embed(11)), _request([11]))
|
||||
assert chunk0 is not None # at the prefill step, not after the next decode step
|
||||
assert torch.equal(chunk0.embed.prefill, torch.cat([_embed(t) for t in [*PROMPT, 11]]))
|
||||
assert chunk0.hidden_states.output.shape[0] == len(PROMPT)
|
||||
assert chunk0.ids.all == [*PROMPT, 11]
|
||||
# Decode step 1 processes 11 and samples 12: it carries 12.
|
||||
chunk1 = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False))
|
||||
assert torch.equal(chunk1.embed.decode, _embed(12))
|
||||
|
||||
|
||||
def test_chunked_prefill_uses_sample_of_final_chunk_only():
|
||||
early, chunks, _manager = _run_turn([11, 12, EOS], publish=True, prefill_chunks=(5, 3))
|
||||
assert early == [11, 12] and chunks == 2
|
||||
|
||||
|
||||
def test_terminal_first_token_keeps_captured_stream():
|
||||
manager = _Manager()
|
||||
assert manager.process(_output(PROMPT, sampled=_embed(EOS)), _request([EOS])) is None
|
||||
assert q3._SAMPLED_TEXT_STREAM not in manager.request_payload["r"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("prefill_rows", [0, 1, 3])
|
||||
def test_missing_sample_on_sampled_stream(prefill_rows):
|
||||
manager = _Manager()
|
||||
manager.process(_output(PROMPT, sampled=_embed(11)), _request([11]))
|
||||
chunk = manager.process(_output([11], sampled=_embed(12)), _request([11, 12], history=False))
|
||||
assert torch.equal(chunk.embed.decode, _embed(12))
|
||||
if not prefill_rows:
|
||||
with pytest.raises(RuntimeError, match="sample embedding missing"):
|
||||
manager.process(_output([12]), _request([11, 12], history=False))
|
||||
else:
|
||||
before = manager.put_req_chunk["r"]
|
||||
for tokens in (PROMPT[:prefill_rows], [11]):
|
||||
partial = _output(tokens, sampled=torch.empty(0))
|
||||
assert manager.process(partial, _request([11, 12], history=False)) is None
|
||||
assert manager.put_req_chunk["r"] == before
|
||||
chunk = manager.process(_output([12], sampled=_embed(13)), _request([11, 12, 13], history=False))
|
||||
assert torch.equal(chunk.embed.decode, _embed(13))
|
||||
@@ -1047,3 +1047,80 @@ def test_qwen2_5_omni_thinker2talker_full_payload_noop() -> None:
|
||||
|
||||
payload = thinker2talker_full_payload(None, {"any": "thing"}, None)
|
||||
assert payload is None
|
||||
|
||||
|
||||
def _codec_transfer_manager(request_id: str, *, initial_frames: int = 2, chunk_frames: int = 4):
|
||||
return SimpleNamespace(
|
||||
code_prompt_token_ids=defaultdict(list),
|
||||
put_req_chunk=defaultdict(int, {request_id: 0}),
|
||||
connector=SimpleNamespace(
|
||||
config={
|
||||
"extra": {
|
||||
"initial_codec_chunk_frames": initial_frames,
|
||||
"codec_chunk_frames": chunk_frames,
|
||||
"codec_left_context_frames": 4,
|
||||
}
|
||||
}
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def _codec_request(request_id: str) -> SimpleNamespace:
|
||||
return SimpleNamespace(
|
||||
external_req_id=request_id,
|
||||
sampling_params=SimpleNamespace(stop_token_ids=[2150], stop_token_id=None),
|
||||
)
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_reads_last_valid_row_of_multi_row_span() -> None:
|
||||
# MRv2 eager frames: a prefill span carries its first frame in its last row.
|
||||
request_id = "eager_prefill"
|
||||
manager = _codec_transfer_manager(request_id)
|
||||
codes = torch.zeros((5, 3), dtype=torch.long)
|
||||
codes[-1] = torch.tensor([11, 12, 13])
|
||||
valid = torch.tensor([0, 0, 0, 0, 1], dtype=torch.int8)
|
||||
|
||||
payload = q3.talker2code2wav_async_chunk(
|
||||
manager, {"codes": {"audio": codes}, "meta": {"codec_frame_valid": valid}}, _codec_request(request_id)
|
||||
)
|
||||
|
||||
assert payload is None # below the initial chunk size
|
||||
assert [frame.tolist() for frame in manager.code_prompt_token_ids[request_id]] == [[[11, 12, 13]]]
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_skips_invalid_frame_and_flushes_on_finish() -> None:
|
||||
request_id = "eager_eos"
|
||||
manager = _codec_transfer_manager(request_id)
|
||||
manager.code_prompt_token_ids[request_id] = [torch.tensor([[7, 8, 9]])]
|
||||
# The step that sampled codec EOS carries an invalid (placeholder) frame.
|
||||
output = {
|
||||
"codes": {"audio": torch.tensor([[2150, 0, 0]])},
|
||||
"meta": {"codec_frame_valid": torch.tensor([0], dtype=torch.int8)},
|
||||
}
|
||||
|
||||
payload = q3.talker2code2wav_async_chunk(manager, output, _codec_request(request_id), is_finished=True)
|
||||
|
||||
assert len(manager.code_prompt_token_ids[request_id]) == 1
|
||||
assert payload is not None
|
||||
assert payload.codes.audio.tolist() == [7, 8, 9]
|
||||
assert payload.meta.finished.item() is True
|
||||
|
||||
|
||||
def test_talker2code2wav_async_chunk_without_validity_keeps_one_row_contract() -> None:
|
||||
# V1 / MRv2 deferred frames: one row per step, zero rows are prefill
|
||||
# placeholders and a stop-token CB0 marks the post-EOS step.
|
||||
request_id = "legacy_rows"
|
||||
manager = _codec_transfer_manager(request_id)
|
||||
request = _codec_request(request_id)
|
||||
|
||||
assert (
|
||||
q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.zeros((4, 3), dtype=torch.long)}}, request)
|
||||
is None
|
||||
)
|
||||
assert q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[2150, 1, 1]])}}, request) is None
|
||||
assert manager.code_prompt_token_ids[request_id] == []
|
||||
q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[5, 6, 7]])}}, request)
|
||||
payload = q3.talker2code2wav_async_chunk(manager, {"codes": {"audio": torch.tensor([[8, 9, 10]])}}, request)
|
||||
|
||||
assert payload is not None # two frames fill the initial chunk
|
||||
assert payload.codes.audio.tolist() == [5, 8, 6, 9, 7, 10]
|
||||
Reference in new issue
Block a user