From 5f7f62e13cfb88bc1add00db5d3ac2e9ef1328f5 Mon Sep 17 00:00:00 2001 From: y-null Date: Wed, 16 Sep 2026 16:24:31 +0800 Subject: [PATCH] [Bugfix][MiniCPM-o] Stabilize code2wav CFM CUDA graph caching with reference audio normalization and mel-frame bucketing (#7416) Signed-off-by: y_null Co-authored-by: amy-why-3459 --- tests/config/test_config_factory.py | 6 +- .../minicpmo_4_5/test_cfm_graph_bucketing.py | 162 ++++++++++++++ .../minicpmo_4_5/test_code2wav_batching.py | 102 ++++++++- .../minicpmo_4_5/test_cuda_graph_wrapper.py | 44 +++- .../test_reference_audio_normalization.py | 204 ++++++++++++++++++ vllm_omni/deploy/minicpmo_4_5.yaml | 2 +- vllm_omni/deploy/minicpmo_4_5_2gpu.yaml | 2 +- vllm_omni/deploy/minicpmo_4_5_3gpu.yaml | 2 +- vllm_omni/deploy/minicpmo_4_5_8x4090.yaml | 2 +- .../models/minicpmo_4_5/batched_token2wav.py | 147 +++++++++++++ .../models/minicpmo_4_5/cuda_graph_wrapper.py | 33 ++- .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 143 +++++++++++- .../npu/models/minicpmo_4_5_code2wav.py | 14 ++ 13 files changed, 838 insertions(+), 25 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_reference_audio_normalization.py diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index d1dc0fd6d..3699eeda7 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1328,7 +1328,11 @@ class TestDeployConfigLoading: stages = merge_pipeline_deploy(pipeline, deploy) assert deploy.session_mode == "duplex" - assert deploy.active_stream_window == max_sessions + # ``0`` disables the limiter: the configured session/stage capacity is + # what bounds the pipeline. Any positive window serializes concurrent + # audio first-packet generation (measured: TTFP -23% at 0 vs 4 on + # L20X, 128 requests / concurrency 8). + assert deploy.active_stream_window == 0 assert deploy.duplex_session.max_sessions == max_sessions assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"] assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"] diff --git a/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py new file mode 100644 index 000000000..c88fc6d5b --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py @@ -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) diff --git a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py index a0e7ed4db..edb813af3 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py +++ b/tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py @@ -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) diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index 26daeff7b..959156b8f 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -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() diff --git a/tests/model_executor/models/minicpmo_4_5/test_reference_audio_normalization.py b/tests/model_executor/models/minicpmo_4_5/test_reference_audio_normalization.py new file mode 100644 index 000000000..9b9bbf570 --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_reference_audio_normalization.py @@ -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) diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index ca772afbd..5c26fc2c6 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -6,7 +6,6 @@ async_chunk: true session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. -active_stream_window: 4 duplex_session: # Matches Stage 0/1 max_num_seqs below. max_sessions: 4 @@ -30,6 +29,7 @@ connectors: enable_hift_graph: true enable_cfm_graph: true cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 9600e5959..136dace7e 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -9,7 +9,6 @@ async_chunk: true session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. -active_stream_window: 4 duplex_session: # Matches Stage 0/1 max_num_seqs below. max_sessions: 4 @@ -33,6 +32,7 @@ connectors: enable_hift_graph: true enable_cfm_graph: true cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml index aae531218..9669a2768 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -9,7 +9,6 @@ async_chunk: true session_mode: duplex # Keep the active window aligned with the session and stage capacities below. # A smaller window serializes first-packet generation for concurrent chat turns. -active_stream_window: 4 duplex_session: # Matches Stage 0/1 max_num_seqs below. max_sessions: 4 @@ -33,6 +32,7 @@ connectors: enable_hift_graph: true enable_cfm_graph: true cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 stages: - stage_id: 0 diff --git a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml index 00b1090f8..92065dd4f 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -19,7 +19,6 @@ async_chunk: true # (/v1/realtime?duplex=1). The default duplex_session caps sessions at 1, # matching Stage 0/1 max_num_seqs below. session_mode: duplex -active_stream_window: 1 trust_remote_code: true enable_prefix_caching: false @@ -40,6 +39,7 @@ connectors: enable_hift_graph: true enable_cfm_graph: true cfm_max_graphs: 32 + cfm_graph_bucket_frames: 16 stages: - stage_id: 0 diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py index 2bffeb49e..bec442282 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/batched_token2wav.py @@ -46,6 +46,80 @@ def relpos_encode_token_budget( return max(lookahead + 1, min(int(cap), room)) +def _cfm_pad_frames( + *, + mel_frames: int, + offset: int, + noise_capacity: int, + bucket_frames: int, + disabled: bool, +) -> int: + """Frame padding that aligns one CFM call onto the capture-shape grid. + + Chunk lengths vary per request: the steady-state chunk is 50 mel frames, + but the first/last chunk and the ``plan_token2wav_encode_slices`` splits + land on arbitrary lengths, and every distinct length is its own CUDA-graph + capture shape. Padding the frame axis up to ``bucket_frames`` collapses + them onto one grid (the caller trims the output back). Measured on the + shipped config, 128 requests / concurrency 8: 61 captures and 1 cache + flush without bucketing versus 15 captures and 0 flushes with it, audio + RTF 1.317 -> 1.191 under otherwise identical conditions. + + The attention and CNN caches also advance by the padded width so their + shapes stay on the same grid; in practice that is bounded by the trim in + ``_decode_batch_once`` (the cache width saturates at ``prompt_len + + 100``) rather than by the decoder's ``rand_noise`` capacity. + + Returns 0 when bucketing does not apply: the ragged valid-lengths path, + graphs unavailable or already disabled by ``CFMGraphWrapper._disable``, + or padding that would overflow the decoder's noise buffer. + """ + if disabled or bucket_frames <= 1: + return 0 + pad = (bucket_frames - mel_frames % bucket_frames) % bucket_frames + if offset + mel_frames + pad > noise_capacity: + return 0 + return pad + + +def _zero_padded_frames(tensor: torch.Tensor, valid_frames: int | None) -> None: + """Keep the padded columns of ``tensor`` at zero, in place. + + Bucketing pads the frame axis so the capture shape stays fixed. The padded + columns carry no valid content: they are excluded from the attention by the + ``attn_mask`` built in ``_decode_batch_once``, and zeroing keeps them from + leaking back in. The integration step ``x = x + dt * velocity`` would + otherwise make the padded region non-zero again after the initial zeroing, + and the cross-chunk caches would read those values back. + """ + if valid_frames is None or valid_frames >= int(tensor.shape[-1]): + return + tensor[..., valid_frames:] = 0.0 + + +def _zero_padded_cnn_cache( + cnn_cache: torch.Tensor, + estimator: nn.Module, + pad_frames: int, +) -> None: + """Clear the cache positions that come from padded frames, in place. + + Each block's CNN cache holds the tail of its convolution output + (``new_cnn_cache = x[..., -causal_padding[0]:]``, inside ``stepaudio2``), + so when the padding sits at the chunk tail those positions are + padding-derived and would otherwise become the next chunk's left context. + Padding is at most ``pad_frames`` wide, so only the trailing positions that + can come from it are cleared; the valid part of the window is kept. + """ + for index, block in enumerate(estimator.blocks): + width = int(block.conv.block[1].causal_padding[0]) + if width <= 0: + continue + zero_from = max(0, width - pad_frames) + if zero_from < width: + cnn_cache[index][..., zero_from:] = 0.0 + + def plan_token2wav_encode_slices( num_frames: int, *, @@ -249,6 +323,17 @@ class BatchedToken2Wav(nn.Module): "CFM CUDA Graph is disabled on device type %s", flow_parameter.device.type if flow_parameter is not None else "unknown", ) + # mel-frame bucket size for the CFM CUDA Graph path. Pad each decode + # chunk up to a multiple of this many frames so the graph cache key + # space stays small (0 disables bucketing, e.g. when graphs are off). + self._cfm_graph_bucket_frames = ( + int(cfm_graph_cfg.get("bucket_frames", 0)) if self._cfm_graph_wrapper is not None else 0 + ) + if self._cfm_graph_bucket_frames > 1: + logger.info( + "CFM CUDA Graph bucketing enabled (bucket_frames=%d)", + self._cfm_graph_bucket_frames, + ) self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} def _hift_inference( @@ -398,6 +483,7 @@ class BatchedToken2Wav(nn.Module): att_cache: torch.Tensor | None, attn_mask: torch.Tensor | None = None, valid_lengths: list[int] | None = None, + valid_frames: int | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: if self._trt_stepper is not None and valid_lengths is None: out, new_cnn, new_att = self._trt_stepper.step( @@ -413,6 +499,13 @@ class BatchedToken2Wav(nn.Module): time_embedding = estimator.t_embedder(time).unsqueeze(1) width = int(x.shape[-1]) speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) + if valid_frames is not None and valid_frames < width: + # Every frame gets the speaker vector, so the padded columns would + # stay non-zero even with x/mu/cond zeroed -- and the estimator + # would then emit non-zero features there, leaving the + # chunk-boundary attention/CNN caches padding-derived. Zero them. + speaker_features = speaker_features.clone() + _zero_padded_frames(speaker_features, valid_frames) estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) cnn_out, att_out = self._estimator_buffers(estimator, estimator_input, att_cache) old_cnn: Any = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) @@ -433,6 +526,7 @@ class BatchedToken2Wav(nn.Module): graph_att, cnn_out, att_out, + attn_mask, ) if valid_lengths is not None: if not hasattr(estimator, "in_proj"): @@ -566,6 +660,26 @@ class BatchedToken2Wav(nn.Module): estimator = decoder.estimator batch_size = int(mu.shape[0]) offset = int(att_cache.shape[4]) if att_cache is not None else 0 + mel_frames = int(mu.shape[2]) + pad_frames = _cfm_pad_frames( + mel_frames=mel_frames, + offset=offset, + noise_capacity=int(decoder.rand_noise.shape[2]), + bucket_frames=self._cfm_graph_bucket_frames, + disabled=( + valid_lengths is not None + or self._cfm_graph_wrapper is None + # `_disable` keeps the wrapper object alive, so check the flag + # too: padding under a disabled wrapper is pure overhead. + or not self._cfm_graph_wrapper.enabled + ), + ) + if pad_frames: + # Replicate rather than zero: a repeated last frame is a closer + # continuation of the chunk than silence, so the padded positions + # carry less of a step change into the attention and CNN caches. + mu = torch.nn.functional.pad(mu, (0, pad_frames), mode="replicate") + cond = torch.nn.functional.pad(cond, (0, pad_frames), mode="replicate") end = offset + int(mu.shape[2]) if end > int(decoder.rand_noise.shape[2]): raise RuntimeError( @@ -574,6 +688,13 @@ class BatchedToken2Wav(nn.Module): f'"available":{int(decoder.rand_noise.shape[2])}}}' ) x = decoder.rand_noise[:, :, offset:end].expand(batch_size, -1, -1).clone() + if pad_frames: + # The padded columns would otherwise carry real noise values. They + # are excluded from this chunk's attention by ``attn_mask`` below + # and held at zero here, so the padded region stays inert. Masking + # is what removes them; zeroing alone would not, since a zero row + # still occupies part of the softmax denominator. + x[:, :, mel_frames:] = 0.0 timeline = torch.linspace( 0, 1, @@ -603,6 +724,20 @@ class BatchedToken2Wav(nn.Module): device=mu.device, ) attn_mask = valid_queries.unsqueeze(2) & torch.cat((current_keys, old_keys), dim=2) + elif pad_frames: + # Mask the padded keys instead of only zeroing their content: a + # zero-valued key/value pair still takes probability mass out of the + # softmax denominator, so the real frames keep attending to the + # padding unless it is explicitly excluded. + kv_len = int(mu.shape[2]) + offset + attn_mask = torch.ones( + 2 * batch_size, + int(mu.shape[2]), + kv_len, + dtype=torch.bool, + device=mu.device, + ) + attn_mask[:, :, mel_frames : mel_frames + pad_frames] = False next_cnn: list[torch.Tensor] = [] next_att_cache: torch.Tensor | None = None ragged_att_cache: list[torch.Tensor] | None = None @@ -622,10 +757,14 @@ class BatchedToken2Wav(nn.Module): att_cache=old_att, attn_mask=attn_mask, valid_lengths=valid_lengths, + valid_frames=mel_frames if pad_frames else None, ) + if pad_frames: + _zero_padded_cnn_cache(step_cnn, estimator, pad_frames) conditional, unconditional = estimate.split(batch_size, dim=0) velocity = (1.0 + decoder.inference_cfg_rate) * conditional - decoder.inference_cfg_rate * unconditional x = x + dt * velocity + _zero_padded_frames(x, mel_frames if pad_frames else None) time = time + dt if step + 1 < self.n_timesteps: dt = timeline[step + 2] - time[0] @@ -666,7 +805,15 @@ class BatchedToken2Wav(nn.Module): dtype=self._estimator_att_cache_dtype, ) next_att_cache[step].copy_(step_att) + if pad_frames: + # The padded steps ran but hold no valid content, and each + # cache entry becomes the next chunk's keys. Clear the + # padded columns so the boundary the next chunk reads is + # the valid one, in step with the mask built above. + next_att_cache[step][..., mel_frames : mel_frames + pad_frames, :] = 0.0 del step_att + if pad_frames: + x = x[:, :, :mel_frames] if ragged_att_cache is not None: return x, torch.stack(next_cnn), ragged_att_cache assert next_att_cache is not None diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py index d737905fd..4db0c39bf 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/cuda_graph_wrapper.py @@ -135,11 +135,13 @@ class HiFTGraphWrapper: return speech, cache_source -def _tensor_signature(value: torch.Tensor) -> tuple: +def _tensor_signature(value: torch.Tensor | None) -> tuple: + if value is None: + return (None,) return tuple(value.shape), str(value.dtype), str(value.device) -_DTYPE_MAP = {str(dtype): dtype for dtype in (torch.float32, torch.float16, torch.bfloat16, torch.float64)} +_DTYPE_MAP = {str(dtype): dtype for dtype in (torch.float32, torch.float16, torch.bfloat16, torch.float64, torch.bool)} def _memory_snapshot(device: torch.device) -> tuple[int, int] | None: @@ -231,7 +233,7 @@ class CFMGraphWrapper: return {**self._stats, "cache_size": len(self._cache)} def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor: - return self.graph_fn(args[0], args[1], None, args[2], args[3], args[4], args[5]) + return self.graph_fn(args[0], args[1], args[6], args[2], args[3], args[4], args[5]) def _eager(self, inputs: tuple[torch.Tensor, ...]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: self._stats["eager"] += 1 @@ -261,10 +263,21 @@ class CFMGraphWrapper: self.enabled = False self._flush() - def _capture(self, key: tuple) -> tuple | None: - """Capture a CUDA graph for the given key. Returns None on failure.""" + def _capture(self, key: tuple, inputs: tuple | None = None) -> tuple | None: + """Capture a CUDA graph for the given key. Returns None on failure. + + ``inputs`` should be the real tensors for this key whenever they are + available: the static buffers must be built from the values the graph + will actually run with. Zero-filled placeholders are only a fallback -- + capturing with them bakes the placeholder branch into the graph (e.g. an + attention mask that looks like "nothing is masked"). + """ try: - static_inputs = _tensors_from_key(key) + static_inputs = ( + tuple(None if tensor is None else tensor.detach().clone() for tensor in inputs) + if inputs is not None + else _tensors_from_key(key) + ) except KeyError: # An unsupported dtype is a property of this shape alone, so run it # eager and keep the other shapes on graphs. @@ -315,8 +328,9 @@ class CFMGraphWrapper: att_cache: torch.Tensor, cnn_out: torch.Tensor, att_out: torch.Tensor, + attn_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out) + inputs = (estimator_input, time_emb, cnn_cache, att_cache, cnn_out, att_out, attn_mask) self._stats["calls"] += 1 if not self.enabled or torch.cuda.is_current_stream_capturing() or estimator_input.device.type != "cuda": @@ -330,7 +344,7 @@ class CFMGraphWrapper: if entry is None: if len(self._cache) >= self.max_graphs: self._flush() - entry = self._capture(key) + entry = self._capture(key, inputs) if entry is None: return self._eager(inputs) self._cache[key] = entry @@ -339,7 +353,8 @@ class CFMGraphWrapper: static_inputs, static_output, graph = entry for static, current in zip(static_inputs, inputs, strict=True): - static.copy_(current) + if static is not None: + static.copy_(current) graph.replay() return ( static_output.detach().clone(), diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py index f701e8b26..da2c0eb72 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_code2wav.py @@ -9,6 +9,7 @@ import os import tempfile from collections.abc import Iterable, Mapping, Sequence from dataclasses import dataclass +from functools import lru_cache from hashlib import sha256 from pathlib import Path from typing import Any @@ -57,6 +58,92 @@ def _scalar(value: Any, default: Any = None) -> Any: return default if value is None else value +_REF_TARGET_SAMPLE_RATE = 24000 +# Overridable per deployment via ``ref_audio_max_seconds``. +_REF_MAX_SECONDS = 6.0 +# Channel ceiling for the optional (channels, samples) downmix guard. +_REF_MAX_CHANNELS = 8 + + +@lru_cache(maxsize=8) +def _get_resampler(orig_freq: int, new_freq: int): + """Cached anti-aliased resampler matching token2wav's own loader. + + ``torchaudio`` is imported lazily (it is not a hard runtime requirement of + this model on every platform) and the transform is built once per rate pair + instead of once per request, keeping it off the first-packet path. + """ + import torchaudio + + return torchaudio.transforms.Resample(orig_freq=orig_freq, new_freq=new_freq) + + +def _read_reference_wav(path: str) -> tuple[Any, int]: + """Read a WAV file for :func:`_normalize_reference`. + + ``soundfile`` returns ``(samples, channels)`` for multi-channel files + while the normalizer expects ``(channels, samples)``. Without the + transpose a stereo prompt looks like a huge channel count, the downmix + guard rejects it, and the caller silently falls back to the raw file -- + a second L0 value for the same content. + """ + waveform, sample_rate_hz = sf.read(path, dtype="float32", always_2d=False) + if getattr(waveform, "ndim", 1) > 1: + waveform = waveform.T + return waveform, int(sample_rate_hz) + + +def _normalize_reference( + ref_audio: Any, + sample_rate_hz: int, + *, + max_seconds: float = _REF_MAX_SECONDS, +) -> tuple[torch.Tensor, int]: + """Normalize a per-request reference waveform onto a shared grid. + + MiniCPM-o streaming uses the reference-audio length as the CFM attention + cache origin (L0). Unnormalized references give every request a distinct + L0, which explodes the CFM CUDA-graph cache key space (see #6628). Fold + every reference onto the same sample rate and a fixed length (truncate + long ones, zero-pad short ones) so L0 becomes one constant value. + + References longer than ``max_seconds`` are truncated with a warning; + the window is configurable via ``ref_audio_max_seconds``. + """ + tensor = torch.as_tensor(ref_audio, dtype=torch.float32) + if tensor.dim() > 1: + # (channels, samples) -> mono; plain reshape(-1) would interleave. + # Upstream flattens request references to 1-D, so anything else is a + # non-standard caller: reject layouts we cannot downmix unambiguously + # instead of silently averaging the waveform away. + if tensor.shape[0] > _REF_MAX_CHANNELS: + raise ValueError(f"reference audio must be 1-D or (channels, samples); got shape {tuple(tensor.shape)}") + tensor = tensor.mean(dim=0) + waveform = tensor.reshape(-1).cpu().contiguous() + if waveform.numel() == 0: + # Keep the caller's "empty_ref_audio" error path intact: an empty + # waveform must not be zero-padded into a valid-length reference. + return waveform, _REF_TARGET_SAMPLE_RATE + if sample_rate_hz != _REF_TARGET_SAMPLE_RATE: + # Match token2wav's own loader (torchaudio, anti-aliased): it only + # resamples when the stored rate differs, so this output is what the + # model consumes. + waveform = _get_resampler(sample_rate_hz, _REF_TARGET_SAMPLE_RATE)(waveform.view(1, -1)).view(-1).contiguous() + max_samples = max(1, int(max_seconds * _REF_TARGET_SAMPLE_RATE)) + if waveform.numel() > max_samples: + logger.warning( + "Reference audio is %.2fs, truncating to the %.2fs window (set ``ref_audio_max_seconds`` to keep more).", + waveform.numel() / _REF_TARGET_SAMPLE_RATE, + max_seconds, + ) + waveform = waveform[:max_samples].contiguous() + elif waveform.numel() < max_samples: + # Zero-pad short references so every request shares one L0; trailing + # silence has minimal style impact (verified via E3 WER/SIM). + waveform = torch.nn.functional.pad(waveform, (0, max_samples - waveform.numel())) + return waveform, _REF_TARGET_SAMPLE_RATE + + def _codec_tensor(value: Any, fallback: torch.Tensor) -> torch.Tensor: if isinstance(value, torch.Tensor): return value.reshape(-1).to(device=fallback.device, dtype=torch.long) @@ -181,7 +268,11 @@ class MiniCPMO45Code2Wav(nn.Module): self._cfm_graph_config = { "enabled": bool(extra.get("enable_cfm_graph", False)), "max_graphs": int(extra.get("cfm_max_graphs", 32)), + "bucket_frames": int(extra.get("cfm_graph_bucket_frames", 0)), } + self._ref_max_seconds = float(extra.get("ref_audio_max_seconds", _REF_MAX_SECONDS)) + if self._ref_max_seconds <= 0: + raise ValueError("MiniCPM-o Code2Wav ref_audio_max_seconds must be > 0") self._min_batch_size = int(extra.get("code2wav_min_batch_size", 1)) if self._min_batch_size < 1: raise ValueError("MiniCPM-o Code2Wav code2wav_min_batch_size must be >= 1") @@ -192,6 +283,7 @@ class MiniCPMO45Code2Wav(nn.Module): raise ValueError("MiniCPM-o Code2Wav code2wav_initial_batch_size must be 0 or >= code2wav_min_batch_size") self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") + self._default_prompt_normalized: tuple[str, str] | None = None @property def _default_prompt_wav(self) -> str: @@ -199,6 +291,46 @@ class MiniCPMO45Code2Wav(nn.Module): return str(self._prompt_wav_override) return str(Path(self.model_path) / "assets" / "HT_ref_audio.wav") + def _normalized_default_prompt(self) -> tuple[str, str]: + """Fold the shipped default prompt onto the request-reference grid. + + The shipped asset is 6.016 s and is loaded directly by token2wav, so a + request without a reference would otherwise keep a second L0 value a + couple of frames away from the normalized request references -- two + graph-key families instead of one (both use ``ref_audio_max_seconds``). + Returns ``(prompt_wav, + prompt_cache_id)``; if the asset cannot be read, the shipped path is + returned unchanged. + """ + if self._default_prompt_normalized is not None: + return self._default_prompt_normalized + source = self._default_prompt_wav + fallback = (source, self._default_prompt_id) + try: + waveform, sample_rate_hz = _read_reference_wav(source) + normalized, target_sr = _normalize_reference( + waveform, + sample_rate_hz, + max_seconds=self._ref_max_seconds, + ) + except Exception: + logger.warning("Could not normalize the default prompt %s; using it as-is", source) + self._default_prompt_normalized = fallback + return self._default_prompt_normalized + if normalized.numel() == 0: + self._default_prompt_normalized = fallback + return self._default_prompt_normalized + digest = sha256() + digest.update(normalized.numpy().tobytes()) + digest.update(str(target_sr).encode()) + cache_id = f"default-ref-{digest.hexdigest()[:24]}-{target_sr}" + path = Path(self._runtime_prompt_dir.name) / f"{cache_id}.wav" + if not path.is_file(): + sf.write(path, normalized.numpy(), target_sr, format="WAV") + logger.info("Default prompt normalized onto the %.2fs window: %s", self._ref_max_seconds, path) + self._default_prompt_normalized = (str(path), cache_id) + return self._default_prompt_normalized + def _extra_config(self) -> dict[str, Any]: model_config = getattr(self.vllm_config, "model_config", None) connector = getattr(model_config, "stage_connector_config", None) @@ -220,9 +352,13 @@ class MiniCPMO45Code2Wav(nn.Module): sample_rate: Any, ) -> tuple[str, _RuntimePrompt]: sample_rate_hz = int(_scalar(sample_rate, 0)) - waveform = torch.as_tensor(ref_audio, dtype=torch.float32).reshape(-1).cpu().contiguous() if sample_rate_hz <= 0: raise _batch_error("invalid_ref_audio_sample_rate", sample_rate=sample_rate_hz) + waveform, sample_rate_hz = _normalize_reference( + ref_audio, + sample_rate_hz, + max_seconds=self._ref_max_seconds, + ) if waveform.numel() == 0: raise _batch_error("empty_ref_audio") if not bool(torch.isfinite(waveform).all().item()): @@ -288,9 +424,10 @@ class MiniCPMO45Code2Wav(nn.Module): if entry is not None: return entry.cache_id, entry.path, cache_key + default_wav, default_cache_id = self._normalized_default_prompt() return ( - str(_scalar(meta.get("prompt_cache_id"), self._default_prompt_id)), - str(_scalar(meta.get("prompt_wav"), self._default_prompt_wav)), + str(_scalar(meta.get("prompt_cache_id"), default_cache_id)), + str(_scalar(meta.get("prompt_wav"), default_wav)), None, ) diff --git a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py index 0d34c93a5..4346e83b6 100644 --- a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py +++ b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py @@ -75,10 +75,15 @@ def _graphable_estimator_step( cond, cnn_cache, att_cache, + valid_frames=None, ): """Run the CFM estimator body after host-backed timestep embedding.""" width = int(x.shape[-1]) speaker_features = speakers.unsqueeze(-1).expand(-1, -1, width) + if valid_frames is not None and valid_frames < width: + # Match the CUDA path: padded columns must not carry the speaker vector. + speaker_features = speaker_features.clone() + speaker_features[:, :, valid_frames:] = 0.0 estimator_input = torch.cat((x, mu, speaker_features, cond), dim=1) cnn_out, att_out = backend._estimator_buffers(estimator, estimator_input, att_cache) old_cnn = cnn_cache if cnn_cache is not None else [None] * len(estimator.blocks) @@ -108,6 +113,7 @@ def _patched_estimator_step( att_cache, attn_mask=None, valid_lengths=None, + valid_frames=None, ): assert _original_estimator_step is not None graph_runner = _backend_graph_runners.get(self) @@ -130,10 +136,16 @@ def _patched_estimator_step( att_cache=att_cache, attn_mask=attn_mask, valid_lengths=valid_lengths, + valid_frames=valid_frames, ) if (cnn_cache is None) != (att_cache is None): raise ValueError("estimator CNN and attention caches must both be present or absent") + # valid_frames only matters for padded inputs; skip forwarding it for + # unpadded calls so platform-owned graphable bodies that predate the + # padding feature keep working with the old signature. + graphable_kwargs = {"valid_frames": valid_frames} if valid_frames is not None else {} + # The upstream embedder creates a frequency tensor on the host. Keep it # outside capture while retaining the tensor-only estimator body in graph. time_embedding = estimator.t_embedder(time).unsqueeze(1) @@ -152,6 +164,7 @@ def _patched_estimator_step( cond=step_cond, cnn_cache=None, att_cache=None, + **graphable_kwargs, ), ) @@ -169,6 +182,7 @@ def _patched_estimator_step( cond=step_cond, cnn_cache=step_cnn, att_cache=step_att, + **graphable_kwargs, ), )