[Bugfix][MiniCPM-o] Stabilize code2wav CFM CUDA graph caching with reference audio normalization and mel-frame bucketing (#7416)

Signed-off-by: y_null <y_null@qq.com>
Co-authored-by: amy-why-3459 <wuhaiyan17@huawei.com>
This commit is contained in:
y-nullandamy-why-3459 authored and GitHub committed 2026-09-16 01:24:31 -07:00
1 parent 507cb1d834
commit 5f7f62e13c
13 files changed
+838 -25

No files matched your search

@@ -0,0 +1,162 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the CFM mel-frame bucketing decision.
Steady-state chunk lengths vary per request; without bucketing every length
becomes its own CUDA-graph capture shape (428 captures / 13 flushes in the
#6628 regression). ``_cfm_pad_frames`` aligns the frame axis onto a grid so
steady-state calls share one capture shape; ``_decode_cfm`` trims the output
back. These tests pin the decision math on CPU, no CUDA required.
"""
import pytest
import torch
from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
_cfm_pad_frames,
_zero_padded_frames,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_pads_partial_chunk_up_to_bucket():
# 50 frames on a 16-frame grid -> 14 padding frames.
assert _cfm_pad_frames(mel_frames=50, offset=0, noise_capacity=1000, bucket_frames=16, disabled=False) == 14
def test_aligned_chunk_needs_no_padding():
assert _cfm_pad_frames(mel_frames=64, offset=0, noise_capacity=1000, bucket_frames=16, disabled=False) == 0
def test_ragged_valid_lengths_path_disables_bucketing():
"""The ragged per-request cache path must never be padded."""
for ragged in (True,):
for mel in (50, 64, 96):
assert (
_cfm_pad_frames(
mel_frames=mel,
offset=0,
noise_capacity=1000,
bucket_frames=16,
disabled=ragged,
)
== 0
)
def test_disabled_bucketing_returns_zero():
"""bucket_frames <= 1 (or wrapper absent) means plain eager behavior."""
for bucket in (0, 1):
assert (
_cfm_pad_frames(
mel_frames=50,
offset=0,
noise_capacity=1000,
bucket_frames=bucket,
disabled=False,
)
== 0
)
def test_padding_never_overflows_noise_buffer():
"""Padding past the decoder's noise buffer must fall back to zero.
The padded call slices x from decoder.rand_noise[offset:end]; exceeding
its width would crash. Bucketing is best-effort, never required.
"""
# 96 frames + 16 pad = 112 > 100 capacity -> give up, pad 0.
assert _cfm_pad_frames(mel_frames=96, offset=0, noise_capacity=100, bucket_frames=16, disabled=False) == 0
# Same mel length fits when un-padded: 96 <= 100.
assert _cfm_pad_frames(mel_frames=96, offset=0, noise_capacity=100, bucket_frames=1, disabled=False) == 0
# Cache offset eats into the capacity: 64 + 0 pad would fit, 64 + 48 pad
# would not, so the pad must be dropped.
assert _cfm_pad_frames(mel_frames=64, offset=50, noise_capacity=110, bucket_frames=16, disabled=False) == 0
# The original chunk fits while the padded one does not: 50 + 50 = 100
# <= 110, but 100 + 14 = 114 > 110, so the pad must be dropped.
assert _cfm_pad_frames(mel_frames=50, offset=50, noise_capacity=110, bucket_frames=16, disabled=False) == 0
# Same numbers with room for the padding: 114 <= 120, so it stands.
assert _cfm_pad_frames(mel_frames=50, offset=50, noise_capacity=120, bucket_frames=16, disabled=False) == 14
def test_padding_applies_with_cache_offset_when_it_fits():
# offset 50 + mel 50 + pad 14 = 114 <= 200 -> pad stands.
assert _cfm_pad_frames(mel_frames=50, offset=50, noise_capacity=200, bucket_frames=16, disabled=False) == 14
def test_steady_state_cache_width_saturates_on_bucket_grid():
"""Pin the real recurrence: grow, then trim to ``prompt_len + 100``.
``_decode_batch_once`` trims the estimator attention cache after every
decode, so the width saturates instead of growing without bound. What
matters for the graph cache is that the steady-state
``(chunk_width, cache_width)`` pair settles on a single grid point.
"""
bucket, prompt_len = 16, 304
cache_cap = prompt_len + 100
width = prompt_len
shapes = set()
for _ in range(20):
pad = _cfm_pad_frames(
mel_frames=50,
offset=width,
noise_capacity=30000,
bucket_frames=bucket,
disabled=False,
)
assert pad == 14
shapes.add((50 + pad, width)) # (chunk width, cache width) as used
width = min(width + 50 + pad, cache_cap)
assert width == cache_cap
# (64, 304) -> (64, 368) -> (64, 404), then stable.
assert shapes == {(64, 304), (64, 368), (64, 404)}
def test_varied_chunk_lengths_collapse_onto_few_widths():
"""Bucketing exists for the varied-length calls, not the steady 50-frame one.
The first/last chunk and the ``plan_token2wav_encode_slices`` splits land on
arbitrary lengths; those are the calls that would each become their own
capture shape. Pin that a realistic mix collapses onto a few padded widths.
"""
varied = (7, 12, 17, 25, 33, 50, 51, 64)
padded = {
mel
+ _cfm_pad_frames(
mel_frames=mel,
offset=300,
noise_capacity=30000,
bucket_frames=16,
disabled=False,
)
for mel in varied
}
assert len(padded) == 4, sorted(padded) # 16 / 32 / 48 / 64
assert len(padded) < len(varied)
def test_zero_padded_frames_survives_an_integration_step():
"""The padded region must stay zero after ``x = x + dt * velocity``.
Zeroing once before the CFM loop is not enough: the update touches every
column, so the padded region becomes non-zero again (0 -> 0.010 -> 0.020
over the first steps). Pin the helper that re-zeroes it each step.
"""
mel, pad = 50, 14
x = torch.zeros(2, 4, mel + pad)
x[..., :mel] = 0.05
_zero_padded_frames(x, mel)
assert torch.all(x[..., mel:] == 0.0)
for _ in range(3):
x = x + 0.02 * torch.ones_like(x) # the integration step
_zero_padded_frames(x, mel)
assert torch.all(x[..., mel:] == 0.0)
assert torch.all(x[..., :mel] > 0.0)
def test_zero_padded_frames_is_a_noop_without_padding():
x = torch.ones(1, 2, 32)
_zero_padded_frames(x, None)
_zero_padded_frames(x, 32)
assert torch.all(x == 1.0)
@@ -523,7 +523,7 @@ def test_cfm_graph_receives_bfloat16_cache_in_compute_dtype():
seen_dtypes: list[torch.dtype] = []
class RecordingGraph:
def replay(self, estimator_input, time_embedding, cnn_cache, att_cache, cnn_out, att_out):
def replay(self, estimator_input, time_embedding, cnn_cache, att_cache, cnn_out, att_out, attn_mask=None):
del time_embedding, cnn_cache
seen_dtypes.append(att_cache.dtype)
return estimator_input[:, :1], cnn_out, att_out
@@ -1313,3 +1313,103 @@ def test_reference_voice_and_duplex_metadata_follow_request_lifecycle():
assert prompt_key not in model._runtime_prompts
assert not Path(prompt_wav).exists()
assert (prompt_cache_id, prompt_wav) not in model.backend._prompt_features
def _eager_cfm_wrapper(estimator: nn.Module) -> SimpleNamespace:
"""A CFM graph wrapper that honours the contract but runs eagerly.
``_decode_cfm`` only takes the padded path when a wrapper is present, so
CPU tests that need that path (CUDA graphs are unavailable here) supply
this in place of the real wrapper.
"""
def replay(estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out, attn_mask=None):
output = estimator.blocks_forward_chunk(
estimator_input,
time_emb,
attn_mask,
cnn_cache,
att_cache,
cnn_out,
att_out,
)
return output, cnn_out, att_out
return SimpleNamespace(enabled=True, replay=replay)
def _padded_decode_adapter(bucket_frames: int) -> tuple[BatchedToken2Wav, int, int]:
"""Build an adapter over a real ``DiT`` that exercises the CFM path.
Returns the adapter, the valid width, and the padding width. Every call
seeds the same way, so two adapters built here agree on weights and noise.
"""
from cosyvoice2.flow.decoder_dit import DiT
torch.manual_seed(23)
estimator = DiT(
in_channels=4,
out_channels=1,
depth=2,
num_heads=2,
head_dim=2,
hidden_size=4,
).eval()
token2wav = _FakeToken2Wav()
token2wav.flow.decoder.estimator = estimator
token2wav.flow.decoder.rand_noise = torch.randn(1, 1, 160)
adapter = BatchedToken2Wav(token2wav)
adapter._cfm_graph_wrapper = _eager_cfm_wrapper(estimator)
adapter._cfm_graph_bucket_frames = bucket_frames
mel_frames = 51 # the 16-frame grid pads this to 64
pad_frames = (-(-mel_frames // 16) * 16 - mel_frames) if bucket_frames else 0
return adapter, mel_frames, pad_frames
def test_padded_chunk_keeps_the_cross_chunk_caches_on_the_valid_boundary():
"""Caches the next chunk reads must not carry padding.
The padded steps still run, so ``step_cnn``/``step_att`` come back with
padding-derived values in them. Unless those positions are cleared, the
following chunk consumes them as its keys and as its CNN left context.
This drives a padded chunk through the real ``DiT`` and checks both caches
at the boundary.
"""
adapter, mel_frames, pad_frames = _padded_decode_adapter(bucket_frames=16)
assert pad_frames == 13
mu = torch.randn(1, 1, mel_frames)
cond = torch.zeros(1, 1, mel_frames)
speakers = torch.randn(1, 1)
with torch.no_grad():
x, cnn, att = adapter._decode_cfm(mu, speakers, cond, cnn_cache=None, att_cache=None)
assert int(x.shape[2]) == mel_frames
assert int(att.shape[4]) == mel_frames + pad_frames
# These columns become the next chunk's keys.
assert torch.count_nonzero(att[..., mel_frames : mel_frames + pad_frames, :]) == 0
# This tail becomes the next chunk's CNN left context.
assert torch.count_nonzero(cnn[..., -pad_frames:]) == 0
def test_padding_does_not_change_the_valid_frames():
"""A padded chunk returns the same valid frames as an exact-width one.
Masking takes the padded keys out of the attention and the padded CNN tail
never reaches the valid positions, so the only difference between the two
runs is the capture shape -- not the audio the chunk produces.
"""
padded, mel_frames, pad_frames = _padded_decode_adapter(bucket_frames=16)
exact, _, exact_pad = _padded_decode_adapter(bucket_frames=0)
assert pad_frames == 13
assert exact_pad == 0
mu = torch.randn(1, 1, mel_frames, generator=torch.Generator().manual_seed(5))
cond = torch.zeros(1, 1, mel_frames)
speakers = torch.randn(1, 1, generator=torch.Generator().manual_seed(7))
with torch.no_grad():
padded_x, _, _ = padded._decode_cfm(mu, speakers, cond, cnn_cache=None, att_cache=None)
exact_x, _, _ = exact._decode_cfm(mu, speakers, cond, cnn_cache=None, att_cache=None)
assert int(padded_x.shape[2]) == int(exact_x.shape[2]) == mel_frames
assert torch.allclose(padded_x, exact_x, atol=1e-5)
@@ -252,6 +252,37 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes(
wrapper._flush()
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
def test_cfm_capture_keeps_the_real_mask(monkeypatch: pytest.MonkeyPatch) -> None:
"""Static buffers must be built from the real inputs, mask included.
Capturing from zero-filled placeholders would bake the "nothing is masked"
branch into the graph, while replay copies the real mask into those very
buffers -- so the captured branch has to come from the real value.
"""
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
torch.manual_seed(0)
estimator = _MiniDiT().eval().cuda()
wrapper = CFMGraphWrapper(graph_fn=estimator.blocks_forward_chunk, max_graphs=4)
inputs = _cfm_inputs(2, 12, 0)
frames = int(inputs[0].shape[2])
mask = torch.ones(2, frames, frames, dtype=torch.bool, device="cuda")
mask[:, :, -2:] = False # two padded keys are masked out
with torch.inference_mode():
wrapper.replay(*inputs, mask)
assert wrapper._stats["captures"] == 1
static_inputs, _, _ = next(iter(wrapper._cache.values()))
static_mask = static_inputs[6]
assert static_mask is not None
assert torch.equal(static_mask, mask)
wrapper._flush()
def test_importing_wrapper_does_not_resolve_platform() -> None:
"""Importing this module must not build the OmniPlatform singleton.
@@ -354,7 +385,8 @@ def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.Monkey
"""A key that cannot round-trip is a property of one shape, not of the GPU.
It must not disable the wrapper or retire the generation the way a capture
failure does.
failure does. `_tensors_from_key` is the fallback used when no real inputs
are available, so the probe drives `_capture` without them.
"""
pool = torch.cuda.graph_pool_handle()
monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool)
@@ -374,7 +406,10 @@ def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.Monkey
monkeypatch.setattr(wrapper_module, "_tensors_from_key", _reject_one_shape)
with torch.inference_mode():
wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0))
key = ("estimator_step",) + tuple(
wrapper_module._tensor_signature(t) for t in _cfm_inputs(2, unbuildable_width, 0)
)
assert wrapper._capture(key, None) is None
assert wrapper.enabled is True
assert wrapper._stats["captures"] == 0
@@ -382,11 +417,6 @@ def test_cfm_unsupported_dtype_eagers_only_that_shape(monkeypatch: pytest.Monkey
wrapper.replay(*_cfm_inputs(2, 12, 0))
assert wrapper._stats["captures"] == 1
# and the rejected shape stays eager without retrying the capture
wrapper.replay(*_cfm_inputs(2, unbuildable_width, 0))
assert wrapper._stats["captures"] == 1
assert wrapper._stats["eager"] == 2
wrapper._flush()
@@ -0,0 +1,204 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
"""Unit tests for the reference-audio normalization grid.
Every request's reference audio is folded onto one sample rate and one fixed
length so the CFM attention cache origin (L0) is a single constant; that is
what collapsed the CFM CUDA-graph key space back from 58 shapes to a handful
and cut the capture storm behind the #6628 regression. These tests pin the
pure math so future edits cannot silently widen L0 again.
"""
import logging
from types import SimpleNamespace
import pytest
import torch
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import (
_REF_MAX_SECONDS,
_REF_TARGET_SAMPLE_RATE,
_normalize_reference,
)
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
TARGET = _REF_TARGET_SAMPLE_RATE
MAX_SAMPLES = int(_REF_MAX_SECONDS * TARGET) # 6 s @ 24 kHz
def test_mono_downmix_averages_channels_without_interleaving():
"""(2, T) stereo must become channel-mean mono, not a reshape interleave.
reshape(-1) on (2, T) would alternate 1/0 samples; mean(dim=0) yields a
constant 0.5 waveform. The whole-output equality fails on any interleave.
"""
stereo = torch.zeros(2, 16)
stereo[0] = 1.0
waveform, sr = _normalize_reference(stereo, TARGET)
assert sr == TARGET
assert waveform.shape == (MAX_SAMPLES,)
assert torch.all(waveform[:16] == 0.5) # mono body
assert torch.all(waveform[16:] == 0.0) # zero-padded tail
def test_resample_48k_to_24k_halves_sample_count():
"""Off-grid sample rates are resampled onto 24 kHz before padding."""
wav_48k = torch.arange(24000, dtype=torch.float32)
waveform, sr = _normalize_reference(wav_48k, 48000)
assert sr == TARGET
assert waveform.shape == (MAX_SAMPLES,)
assert torch.all(waveform[:12000] > 0) # resampled body, then silence
assert torch.all(waveform[12000:] == 0)
def test_long_reference_is_truncated_to_6s():
samples_7s = torch.arange(7 * TARGET, dtype=torch.float32)
waveform, sr = _normalize_reference(samples_7s, TARGET)
assert sr == TARGET
assert waveform.shape == (MAX_SAMPLES,)
assert torch.equal(waveform, samples_7s[:MAX_SAMPLES])
def test_short_reference_is_zero_padded_to_6s():
"""The zero-pad is the load-bearing half of S1b.
Truncation alone only caps long references; short ones would keep
distinct L0 values (verified experimentally: rtf 3.375 with truncation
only, 1.2 with zero-padding). Pin the padding behavior.
"""
samples_1s = torch.arange(TARGET, dtype=torch.float32)
waveform, sr = _normalize_reference(samples_1s, TARGET)
assert sr == TARGET
assert waveform.shape == (MAX_SAMPLES,)
assert torch.equal(waveform[:TARGET], samples_1s)
assert torch.all(waveform[TARGET:] == 0)
def test_exact_length_reference_passes_through_unchanged():
samples_6s = torch.arange(MAX_SAMPLES, dtype=torch.float32)
waveform, sr = _normalize_reference(samples_6s, TARGET)
assert sr == TARGET
assert torch.equal(waveform, samples_6s)
def test_every_output_shares_one_length():
"""Whatever comes in, one request's L0 equals any other's."""
for length in (TARGET // 4, TARGET, 3 * TARGET, MAX_SAMPLES + 1):
waveform, _ = _normalize_reference(torch.ones(length), TARGET)
assert waveform.numel() == MAX_SAMPLES
def test_empty_reference_stays_empty():
"""An empty waveform must not be padded into a valid-length reference.
``_materialize_runtime_prompt`` relies on numel() == 0 to raise
"empty_ref_audio" and reject the request; zero-padding would turn a
broken input into a silently accepted all-silence reference.
"""
waveform, sr = _normalize_reference(torch.zeros(0), TARGET)
assert waveform.numel() == 0
assert sr == TARGET
def test_window_length_is_configurable():
"""Deployments can widen the window via ``ref_audio_max_seconds``."""
samples_5s = torch.arange(5 * TARGET, dtype=torch.float32)
waveform, sr = _normalize_reference(samples_5s, TARGET, max_seconds=3.0)
assert sr == TARGET
assert waveform.numel() == 3 * TARGET
# A longer window keeps the whole 5s clip instead of truncating it.
waveform, _ = _normalize_reference(samples_5s, TARGET, max_seconds=8.0)
assert waveform.numel() == 8 * TARGET
assert torch.equal(waveform[: 5 * TARGET], samples_5s)
def test_truncation_emits_warning(caplog):
"""Silent truncation would hide a user-visible behavior change.
A bare ``caplog.at_level`` does not capture vLLM loggers (they do not
propagate to the root handlers the fixture taps), so attach the fixture
handler to the module logger directly -- same idiom as
``tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py``.
"""
samples_7s = torch.arange(7 * TARGET, dtype=torch.float32)
target_logger = logging.getLogger("vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav")
target_logger.addHandler(caplog.handler)
prev_level = target_logger.level
target_logger.setLevel(logging.WARNING)
try:
_normalize_reference(samples_7s, TARGET)
finally:
target_logger.removeHandler(caplog.handler)
target_logger.setLevel(prev_level)
assert any("truncating" in record.getMessage() for record in caplog.records)
def test_soundfile_2d_output_is_transposed_before_the_normalizer(tmp_path):
"""``sf.read`` gives (samples, channels); the normalizer wants the transpose.
Without it a stereo default prompt looks like a huge channel count, the
downmix guard rejects it, and the caller falls back to the raw file -- so
that prompt keeps a second L0 value next to the normalized request
references.
"""
import soundfile as sf
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import (
_read_reference_wav,
)
frames = TARGET
# Opposite-signed channels: the downmix is a true mean only when the two
# channels were aligned as rows, so a wrong layout cannot look correct.
stereo = torch.stack((torch.full((frames,), 0.5), torch.full((frames,), -0.5)), dim=1)
path = tmp_path / "stereo.wav"
sf.write(str(path), stereo.numpy(), TARGET, format="WAV")
waveform, sample_rate_hz = _read_reference_wav(str(path))
assert sample_rate_hz == TARGET
assert waveform.shape == (2, frames)
normalized, target_sr = _normalize_reference(waveform, sample_rate_hz, max_seconds=1.0)
assert target_sr == TARGET
assert normalized.shape[-1] == TARGET
assert torch.allclose(normalized, torch.zeros_like(normalized), atol=1e-6)
def test_default_prompt_uses_the_configured_window(tmp_path):
"""The default prompt has to honour ``ref_audio_max_seconds``.
With a non-default window, a default prompt normalized at the fixed 6 s
constant would keep a different L0 from the request references, which is
the mismatch this normalization exists to remove.
"""
import tempfile
import soundfile as sf
from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import (
MiniCPMO45Code2Wav,
_read_reference_wav,
)
frames = 2 * TARGET # shorter than the window, so it gets zero-padded
samples = torch.full((frames, 2), 0.25)
source = tmp_path / "default.wav"
sf.write(str(source), samples.numpy(), TARGET, format="WAV")
window_seconds = 8.0
with tempfile.TemporaryDirectory() as runtime_dir:
stub = SimpleNamespace(
_default_prompt_wav=str(source),
_default_prompt_id="HT_ref_audio",
_default_prompt_normalized=None,
_ref_max_seconds=window_seconds,
_runtime_prompt_dir=SimpleNamespace(name=runtime_dir),
)
prompt_wav, _ = MiniCPMO45Code2Wav._normalized_default_prompt(stub)
# Read inside the block: the temporary directory is removed on exit.
normalized, sample_rate_hz = _read_reference_wav(prompt_wav)
assert sample_rate_hz == TARGET
# The configured window, not the fixed 6 s default.
assert normalized.shape[-1] == int(window_seconds * TARGET)