Optimize CosyVoice3 Stage1 flow batching (#4876)
Signed-off-by: gerayking <399geray@gmail.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
This commit is contained in:
gerayandHongsheng Liu authored and GitHub committed
1 parent
28291df179
commit
9edc1b4667
10 files changed
+825
-198
No files matched your search
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
"""Unit tests for CosyVoice3 components."""
|
||||
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
from types import SimpleNamespace
|
||||
|
||||
@@ -9,6 +10,7 @@ import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from vllm_omni.diffusion.config import set_current_diffusion_config
|
||||
from vllm_omni.diffusion.data import AttentionConfig
|
||||
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import (
|
||||
@@ -29,6 +31,23 @@ def _force_torch_sdpa():
|
||||
yield
|
||||
|
||||
|
||||
def test_batch_flow_runtime_flags_are_opt_in(monkeypatch):
|
||||
from vllm_omni.model_executor.models.cosyvoice3.runtime import (
|
||||
cosyvoice3_batch_flow_debug,
|
||||
cosyvoice3_batch_flow_enabled,
|
||||
)
|
||||
|
||||
monkeypatch.delenv("COSYVOICE3_BATCH_FLOW", raising=False)
|
||||
monkeypatch.delenv("COSYVOICE3_BATCH_FLOW_DEBUG", raising=False)
|
||||
assert not cosyvoice3_batch_flow_enabled()
|
||||
assert not cosyvoice3_batch_flow_debug()
|
||||
|
||||
monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1")
|
||||
monkeypatch.setenv("COSYVOICE3_BATCH_FLOW_DEBUG", "1")
|
||||
assert cosyvoice3_batch_flow_enabled()
|
||||
assert cosyvoice3_batch_flow_debug()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def causal_hift():
|
||||
return CausalHiFTGenerator(
|
||||
@@ -139,6 +158,39 @@ class TestDiTAttention:
|
||||
assert attention.to_k.out_features == 512
|
||||
assert attention.to_v.out_features == 512
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_forward_passes_2d_mask_without_casting_qkv(self):
|
||||
"""The float32 fallback receives the per-row validity mask."""
|
||||
from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiTAttention
|
||||
|
||||
class CaptureAttention(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.last_metadata = None
|
||||
self.last_dtype = None
|
||||
|
||||
def forward(self, q, k, v, attn_metadata=None):
|
||||
self.last_metadata = attn_metadata
|
||||
self.last_dtype = q.dtype
|
||||
return torch.zeros_like(q)
|
||||
|
||||
attention = DiTAttention(dim=16, heads=2, dim_head=8, dropout=0.0)
|
||||
capture = CaptureAttention()
|
||||
attention.attn = capture
|
||||
|
||||
x = torch.randn(2, 5, 16)
|
||||
mask = torch.tensor([[True, True, True, False, False], [True, True, True, True, False]])
|
||||
|
||||
out = attention(x, mask=mask)
|
||||
|
||||
assert out.shape == x.shape
|
||||
assert capture.last_metadata is not None
|
||||
assert torch.equal(capture.last_metadata.attn_mask, mask)
|
||||
assert capture.last_dtype == x.dtype
|
||||
assert out.dtype == x.dtype
|
||||
assert torch.allclose(out[~mask], torch.zeros_like(out[~mask]))
|
||||
|
||||
|
||||
class TestDiTBlock:
|
||||
"""Tests for DiTBlock."""
|
||||
@@ -422,6 +474,52 @@ class TestCFM:
|
||||
assert reused_context is context
|
||||
assert reused_stream is stream
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.cpu
|
||||
def test_causal_conditional_cfm_batches_cfg_estimator(self):
|
||||
"""Batched flow calls should invoke the CFG estimator with 2B rows."""
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import CausalConditionalCFM
|
||||
|
||||
class RecordingEstimator(nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.shapes: list[tuple[int, ...]] = []
|
||||
|
||||
def forward(self, x, mask, mu, t, spks=None, cond=None):
|
||||
self.shapes.append(tuple(x.shape))
|
||||
return torch.zeros_like(x)
|
||||
|
||||
estimator = RecordingEstimator()
|
||||
cfm_params = DictConfig(
|
||||
{
|
||||
"sigma_min": 1e-6,
|
||||
"solver": "euler",
|
||||
"t_scheduler": "cosine",
|
||||
"training_cfg_rate": 0.2,
|
||||
"inference_cfg_rate": 0.7,
|
||||
}
|
||||
)
|
||||
cfm = CausalConditionalCFM(
|
||||
in_channels=80,
|
||||
cfm_params=cfm_params,
|
||||
n_spks=1,
|
||||
spk_emb_dim=80,
|
||||
estimator=estimator,
|
||||
)
|
||||
|
||||
batch, mel_dim, seq_len = 3, 80, 16
|
||||
mu = torch.randn(batch, mel_dim, seq_len)
|
||||
mask = torch.ones(batch, 1, seq_len)
|
||||
spks = torch.randn(batch, 80)
|
||||
cond = torch.randn(batch, mel_dim, seq_len)
|
||||
|
||||
out, _ = cfm(mu, mask, n_timesteps=2, spks=spks, cond=cond)
|
||||
|
||||
assert out.shape == mu.shape
|
||||
assert estimator.shapes == [(2 * batch, mel_dim, seq_len)] * 2
|
||||
|
||||
|
||||
class TestSDPAFallback:
|
||||
"""Test SDPA fallback for float32 inputs."""
|
||||
@@ -484,3 +582,157 @@ def test_code2wav_forward_finalizes_hift_tail():
|
||||
assert out.shape == (1, 1, 8)
|
||||
assert model.hift.finalize_calls == [True]
|
||||
assert forward_mel_calls[0]["token_offset_tokens"] == 0
|
||||
|
||||
|
||||
def test_code2wav_streaming_batch_pads_codec_tokens_and_preserves_lengths():
|
||||
from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
|
||||
|
||||
model = object.__new__(CosyVoice3Code2Wav)
|
||||
model.flow_model = SimpleNamespace(pre_lookahead_len=1, token_mel_ratio=2)
|
||||
forward_mel_calls = []
|
||||
|
||||
def fake_forward_mel(self, **kwargs):
|
||||
forward_mel_calls.append(kwargs)
|
||||
return torch.arange(2 * 80 * 8, dtype=torch.float32).reshape(2, 80, 8)
|
||||
|
||||
def fake_stream_hift(self, feat, *, cache_state=None, finalize=False):
|
||||
return feat, None
|
||||
|
||||
model._forward_mel = types.MethodType(fake_forward_mel, model)
|
||||
model._stream_hift_from_feat = types.MethodType(fake_stream_hift, model)
|
||||
items = [
|
||||
{
|
||||
"token": torch.ones(1, 3, dtype=torch.int32),
|
||||
"prompt_token": torch.ones(1, 4, dtype=torch.int32),
|
||||
"prompt_feat": torch.ones(1, 8, 80),
|
||||
"embedding": torch.ones(1, 192),
|
||||
"finalize": False,
|
||||
},
|
||||
{
|
||||
"token": torch.ones(1, 5, dtype=torch.int32),
|
||||
"prompt_token": torch.ones(1, 4, dtype=torch.int32),
|
||||
"prompt_feat": torch.ones(1, 8, 80),
|
||||
"embedding": torch.ones(1, 192),
|
||||
"finalize": False,
|
||||
},
|
||||
]
|
||||
|
||||
results = model.forward_streaming_batch(items)
|
||||
|
||||
assert len(forward_mel_calls) == 1
|
||||
call = forward_mel_calls[0]
|
||||
assert call["token"].shape == (2, 5)
|
||||
assert torch.equal(call["token_lens"], torch.tensor([3, 5], dtype=torch.int32))
|
||||
assert results[0][0].shape[-1] == 4
|
||||
assert results[1][0].shape[-1] == 8
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_code2wav_streaming_batch_matches_ragged_flow_numerics(monkeypatch):
|
||||
"""A padded flow batch must match individual flow calls on valid mels."""
|
||||
from omegaconf import DictConfig
|
||||
|
||||
from vllm_omni.diffusion.models.cosyvoice3_audio.cosyvoice3_dit import DiT
|
||||
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm import (
|
||||
CausalConditionalCFM,
|
||||
CausalMaskedDiffWithDiT,
|
||||
)
|
||||
from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.layers import PreLookaheadLayer
|
||||
from vllm_omni.model_executor.models.cosyvoice3.cosyvoice3_code2wav import CosyVoice3Code2Wav
|
||||
|
||||
torch.manual_seed(0)
|
||||
estimator = DiT(
|
||||
dim=32,
|
||||
depth=1,
|
||||
heads=4,
|
||||
dim_head=8,
|
||||
dropout=0.0,
|
||||
ff_mult=2,
|
||||
mel_dim=80,
|
||||
mu_dim=80,
|
||||
spk_dim=80,
|
||||
out_channels=80,
|
||||
)
|
||||
decoder = CausalConditionalCFM(
|
||||
in_channels=80,
|
||||
cfm_params=DictConfig(
|
||||
{
|
||||
"sigma_min": 1e-6,
|
||||
"solver": "euler",
|
||||
"t_scheduler": "cosine",
|
||||
"training_cfg_rate": 0.2,
|
||||
"inference_cfg_rate": 0.7,
|
||||
}
|
||||
),
|
||||
n_spks=1,
|
||||
spk_emb_dim=80,
|
||||
estimator=estimator,
|
||||
)
|
||||
flow_model = (
|
||||
CausalMaskedDiffWithDiT(
|
||||
input_size=80,
|
||||
output_size=80,
|
||||
spk_embed_dim=192,
|
||||
vocab_size=64,
|
||||
input_frame_rate=25,
|
||||
only_mask_loss=True,
|
||||
token_mel_ratio=2,
|
||||
pre_lookahead_len=1,
|
||||
pre_lookahead_layer=PreLookaheadLayer(in_channels=80, channels=80, pre_lookahead_len=1),
|
||||
decoder=decoder,
|
||||
)
|
||||
.to(device="cuda", dtype=torch.bfloat16)
|
||||
.eval()
|
||||
)
|
||||
|
||||
model = object.__new__(CosyVoice3Code2Wav)
|
||||
nn.Module.__init__(model)
|
||||
model.flow_model = flow_model
|
||||
|
||||
def return_mel(self, feat, *, cache_state=None, finalize=False):
|
||||
return feat, None
|
||||
|
||||
model._stream_hift_from_feat = types.MethodType(return_mel, model)
|
||||
|
||||
original_randn = torch.randn
|
||||
|
||||
def length_consistent_randn(*size, **kwargs):
|
||||
shape = tuple(size[0]) if len(size) == 1 and isinstance(size[0], (tuple, list)) else tuple(size)
|
||||
if len(shape) == 3 and shape[1] == 80:
|
||||
device = kwargs.get("device")
|
||||
dtype = kwargs.get("dtype", torch.float32)
|
||||
channels = torch.arange(shape[1], device=device, dtype=torch.float32).view(1, -1, 1)
|
||||
positions = torch.arange(shape[2], device=device, dtype=torch.float32).view(1, 1, -1)
|
||||
noise = torch.sin(channels * 0.17 + positions * 0.31)
|
||||
return noise.expand(shape[0], -1, -1).to(dtype=dtype).clone()
|
||||
return original_randn(*size, **kwargs)
|
||||
|
||||
monkeypatch.setattr(torch, "randn", length_consistent_randn)
|
||||
common = {
|
||||
"prompt_token": torch.tensor([[7, 8]], dtype=torch.int32),
|
||||
"prompt_feat": torch.linspace(-0.5, 0.5, 4 * 80).reshape(1, 4, 80),
|
||||
"embedding": torch.linspace(-1.0, 1.0, 192).reshape(1, 192),
|
||||
"finalize": False,
|
||||
}
|
||||
items = [
|
||||
{**common, "token": torch.tensor([[1, 2, 3]], dtype=torch.int32)},
|
||||
{**common, "token": torch.tensor([[1, 2, 3, 4, 5]], dtype=torch.int32)},
|
||||
]
|
||||
|
||||
batched = model.forward_streaming_batch(items, n_timesteps=2)
|
||||
individual = [
|
||||
model.forward_streaming(
|
||||
token=item["token"],
|
||||
prompt_token=item["prompt_token"],
|
||||
prompt_feat=item["prompt_feat"],
|
||||
embedding=item["embedding"],
|
||||
n_timesteps=2,
|
||||
)
|
||||
for item in items
|
||||
]
|
||||
|
||||
for (batched_mel, _), (individual_mel, _) in zip(batched, individual):
|
||||
assert batched_mel.shape == individual_mel.shape
|
||||
rel_mean = (batched_mel - individual_mel).abs().mean() / individual_mel.abs().mean().clamp_min(1e-6)
|
||||
assert rel_mean.item() < 0.05
|
||||
@@ -42,6 +42,7 @@ class _DummyCode2Wav:
|
||||
self.outputs = list(outputs or [])
|
||||
self.forward_calls: list[dict[str, object]] = []
|
||||
self.forward_streaming_calls: list[dict[str, object]] = []
|
||||
self.forward_streaming_batch_calls: list[list[dict[str, object]]] = []
|
||||
|
||||
def forward(self, **kwargs):
|
||||
self.forward_calls.append(kwargs)
|
||||
@@ -65,6 +66,22 @@ class _DummyCode2Wav:
|
||||
}
|
||||
return audio, new_state
|
||||
|
||||
def forward_streaming_batch(self, items, *, n_timesteps: int = 10):
|
||||
self.forward_streaming_batch_calls.append(items)
|
||||
return [
|
||||
self.forward_streaming(
|
||||
token=item["token"],
|
||||
prompt_token=item["prompt_token"],
|
||||
prompt_feat=item["prompt_feat"],
|
||||
embedding=item["embedding"],
|
||||
cache_state=item.get("cache_state"),
|
||||
n_timesteps=n_timesteps,
|
||||
token_offset_tokens=int(item.get("token_offset_tokens", 0)),
|
||||
finalize=bool(item.get("finalize", False)),
|
||||
)
|
||||
for item in items
|
||||
]
|
||||
|
||||
|
||||
def _make_code2wav_model(
|
||||
*,
|
||||
@@ -373,6 +390,53 @@ def test_forward_clears_streaming_cache_on_terminal_chunk():
|
||||
assert "rid-stream" not in model._stream_vocoder_cache_by_req
|
||||
|
||||
|
||||
def test_forward_batches_streaming_flow_items(monkeypatch):
|
||||
monkeypatch.setenv("COSYVOICE3_BATCH_FLOW", "1")
|
||||
model = _make_code2wav_model()
|
||||
runtime_info = [
|
||||
{
|
||||
"embed": {
|
||||
"speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
|
||||
"embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
|
||||
},
|
||||
"meta": {
|
||||
"req_id": ["rid-a"],
|
||||
"stream_finished": torch.tensor(False),
|
||||
"left_context_size": 0,
|
||||
},
|
||||
},
|
||||
{
|
||||
"embed": {
|
||||
"speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
|
||||
"speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
|
||||
"embedding": torch.tensor([[0.7, 0.8]], dtype=torch.float32),
|
||||
},
|
||||
"meta": {
|
||||
"req_id": ["rid-b"],
|
||||
"stream_finished": torch.tensor(False),
|
||||
"left_context_size": 1,
|
||||
},
|
||||
},
|
||||
]
|
||||
|
||||
out = model.forward(
|
||||
input_ids=torch.tensor([0, 1, 2, 0, 1, 2], dtype=torch.long),
|
||||
positions=torch.arange(6, dtype=torch.long),
|
||||
model_intermediate_buffer=runtime_info,
|
||||
seq_token_counts=[3, 3],
|
||||
)
|
||||
|
||||
assert len(out.multimodal_outputs["audio"]) == 2
|
||||
assert len(model.code2wav.forward_streaming_batch_calls) == 1
|
||||
batch_items = model.code2wav.forward_streaming_batch_calls[0]
|
||||
assert [item["index"] for item in batch_items] == [0, 1]
|
||||
assert torch.equal(batch_items[0]["token"], torch.tensor([[0, 1, 2]]))
|
||||
assert batch_items[1]["token_offset_tokens"] == 1
|
||||
assert "rid-a" in model._stream_vocoder_cache_by_req
|
||||
assert "rid-b" in model._stream_vocoder_cache_by_req
|
||||
|
||||
|
||||
def test_sample_uses_ras_rejection_for_recent_repetition():
|
||||
model = _make_talker_model()
|
||||
metadata = _make_sampling_metadata(output_token_ids=[[1] * 10])
|
||||
|
||||
Reference in new issue
Block a user