diff --git a/docs/serving/realtime_duplex_api.md b/docs/serving/realtime_duplex_api.md index 7454d15f4..5095d1ca7 100644 --- a/docs/serving/realtime_duplex_api.md +++ b/docs/serving/realtime_duplex_api.md @@ -54,7 +54,7 @@ vllm-omni serve openbmb/MiniCPM-o-4_5 \ ``` `vllm_omni/deploy/minicpmo_4_5.yaml` declares `session_mode: duplex` and -`duplex_session.max_sessions: 4`. Because the MiniCPM-o 4.5 pipeline declares a +`duplex_session.max_sessions: 16`. Because the MiniCPM-o 4.5 pipeline declares a `duplex_plugin`, `vllm-omni serve` runs it through `DuplexOmni`: the server mounts `ws://:8099/v1/realtime?duplex=1` (this page; `ws://:8099/v1/duplex` is an alias of the same route), `POST /v1/chat/completions`, `/v1/models` and diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index 6c4b5c7ab..df1a3a036 100644 --- a/tests/config/test_config_factory.py +++ b/tests/config/test_config_factory.py @@ -1408,12 +1408,12 @@ class TestDeployConfigLoading: @pytest.mark.parametrize( ("filename", "max_sessions"), [ - ("minicpmo_4_5.yaml", 4), + ("minicpmo_4_5.yaml", 16), ("minicpmo_4_5_2gpu.yaml", 4), ("minicpmo_4_5_3gpu.yaml", 4), ("minicpmo_4_5_8x4090.yaml", 1), - ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4), - ("minicpmo_4_5_4gpu_stage1_replicas.yaml", 4), + ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 16), + ("minicpmo_4_5_4gpu_stage1_replicas.yaml", 16), ("minicpmo_4_5_8x4090_stage1_replicas.yaml", 1), ], ) @@ -2886,7 +2886,7 @@ class TestPlatformOverrides: cuda = _apply_platform_overrides(load_deploy_config(deploy_path), platform="cuda") cuda_stages = merge_pipeline_deploy(pipeline, cuda) - assert cuda_stages[1].yaml_engine_args["kv_cache_memory_bytes"] == 2 * 1024**3 + assert cuda_stages[1].yaml_engine_args["kv_cache_memory_bytes"] == 4 * 1024**3 # The CUDA memory budget must not leak into the existing NPU profile. npu = _apply_platform_overrides(load_deploy_config(deploy_path), platform="npu") @@ -2903,7 +2903,7 @@ class TestPlatformOverrides: load_deploy_config(Path(get_deploy_config_path(filename))), platform="cuda" ) replica_stages = merge_pipeline_deploy(pipeline, replica) - # Explicit null clears the inherited single-GPU 2 GiB CUDA cap. + # Explicit null clears the inherited single-GPU 4 GiB CUDA cap. assert replica_stages[1].yaml_engine_args.get("kv_cache_memory_bytes") is None def test_fish_speech_npu_uses_ascend_kv_block_size(self): diff --git a/tests/engine/duplex/test_duplex_deploy_config.py b/tests/engine/duplex/test_duplex_deploy_config.py index e9168729c..f7af09d2b 100644 --- a/tests/engine/duplex/test_duplex_deploy_config.py +++ b/tests/engine/duplex/test_duplex_deploy_config.py @@ -96,9 +96,9 @@ def test_duplex_session_runtime_rejects_non_positive_values(tmp_path, name: str, @pytest.mark.parametrize( ("deploy_yaml", "expected"), [ - ("minicpmo_4_5.yaml", 4), + ("minicpmo_4_5.yaml", 16), ("minicpmo_4_5_8x4090.yaml", 1), - ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 4), + ("minicpmo_4_5_3gpu_stage1_replicas.yaml", 16), ], ) def test_deploy_duplex_max_sessions_tracks_the_deploy_config(deploy_yaml: str, expected: int) -> None: 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 index 39ff8ebdc..c72e947c1 100644 --- 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 @@ -16,9 +16,15 @@ import torch from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( _cfm_pad_frames, - _zero_padded_cnn_cache, _zero_padded_frames, ) +from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( + _att_keep_ranges, + _build_capture_mask, + _capture_query_width, + _whole_euler_att_segments, + _zero_padded_cnn_cache, +) pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -139,6 +145,75 @@ def test_varied_chunk_lengths_collapse_onto_few_widths(): assert len(padded) < len(varied) +def test_capture_query_width_collapses_decode_grid_onto_one_bucket(): + """Whole-Euler capture bucket is the #7416 pad helper with a 64-wide decode grid.""" + assert _capture_query_width(16, 0) == 16 + assert _capture_query_width(16, 1) == 16 + assert {_capture_query_width(w, 64) for w in (16, 32, 48, 50, 64)} == {64} + assert _capture_query_width(304, 64) == 320 + assert _capture_query_width(320, 64) == 320 + + +def test_att_keep_ranges_match_the_streaming_trim(): + """Past ``prompt_len + 100`` frames the cache keeps its first ``prompt_len`` and last 100.""" + assert _att_keep_ranges(368, (300, 100)) == [(0, 368)] + assert _att_keep_ranges(400, (300, 100)) == [(0, 400)] + assert _att_keep_ranges(464, (300, 100)) == [(0, 300), (364, 100)] + assert _att_keep_ranges(464, None) == [(0, 464)] + + +def test_whole_euler_att_segments_skip_capture_padding(): + # Steady chunk: capture width == chunk width, so only the trim applies. + assert _whole_euler_att_segments(mel_width=64, query_cap=64, offset=400, keep=(300, 100)) == [ + (0, 300), + (364, 100), + ] + # A 32-wide chunk on the 64 capture: the graph wrote [current 64 | cache 400]; + # frames 32:64 are capture padding and the cache starts at 64. + assert _whole_euler_att_segments(mel_width=32, query_cap=64, offset=400, keep=None) == [(0, 32), (64, 400)] + assert _whole_euler_att_segments(mel_width=32, query_cap=64, offset=400, keep=(300, 100)) == [ + (0, 32), + (64, 268), + (364, 100), + ] + # First chunk: no cache, only the tail of the capture width is dropped. + assert _whole_euler_att_segments(mel_width=304, query_cap=320, offset=0, keep=None) == [(0, 304)] + + +def test_capture_mask_moves_the_cache_block_past_capture_padding(): + width, cap, offset, valid = 32, 64, 8, 30 + caller = torch.ones(2, width, width + offset, dtype=torch.bool) + caller[:, :, valid:width] = False + mask = _build_capture_mask( + attn_mask=caller, + batch_size=1, + query_cap=cap, + offset=offset, + mel_width=width, + mel_frames=valid, + device=torch.device("cpu"), + ) + assert mask.shape == (2, cap, cap + offset) + assert mask[:, :, :valid].all() + assert not mask[:, :, valid:cap].any() + assert mask[:, :, cap:].all() + # Capture-padding query rows are never empty, so SDPA cannot produce NaN there. + assert mask.any(dim=-1).all() + + implicit = _build_capture_mask( + attn_mask=None, + batch_size=1, + query_cap=cap, + offset=offset, + mel_width=width, + mel_frames=width, + device=torch.device("cpu"), + ) + assert implicit[:, :, :width].all() + assert not implicit[:, :, width:cap].any() + assert implicit[:, :, cap:].all() + + def test_zero_padded_frames_survives_an_integration_step(): """The padded region must stay zero after ``x = x + dt * velocity``. 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 b74a8aaaa..e97eb8e5a 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 @@ -172,6 +172,19 @@ def _config(minimum: int = 1, initial: int = 0, *, runtime_prompt_cache_size: in ) +@pytest.mark.parametrize( + ("extra", "max_num_seqs", "micro"), [({}, 6, 6), ({}, 64, 16), ({"micro_batch_size": 2}, 6, 2)] +) +def test_whole_euler_micro_batch_defaults_to_stage_concurrency(extra, max_num_seqs, micro): + config = _config() + config.model_config.stage_connector_config["extra"].update(extra) + config.scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs) + + model = MiniCPMO45Code2Wav(vllm_config=config) + + assert model._cfm_graph_config["micro_batch_size"] == micro + + def _model(initial: int = 0, minimum: int = 1, *, runtime_prompt_cache_size: int = 4, setup_cache_size: int = 1): token2wav = _FakeToken2Wav() backend = BatchedToken2Wav(token2wav, setup_cache_size=setup_cache_size) @@ -213,17 +226,43 @@ def test_setup_cache_reuses_read_only_state_for_the_same_exact_batch(): torch.testing.assert_close(cache[name], expected) -def test_setup_cache_keys_batch_size_and_evicts_least_recent_entry(): +def test_setup_solves_the_prompt_once_for_every_batch_size(): token2wav = _FakeToken2Wav() adapter = BatchedToken2Wav(token2wav, setup_cache_size=1) prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") - adapter.setup_batch(prompt, 1) - adapter.setup_batch(prompt, 2) - adapter.setup_batch(prompt, 1) + single = adapter.setup_batch(prompt, 1) + pair = adapter.setup_batch(prompt, 2) - assert token2wav.flow.encoder.calls == [1, 2, 1] - assert list(adapter._setup_cache) == [(prompt.cache_key, 1, 0)] + assert token2wav.flow.encoder.calls == [1] + assert pair == [single[0], single[0]] + assert list(adapter._setup_cache) == [(prompt.cache_key, 0)] + + +def test_setup_cache_evicts_least_recent_prompt(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav, setup_cache_size=1) + first = adapter.prepare_prompt("first", "/fake/first.wav") + second = adapter.prepare_prompt("second", "/fake/second.wav") + + adapter.setup_batch(first, 1) + adapter.setup_batch(second, 2) + adapter.setup_batch(first, 1) + + assert token2wav.flow.encoder.calls == [1, 1, 1] + assert list(adapter._setup_cache) == [(first.cache_key, 0)] + + +def test_setup_without_cache_still_solves_one_shared_row(): + token2wav = _FakeToken2Wav() + adapter = BatchedToken2Wav(token2wav, setup_cache_size=0) + prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav") + + states = adapter.setup_batch(prompt, 3) + + assert token2wav.flow.encoder.calls == [1] + assert states[0] is states[1] is states[2] + assert not adapter._setup_cache def test_evict_prompt_clears_features_and_setup_state(): @@ -481,9 +520,12 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches(): ) assert token2wav.prompt_calls == 1 - assert token2wav.flow.encoder.calls == [2, 2] - assert token2wav.flow.decoder.estimator.cfg_batches == [4, 4, 4, 4] - assert all(order == [1.0, 1.0, 0.0, 0.0] for order in token2wav.flow.decoder.estimator.speaker_order) + # Setup solves the shared prompt once; the chunk runs both requests together. + assert token2wav.flow.encoder.calls == [1, 2] + assert token2wav.flow.decoder.estimator.cfg_batches == [2, 2, 4, 4] + speaker_order = token2wav.flow.decoder.estimator.speaker_order + assert all(order == [1.0, 0.0] for order in speaker_order[:2]) + assert all(order == [1.0, 1.0, 0.0, 0.0] for order in speaker_order[2:]) assert token2wav.hift.calls == [2] assert len(audios) == 2 cache0 = states[0].flow_cache["estimator_cnn_cache"] @@ -706,6 +748,40 @@ def test_decode_cfm_enters_platform_sdpa_context(monkeypatch): assert entered == ["enter", "exit"] +def test_eager_cfm_reads_the_step_sizes_to_the_host_once(monkeypatch): + """The eager Euler loop must not sync the stream for ``dt`` on every decode.""" + adapter = BatchedToken2Wav(_FakeToken2Wav()) + reads: list[str] = [] + + def recording(name): + original = getattr(torch.Tensor, name) + + def read(self, *args, **kwargs): + # The fake estimator records its inputs with ``tolist``; skip those. + if sys._getframe(1).f_code.co_filename != __file__: + reads.append(name) + return original(self, *args, **kwargs) + + return read + + for name in ("item", "tolist"): + monkeypatch.setattr(torch.Tensor, name, recording(name)) + + def decode(): + adapter._decode_cfm( + torch.ones((1, 1, 2)), + torch.ones((1, 1)), + torch.zeros((1, 1, 2)), + cnn_cache=None, + att_cache=None, + ) + + decode() + assert reads == ["tolist"] + decode() + assert reads == ["tolist"] + + def test_ragged_decode_bypasses_exact_shape_accelerators(): adapter = BatchedToken2Wav(_FakeToken2Wav()) _enable_fake_ragged_kernel(adapter) @@ -963,6 +1039,44 @@ def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows(): ) +@pytest.mark.parametrize("attention_dtype", [torch.float32, torch.bfloat16]) +def test_split_flow_cache_trims_each_request_like_the_batch_trim(attention_dtype): + """Trimming per request while splitting equals trimming the stacked batch first.""" + adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=attention_dtype == torch.bfloat16) + batch_size, prompt_len, frames = 3, 5, 12 + torch.manual_seed(0) + stacked = { + "conformer_cnn_cache": torch.randn(batch_size, 2, 3), + "conformer_att_cache": torch.randn(2, batch_size, 2, 4, 6), + "estimator_cnn_cache": torch.randn(4, 2, 2 * batch_size, 3, 2), + "estimator_att_cache": torch.randn(4, 2, 2 * batch_size, 2, frames, 6), + } + + restored = adapter._split_flow_cache(stacked, batch_size, estimator_att_keep=(prompt_len, 4)) + + att = stacked["estimator_att_cache"] + trimmed = torch.cat((att[..., :prompt_len, :], att[..., -4:, :]), dim=4) + for row, request in enumerate(restored): + expected = torch.cat( + (trimmed[:, :, row : row + 1], trimmed[:, :, batch_size + row : batch_size + row + 1]), + dim=2, + ).to(attention_dtype) + assert request["estimator_att_cache"].dtype == attention_dtype + torch.testing.assert_close(request["estimator_att_cache"], expected) + torch.testing.assert_close( + request["estimator_cnn_cache"], + torch.cat( + ( + stacked["estimator_cnn_cache"][:, :, row : row + 1], + stacked["estimator_cnn_cache"][:, :, batch_size + row : batch_size + row + 1], + ), + dim=2, + ), + ) + untrimmed = adapter._split_flow_cache(stacked, batch_size, estimator_att_keep=(frames, 4)) + assert untrimmed[0]["estimator_att_cache"].shape[4] == frames + + def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep(): token2wav = _FakeToken2Wav() adapter = BatchedToken2Wav(token2wav, bfloat16_attention_cache=True) @@ -970,9 +1084,8 @@ def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep() states = adapter.setup_batch(prompt, 2) assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in states) - assert ( - states[0].flow_cache["estimator_att_cache"].data_ptr() != states[1].flow_cache["estimator_att_cache"].data_ptr() - ) + # Both requests start from the one shared prompt state. + assert states[0] is states[1] stacked = adapter._stack_flow_cache(states) assert stacked["estimator_att_cache"].dtype == torch.bfloat16 @@ -991,6 +1104,10 @@ def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep() last_chunk=False, ) assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in next_states) + assert ( + next_states[0].flow_cache["estimator_att_cache"].data_ptr() + != next_states[1].flow_cache["estimator_att_cache"].data_ptr() + ) assert token2wav.flow.decoder.estimator.attention_cache_dtypes[-2:] == [torch.float32, torch.float32] @@ -1117,9 +1234,17 @@ def test_initial_empty_segment_marker_initializes_stream_without_audio(): assert "duplex" in model._states -@pytest.mark.parametrize("setup_cache_size,expected_calls", [(0, [2, 2, 1]), (1, [2, 1])]) +@pytest.mark.parametrize("setup_cache_size,expected_calls", [(0, [1, 1, 1]), (1, [1])]) def test_initial_empty_segment_markers_respect_initial_batch_limit(setup_cache_size, expected_calls): model, token2wav = _model(initial=2, setup_cache_size=setup_cache_size) + setup_batches: list[int] = [] + setup_batch = model.backend.setup_batch + + def record_setup(features, batch_size): + setup_batches.append(batch_size) + return setup_batch(features, batch_size) + + model.backend.setup_batch = record_setup names = ["a", "b", "c", "d", "e"] boundaries = [] for name in names: @@ -1135,6 +1260,8 @@ def test_initial_empty_segment_markers_respect_initial_batch_limit(setup_cache_s output = _forward(model, boundaries) + assert setup_batches == [2, 2, 1] + # One shared prompt solve per setup, or one in all with the setup cache. assert token2wav.flow.encoder.calls == expected_calls assert [audio.numel() for audio in output.multimodal_outputs["model_outputs"]] == [0] * len(names) assert set(model._states) == set(names) @@ -1631,8 +1758,8 @@ def test_setup_cache_misses_when_cfm_graph_padding_is_disabled(monkeypatch, disa assert unpadded[0] is not padded[0] assert reused[0] is unpadded[0] assert adapter.flow.encoder.calls == [1, 1] - assert padded_keys == [(prompt.cache_key, 1, 16)] - assert list(adapter._setup_cache) == [(prompt.cache_key, 1, 0)] + assert padded_keys == [(prompt.cache_key, 16)] + assert list(adapter._setup_cache) == [(prompt.cache_key, 0)] def test_padded_chunk_keeps_the_cross_chunk_caches_on_the_valid_boundary(): 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 555e72e46..8283bf635 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 @@ -8,15 +8,18 @@ from unittest.mock import Mock import pytest import torch import torch.nn as nn +import torch.nn.functional as F from vllm.platforms import current_platform import vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper as wrapper_module from vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan import ( HiFTGenerator, ) +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import ( CFMGraphWrapper, HiFTGraphWrapper, + WholeEulerCFMGraphWrapper, ) pytestmark = [pytest.mark.core_model, pytest.mark.cuda] @@ -108,6 +111,7 @@ def _fake_wrapper(monkeypatch: pytest.MonkeyPatch) -> HiFTGraphWrapper: wrapper.static_cache_source_outputs = {} wrapper.lazy_graph_count = 0 wrapper.max_lazy_graphs = 1 + wrapper.max_serial_batch = 8 wrapper.decode_fn = Mock(return_value=(torch.tensor([[99.0]]), torch.tensor([[[98.0]]]))) wrapper.finalize_fn = lambda magnitude, phase: magnitude + phase @@ -148,6 +152,20 @@ def test_lazy_capture_limit_falls_back_to_eager(monkeypatch: pytest.MonkeyPatch) assert result is wrapper.decode_fn.return_value +def test_multibatch_serial_replay_with_batch1_graph(monkeypatch: pytest.MonkeyPatch) -> None: + wrapper = _fake_wrapper(monkeypatch) + # 3-batch input with only batch size 1 supported + speech_feat = torch.randn(3, 80, 7) + cache_source = torch.zeros(3, 1, 0) + speech, source = wrapper.replay(speech_feat, cache_source) + + # Lazily captured once for shape (1, 7, 0) + wrapper._capture.assert_called_once_with(1, 7, 0) + assert speech.shape == (3, 1, 7) + assert source.shape == (3, 1, 7) + wrapper.decode_fn.assert_not_called() + + # --------------------------------------------------------------------------- # CFMGraphWrapper tests # --------------------------------------------------------------------------- @@ -315,6 +333,7 @@ def _cfm_mock_wrapper(monkeypatch: pytest.MonkeyPatch, *, max_graphs: int = 1) - monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) wrapper = object.__new__(CFMGraphWrapper) wrapper.max_graphs = max_graphs + wrapper.max_serial_batch = 8 wrapper.enabled = True wrapper.graph_fn = Mock(return_value=torch.tensor([42.0])) wrapper.device = torch.device("cuda") @@ -582,3 +601,1721 @@ def test_cfm_none_cache_parity_between_graph_and_eager( torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5) torch.testing.assert_close(graph_cnn, eager_cnn_out, rtol=1e-4, atol=1e-5) torch.testing.assert_close(graph_att, eager_att_out, rtol=1e-4, atol=1e-5) + + +# --------------------------------------------------------------------------- +# WholeEulerCFMGraphWrapper tests +# --------------------------------------------------------------------------- + + +class _WholeEulerDiT(nn.Module): + """DiT module with t_embedder and buffer dimensions for Whole-Euler testing.""" + + def __init__(self, x_dim: int = 4, hidden: int = 8, depth: int = 2) -> None: + super().__init__() + self.x_dim = x_dim + self.in_proj = nn.Linear(x_dim * 4, hidden) + self.blocks = nn.ModuleList([nn.LayerNorm(hidden) for _ in range(depth)]) + self.final_layer = nn.Linear(hidden, x_dim) + for b in self.blocks: + b.conv = SimpleNamespace( + in_channels=4, + out_channels=4, + block=[None, SimpleNamespace(causal_padding=[2])], + ) + b.attn = SimpleNamespace(num_heads=2, head_dim=4) + + def t_embedder(self, t: torch.Tensor) -> torch.Tensor: + return t[:, None].expand(-1, 8) + + def blocks_forward_chunk( + self, + x: torch.Tensor, + t: torch.Tensor, + mask: torch.Tensor | None, + cnn_cache: torch.Tensor | None = None, + att_cache: torch.Tensor | None = None, + cnn_cache_buffer: torch.Tensor | None = None, + att_cache_buffer: torch.Tensor | None = None, + ) -> torch.Tensor: + assert cnn_cache_buffer is not None + assert att_cache_buffer is not None + x = x.transpose(1, 2) + x = self.in_proj(x) + for b_idx in range(len(self.blocks)): + cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None + if cnn_b is not None: + x[:, : cnn_b.shape[2], :] += cnn_b.transpose(1, 2) + att_b = att_cache[b_idx] if att_cache is not None else None + old_len = 0 + if att_b is not None and att_b.shape[2] > 0: + old_len = att_b.shape[2] + x += att_b.sum(dim=(1, 2), keepdim=False).unsqueeze(1) + att_cache_buffer[b_idx][:, :, :old_len, :] = att_b + x = self.blocks[b_idx](x) + x = x + t + cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous() + dt = x.shape[1] + att_cache_buffer[b_idx][:, :, old_len : old_len + dt, :] = x.unsqueeze(1) + x = self.final_layer(x) + return x.transpose(1, 2) + + +def _eager_solve_euler( + estimator: nn.Module, + x: torch.Tensor, + mu_cfg: torch.Tensor, + speakers_cfg: torch.Tensor, + cond_cfg: torch.Tensor, + cnn_cache: torch.Tensor | None, + att_cache: torch.Tensor | None, + attn_mask: torch.Tensor | None, + timeline: torch.Tensor, + *, + n_timesteps: int = 10, + inference_cfg_rate: float = 0.7, + mel_frames: int | None = None, + pad_frames: int = 0, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + cur_x = x.clone() + batch_size = int(x.shape[0]) + width = int(mu_cfg.shape[2]) + if mel_frames is None: + mel_frames = width - pad_frames + speaker_features = speakers_cfg.unsqueeze(-1).expand(-1, -1, width) + if pad_frames > 0: + speaker_features = speaker_features.clone() + speaker_features[..., mel_frames:] = 0.0 + + depth = len(estimator.blocks) + block0 = estimator.blocks[0] + cnn_channels = int(block0.conv.in_channels + block0.conv.out_channels) + cnn_width = int(block0.conv.block[1].causal_padding[0]) + heads = int(block0.attn.num_heads) + att_width = int(block0.attn.head_dim * 2) + offset = int(att_cache.shape[4]) if att_cache is not None else 0 + + next_cnns = [] + next_atts = [] + + for step in range(n_timesteps): + t_val = timeline[step].expand(2 * batch_size) + dt = timeline[step + 1] - timeline[step] + time_embedding = estimator.t_embedder(t_val).unsqueeze(1) + x_cfg = torch.cat((cur_x, cur_x), dim=0) + estimator_input = torch.cat((x_cfg, mu_cfg, speaker_features, cond_cfg), dim=1) + + c_out = torch.empty(depth, 2 * batch_size, cnn_channels, cnn_width, device=x.device, dtype=x.dtype) + a_out = torch.empty(depth, 2 * batch_size, heads, offset + width, att_width, device=x.device, dtype=x.dtype) + old_c = cnn_cache[step] if cnn_cache is not None else [None] * depth + old_a = att_cache[step] if att_cache is not None else [None] * depth + + est = estimator.blocks_forward_chunk( + estimator_input, + time_embedding, + attn_mask, + old_c, + old_a, + c_out, + a_out, + ) + + if pad_frames > 0: + wrapper_module._zero_padded_cnn_cache(c_out, estimator, pad_frames) + a_out[..., mel_frames : mel_frames + pad_frames, :] = 0.0 + + conditional, unconditional = est.split(batch_size, dim=0) + velocity = (1.0 + inference_cfg_rate) * conditional - inference_cfg_rate * unconditional + cur_x = cur_x + dt * velocity + if pad_frames > 0: + cur_x[..., mel_frames:] = 0.0 + + next_cnns.append(c_out) + next_atts.append(a_out) + + return cur_x[:, :, :mel_frames], torch.stack(next_cnns), torch.stack(next_atts) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_whole_euler_graph_replay_matches_eager_for_uncached_and_cached_shapes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=32) + + batch_size = 2 + chunk_size = 10 + x = torch.randn(batch_size, 4, chunk_size, device="cuda") + mu = torch.randn(batch_size, 4, chunk_size, device="cuda") + speakers = torch.randn(batch_size, 4, device="cuda") + cond = torch.randn(batch_size, 4, chunk_size, device="cuda") + + mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=0) + speakers_cfg = torch.cat((speakers, torch.zeros_like(speakers)), dim=0) + cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=0) + + # 1. Uncached case + eager_x, eager_cnn, eager_att = _eager_solve_euler( + estimator, x, mu_cfg, speakers_cfg, cond_cfg, None, None, None, wrapper.timeline + ) + graph_res = wrapper.replay( + x=x.clone(), + mu_cfg=mu_cfg.clone(), + speakers_cfg=speakers_cfg.clone(), + cond_cfg=cond_cfg.clone(), + cnn_cache=None, + att_cache=None, + attn_mask=None, + ) + assert graph_res is not None + graph_x, graph_cnn, graph_att = graph_res + + torch.testing.assert_close(graph_x, eager_x, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_cnn, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_att, rtol=1e-4, atol=1e-5) + + # 2. Cached case (carry forward previous chunk outputs) + x2 = torch.randn(batch_size, 4, chunk_size, device="cuda") + eager_x2, eager_cnn2, eager_att2 = _eager_solve_euler( + estimator, x2, mu_cfg, speakers_cfg, cond_cfg, eager_cnn, eager_att, None, wrapper.timeline + ) + graph_res2 = wrapper.replay( + x=x2.clone(), + mu_cfg=mu_cfg.clone(), + speakers_cfg=speakers_cfg.clone(), + cond_cfg=cond_cfg.clone(), + cnn_cache=graph_cnn, + att_cache=graph_att, + attn_mask=None, + ) + assert graph_res2 is not None + graph_x2, graph_cnn2, graph_att2 = graph_res2 + + torch.testing.assert_close(graph_x2, eager_x2, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn2, eager_cnn2, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att2, eager_att2, rtol=1e-4, atol=1e-5) + + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_whole_euler_graph_with_padding_matches_eager( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(42) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=16) + + batch_size = 2 + mel_frames = 8 + pad_frames = 2 + chunk_size = mel_frames + pad_frames + + x = torch.randn(batch_size, 4, chunk_size, device="cuda") + x[:, :, mel_frames:] = 0.0 + mu = torch.randn(batch_size, 4, chunk_size, device="cuda") + speakers = torch.randn(batch_size, 4, device="cuda") + cond = torch.randn(batch_size, 4, chunk_size, device="cuda") + + mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=0) + speakers_cfg = torch.cat((speakers, torch.zeros_like(speakers)), dim=0) + cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=0) + + attn_mask = torch.ones(2 * batch_size, chunk_size, chunk_size, dtype=torch.bool, device="cuda") + attn_mask[:, :, mel_frames : mel_frames + pad_frames] = False + + eager_x, eager_cnn, eager_att = _eager_solve_euler( + estimator, + x, + mu_cfg, + speakers_cfg, + cond_cfg, + None, + None, + attn_mask, + wrapper.timeline, + mel_frames=mel_frames, + pad_frames=pad_frames, + ) + + graph_res = wrapper.replay( + x=x.clone(), + mu_cfg=mu_cfg.clone(), + speakers_cfg=speakers_cfg.clone(), + cond_cfg=cond_cfg.clone(), + cnn_cache=None, + att_cache=None, + attn_mask=attn_mask.clone(), + mel_frames=mel_frames, + pad_frames=pad_frames, + ) + assert graph_res is not None + graph_x, graph_cnn, graph_att = graph_res + + assert graph_x.shape[2] == mel_frames + assert eager_x.shape[2] == mel_frames + torch.testing.assert_close(graph_x, eager_x, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_cnn, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_att, rtol=1e-4, atol=1e-5) + + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_whole_euler_cache_flushes_whole_generation( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: torch.cuda.graph_pool_handle()) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=2) + + def _call(w: int): + x = torch.randn(1, 4, w, device="cuda") + mu_cfg = torch.randn(2, 4, w, device="cuda") + spk_cfg = torch.randn(2, 4, device="cuda") + cond_cfg = torch.randn(2, 4, w, device="cuda") + return wrapper.replay( + x=x, + mu_cfg=mu_cfg, + speakers_cfg=spk_cfg, + cond_cfg=cond_cfg, + cnn_cache=None, + att_cache=None, + ) + + _call(10) + assert len(wrapper._cache) == 1 + assert wrapper._stats["captures"] == 1 + + # Same shape hits + _call(10) + assert wrapper._stats["hits"] == 1 + + # Second distinct shape fills cache to max_graphs=2 + _call(12) + assert len(wrapper._cache) == 2 + assert wrapper._stats["captures"] == 2 + assert wrapper._stats["flushes"] == 0 + + # Third distinct shape exceeds max_graphs and triggers whole-generation flush + _call(14) + assert wrapper._stats["flushes"] == 1 + assert len(wrapper._cache) == 1 + assert wrapper._stats["captures"] == 3 + + # Replay evicted shape 10 and assert numerical match with eager + torch.manual_seed(42) + x10 = torch.randn(1, 4, 10, device="cuda") + mu10 = torch.randn(2, 4, 10, device="cuda") + spk10 = torch.randn(2, 4, device="cuda") + cond10 = torch.randn(2, 4, 10, device="cuda") + res10 = wrapper.replay( + x=x10.clone(), + mu_cfg=mu10.clone(), + speakers_cfg=spk10.clone(), + cond_cfg=cond10.clone(), + cnn_cache=None, + att_cache=None, + ) + assert res10 is not None + eager_x10, _, _ = _eager_solve_euler( + estimator, + x10.clone(), + mu10.clone(), + spk10.clone(), + cond10.clone(), + None, + None, + None, + wrapper.timeline, + mel_frames=10, + pad_frames=0, + ) + torch.testing.assert_close(res10[0], eager_x10, rtol=1e-4, atol=1e-5) + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_whole_euler_arena_cleaned_when_first_capture_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=2) + + # Force failure during capture + def mock_fail(*args, **kwargs): + raise RuntimeError("Simulated first capture failure") + + monkeypatch.setattr(wrapper, "_run_euler_loop", mock_fail) + + x = torch.randn(1, 4, 10, device="cuda") + mu_cfg = torch.randn(2, 4, 10, device="cuda") + spk_cfg = torch.randn(2, 4, device="cuda") + cond_cfg = torch.randn(2, 4, 10, device="cuda") + + res = wrapper.replay( + x=x, + mu_cfg=mu_cfg, + speakers_cfg=spk_cfg, + cond_cfg=cond_cfg, + cnn_cache=None, + att_cache=None, + ) + assert res is None + assert wrapper.enabled is False + assert len(wrapper._cache) == 0 + # Arena must be cleared even when cache has zero entries + assert not wrapper.arena._buffers + + +@pytest.mark.parametrize( + ("micro", "pad_max_rows", "plans"), + [ + (4, 1, {1: [(1, 1)], 2: [(2, 2)], 3: [(4, 3)], 6: [(4, 4), (2, 2)], 7: [(4, 4), (4, 3)]}), + (4, 0, {3: [(2, 2), (1, 1)], 7: [(4, 4), (2, 2), (1, 1)]}), + ( + 8, + 2, + { + 5: [(4, 4), (1, 1)], + 6: [(8, 6)], + 12: [(8, 8), (4, 4)], + 13: [(8, 8), (4, 4), (1, 1)], + 16: [(8, 8), (8, 8)], + }, + ), + ], +) +def test_whole_euler_plan_groups_uses_power_of_two_graph_batches( + micro: int, + pad_max_rows: int, + plans: dict[int, list[tuple[int, int]]], +) -> None: + wrapper = WholeEulerCFMGraphWrapper( + estimator=_WholeEulerDiT(), + n_timesteps=10, + max_graphs=8, + max_graph_batch=16, + micro_batch_size=micro, + pad_max_rows=pad_max_rows, + ) + for batch_size, plan in plans.items(): + assert wrapper._plan_groups(batch_size) == plan + assert sum(rows for _, rows in plan) == batch_size + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +@pytest.mark.parametrize(("pad_max_rows", "graph_batches"), [(1, {4}), (0, {1, 2})]) +def test_whole_euler_multibatch_remainder_replay( + monkeypatch: pytest.MonkeyPatch, + pad_max_rows: int, + graph_batches: set[int], +) -> None: + """B=3 runs one padded B=4 replay by default, or B=2 + B=1 replays; both match eager.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=4, pad_max_rows=pad_max_rows) + + batch_size = 3 + w = 8 + x = torch.randn(batch_size, 4, w, device="cuda") + mu_cfg = torch.randn(2 * batch_size, 4, w, device="cuda") + spk_cfg = torch.randn(2 * batch_size, 4, device="cuda") + cond_cfg = torch.randn(2 * batch_size, 4, w, device="cuda") + + out_mel, out_cnn, out_att = wrapper.replay( + x=x.clone(), + mu_cfg=mu_cfg, + speakers_cfg=spk_cfg, + cond_cfg=cond_cfg, + cnn_cache=None, + att_cache=None, + ) + assert out_mel.shape == (batch_size, 4, w) + assert out_cnn.shape[2] == 2 * batch_size + assert out_att.shape[2] == 2 * batch_size + assert {key[1] for key in wrapper._cache} == graph_batches + + eager_x, eager_cnn, eager_att = _eager_solve_euler( + estimator, x, mu_cfg, spk_cfg, cond_cfg, None, None, None, wrapper.timeline + ) + torch.testing.assert_close(out_mel, eager_x, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(out_cnn, eager_cnn, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(out_att, eager_att, rtol=1e-4, atol=1e-5) + + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_hierarchical_microbatch_b4_and_b8(monkeypatch: pytest.MonkeyPatch) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=8, max_graph_batch=8) + + w = 8 + # 1. Native B=4 capture + batch_size = 4 + x4 = torch.randn(batch_size, 4, w, device="cuda") + mu4 = torch.randn(2 * batch_size, 4, w, device="cuda") + spk4 = torch.randn(2 * batch_size, 4, device="cuda") + cond4 = torch.randn(2 * batch_size, 4, w, device="cuda") + + out_mel4, out_cnn4, out_att4 = wrapper.replay( + x=x4, + mu_cfg=mu4, + speakers_cfg=spk4, + cond_cfg=cond4, + cnn_cache=None, + att_cache=None, + ) + assert out_mel4 is not None + assert out_mel4.shape == (4, 4, w) + assert out_cnn4.shape[2] == 8 + assert out_att4.shape[2] == 8 + assert any(key[1] == 4 for key in wrapper._cache) + + # 2. B=8 partitions into [4, 4], reusing the native B=4 graph + batch_size = 8 + x8 = torch.randn(batch_size, 4, w, device="cuda") + mu8 = torch.randn(2 * batch_size, 4, w, device="cuda") + spk8 = torch.randn(2 * batch_size, 4, device="cuda") + cond8 = torch.randn(2 * batch_size, 4, w, device="cuda") + + out_mel8, out_cnn8, out_att8 = wrapper.replay( + x=x8, + mu_cfg=mu8, + speakers_cfg=spk8, + cond_cfg=cond8, + cnn_cache=None, + att_cache=None, + ) + assert out_mel8 is not None + assert out_mel8.shape == (8, 4, w) + assert out_cnn8.shape[2] == 16 + assert out_att8.shape[2] == 16 + + # 3. B=5 partitions into [4, 1], capturing B=1 graph as well + batch_size = 5 + x5 = torch.randn(batch_size, 4, w, device="cuda") + mu5 = torch.randn(2 * batch_size, 4, w, device="cuda") + spk5 = torch.randn(2 * batch_size, 4, device="cuda") + cond5 = torch.randn(2 * batch_size, 4, w, device="cuda") + + out_mel5, out_cnn5, out_att5 = wrapper.replay( + x=x5, + mu_cfg=mu5, + speakers_cfg=spk5, + cond_cfg=cond5, + cnn_cache=None, + att_cache=None, + ) + assert out_mel5 is not None + assert out_mel5.shape == (5, 4, w) + assert out_cnn5.shape[2] == 10 + assert out_att5.shape[2] == 10 + cached_batches = {key[1] for key in wrapper._cache} + assert 4 in cached_batches + assert 1 in cached_batches + + # 4. B=16 partitions into [4, 4, 4, 4] with max_graph_batch=16 + wrapper_16 = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=8, max_graph_batch=16) + batch_size = 16 + x16 = torch.randn(batch_size, 4, w, device="cuda") + mu16 = torch.randn(2 * batch_size, 4, w, device="cuda") + spk16 = torch.randn(2 * batch_size, 4, device="cuda") + cond16 = torch.randn(2 * batch_size, 4, w, device="cuda") + + out_mel16, out_cnn16, out_att16 = wrapper_16.replay( + x=x16, + mu_cfg=mu16, + speakers_cfg=spk16, + cond_cfg=cond16, + cnn_cache=None, + att_cache=None, + ) + assert out_mel16 is not None + assert out_mel16.shape == (16, 4, w) + assert out_cnn16.shape[2] == 32 + assert out_att16.shape[2] == 32 + + # B > 16 returns None (eager fallback) + x17 = torch.randn(17, 4, w, device="cuda") + mu17 = torch.randn(34, 4, w, device="cuda") + spk17 = torch.randn(34, 4, device="cuda") + cond17 = torch.randn(34, 4, w, device="cuda") + assert ( + wrapper_16.replay( + x=x17, + mu_cfg=mu17, + speakers_cfg=spk17, + cond_cfg=cond17, + cnn_cache=None, + att_cache=None, + ) + is None + ) + + wrapper._flush() + wrapper_16._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_same_bucket_different_padding_hits_cache(monkeypatch: pytest.MonkeyPatch) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=4) + + batch_size = 1 + mel_width = 16 + + # Call 1: mel_frames=12, pad_frames=4 -> captures graph 1 + x1 = torch.randn(batch_size, 4, mel_width, device="cuda") + mu1 = torch.randn(2 * batch_size, 4, mel_width, device="cuda") + spk1 = torch.randn(2 * batch_size, 4, device="cuda") + cond1 = torch.randn(2 * batch_size, 4, mel_width, device="cuda") + mask1 = torch.ones(2 * batch_size, mel_width, mel_width, dtype=torch.bool, device="cuda") + mask1[:, :, 12:] = False + + res1 = wrapper.replay( + x=x1, + mu_cfg=mu1, + speakers_cfg=spk1, + cond_cfg=cond1, + cnn_cache=None, + att_cache=None, + attn_mask=mask1, + mel_frames=12, + pad_frames=4, + ) + assert res1 is not None + assert res1[0].shape[-1] == 12 + assert wrapper._stats["captures"] == 1 + assert wrapper._stats["hits"] == 0 + + # Call 2: mel_frames=10, pad_frames=6 (different unpadded length, same bucket mel_width) + x2 = torch.randn(batch_size, 4, mel_width, device="cuda") + mu2 = torch.randn(2 * batch_size, 4, mel_width, device="cuda") + spk2 = torch.randn(2 * batch_size, 4, device="cuda") + cond2 = torch.randn(2 * batch_size, 4, mel_width, device="cuda") + mask2 = torch.ones(2 * batch_size, mel_width, mel_width, dtype=torch.bool, device="cuda") + mask2[:, :, 10:] = False + + res2 = wrapper.replay( + x=x2, + mu_cfg=mu2, + speakers_cfg=spk2, + cond_cfg=cond2, + cnn_cache=None, + att_cache=None, + attn_mask=mask2, + mel_frames=10, + pad_frames=6, + ) + assert res2 is not None + assert res2[0].shape[-1] == 10 + # Must be a cache hit! + assert wrapper._stats["captures"] == 1 + assert wrapper._stats["hits"] == 1 + assert len(wrapper._cache) == 1 + + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_mask_presence_does_not_split_key(monkeypatch: pytest.MonkeyPatch) -> None: + """#7416 fills static buffers every replay; mask presence is data, not a key.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=4) + + batch_size = 1 + w = 16 + x = torch.randn(batch_size, 4, w, device="cuda") + mu = torch.randn(2 * batch_size, 4, w, device="cuda") + spk = torch.randn(2 * batch_size, 4, device="cuda") + cond = torch.randn(2 * batch_size, 4, w, device="cuda") + mask = torch.ones(2 * batch_size, w, w, dtype=torch.bool, device="cuda") + mask[:, :, 12:] = False + + assert ( + wrapper.replay( + x=x, + mu_cfg=mu, + speakers_cfg=spk, + cond_cfg=cond, + cnn_cache=None, + att_cache=None, + attn_mask=None, + ) + is not None + ) + assert wrapper._stats["captures"] == 1 + + assert ( + wrapper.replay( + x=x, + mu_cfg=mu, + speakers_cfg=spk, + cond_cfg=cond, + cnn_cache=None, + att_cache=None, + attn_mask=mask, + ) + is not None + ) + assert wrapper._stats["captures"] == 1 + assert wrapper._stats["hits"] == 1 + assert len(wrapper._cache) == 1 + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_query_capture_bucket_reuses_decode_widths(monkeypatch: pytest.MonkeyPatch) -> None: + """Same capture grid as _cfm_pad_frames, but 64 so 16/32/48/64 share one graph.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=4, + query_bucket_frames=64, + ) + + def _call(width: int): + x = torch.randn(1, 4, width, device="cuda") + mu = torch.randn(2, 4, width, device="cuda") + spk = torch.randn(2, 4, device="cuda") + cond = torch.randn(2, 4, width, device="cuda") + return wrapper.replay( + x=x, + mu_cfg=mu, + speakers_cfg=spk, + cond_cfg=cond, + cnn_cache=None, + att_cache=None, + ) + + # The wide offset-0 prompt solve comes first, as in serving; it reserves + # the storage the decode widths then share. + assert _call(304) is not None + assert _call(16) is not None + assert _call(48) is not None + assert _call(64) is not None + assert wrapper._stats["captures"] == 2 + assert wrapper._stats["hits"] == 2 + assert wrapper._stats["flushes"] == 0 + query_caps = {key[2] for key in wrapper._cache} + assert query_caps == {64, 320} + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_max_serial_batch_dispatch(monkeypatch: pytest.MonkeyPatch) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + # Configure threshold = 2 + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=4, + max_serial_batch=2, + ) + + w = 8 + # Batch = 2 <= max_serial_batch: uses CUDA graph serial replay + x_small = torch.randn(2, 4, w, device="cuda") + mu_small = torch.randn(4, 4, w, device="cuda") + spk_small = torch.randn(4, 4, device="cuda") + cond_small = torch.randn(4, 4, w, device="cuda") + + res_small = wrapper.replay( + x=x_small, + mu_cfg=mu_small, + speakers_cfg=spk_small, + cond_cfg=cond_small, + cnn_cache=None, + att_cache=None, + ) + assert res_small is not None + assert wrapper._stats["captures"] == 1 + + # Batch = 4 > max_serial_batch: returns None to fall back to batched eager + x_large = torch.randn(4, 4, w, device="cuda") + mu_large = torch.randn(8, 4, w, device="cuda") + spk_large = torch.randn(8, 4, device="cuda") + cond_large = torch.randn(8, 4, w, device="cuda") + + res_large = wrapper.replay( + x=x_large, + mu_cfg=mu_large, + speakers_cfg=spk_large, + cond_cfg=cond_large, + cnn_cache=None, + att_cache=None, + ) + assert res_large is None + # No new graph captures triggered for large batch + assert wrapper._stats["captures"] == 1 + + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_cfm_step_graph_captures_distinct_batch_shapes(monkeypatch: pytest.MonkeyPatch) -> None: + """Step-level CFM keys by tensor signature (upstream #7416): a larger CFG batch is another shape.""" + 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) + + wrapper.replay(*_cfm_inputs(batch_size=4, chunk_size=8, old_att_len=0)) + assert wrapper._stats["captures"] == 1 + assert wrapper._stats["eager"] == 0 + + wrapper.replay(*_cfm_inputs(batch_size=8, chunk_size=8, old_att_len=0)) + assert wrapper._stats["captures"] == 2 + assert wrapper._stats["eager"] == 0 + + wrapper._flush() + + +def test_zero_padded_cnn_cache_clears_every_timestep_of_a_whole_euler_cache() -> None: + # Whole-Euler keeps every timestep's CNN cache as (n_timesteps, depth, 2B, C, W) + # and hands the helper a block-first view of it. + estimator = _WholeEulerDiT() + width = int(estimator.blocks[0].conv.block[1].causal_padding[0]) + cache = torch.randn(10, len(estimator.blocks), 4, 8, width) + expected = cache.clone() + expected[..., width - 1 :] = 0.0 + + wrapper_module._zero_padded_cnn_cache(cache.transpose(0, 1), estimator, 1) + + torch.testing.assert_close(cache, expected) + + +def test_euler_step_applies_cfg_guidance() -> None: + torch.manual_seed(42) + B, C, T = 2, 8, 16 + cur_x = torch.randn(B, C, T) + estimate = torch.randn(2 * B, C, T) + dt = 0.1 + cfg = 0.7 + + cond, uncond = estimate.split(B, dim=0) + v = (1.0 + cfg) * cond - cfg * uncond + expected = cur_x + dt * v + + actual = wrapper_module._euler_step(cur_x.clone(), estimate, dt, cfg, B) + torch.testing.assert_close(actual, expected) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_execution_arena_buffer_reuse(monkeypatch: pytest.MonkeyPatch) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=4) + + w = 10 + x1 = torch.randn(1, 4, w, device="cuda") + mu1 = torch.randn(2, 4, w, device="cuda") + spk1 = torch.randn(2, 4, device="cuda") + cond1 = torch.randn(2, 4, w, device="cuda") + + # First call captures shape 1 (offset=0); att_keep reserves both offsets. + res1 = wrapper.replay( + x=x1, + mu_cfg=mu1, + speakers_cfg=spk1, + cond_cfg=cond1, + cnn_cache=None, + att_cache=None, + att_keep=(w, w), + ) + assert res1 is not None + assert wrapper._stats["captures"] == 1 + arena = wrapper.arena + shared = {key: buf for key, buf in arena._buffers.items() if key[0] in ("time_embeddings", "cnn_in", "cnn_out")} + assert sorted(key[0] for key in shared) == ["cnn_in", "cnn_out", "time_embeddings"] + + # Second call with cache (offset=10) + res2 = wrapper.replay( + x=x1, + mu_cfg=mu1, + speakers_cfg=spk1, + cond_cfg=cond1, + cnn_cache=res1[1], + att_cache=res1[2], + ) + assert res2 is not None + assert wrapper._stats["captures"] == 2 + + # Verify that the shared buffers are the exact same tensor instances (no duplicate allocation) + assert all(arena._buffers[key] is buf for key, buf in shared.items()) + + wrapper._flush() + assert not arena._buffers + + +def _whole_euler_chunk(batch_size: int, width: int) -> dict[str, torch.Tensor]: + return { + "x": torch.randn(batch_size, 4, width, device="cuda"), + "mu_cfg": torch.randn(2 * batch_size, 4, width, device="cuda"), + "speakers_cfg": torch.randn(2 * batch_size, 4, device="cuda"), + "cond_cfg": torch.randn(2 * batch_size, 4, width, device="cuda"), + } + + +def _split_cfg_rows(stacked: torch.Tensor, batch_size: int) -> list[torch.Tensor]: + return [ + torch.cat((stacked[:, :, row : row + 1], stacked[:, :, batch_size + row : batch_size + row + 1]), dim=2) + for row in range(batch_size) + ] + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_graphs_share_one_attention_storage(monkeypatch: pytest.MonkeyPatch) -> None: + """Every offset and batch size reads and writes prefix views of one storage, not a copy each.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=8) + + w = 8 + cnn = att = None + eager_cnn = eager_att = None + for _ in range(3): + chunk = _whole_euler_chunk(1, w) + # att_keep sums to the third chunk's length: it reserves that much, and trims nothing. + graph_x, cnn, att = wrapper.replay(**chunk, cnn_cache=cnn, att_cache=att, att_keep=(w, 2 * w)) + eager_x, eager_cnn, eager_att = _eager_solve_euler( + estimator, + chunk["x"], + chunk["mu_cfg"], + chunk["speakers_cfg"], + chunk["cond_cfg"], + eager_cnn, + eager_att, + None, + wrapper.timeline, + ) + torch.testing.assert_close(graph_x, eager_x, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(att, eager_att, rtol=1e-4, atol=1e-5) + + # A B=2 stream at a known offset maps onto the same storage. + _, pair_cnn, pair_att = wrapper.replay(**_whole_euler_chunk(2, w), cnn_cache=None, att_cache=None) + wrapper.replay(**_whole_euler_chunk(2, w), cnn_cache=pair_cnn, att_cache=pair_att) + + assert wrapper._stats["captures"] == 5 + assert {(key[1], key[3]) for key in wrapper._cache} == {(1, 0), (1, w), (1, 2 * w), (2, 0), (2, w)} + storage = wrapper.arena._att + assert storage is not None + assert int(storage.shape[2]) == 2 * wrapper.micro_batch_size + for statics, _final_x, _out_cnn, out_att, _graph in wrapper._cache.values(): + assert statics.att_cache.untyped_storage().data_ptr() == storage.untyped_storage().data_ptr() + assert out_att.untyped_storage().data_ptr() == storage.untyped_storage().data_ptr() + wrapper._flush() + assert wrapper.arena._att is None + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_storage_growth_retires_the_graphs_on_the_old_storage(monkeypatch: pytest.MonkeyPatch) -> None: + """A capture that outgrows the attention storage must not leave the old one alive beside it.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=8) + + w = 8 + first = _whole_euler_chunk(1, w) + # att_keep reserves 2 * w frames: the next w-wide chunk fits, a 2w-wide one does not. + _, cnn, att = wrapper.replay(**first, cnn_cache=None, att_cache=None, att_keep=(w, w)) + _, eager_cnn, eager_att = _eager_solve_euler( + estimator, + first["x"], + first["mu_cfg"], + first["speakers_cfg"], + first["cond_cfg"], + None, + None, + None, + wrapper.timeline, + ) + wrapper.replay(**_whole_euler_chunk(1, w), cnn_cache=cnn, att_cache=att) + assert wrapper._stats["flushes"] == 0 + assert len(wrapper._cache) == 2 + + wide = _whole_euler_chunk(1, 2 * w) + graph_x, _, graph_att = wrapper.replay(**wide, cnn_cache=cnn, att_cache=att) + + assert wrapper._stats["flushes"] == 1 + assert len(wrapper._cache) == 1 + storage = wrapper.arena._att + assert storage is not None and int(storage.shape[4]) >= 3 * w + for statics, _final_x, _out_cnn, out_att, _graph in wrapper._cache.values(): + assert statics.att_cache.untyped_storage().data_ptr() == storage.untyped_storage().data_ptr() + assert out_att.untyped_storage().data_ptr() == storage.untyped_storage().data_ptr() + eager_x, _, eager_wide_att = _eager_solve_euler( + estimator, + wide["x"], + wide["mu_cfg"], + wide["speakers_cfg"], + wide["cond_cfg"], + eager_cnn, + eager_att, + None, + wrapper.timeline, + ) + torch.testing.assert_close(graph_x, eager_x, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_att, eager_wide_att, rtol=1e-4, atol=1e-5) + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_per_request_rows_match_stacked(monkeypatch: pytest.MonkeyPatch) -> None: + """Per-request cache in/out (with the streaming trim) equals the stacked path, row for row.""" + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(0) + estimator = _WholeEulerDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=estimator, n_timesteps=10, max_graphs=8, max_graph_batch=8) + + batch_size = 6 # one B=4 replay plus one B=2 replay + w = 8 + _, cnn, att = wrapper.replay(**_whole_euler_chunk(batch_size, w), cnn_cache=None, att_cache=None) + chunk = _whole_euler_chunk(batch_size, w) + stacked_x, stacked_cnn, stacked_att = wrapper.replay(**chunk, cnn_cache=cnn, att_cache=att) + keep = (4, 6) + rows_x, rows_cnn, rows_att = wrapper.replay( + **chunk, + cnn_cache=cnn, + att_cache=_split_cfg_rows(att, batch_size), + att_keep=keep, + ) + + torch.testing.assert_close(rows_x, stacked_x) + torch.testing.assert_close(rows_cnn, stacked_cnn) + assert isinstance(rows_att, list) and len(rows_att) == batch_size + kept = wrapper_module._att_keep_ranges(int(stacked_att.shape[4]), keep) + assert sum(length for _, length in kept) == 10 + for row, expected in enumerate(_split_cfg_rows(stacked_att, batch_size)): + trimmed = torch.cat([expected[..., start : start + length, :] for start, length in kept], dim=-2) + torch.testing.assert_close(rows_att[row], trimmed) + wrapper._flush() + + +class _RealCausalDiTBlock(nn.Module): + """Causal convolution and self-attention block for multi-chunk numerical verification.""" + + def __init__(self, channels: int = 4, hidden: int = 8, causal_padding: int = 2, num_heads: int = 2) -> None: + super().__init__() + self.causal_padding = causal_padding + self.num_heads = num_heads + self.head_dim = hidden // num_heads + self.conv1d = nn.Conv1d(hidden, hidden, kernel_size=causal_padding + 1, padding=0) + self.q_proj = nn.Linear(hidden, hidden) + self.k_proj = nn.Linear(hidden, hidden) + self.v_proj = nn.Linear(hidden, hidden) + self.out_proj = nn.Linear(hidden, hidden) + self.norm1 = nn.LayerNorm(hidden) + self.norm2 = nn.LayerNorm(hidden) + self.conv = SimpleNamespace( + in_channels=channels, + out_channels=channels, + block=[None, SimpleNamespace(causal_padding=[causal_padding])], + ) + self.attn = SimpleNamespace(num_heads=num_heads, head_dim=self.head_dim) + + def forward_chunk( + self, + x: torch.Tensor, + mask: torch.Tensor | None, + cnn_cache: torch.Tensor | None, + att_cache: torch.Tensor | None, + cnn_buf: torch.Tensor, + att_buf: torch.Tensor, + ) -> torch.Tensor: + B, C, T = x.shape + # 1. Causal Conv1D + if cnn_cache is not None and cnn_cache.shape[-1] > 0: + conv_in = torch.cat([cnn_cache, x], dim=-1) + else: + conv_in = F.pad(x, (self.causal_padding, 0)) + cnn_buf.copy_(conv_in[..., -self.causal_padding :]) + h_conv = self.conv1d(conv_in) + + # 2. Multi-head self-attention with attn_mask + x_time = x.transpose(1, 2) + Q = self.q_proj(x_time).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + K = self.k_proj(x_time).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + V = self.v_proj(x_time).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + + if att_cache is not None and att_cache.shape[2] > 0: + past_k = att_cache[..., : self.head_dim] + past_v = att_cache[..., self.head_dim :] + K_all = torch.cat([past_k, K], dim=2) + V_all = torch.cat([past_v, V], dim=2) + else: + K_all = K + V_all = V + + att_buf.copy_(torch.cat([K_all, V_all], dim=-1)) + + scores = torch.matmul(Q, K_all.transpose(-2, -1)) / (self.head_dim**0.5) + if mask is not None: + scores = scores.masked_fill(~mask.unsqueeze(1), -1e4) + probs = F.softmax(scores, dim=-1) + h_attn = torch.matmul(probs, V_all).transpose(1, 2).contiguous().view(B, T, C).transpose(1, 2) + h_attn = self.out_proj(h_attn.transpose(1, 2)).transpose(1, 2) + + x = x + self.norm1((x + h_attn).transpose(1, 2)).transpose(1, 2) + x = x + self.norm2((x + h_conv).transpose(1, 2)).transpose(1, 2) + return x + + +class _RealAttentionCausalConvDiT(nn.Module): + """Estimator with real attention consuming attn_mask and causal convolutions for cross-chunk parity testing.""" + + def __init__(self, x_dim: int = 4, hidden: int = 8, depth: int = 2, causal_pad: int = 2) -> None: + super().__init__() + self.x_dim = x_dim + self.in_proj = nn.Linear(x_dim * 4, hidden) + self.blocks = nn.ModuleList( + [_RealCausalDiTBlock(channels=x_dim, hidden=hidden, causal_padding=causal_pad) for _ in range(depth)] + ) + self.final_layer = nn.Linear(hidden, x_dim) + + def t_embedder(self, t: torch.Tensor) -> torch.Tensor: + return t[:, None].expand(-1, 8) + + def blocks_forward_chunk( + self, + x: torch.Tensor, + t: torch.Tensor, + mask: torch.Tensor | None, + cnn_cache: torch.Tensor | None = None, + att_cache: torch.Tensor | None = None, + cnn_cache_buffer: torch.Tensor | None = None, + att_cache_buffer: torch.Tensor | None = None, + ) -> torch.Tensor: + assert cnn_cache_buffer is not None and att_cache_buffer is not None + x = self.in_proj(x.transpose(1, 2)).transpose(1, 2) + t_feat = t.transpose(1, 2) + for b_idx, block in enumerate(self.blocks): + cnn_b = cnn_cache[b_idx] if cnn_cache is not None else None + att_b = att_cache[b_idx] if att_cache is not None else None + x = block.forward_chunk(x, mask, cnn_b, att_b, cnn_cache_buffer[b_idx], att_cache_buffer[b_idx]) + x = x + t_feat + x = self.final_layer(x.transpose(1, 2)).transpose(1, 2) + return x + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_real_attention_causal_conv_parity_across_chunks(monkeypatch: pytest.MonkeyPatch) -> None: + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(42) + dit = _RealAttentionCausalConvDiT().eval().cuda() + wrapper = WholeEulerCFMGraphWrapper(estimator=dit, n_timesteps=10, max_graphs=4) + + B = 1 + W = 16 + + # Call 1 (bucket=16, mel_frames=12, pad_frames=4, offset=0) + x1 = torch.randn(B, 4, W, device="cuda") + mu1 = torch.randn(2 * B, 4, W, device="cuda") + spk1 = torch.randn(2 * B, 4, device="cuda") + cond1 = torch.randn(2 * B, 4, W, device="cuda") + mask1 = torch.ones(2 * B, W, W, dtype=torch.bool, device="cuda") + mask1[:, :, 12:] = False + + g_x1, g_cnn1, g_att1 = wrapper.replay( + x=x1, + mu_cfg=mu1, + speakers_cfg=spk1, + cond_cfg=cond1, + cnn_cache=None, + att_cache=None, + attn_mask=mask1, + mel_frames=12, + pad_frames=4, + # Reserves chunk 2's 12 + 16 frames up front, as serving does; trims nothing. + att_keep=(12, 16), + ) + e_x1, e_cnn1, e_att1 = _eager_solve_euler( + dit, + x1, + mu1, + spk1, + cond1, + None, + None, + mask1, + wrapper.timeline, + mel_frames=12, + pad_frames=4, + ) + + torch.testing.assert_close(g_x1, e_x1, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(g_cnn1, e_cnn1, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(g_att1[..., :12, :], e_att1[..., :12, :], rtol=1e-5, atol=1e-5) + + # Chunk 2 consuming Chunk 1 caches (offset=12, mel_frames=16, pad_frames=0) + x1_c2 = torch.randn(B, 4, W, device="cuda") + mu1_c2 = torch.randn(2 * B, 4, W, device="cuda") + spk1_c2 = torch.randn(2 * B, 4, device="cuda") + cond1_c2 = torch.randn(2 * B, 4, W, device="cuda") + mask1_c2 = torch.ones(2 * B, W, W + 12, dtype=torch.bool, device="cuda") + + g_x1_c2, g_cnn1_c2, g_att1_c2 = wrapper.replay( + x=x1_c2, + mu_cfg=mu1_c2, + speakers_cfg=spk1_c2, + cond_cfg=cond1_c2, + cnn_cache=g_cnn1, + att_cache=g_att1[..., :12, :], + attn_mask=mask1_c2, + mel_frames=16, + pad_frames=0, + ) + e_x1_c2, e_cnn1_c2, e_att1_c2 = _eager_solve_euler( + dit, + x1_c2, + mu1_c2, + spk1_c2, + cond1_c2, + e_cnn1, + e_att1[..., :12, :], + mask1_c2, + wrapper.timeline, + mel_frames=16, + pad_frames=0, + ) + + torch.testing.assert_close(g_x1_c2, e_x1_c2, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(g_cnn1_c2, e_cnn1_c2, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(g_att1_c2, e_att1_c2, rtol=1e-5, atol=1e-5) + + # Call 2 on same bucket (bucket=16, mel_frames=15, pad_frames=1, offset=0) + x2 = torch.randn(B, 4, W, device="cuda") + mu2 = torch.randn(2 * B, 4, W, device="cuda") + spk2 = torch.randn(2 * B, 4, device="cuda") + cond2 = torch.randn(2 * B, 4, W, device="cuda") + mask2 = torch.ones(2 * B, W, W, dtype=torch.bool, device="cuda") + mask2[:, :, 15:] = False + + g_x2, g_cnn2, g_att2 = wrapper.replay( + x=x2, + mu_cfg=mu2, + speakers_cfg=spk2, + cond_cfg=cond2, + cnn_cache=None, + att_cache=None, + attn_mask=mask2, + mel_frames=15, + pad_frames=1, + ) + e_x2, e_cnn2, e_att2 = _eager_solve_euler( + dit, + x2, + mu2, + spk2, + cond2, + None, + None, + mask2, + wrapper.timeline, + mel_frames=15, + pad_frames=1, + ) + + torch.testing.assert_close(g_x2, e_x2, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(g_cnn2, e_cnn2, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(g_att2[..., :15, :], e_att2[..., :15, :], rtol=1e-5, atol=1e-5) + + # Subsequent chunk consuming Call 2 caches (offset=15, mel_frames=16, pad_frames=0) + mask2_c2 = torch.ones(2 * B, W, W + 15, dtype=torch.bool, device="cuda") + g_x2_c2, g_cnn2_c2, g_att2_c2 = wrapper.replay( + x=x1_c2, + mu_cfg=mu1_c2, + speakers_cfg=spk1_c2, + cond_cfg=cond1_c2, + cnn_cache=g_cnn2, + att_cache=g_att2[..., :15, :], + attn_mask=mask2_c2, + mel_frames=16, + pad_frames=0, + ) + e_x2_c2, e_cnn2_c2, e_att2_c2 = _eager_solve_euler( + dit, + x1_c2, + mu1_c2, + spk1_c2, + cond1_c2, + e_cnn2, + e_att2[..., :15, :], + mask2_c2, + wrapper.timeline, + mel_frames=16, + pad_frames=0, + ) + torch.testing.assert_close(g_x2_c2, e_x2_c2, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(g_cnn2_c2, e_cnn2_c2, rtol=1e-5, atol=1e-5) + torch.testing.assert_close(g_att2_c2, e_att2_c2, rtol=1e-5, atol=1e-5) + + # Cache hit check: Call 2 must hit the graph captured in Call 1 + assert wrapper._stats["hits"] >= 1 + wrapper._flush() + + +class _CurrentFirstCausalDiTBlock(_RealCausalDiTBlock): + """``_RealCausalDiTBlock`` with stepaudio2's cache layout and masked SDPA. + + stepaudio2 ``Attention.forward_chunk`` concatenates ``[k, k_cache]``: the + current chunk comes first and the cache after it. Query capture padding + has to move the cache block, which an old-first fake cannot catch. + """ + + def forward_chunk(self, x, mask, cnn_cache, att_cache, cnn_buf, att_buf): + B, C, T = x.shape + if cnn_cache is not None and cnn_cache.shape[-1] > 0: + conv_in = torch.cat([cnn_cache, x], dim=-1) + else: + conv_in = F.pad(x, (self.causal_padding, 0)) + cnn_buf.copy_(conv_in[..., -self.causal_padding :]) + h_conv = self.conv1d(conv_in) + + x_time = x.transpose(1, 2) + Q = self.q_proj(x_time).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + K = self.k_proj(x_time).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + V = self.v_proj(x_time).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) + if att_cache is not None and att_cache.shape[2] > 0: + K = torch.cat([K, att_cache[..., : self.head_dim]], dim=2) + V = torch.cat([V, att_cache[..., self.head_dim :]], dim=2) + att_buf.copy_(torch.cat([K, V], dim=-1)) + h_attn = F.scaled_dot_product_attention(Q, K, V, attn_mask=None if mask is None else mask.unsqueeze(1)) + h_attn = self.out_proj(h_attn.transpose(1, 2).reshape(B, T, C)).transpose(1, 2) + + x = x + self.norm1((x + h_attn).transpose(1, 2)).transpose(1, 2) + x = x + self.norm2((x + h_conv).transpose(1, 2)).transpose(1, 2) + return x + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_query_bucket_keeps_current_first_cache_layout(monkeypatch: pytest.MonkeyPatch) -> None: + """A narrow padded chunk replayed on a wider capture keeps the cache aligned. + + The capture mask has to put the cache block after the full capture width, + and the returned cache has to drop the capture padding -- otherwise the + chunk attends to padding, loses its newest cache frames, and hands a + shifted cache to the next chunk. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + + torch.manual_seed(7) + dit = _RealAttentionCausalConvDiT().eval().cuda() + dit.blocks = nn.ModuleList([_CurrentFirstCausalDiTBlock() for _ in range(2)]).eval().cuda() + exact = WholeEulerCFMGraphWrapper(estimator=dit, n_timesteps=10, max_graphs=8, query_bucket_frames=0) + bucketed = WholeEulerCFMGraphWrapper(estimator=dit, n_timesteps=10, max_graphs=8, query_bucket_frames=16) + + def _mask(width: int, valid: int, offset: int) -> torch.Tensor: + mask = torch.ones(2, width, width + offset, dtype=torch.bool, device="cuda") + mask[:, :, valid:width] = False + return mask + + chunks = [ + (_whole_euler_chunk(1, 16), 16, 0), # full-width chunk builds a cache + (_whole_euler_chunk(1, 8), 6, 2), # narrow padded chunk: capture pads 8 -> 16 + (_whole_euler_chunk(1, 16), 16, 0), # consumes the narrow chunk's cache + ] + caches = {"exact": (None, None), "bucketed": (None, None)} + for chunk, mel_frames, pad_frames in chunks: + chunk["x"][:, :, mel_frames:] = 0.0 + results = {} + for name, wrapper in (("exact", exact), ("bucketed", bucketed)): + cnn, att = caches[name] + offset = 0 if att is None else int(att.shape[4]) + width = int(chunk["mu_cfg"].shape[2]) + results[name] = wrapper.replay( + **{key: value.clone() for key, value in chunk.items()}, + cnn_cache=cnn, + att_cache=att, + attn_mask=_mask(width, mel_frames, offset) if pad_frames else None, + mel_frames=mel_frames, + pad_frames=pad_frames, + # Reserves all three chunks (16 + 6 + 16 frames) up front; trims nothing. + att_keep=(16, 48), + ) + assert results[name] is not None + assert all(torch.isfinite(t).all() for t in results[name]) + caches[name] = (results[name][1], results[name][2]) + for exact_value, bucketed_value in zip(results["exact"], results["bucketed"], strict=True): + torch.testing.assert_close(bucketed_value, exact_value, rtol=1e-4, atol=1e-5) + + assert {key[2] for key in bucketed._cache} == {16} + assert {key[2] for key in exact._cache} == {8, 16} + exact._flush() + bucketed._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_skipped_when_trt_stepper_configured() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( + BatchedToken2Wav, + ) + + mock_stepper = Mock() + mock_stepper.step.side_effect = lambda x, mu, t, spks, cond, cnn_cache, att_cache: ( + x, + torch.zeros(2, x.shape[0], 4, 2, device=x.device), + torch.zeros(2, x.shape[0], 2, x.shape[2], 4, device=x.device), + ) + + class _MockDecoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.estimator = _WholeEulerDiT().eval().cuda() + self.inference_cfg_rate = 0.7 + self.register_buffer("rand_noise", torch.zeros(1, 4, 100, device="cuda"), persistent=False) + + class _MockFlow(nn.Module): + def __init__(self) -> None: + super().__init__() + self.decoder = _MockDecoder() + self.spk_embed_affine_layer = nn.Identity() + + class _MockToken2Wav: + def __init__(self) -> None: + self.flow = _MockFlow() + self.hift = nn.Module() + self.float16 = False + self.n_timesteps = 10 + self.mel_cache_len = 1 + self.source_cache_len = 2 + self.speech_window = torch.hamming_window(4, periodic=False) + + # Both TRT stepper and enable_cfm_graph (with enable_whole_euler=True) configured + adapter = BatchedToken2Wav( + _MockToken2Wav(), + trt_stepper=mock_stepper, + cfm_graph_config={"enabled": True, "enable_whole_euler": True, "bucket_frames": 16}, + ) + + # 1. Whole-Euler graph wrapper must not be initialized + assert adapter._whole_euler_graph_wrapper is None + + # 2. Replay/decode must route to TRT stepper rather than Whole-Euler + mu = torch.randn(1, 4, 16, device="cuda") + spk = torch.randn(1, 4, device="cuda") + cond = torch.randn(1, 4, 16, device="cuda") + adapter._decode_cfm( + mu=mu, + speakers=spk, + cond=cond, + cnn_cache=None, + att_cache=None, + ) + assert mock_stepper.step.called + assert mock_stepper.step.call_count == adapter.n_timesteps + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_disabled_via_serving_config() -> None: + from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import ( + BatchedToken2Wav, + ) + + class _MockDecoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.estimator = _WholeEulerDiT().eval().cuda() + self.inference_cfg_rate = 0.7 + self.register_buffer("rand_noise", torch.zeros(1, 4, 100, device="cuda"), persistent=False) + + class _MockFlow(nn.Module): + def __init__(self) -> None: + super().__init__() + self.decoder = _MockDecoder() + self.spk_embed_affine_layer = nn.Identity() + + class _MockToken2Wav: + def __init__(self) -> None: + self.flow = _MockFlow() + self.hift = nn.Module() + self.float16 = False + self.n_timesteps = 10 + self.mel_cache_len = 1 + self.source_cache_len = 2 + self.speech_window = torch.hamming_window(4, periodic=False) + + adapter_disabled = BatchedToken2Wav( + _MockToken2Wav(), + cfm_graph_config={"enabled": True, "enable_whole_euler": False}, + ) + assert adapter_disabled._whole_euler_graph_wrapper is None + assert adapter_disabled._cfm_graph_wrapper is not None + + adapter_enabled = BatchedToken2Wav( + _MockToken2Wav(), + cfm_graph_config={"enabled": True, "enable_whole_euler": True}, + ) + assert adapter_enabled._whole_euler_graph_wrapper is not None + assert adapter_enabled._cfm_graph_wrapper is not None + + # Unconfigured, the capture width is the steady chunk's mel width. + adapter_derived = BatchedToken2Wav( + _MockToken2Wav(), + connector_config={"codec_chunk_frames": 25, "codec_left_context_frames": 3}, + cfm_graph_config={"enabled": True, "enable_whole_euler": True}, + ) + assert adapter_derived._whole_euler_graph_wrapper.query_bucket_frames == 50 + + +def _tiny_upstream_dit() -> nn.Module: + """The shipped DiT architecture at toy width, so ``_blocks_forward_chunk_ragged`` runs as in serving.""" + decoder_dit = pytest.importorskip("stepaudio2.cosyvoice2.flow.decoder_dit") + torch.manual_seed(0) + estimator = decoder_dit.DiT(in_channels=16, out_channels=4, depth=2, num_heads=2, head_dim=8, hidden_size=16) + with torch.no_grad(): + # The adaLN-Zero init makes every block an identity; any weights will do here. + for parameter in estimator.parameters(): + parameter.normal_(0.0, 0.1) + return estimator.eval().cuda() + + +def _ragged_mask(lengths: list[int], width: int, offset: int) -> torch.Tensor: + """The attention mask ``_decode_cfm`` builds for ``valid_lengths``.""" + cfg_lengths = torch.tensor((*lengths, *lengths), device="cuda") + valid = torch.arange(width, device="cuda").unsqueeze(0) < cfg_lengths.unsqueeze(1) + current = valid.unsqueeze(1).expand(-1, width, -1) + old = torch.ones((2 * len(lengths), width, offset), dtype=torch.bool, device="cuda") + return valid.unsqueeze(2) & torch.cat((current, old), dim=2) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_ragged_rows_match_per_row_exact_solves(monkeypatch: pytest.MonkeyPatch) -> None: + """One graph solves rows of different lengths exactly as each row's own unpadded solve. + + Chunk 1 runs through a 16-frame query bucket at width 8: its CNN cache must + be the unpadded solve's, not zeros. Chunk 2 is ragged (8/5/3 valid frames) + with garbage in every padded column. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + estimator = _tiny_upstream_dit() + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=8, + query_bucket_frames=16, + ragged_body=BatchedToken2Wav._blocks_forward_chunk_ragged, + ) + batch_size, width = 3, 8 + + first = _whole_euler_chunk(batch_size, width) + graph_x, graph_cnn, graph_att = wrapper.replay(**first, cnn_cache=None, att_cache=None) + eager_x, eager_cnn, eager_att = _eager_solve_euler( + estimator, + first["x"], + first["mu_cfg"], + first["speakers_cfg"], + first["cond_cfg"], + None, + None, + None, + wrapper.timeline, + ) + torch.testing.assert_close(graph_x, eager_x, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(graph_cnn, eager_cnn, rtol=1e-4, atol=1e-5) + assert eager_cnn.abs().amax() > 0 + torch.testing.assert_close(graph_att, eager_att, rtol=1e-4, atol=1e-5) + + lengths = [8, 5, 3] + second = _whole_euler_chunk(batch_size, width) + rows_att = _split_cfg_rows(eager_att, batch_size) + out_x, out_cnn, out_att = wrapper.replay( + **second, + cnn_cache=eager_cnn, + att_cache=rows_att, + attn_mask=_ragged_mask(lengths, width, width), + valid_lengths=lengths, + ) + assert isinstance(out_att, list) and len(out_att) == batch_size + cnn_rows = _split_cfg_rows(out_cnn, batch_size) + eager_cnn_rows = _split_cfg_rows(eager_cnn, batch_size) + for row, length in enumerate(lengths): + cfg = [row, batch_size + row] + ref_x, ref_cnn, ref_att = _eager_solve_euler( + estimator, + second["x"][row : row + 1, :, :length], + second["mu_cfg"][cfg, :, :length], + second["speakers_cfg"][cfg], + second["cond_cfg"][cfg, :, :length], + eager_cnn_rows[row], + rows_att[row], + None, + wrapper.timeline, + ) + torch.testing.assert_close(out_x[row : row + 1, :, :length], ref_x, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(cnn_rows[row], ref_cnn, rtol=1e-4, atol=1e-5) + assert out_att[row].shape[4] == length + width + torch.testing.assert_close(out_att[row], ref_att, rtol=1e-4, atol=1e-5) + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_padded_ragged_replay_after_a_wider_replay(monkeypatch: pytest.MonkeyPatch) -> None: + """A padded ragged replay is exact after a wider replay of the same graph batch. + + Every graph of one batch size shares the lengths buffer. The 16-frame + replay leaves 16 in the padded row, past the 8-frame history the ragged + body gathers from. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + estimator = _tiny_upstream_dit() + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=8, + micro_batch_size=4, + ragged_body=BatchedToken2Wav._blocks_forward_chunk_ragged, + ) + assert wrapper.replay(**_whole_euler_chunk(4, 16), cnn_cache=None, att_cache=None) is not None + + lengths, width = [8, 5, 3], 8 + assert wrapper._plan_groups(len(lengths)) == [(4, 3)] + chunk = _whole_euler_chunk(len(lengths), width) + out_x, out_cnn, out_att = wrapper.replay( + **chunk, + cnn_cache=None, + att_cache=None, + attn_mask=_ragged_mask(lengths, width, 0), + valid_lengths=lengths, + ) + cnn_rows = _split_cfg_rows(out_cnn, len(lengths)) + for row, length in enumerate(lengths): + cfg = [row, len(lengths) + row] + ref_x, ref_cnn, ref_att = _eager_solve_euler( + estimator, + chunk["x"][row : row + 1, :, :length], + chunk["mu_cfg"][cfg, :, :length], + chunk["speakers_cfg"][cfg], + chunk["cond_cfg"][cfg, :, :length], + None, + None, + None, + wrapper.timeline, + ) + torch.testing.assert_close(out_x[row : row + 1, :, :length], ref_x, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(cnn_rows[row], ref_cnn, rtol=1e-4, atol=1e-5) + torch.testing.assert_close(out_att[row], ref_att, rtol=1e-4, atol=1e-5) + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_whole_euler_request_caches_grow_and_update_in_place(monkeypatch: pytest.MonkeyPatch) -> None: + """A request's cache is allocated once, with room for the steady length, then updated in place. + + Both requests start from one shared prompt cache, which is never written. + Each chunk must match a solve whose input caches cannot be updated in place. + """ + pool = torch.cuda.graph_pool_handle() + monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: pool) + wrapper = WholeEulerCFMGraphWrapper( + estimator=_tiny_upstream_dit(), + n_timesteps=10, + max_graphs=8, + query_bucket_frames=8, + ragged_body=BatchedToken2Wav._blocks_forward_chunk_ragged, + ) + batch_size, width, keep = 2, 8, (8, 16) + _, _, prompt_att = wrapper.replay(**_whole_euler_chunk(1, width), cnn_cache=None, att_cache=None, att_keep=keep) + (shared,) = _split_cfg_rows(prompt_att, 1) + shared_before = shared.clone() + + def fixed_copy(cache: torch.Tensor) -> torch.Tensor: + # Same values in another memory layout, which is never resized in place. + return cache.transpose(0, 1).contiguous().transpose(0, 1) + + rows = [shared, shared] + storages: list[int] | None = None + for frames in (16, 24, 24): + chunk = _whole_euler_chunk(batch_size, width) + _, _, expected = wrapper.replay( + **chunk, cnn_cache=None, att_cache=[fixed_copy(row) for row in rows], att_keep=keep + ) + _, _, rows = wrapper.replay(**chunk, cnn_cache=None, att_cache=rows, att_keep=keep) + for row, reference in zip(rows, expected, strict=True): + assert row.shape[4] == frames + assert torch.equal(row, reference) + assert row.untyped_storage().nbytes() == row[..., :1, :].numel() * sum(keep) * row.element_size() + if storages is None: + storages = [row.data_ptr() for row in rows] + assert shared.data_ptr() not in storages and len(set(storages)) == batch_size + else: + assert [row.data_ptr() for row in rows] == storages + assert torch.equal(shared, shared_before) + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +def test_whole_euler_groups_are_acquired_again_after_a_flush(monkeypatch: pytest.MonkeyPatch) -> None: + """Capturing a later group can retire an earlier one; no group replays until all are held at once.""" + wrapper = WholeEulerCFMGraphWrapper(estimator=_tiny_upstream_dit(), n_timesteps=10, max_graphs=8) + calls: list[int] = [] + + def entry(*, graph_batch: int, fill, **_: object) -> tuple: + calls.append(graph_batch) + if len(calls) == 2 or flush_always: + wrapper._stats["flushes"] += 1 + return ("entry", graph_batch, len(calls)) + + monkeypatch.setattr(wrapper, "_entry", entry) + flush_always = False + groups, fills = [(16, 16), (1, 1)], [None, None] + assert wrapper._group_entries(groups, fills) == [("entry", 16, 3), ("entry", 1, 4)] + assert calls == [16, 1, 16, 1] + flush_always = True + assert wrapper._group_entries(groups, fills) is None diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index 92e96b25a..d269a4227 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py @@ -170,7 +170,7 @@ class TestDeployTopology: assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor) assert "hf_overrides" not in stages[1].yaml_engine_args if filename == "minicpmo_4_5.yaml": - assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4] + assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [16, 16, 16] memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages] assert memory_utilizations == [ 0.55, diff --git a/vllm_omni/deploy/minicpmo_4_5.yaml b/vllm_omni/deploy/minicpmo_4_5.yaml index 07090660b..c273567aa 100644 --- a/vllm_omni/deploy/minicpmo_4_5.yaml +++ b/vllm_omni/deploy/minicpmo_4_5.yaml @@ -7,7 +7,7 @@ session_mode: duplex # A smaller window serializes first-packet generation for concurrent chat turns. duplex_session: # Matches Stage 0/1 max_num_seqs below. - max_sessions: 4 + max_sessions: 16 trust_remote_code: true enable_prefix_caching: false @@ -27,12 +27,17 @@ connectors: # HiFTGraphWrapper.derive_capture_bucket_size explicitly includes it. enable_hift_graph: true enable_cfm_graph: true + # Whole-Euler CUDA graphs exist for batch 1, 2, 4, ... up to + # micro_batch_size, which defaults to Stage 2 max_num_seqs: all graphs + # share one attention-cache arena of that many requests (~0.6 GiB each + # in fp32, 10 GiB at 16). Larger batches replay micro-batches; + # B > max_graph_batch runs eager. + max_graph_batch: 16 cfm_max_graphs: 32 - cfm_graph_bucket_frames: 16 stages: - stage_id: 0 - max_num_seqs: 4 + max_num_seqs: 16 # Leave headroom on the shared GPU for Talker KV + concurrent Code2Wav. gpu_memory_utilization: 0.55 enforce_eager: false @@ -62,7 +67,7 @@ stages: repetition_penalty: 1.0 - stage_id: 1 - max_num_seqs: 4 + max_num_seqs: 16 # Preserve the checkpoint's full-attention Talker prompt while it fits; # native duplex rolls over to a one-chunk window at context capacity. # Leave device-level headroom for the colocated HiFi-GAN cuDNN workspace. @@ -92,7 +97,7 @@ stages: detokenize: false - stage_id: 2 - max_num_seqs: 4 + max_num_seqs: 16 gpu_memory_utilization: 0.18 enforce_eager: true enable_chunked_prefill: false @@ -113,10 +118,11 @@ platforms: stages: - stage_id: 1 # Full-attention occupancy is bounded by max_num_seqs * max_model_len - # (4 * 4096 = 16384 tokens). Keep the reported CUDA KV token capacity - # above that bound because physical-preemption replay is unsupported. - # The explicit cap leaves device headroom for concurrent Stage 2. - kv_cache_memory_bytes: 2147483648 # 2 GiB + # (16 * 4096 = 65536 tokens). Keep the reported CUDA KV token capacity + # above that bound because physical-preemption replay is unsupported + # (2 GiB holds ~35k tokens). The explicit cap leaves device headroom + # for concurrent Stage 2. + kv_cache_memory_bytes: 4294967296 # 4 GiB npu: stages: diff --git a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml index 136dace7e..885892664 100644 --- a/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_2gpu.yaml @@ -32,7 +32,6 @@ 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 9669a2768..6002548c9 100644 --- a/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_3gpu.yaml @@ -32,7 +32,6 @@ 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 92065dd4f..7523a9175 100644 --- a/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_8x4090.yaml @@ -39,7 +39,6 @@ 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 c4240c6cd..09f8d0f65 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 @@ -15,7 +15,16 @@ import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger -from .cuda_graph_wrapper import CFMGraphWrapper, HiFTGraphWrapper +from .cuda_graph_wrapper import ( + CFMGraphWrapper, + HiFTGraphWrapper, + WholeEulerCFMGraphWrapper, + _att_keep_ranges, + _copy_frame_segments, + _euler_step, + _euler_timeline, + _zero_padded_cnn_cache, +) logger = init_logger(__name__) @@ -25,6 +34,17 @@ _SILENCE_TOKEN = 4218 # than this many codec tokens overflows the PE slice (matrix_ac vs matrix_bd). _DEFAULT_RELPOS_MAX_POS = 5000 _MAX_ENCODE_TOKEN_CAP = 1024 +# Streaming cache trim (stepaudio2 ``Token2wav.stream``): past ``prompt_len + +# 100`` frames, keep the first ``prompt_len`` frames and the last 100. +_CACHE_TRIM_SUFFIX = 100 + + +def _trim_streaming_cache(cache: torch.Tensor, prompt_len: int) -> torch.Tensor: + """The streaming trim of a cache whose frame axis is ``-2``.""" + ranges = _att_keep_ranges(int(cache.shape[-2]), (prompt_len, _CACHE_TRIM_SUFFIX)) + if len(ranges) == 1: + return cache + return torch.cat([cache[..., start : start + length, :] for start, length in ranges], dim=-2) def relpos_encode_token_budget( @@ -47,6 +67,26 @@ def relpos_encode_token_budget( return max(lookahead + 1, min(int(cap), room)) +def _supports_ragged_kernel(estimator: nn.Module) -> bool: + """Whether ``_blocks_forward_chunk_ragged`` can run this estimator's DiT blocks.""" + blocks = getattr(estimator, "blocks", None) + if not blocks or not hasattr(estimator, "in_proj") or not hasattr(estimator, "final_layer"): + return False + block = blocks[0] + attn = getattr(block, "attn", None) + conv_blocks = getattr(getattr(block, "conv", None), "block", None) + return ( + hasattr(block, "adaLN_modulation") + and all( + hasattr(attn, name) + for name in ("forward_chunk", "to_q", "to_k", "to_v", "to_heads", "q_norm", "k_norm", "proj") + ) + and conv_blocks is not None + and len(conv_blocks) >= 8 + and all(isinstance(conv_blocks[index], nn.Conv1d) for index in (1, 6)) + ) + + def _cfm_pad_frames( *, mel_frames: int, @@ -98,38 +138,6 @@ def _zero_padded_frames(tensor: torch.Tensor, valid_frames: int | None) -> None: 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. - """ - blocks = estimator.blocks - if pad_frames > 0 and blocks and len(blocks) == cnn_cache.shape[0]: - width = int(cnn_cache.shape[-1]) - if width > 0 and all(int(block.conv.block[1].causal_padding[0]) == width for block in blocks): - # _estimator_buffers packs equal-width blocks into one tensor. - # Clear their shared tail with one write instead of one per block. - cnn_cache[..., max(0, width - pad_frames) :] = 0.0 - return - - 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, *, @@ -297,6 +305,7 @@ class BatchedToken2Wav(nn.Module): token2wav.speech_window.detach().clone(), persistent=False, ) + self._timeline_cache: dict[tuple[torch.device, torch.dtype], tuple[torch.Tensor, list[float]]] = {} self.hift_graph_wrapper: HiFTGraphWrapper | None = None graph_config = dict(hift_graph_config or {}) if bool(graph_config.get("enabled", False)): @@ -312,26 +321,79 @@ class BatchedToken2Wav(nn.Module): "MiniCPM-o HiFT CUDA Graph requires source_cache_len to be divisible by mel_cache_len" ) capture_batch_sizes = graph_config.get("capture_batch_sizes", [1]) - logger.info("Enabling HiFT CUDA Graph with batch sizes %s", capture_batch_sizes) + max_serial_batch = int(graph_config.get("max_serial_batch", 4)) + logger.info( + "Enabling HiFT CUDA Graph with batch sizes %s (max_serial_batch=%d)", + capture_batch_sizes, + max_serial_batch, + ) self.hift_graph_wrapper = HiFTGraphWrapper( token2wav=token2wav, connector_config=dict(connector_config), capture_batch_sizes=capture_batch_sizes, + max_serial_batch=max_serial_batch, ) with torch.inference_mode(), _autocast_disabled(hift_parameter.device): self.hift_graph_wrapper.capture() logger.info("HiFT CUDA Graph captured successfully") self._cfm_graph_wrapper: CFMGraphWrapper | None = None + self._whole_euler_graph_wrapper: WholeEulerCFMGraphWrapper | None = None cfm_graph_cfg = dict(cfm_graph_config or {}) if bool(cfm_graph_cfg.get("enabled", False)): flow_parameter = next(self.flow.parameters(), None) if flow_parameter is not None and flow_parameter.device.type == "cuda": estimator = self.flow.decoder.estimator + max_graphs = int(cfm_graph_cfg.get("max_graphs", 32)) + max_serial_batch = int(cfm_graph_cfg.get("max_serial_batch", 4)) + max_graph_batch_cfg = cfm_graph_cfg.get("max_graph_batch") + max_graph_batch = int(max_graph_batch_cfg) if max_graph_batch_cfg is not None else None + micro_batch_size_cfg = cfm_graph_cfg.get("micro_batch_size") + micro_batch_size = int(micro_batch_size_cfg) if micro_batch_size_cfg is not None else 4 + if bool(cfm_graph_cfg.get("enable_whole_euler", True)) and self._trt_stepper is None: + capture_frames = cfm_graph_cfg.get("capture_frames") + if capture_frames is None and connector_config is not None: + # The steady chunk's mel width, derived like HiFTGraphWrapper's + # capture bucket: steady chunks run unpadded and shorter + # chunks share their graph under a mask. + capture_frames = int(connector_config["codec_chunk_frames"]) * int( + getattr(self.flow, "token_mel_ratio", 2) + ) + query_bucket_frames = int(capture_frames if capture_frames is not None else 64) + self._whole_euler_graph_wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=self.n_timesteps, + inference_cfg_rate=getattr(self.flow.decoder, "inference_cfg_rate", 0.7), + att_cache_dtype=self._estimator_att_cache_dtype, + max_graphs=max_graphs, + max_serial_batch=max_serial_batch, + max_graph_batch=max_graph_batch, + micro_batch_size=micro_batch_size, + query_bucket_frames=query_bucket_frames, + pad_max_rows=cfm_graph_cfg.get("pad_max_rows"), + # Looked up at capture time, so a replaced ragged kernel is the one captured. + ragged_body=( + (lambda *args: self._blocks_forward_chunk_ragged(*args)) + if _supports_ragged_kernel(estimator) + else None + ), + ) + logger.info( + "Whole-Euler CFM CUDA Graph enabled " + "(max_graphs=%d, max_serial_batch=%d, max_graph_batch=%s, " + "micro_batch_size=%d, query_bucket_frames=%d)", + max_graphs, + max_serial_batch, + str(max_graph_batch), + micro_batch_size, + query_bucket_frames, + ) + elif self._trt_stepper is not None and bool(cfm_graph_cfg.get("enable_whole_euler", True)): + logger.info("Whole-Euler CFM CUDA Graph disabled because TensorRT stepper is configured") self._cfm_graph_wrapper = CFMGraphWrapper( graph_fn=estimator.blocks_forward_chunk, - max_graphs=int(cfm_graph_cfg.get("max_graphs", 32)), + max_graphs=max_graphs, ) - logger.info("CFM CUDA Graph enabled (max_graphs=%d)", int(cfm_graph_cfg.get("max_graphs", 32))) + logger.info("CFM CUDA Graph enabled (max_graphs=%d)", max_graphs) else: logger.info( "CFM CUDA Graph is disabled on device type %s", @@ -341,7 +403,9 @@ class BatchedToken2Wav(nn.Module): # 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 + int(cfm_graph_cfg.get("bucket_frames", 0)) + if (self._cfm_graph_wrapper is not None or self._whole_euler_graph_wrapper is not None) + else 0 ) if self._cfm_graph_bucket_frames > 1: logger.info( @@ -350,9 +414,7 @@ class BatchedToken2Wav(nn.Module): ) self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} self._setup_cache_size = setup_cache_size - self._setup_cache: OrderedDict[tuple[tuple[str, str], int, int], tuple[BatchedToken2WavState, ...]] = ( - OrderedDict() - ) + self._setup_cache: OrderedDict[tuple[tuple[str, str], int], BatchedToken2WavState] = OrderedDict() def _hift_inference( self, @@ -507,6 +569,7 @@ class BatchedToken2Wav(nn.Module): attn_mask: torch.Tensor | None = None, valid_lengths: list[int] | None = None, valid_frames: int | None = None, + time_embedding: torch.Tensor | 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( @@ -519,7 +582,8 @@ class BatchedToken2Wav(nn.Module): att_cache=att_cache, ) return out.to(mu.dtype), new_cnn, new_att - time_embedding = estimator.t_embedder(time).unsqueeze(1) + if time_embedding is None: + 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: @@ -583,14 +647,95 @@ class BatchedToken2Wav(nn.Module): valid_lengths: torch.Tensor, width: int, ) -> torch.Tensor: + """The ``width`` frames of ``history`` (b, width + t, c) ending each row's valid frames, as (b, c, width).""" indices = valid_lengths[:, None] + torch.arange(width, device=history.device)[None, :] - return history.gather( - 2, - indices[:, None, :].expand(-1, int(history.shape[1]), -1), - ) + return history.gather(1, indices[:, :, None].expand(-1, -1, int(history.shape[2]))).transpose(1, 2) + @staticmethod + def _causal_conv_frames_major(conv: nn.Conv1d, history: torch.Tensor) -> torch.Tensor: + """``conv`` over ``history`` (b, width + t, c_in) as (b, t, c_out), without a layout change. + + The DiT keeps activations as (b, t, c), which already is NHWC for a + height-1 image. Upstream transposes them to (b, c, t) for ``Conv1d``, + and cuDNN converts that to NHWC and back around every convolution. As + a channels-last conv2d it reads the activations where they are. The + channels-last weight is converted once and reused. Other devices run + upstream's ``Conv1d`` layout (Ascend has no channels-last format). + """ + if not history.is_cuda: + output = F.conv1d( + history.transpose(1, 2), + conv.weight, + conv.bias, + conv.stride, + conv.padding, + conv.dilation, + conv.groups, + ) + return output.transpose(1, 2) + weight = conv.weight + version = (weight.data_ptr(), weight._version, weight.dtype) + cached = conv.__dict__.get("_channels_last_weight") + if cached is None or cached[0] != version: + cached = (version, weight.detach().unsqueeze(2).contiguous(memory_format=torch.channels_last)) + conv.__dict__["_channels_last_weight"] = cached + output = F.conv2d( + history.transpose(1, 2).unsqueeze(2), + cached[1], + conv.bias, + (1, conv.stride[0]), + (0, conv.padding[0]), + (1, conv.dilation[0]), + conv.groups, + ) + return output.squeeze(2).transpose(1, 2) + + @staticmethod + def _attend_in_place( + attn: nn.Module, + x: torch.Tensor, + att_cache: torch.Tensor | None, + attn_mask: torch.Tensor | None, + kv: torch.Tensor, + ) -> torch.Tensor: + """``attn.forward_chunk`` that keeps the new keys and values in ``kv``. + + ``kv`` (b, heads, >= t + cache frames, 2 * head_dim) is the block's new + cache, ``[current | cache]`` like upstream's. Upstream concatenates the + whole cache behind the current keys and values, then both into the new + cache, which the caller copied into ``kv``: four passes over Stage 2's + largest tensor. Here only the current chunk is written and attention reads ``kv`` in + place. A cache that already sits behind the current chunk in ``kv`` + (Whole-Euler) is not moved; any other is copied there once. Other + devices keep upstream's attention and copy its cache into ``kv``. + """ + if kv.dtype != x.dtype or not x.is_cuda: + # A reduced-precision cache attends in the activation dtype. + x_att, new_att = attn.forward_chunk(x, att_cache, attn_mask) + kv[:, :, : int(new_att.shape[2])].copy_(new_att) + return x_att + batch, frames, _ = x.shape + q = attn.q_norm(attn.to_heads(attn.to_q(x))) + k = attn.k_norm(attn.to_heads(attn.to_k(x))) + v = attn.to_heads(attn.to_v(x)) + cached = 0 if att_cache is None else int(att_cache.shape[2]) + kv = kv[:, :, : frames + cached] + behind = kv[:, :, frames:] + if cached and (att_cache.data_ptr() != behind.data_ptr() or att_cache.stride() != behind.stride()): + behind.copy_(att_cache) + head_dim = int(k.shape[3]) + kv[..., :frames, :head_dim].copy_(k) + kv[..., :frames, head_dim:].copy_(v) + out = F.scaled_dot_product_attention( + q, + kv[..., :head_dim], + kv[..., head_dim:], + attn_mask=None if attn_mask is None else attn_mask.unsqueeze(1), + ) + return attn.proj(out.transpose(1, 2).reshape(batch, frames, -1)) + + @staticmethod def _blocks_forward_chunk_ragged( - self, estimator: nn.Module, estimator_input: torch.Tensor, time_embedding: torch.Tensor, @@ -599,14 +744,24 @@ class BatchedToken2Wav(nn.Module): att_cache: Any, cnn_cache_buffer: torch.Tensor, att_cache_buffer: torch.Tensor, - valid_lengths: list[int], + valid_lengths: list[int] | torch.Tensor, ) -> torch.Tensor: - """Run one padded DiT batch while capturing exact per-row CNN state.""" - lengths = torch.tensor( - (*valid_lengths, *valid_lengths), - device=estimator_input.device, - dtype=torch.long, - ) + """Run one padded DiT batch while capturing exact per-row CNN state. + + ``valid_lengths`` is one length per request, or the ``(2B,)`` CFG + lengths tensor a CUDA graph captures as a static input. Upstream's + ``DiTBlock.forward_chunk``, with attention over the output cache in + place (``_attend_in_place``) and the causal convolutions on the + frames-major activations (``_causal_conv_frames_major``). + """ + if isinstance(valid_lengths, torch.Tensor): + lengths = valid_lengths + else: + lengths = torch.tensor( + (*valid_lengths, *valid_lengths), + device=estimator_input.device, + dtype=torch.long, + ) x = estimator.in_proj(estimator_input.transpose(1, 2)) for block_index, block in enumerate(estimator.blocks): ( @@ -622,49 +777,55 @@ class BatchedToken2Wav(nn.Module): ) = block.adaLN_modulation(time_embedding).chunk(9, dim=-1) normalized = block.norm1(x) * (1 + scale_msa) + shift_msa - x_att, new_att = block.attn.forward_chunk( + x_att = BatchedToken2Wav._attend_in_place( + block.attn, normalized, att_cache[block_index], attn_mask, + att_cache_buffer[block_index], ) x = x + gate_msa * x_att + # CausalConvBlock: [T, conv1, T, LayerNorm, Mish, T, conv2, T] + # with the transposes (T) dropped. + conv = block.conv conv_input = block.norm3(x) * (1 + scale_conv) + shift_conv old_cnn = cnn_cache[block_index] if old_cnn is None: - width = int(block.conv.block[1].causal_padding[0]) old_cnn = conv_input.new_zeros( ( int(conv_input.shape[0]), - int(block.conv.in_channels + block.conv.out_channels), - width, + int(conv.in_channels + conv.out_channels), + int(conv.block[1].causal_padding[0]), ) ) - old_cnn1, old_cnn2 = old_cnn.split((block.conv.in_channels, block.conv.out_channels), dim=1) - - conv1_input = block.conv.block[0](conv_input) - conv1_output, _ = block.conv.block[1].forward_chunk(conv1_input, old_cnn1) - new_cnn1 = self._gather_causal_cache( - torch.cat((old_cnn1, conv1_input), dim=2), - lengths, - int(old_cnn1.shape[2]), - ) - - conv2_input = block.conv.block[2:6](conv1_output) - conv2_output, _ = block.conv.block[6].forward_chunk(conv2_input, old_cnn2) - new_cnn2 = self._gather_causal_cache( - torch.cat((old_cnn2, conv2_input), dim=2), - lengths, - int(old_cnn2.shape[2]), - ) - x = x + gate_conv * block.conv.block[7](conv2_output) + width = int(old_cnn.shape[2]) + history = torch.cat((old_cnn[:, : conv.in_channels].transpose(1, 2), conv_input), dim=1) + new_cnn1 = BatchedToken2Wav._gather_causal_cache(history, lengths, width) + hidden = conv.block[4](conv.block[3](BatchedToken2Wav._causal_conv_frames_major(conv.block[1], history))) + history = torch.cat((old_cnn[:, conv.in_channels :].transpose(1, 2), hidden), dim=1) + new_cnn2 = BatchedToken2Wav._gather_causal_cache(history, lengths, width) + x = x + gate_conv * BatchedToken2Wav._causal_conv_frames_major(conv.block[6], history) x = x + gate_mlp * block.mlp(block.norm2(x) * (1 + scale_mlp) + shift_mlp) - cnn_cache_buffer[block_index].copy_(torch.cat((new_cnn1, new_cnn2), dim=1)) - att_cache_buffer[block_index][:, :, : int(new_att.shape[2]), :].copy_(new_att) + cnn_cache_buffer[block_index][:, : conv.in_channels].copy_(new_cnn1) + cnn_cache_buffer[block_index][:, conv.in_channels :].copy_(new_cnn2) return estimator.final_layer(x, time_embedding).transpose(1, 2) + def _get_timeline(self, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, list[float]]: + """The Euler timeline and its step sizes, read to the host once. + + ``_euler_step`` takes ``dt`` as a Python float, so reading it + from the device timeline on every decode would sync once per step. + """ + key = (device, dtype) + cached = self._timeline_cache.get(key) + if cached is None: + cached = _euler_timeline(self.n_timesteps, device, dtype) + self._timeline_cache[key] = cached + return cached + def _decode_cfm( self, mu: torch.Tensor, @@ -672,30 +833,44 @@ class BatchedToken2Wav(nn.Module): cond: torch.Tensor, *, cnn_cache: torch.Tensor | None, - att_cache: torch.Tensor | None, + att_cache: torch.Tensor | list[torch.Tensor] | None, valid_lengths: list[int] | None = None, + att_keep: tuple[int, int] | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor | list[torch.Tensor], ]: + """Run the CFM Euler solve for one chunk. + + ``att_cache`` may be the stacked estimator cache or one tensor per + request (the Whole-Euler path). A Whole-Euler solve returns the new + cache in the same form, already trimmed by ``att_keep``; the eager + path returns it stacked and untrimmed (ragged: per request, untrimmed). + """ decoder = self.flow.decoder estimator = decoder.estimator batch_size = int(mu.shape[0]) - offset = int(att_cache.shape[4]) if att_cache is not None else 0 + att_rows = att_cache if isinstance(att_cache, list) else None + if att_rows is not None: + offset = int(att_rows[0].shape[4]) + else: + offset = int(att_cache.shape[4]) if att_cache is not None else 0 mel_frames = int(mu.shape[2]) + graphs_active = valid_lengths is None and ( + self._whole_euler_active() + or ( + self._cfm_graph_wrapper is not None + and getattr(self._cfm_graph_wrapper, "enabled", True) + and self._trt_stepper is None + ) + ) 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 - ), + disabled=(valid_lengths is not None or not graphs_active), ) if pad_frames: # Replicate rather than zero: a repeated last frame is a closer @@ -718,15 +893,7 @@ class BatchedToken2Wav(nn.Module): # 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, - self.n_timesteps + 1, - device=mu.device, - dtype=mu.dtype, - ) - timeline = 1 - torch.cos(timeline * 0.5 * torch.pi) - time = timeline[0].expand(batch_size) + timeline, dt_steps = self._get_timeline(mu.device, mu.dtype) mu_cfg = torch.cat((mu, torch.zeros_like(mu)), dim=0) speakers_cfg = torch.cat((speakers, torch.zeros_like(speakers)), dim=0) cond_cfg = torch.cat((cond, torch.zeros_like(cond)), dim=0) @@ -761,19 +928,42 @@ class BatchedToken2Wav(nn.Module): device=mu.device, ) attn_mask[:, :, mel_frames : mel_frames + pad_frames] = False + + if self._whole_euler_active() and (valid_lengths is None or self._whole_euler_ragged_active()): + whole_euler_result = self._whole_euler_graph_wrapper.replay( + x=x, + mu_cfg=mu_cfg, + speakers_cfg=speakers_cfg, + cond_cfg=cond_cfg, + cnn_cache=cnn_cache, + att_cache=att_cache, + attn_mask=attn_mask, + mel_frames=mel_frames, + pad_frames=pad_frames, + valid_lengths=valid_lengths, + att_keep=att_keep, + ) + if whole_euler_result is not None: + return whole_euler_result + if att_rows is not None: + att_cache = self._stack_estimator_att(att_rows) + next_cnn: list[torch.Tensor] = [] next_att_cache: torch.Tensor | None = None ragged_att_cache: list[torch.Tensor] | None = None - dt = timeline[1] - timeline[0] + time_embeddings = [ + estimator.t_embedder(timeline[s].expand(2 * batch_size)).unsqueeze(1) for s in range(self.n_timesteps) + ] with _token2wav_sdpa_context(mu.device): for step in range(self.n_timesteps): old_cnn = cnn_cache[step] if cnn_cache is not None else None old_att = att_cache[step] if att_cache is not None else None + step_time = timeline[step].expand(2 * batch_size) estimate, step_cnn, step_att = self._estimator_step( estimator, x=torch.cat((x, x), dim=0), mu=mu_cfg, - time=torch.cat((time, time), dim=0), + time=step_time, speakers=speakers_cfg, cond=cond_cfg, cnn_cache=old_cnn, @@ -781,16 +971,13 @@ class BatchedToken2Wav(nn.Module): attn_mask=attn_mask, valid_lengths=valid_lengths, valid_frames=mel_frames if pad_frames else None, + time_embedding=time_embeddings[step], ) 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] + x = _euler_step(x, estimate, dt_steps[step], decoder.inference_cfg_rate, batch_size) + if pad_frames: + _zero_padded_frames(x, mel_frames) next_cnn.append(step_cnn) if valid_lengths is not None: if ragged_att_cache is None: @@ -842,28 +1029,40 @@ class BatchedToken2Wav(nn.Module): assert next_att_cache is not None return x, torch.stack(next_cnn), next_att_cache - def _split_flow_cache(self, cache: dict[str, torch.Tensor], batch_size: int) -> list[dict[str, torch.Tensor]]: + def _split_flow_cache( + self, + cache: dict[str, Any], + batch_size: int, + *, + estimator_att_keep: tuple[int, int] | None = None, + ) -> list[dict[str, torch.Tensor]]: + """Per-request flow caches from a batch's stacked caches. + + ``estimator_att_cache`` is stacked ``(n_t, depth, 2B, heads, L, d)`` + with CFG rows ``[cond x B | uncond x B]``, or already one tensor per + request (Whole-Euler). The streaming trim ``estimator_att_keep`` is + applied while each request is copied out, so no trimmed copy of the + whole batch is built first: at 16 requests that copy alone was ~7 GiB. + """ + estimator_att = cache["estimator_att_cache"] + stacked_att = isinstance(estimator_att, torch.Tensor) + if stacked_att: + keep_ranges = _att_keep_ranges(int(estimator_att.shape[4]), estimator_att_keep) + request_shape = list(estimator_att.shape) + request_shape[2] = 2 + request_shape[4] = sum(length for _, length in keep_ranges) + att_rows = estimator_att.unflatten(2, (2, batch_size)) result: list[dict[str, torch.Tensor]] = [] for row in range(batch_size): - estimator_att = cache["estimator_att_cache"] - if estimator_att.dtype == self._estimator_att_cache_dtype: - request_att = torch.cat( - ( - estimator_att[:, :, row : row + 1], - estimator_att[:, :, batch_size + row : batch_size + row + 1], - ), - dim=2, - ).detach() - else: - request_shape = list(estimator_att.shape) - request_shape[2] = 2 + if stacked_att: request_att = torch.empty( request_shape, device=estimator_att.device, dtype=self._estimator_att_cache_dtype, ) - request_att[:, :, 0:1].copy_(estimator_att[:, :, row : row + 1]) - request_att[:, :, 1:2].copy_(estimator_att[:, :, batch_size + row : batch_size + row + 1]) + _copy_frame_segments(request_att, att_rows[:, :, :, row], keep_ranges) + else: + request_att = estimator_att[row] result.append( { "conformer_cnn_cache": cache["conformer_cnn_cache"][row : row + 1].detach().clone(), @@ -880,20 +1079,45 @@ class BatchedToken2Wav(nn.Module): ) return result - def _stack_flow_cache(self, states: list[BatchedToken2WavState]) -> dict[str, torch.Tensor]: + @staticmethod + def _stack_estimator_att(rows: list[torch.Tensor]) -> torch.Tensor: + """Stack per-request ``(n_t, depth, 2, ...)`` caches into CFG order ``[cond x B | uncond x B]``.""" + return torch.cat((*[row[:, :, 0:1] for row in rows], *[row[:, :, 1:2] for row in rows]), dim=2) + + def _stack_flow_cache( + self, + states: list[BatchedToken2WavState], + *, + include_estimator_att: bool = True, + ) -> dict[str, torch.Tensor | None]: + """Stack per-request flow caches for one batched call. + + The estimator attention cache is by far the largest (~0.5 GiB per + request in fp32). The Whole-Euler path reads it per request, so it + skips stacking it (``include_estimator_att=False``). + """ flows = [state.flow_cache for state in states] conditional_cnn = [flow["estimator_cnn_cache"][:, :, 0:1] for flow in flows] unconditional_cnn = [flow["estimator_cnn_cache"][:, :, 1:2] for flow in flows] - conditional_att = [flow["estimator_att_cache"][:, :, 0:1] for flow in flows] - unconditional_att = [flow["estimator_att_cache"][:, :, 1:2] for flow in flows] - estimator_att = torch.cat((*conditional_att, *unconditional_att), dim=2) return { "conformer_cnn_cache": torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0), "conformer_att_cache": torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1), "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2), - "estimator_att_cache": estimator_att, + "estimator_att_cache": ( + self._stack_estimator_att([flow["estimator_att_cache"] for flow in flows]) + if include_estimator_att + else None + ), } + def _whole_euler_active(self) -> bool: + wrapper = self._whole_euler_graph_wrapper + return wrapper is not None and wrapper.enabled and self._trt_stepper is None + + def _whole_euler_ragged_active(self) -> bool: + """Whether Whole-Euler graphs also solve ragged (per-row length) batches.""" + return self._whole_euler_active() and self._whole_euler_graph_wrapper.ragged_body is not None + def _create_initial_states( self, features: PromptFeatures, @@ -913,12 +1137,15 @@ class BatchedToken2Wav(nn.Module): att_cache=None, ) projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + # The prompt cache never reaches the trim length; ``att_keep`` only + # tells the Whole-Euler arena the steady cache length to reserve. _, estimator_cnn, estimator_att = self._decode_cfm( hidden.transpose(1, 2).contiguous(), projected_speakers, prompt_mels.transpose(1, 2).contiguous(), cnn_cache=None, att_cache=None, + att_keep=(int(prompt_mels.shape[1]), _CACHE_TRIM_SUFFIX), ) flow_cache = { "conformer_cnn_cache": conformer_cnn, @@ -941,7 +1168,17 @@ class BatchedToken2Wav(nn.Module): ] def setup_batch(self, features: PromptFeatures, batch_size: int) -> list[BatchedToken2WavState]: - """Reuse read-only prompt conditioning for the same shape and padding.""" + """Initial states for ``batch_size`` requests sharing one prompt. + + Every request starts from the same prompt tokens, speaker, mels and + noise, so the prompt is solved once and all requests share that one + read-only state. Decoding writes a state only where Whole-Euler updates + an estimator cache in place, which needs a cache held by one request + and allocated with room for its new length; the prompt cache is shared + and allocated exactly, and the first chunk always grows it. Solving it + per batch size recomputed identical rows and cached one ~0.3 GiB + estimator cache copy per request. + """ bucket_frames = ( self._cfm_graph_bucket_frames if self._cfm_graph_wrapper is not None and self._cfm_graph_wrapper.enabled @@ -949,17 +1186,17 @@ class BatchedToken2Wav(nn.Module): ) # Capture the policy before setup: graph capture may disable the wrapper # after padding has already been chosen for these initial states. - cache_key = (features.cache_key, batch_size, bucket_frames) - cached = self._setup_cache.get(cache_key) - if cached is not None: + cache_key = (features.cache_key, bucket_frames) + state = self._setup_cache.get(cache_key) + if state is not None: self._setup_cache.move_to_end(cache_key) - return list(cached) - states = self._create_initial_states(features, batch_size) - if self._setup_cache_size > 0: - self._setup_cache[cache_key] = tuple(states) - while len(self._setup_cache) > self._setup_cache_size: - self._setup_cache.popitem(last=False) - return states + else: + (state,) = self._create_initial_states(features, 1) + if self._setup_cache_size > 0: + self._setup_cache[cache_key] = state + while len(self._setup_cache) > self._setup_cache_size: + self._setup_cache.popitem(last=False) + return [state] * batch_size @staticmethod def _fade_in_out( @@ -1058,7 +1295,12 @@ class BatchedToken2Wav(nn.Module): flush_encoder: bool = False, ) -> tuple[list[torch.Tensor], list[BatchedToken2WavState]]: batch_size = int(tokens.shape[0]) - flow_cache = self._stack_flow_cache(states) + # Whole-Euler reads and writes each request's estimator cache directly, + # so the batch never holds a stacked copy of it. + per_request_att = self._whole_euler_active() + flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_request_att) + prompt_len = int(features.mels.shape[1]) + att_keep = (prompt_len, _CACHE_TRIM_SUFFIX) speakers = features.speaker_embedding.expand(batch_size, -1) with self._autocast(tokens.device): hidden, conformer_cnn, conformer_att = self._encode_chunk( @@ -1074,20 +1316,15 @@ class BatchedToken2Wav(nn.Module): projected_speakers, cond, cnn_cache=flow_cache["estimator_cnn_cache"], - att_cache=flow_cache["estimator_att_cache"], + att_cache=( + [state.flow_cache["estimator_att_cache"] for state in states] + if per_request_att + else flow_cache["estimator_att_cache"] + ), + att_keep=att_keep, ) - prompt_len = int(features.mels.shape[1]) - if estimator_att.shape[4] > prompt_len + 100: - estimator_att = torch.cat( - (estimator_att[..., :prompt_len, :], estimator_att[..., -100:, :]), - dim=4, - ) - if conformer_att.shape[3] > prompt_len + 100: - conformer_att = torch.cat( - (conformer_att[..., :prompt_len, :], conformer_att[..., -100:, :]), - dim=3, - ) + conformer_att = _trim_streaming_cache(conformer_att, prompt_len) new_flow = self._split_flow_cache( { "conformer_cnn_cache": conformer_cnn, @@ -1096,6 +1333,7 @@ class BatchedToken2Wav(nn.Module): "estimator_att_cache": estimator_att, }, batch_size, + estimator_att_keep=att_keep, ) old_mel = torch.cat([state.hift_cache["mel"] for state in states], dim=0) old_source = torch.cat([state.hift_cache["source"] for state in states], dim=0) @@ -1164,10 +1402,17 @@ class BatchedToken2Wav(nn.Module): ) max_frames = self._max_encode_token_frames(states) - if any(int(row_tokens.numel()) > max_frames for row_tokens in tokens): + per_request_att = self._whole_euler_ragged_active() + if (self._whole_euler_active() and not per_request_att) or any( + int(row_tokens.numel()) > max_frames for row_tokens in tokens + ): # decode_batch owns the RelPos-safe encoder slicing used by - # non-async Talker dumps. Preserve that path for overlong rows; - # ordinary short duplex chunks continue through one padded DiT. + # non-async Talker dumps. Preserve that path for overlong rows. + # Whole-Euler graphs that cannot run the ragged kernel also take + # this path: under concurrency most batches hold some request's + # short final chunk, and one padded eager DiT would run the whole + # batch ungraphed on a stacked copy of every cache. Otherwise short + # duplex chunks share one padded DiT solve, graphed when possible. exact_groups: dict[tuple[int, bool], list[int]] = {} for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): exact_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row) @@ -1193,7 +1438,8 @@ class BatchedToken2Wav(nn.Module): conformer_att_rows: list[torch.Tensor | None] = [None] * batch_size for (_, last_chunk), rows in encoder_groups.items(): group_states = [states[row] for row in rows] - group_cache = self._stack_flow_cache(group_states) + # The encoder reads only the conformer caches. + group_cache = self._stack_flow_cache(group_states, include_estimator_att=False) group_tokens = torch.stack([tokens[row] for row in rows], dim=0) with self._autocast(group_tokens.device): hidden, conformer_cnn, conformer_att = self._encode_chunk( @@ -1219,7 +1465,10 @@ class BatchedToken2Wav(nn.Module): for row, hidden in enumerate(resolved_hidden): padded_hidden[row, : int(hidden.shape[1])].copy_(hidden[0]) - flow_cache = self._stack_flow_cache(states) + # A graphed ragged solve reads and writes each request's estimator + # cache directly (trimmed on the way out), like _decode_batch_once. + flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_request_att) + prompt_len = int(features.mels.shape[1]) speakers = features.speaker_embedding.expand(batch_size, -1) with self._autocast(padded_hidden.device): projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) @@ -1229,33 +1478,24 @@ class BatchedToken2Wav(nn.Module): projected_speakers, cond, cnn_cache=flow_cache["estimator_cnn_cache"], - att_cache=flow_cache["estimator_att_cache"], + att_cache=( + [state.flow_cache["estimator_att_cache"] for state in states] + if per_request_att + else flow_cache["estimator_att_cache"] + ), valid_lengths=hidden_lengths, + att_keep=(prompt_len, _CACHE_TRIM_SUFFIX), ) - prompt_len = int(features.mels.shape[1]) assert isinstance(estimator_att, list) new_flow: list[dict[str, torch.Tensor]] = [] for row, row_estimator_att in enumerate(estimator_att): conformer_cnn = conformer_cnn_rows[row] conformer_att = conformer_att_rows[row] assert conformer_cnn is not None and conformer_att is not None - if row_estimator_att.shape[4] > prompt_len + 100: - row_estimator_att = torch.cat( - ( - row_estimator_att[..., :prompt_len, :], - row_estimator_att[..., -100:, :], - ), - dim=4, - ) - if conformer_att.shape[3] > prompt_len + 100: - conformer_att = torch.cat( - ( - conformer_att[..., :prompt_len, :], - conformer_att[..., -100:, :], - ), - dim=3, - ) + # A graphed solve already trimmed the estimator cache; trimming is idempotent. + row_estimator_att = _trim_streaming_cache(row_estimator_att, prompt_len) + conformer_att = _trim_streaming_cache(conformer_att, prompt_len) new_flow.append( { "conformer_cnn_cache": conformer_cnn.detach().clone(), 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 4db0c39bf..a37218fe5 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 @@ -1,7 +1,12 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from collections import Counter +from collections.abc import Callable +from typing import NamedTuple + import torch +import torch.nn.functional as F from torch.cuda import CUDAGraph from vllm.logger import init_logger from vllm.platforms import current_platform @@ -9,8 +14,28 @@ from vllm.platforms import current_platform logger = init_logger(__name__) +def _euler_step( + cur_x: torch.Tensor, + estimate: torch.Tensor, + dt: float, + inference_cfg_rate: float, + batch_size: int, +) -> torch.Tensor: + """One Euler step of the CFG-guided flow: ``cur_x + dt * ((1 + cfg) * cond - cfg * uncond)``.""" + conditional, unconditional = estimate.split(batch_size, dim=0) + velocity = (1.0 + inference_cfg_rate) * conditional - inference_cfg_rate * unconditional + return cur_x + dt * velocity + + +def _euler_timeline(n_timesteps: int, device: torch.device, dtype: torch.dtype) -> tuple[torch.Tensor, list[float]]: + """The cosine Euler timeline and its step sizes, read to the host once.""" + t = torch.linspace(0, 1, n_timesteps + 1, device=device, dtype=dtype) + timeline = 1 - torch.cos(t * 0.5 * torch.pi) + return timeline, (timeline[1:] - timeline[:-1]).tolist() + + class HiFTGraphWrapper: - def __init__(self, token2wav, connector_config, capture_batch_sizes): + def __init__(self, token2wav, connector_config, capture_batch_sizes, max_serial_batch: int | None = None): self.decode_fn = token2wav.hift.inference self.graph_fn = token2wav.hift._inference_pre_istft self.finalize_fn = token2wav.hift._finalize_decode @@ -36,6 +61,7 @@ class HiFTGraphWrapper: self.dtype = parameter.dtype self.max_lazy_graphs = 8 self.lazy_graph_count = 0 + self.max_serial_batch = 4 if max_serial_batch is None else int(max_serial_batch) def derive_capture_bucket_size(self): chunk_mel_frames = ( @@ -108,6 +134,14 @@ class HiFTGraphWrapper: target_b = next((b for b in sorted(self.capture_batch_sizes) if b >= batch_size), None) if target_b is None: + if 1 in self.capture_batch_sizes and 1 < batch_size <= self.max_serial_batch: + speeches = [] + sources = [] + for b in range(batch_size): + sp, src = self.replay(speech_feat[b : b + 1], cache_source[b : b + 1]) + speeches.append(sp) + sources.append(src) + return torch.cat(speeches, dim=0), torch.cat(sources, dim=0) logger.info("Falling back to eager HiFT inference for unsupported batch size %d", batch_size) return self.decode_fn(speech_feat, cache_source) @@ -170,6 +204,179 @@ def _format_memory_delta(before: tuple[int, int] | None, after: tuple[int, int] ) +# Frame granularity of the shared attention cache storage. +_ATT_FRAME_ALIGN = 16 + + +def _align_up(n: int, bucket: int) -> int: + if n <= 0 or bucket <= 1: + return max(0, n) + return ((n + bucket - 1) // bucket) * bucket + + +def _capture_query_width(mel_width: int, bucket: int) -> int: + """Snap the query axis onto the Whole-Euler capture grid. + + ``_cfm_pad_frames`` already aligns onto ``bucket_frames=16``, which is why + 12+4 and 10+6 share a graph (``test_whole_euler_same_bucket_different_padding_hits_cache``) + but 16/32/48/64 stay distinct (``test_varied_chunk_lengths_collapse_onto_few_widths``). + The capture bucket reuses that same pad-and-mask scheme so those + decode widths collapse onto one CUDA graph; first-chunk widths above the + bucket (e.g. 304) align up separately so streaming replay stays on the + small decode graph. + """ + if bucket <= 1: + return int(mel_width) + if mel_width <= bucket: + return int(bucket) + return _align_up(int(mel_width), int(bucket)) + + +def _pad_query_for_capture( + x: torch.Tensor, + mu_cfg: torch.Tensor, + cond_cfg: torch.Tensor, + query_cap: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pad query tensors to ``query_cap`` the same way ``_decode_cfm`` pads to 16. + + ``mu``/``cond`` replicate the last frame (closer continuation than silence); + ``x`` is zeroed on the pad so the Euler state does not integrate noise there. + """ + extra = int(query_cap) - int(mu_cfg.shape[2]) + if extra <= 0: + return x, mu_cfg, cond_cfg + return ( + F.pad(x, (0, extra)), + F.pad(mu_cfg, (0, extra), mode="replicate"), + F.pad(cond_cfg, (0, extra), mode="replicate"), + ) + + +def _build_capture_mask( + *, + attn_mask: torch.Tensor | None, + batch_size: int, + query_cap: int, + offset: int, + mel_width: int, + mel_frames: int, + device: torch.device, +) -> torch.Tensor: + """Always-on attention mask in the capture layout, so presence-of-mask is not a graph key. + + The DiT puts the current chunk ahead of its cache (stepaudio2 + ``Attention.forward_chunk``: ``torch.cat([k, k_cache], dim=2)``). The + caller's keys are therefore ``[current(mel_width) | cache(offset)]`` while + the captured graph sees ``[current(query_cap) | cache(offset)]``: the cache + block moves to ``query_cap`` and the capture padding in between is never + attended. A missing caller mask means "attend to the valid prefix". + + Capture-padding query rows reuse the first row's keys. Their output is + dropped, but an all-False row makes SDPA emit NaN there, and a NaN value in + a masked key still poisons real rows through ``0 * NaN`` in ``P @ V``. + """ + query_cap = int(query_cap) + offset = int(offset) + mel_width = int(mel_width) + mask = torch.zeros( + 2 * int(batch_size), + query_cap, + query_cap + offset, + dtype=torch.bool, + device=device, + ) + if attn_mask is not None: + mask[:, :mel_width, :mel_width] = attn_mask[:, :mel_width, :mel_width] + if offset > 0: + mask[:, :mel_width, query_cap:] = attn_mask[:, :mel_width, mel_width : mel_width + offset] + else: + mask[:, :mel_width, : int(mel_frames)] = True + if offset > 0: + mask[:, :mel_width, query_cap:] = True + if query_cap > mel_width: + mask[:, mel_width:] = mask[:, :1] + return mask + + +def _att_keep_ranges(total: int, keep: tuple[int, int] | None) -> list[tuple[int, int]]: + """``(start, length)`` frame ranges of a new estimator cache that survive the streaming trim. + + Mirrors stepaudio2 ``Token2wav.stream``: once the cache outgrows + ``prompt_len + 100`` frames it keeps the first ``prompt_len`` and the last + 100. ``keep`` is ``(prompt_len, 100)``; ``None`` keeps everything. + """ + total = int(total) + if keep is not None: + prefix, suffix = int(keep[0]), int(keep[1]) + if total > prefix + suffix: + return [(0, prefix), (total - suffix, suffix)] + return [(0, total)] + + +def _whole_euler_att_segments( + *, + mel_width: int, + query_cap: int, + offset: int, + keep: tuple[int, int] | None, +) -> list[tuple[int, int]]: + """Frame segments of a Whole-Euler cache output, in caller layout and trimmed. + + The graph writes ``[current(query_cap) | cache(offset)]`` (current chunk + first, see ``_build_capture_mask``); callers expect + ``[current(mel_width) | cache(offset)]``. The segments skip the capture + padding in between, then apply ``keep`` (``_att_keep_ranges``). + """ + segments: list[tuple[int, int]] = [] + for start, length in _att_keep_ranges(mel_width + offset, keep): + end = start + length + if start < mel_width: + segments.append((start, min(end, mel_width) - start)) + if end > mel_width: + low = max(start, mel_width) + segments.append((query_cap + low - mel_width, end - low)) + merged: list[tuple[int, int]] = [] + for start, length in segments: + if merged and sum(merged[-1]) == start: + merged[-1] = (merged[-1][0], merged[-1][1] + length) + elif length > 0: + merged.append((start, length)) + return merged + + +def _resized_frame_view(cache: torch.Tensor, frames: int) -> torch.Tensor | None: + """``cache`` (..., L, width) as ``frames`` frames of its own allocation, or ``None`` when it has no room. + + Only a view of the first frames of a whole contiguous allocation + qualifies, as ``WholeEulerCFMGraphWrapper.replay`` hands out: a cache + allocated for its steady length then grows into it without a second copy. + """ + width = int(cache.shape[-1]) + room = cache.stride(-3) // max(width, 1) + full = (*cache.shape[:-2], room, width) + strides, step = [], 1 + for size in reversed(full): + strides.append(step) + step *= size + if ( + room < frames + or cache.stride() != tuple(reversed(strides)) + or cache.storage_offset() != 0 + or cache.untyped_storage().nbytes() != step * cache.element_size() + ): + return None + return cache.as_strided((*cache.shape[:-2], frames, width), cache.stride()) + + +def _copy_frame_segments(dst: torch.Tensor, src: torch.Tensor, segments: list[tuple[int, int]]) -> None: + """Write ``src[..., start:start + length, :]`` segments back to back into ``dst``'s frame axis.""" + position = 0 + for start, length in segments: + dst[..., position : position + length, :].copy_(src[..., start : start + length, :]) + position += length + + def _tensors_from_key(key: tuple) -> tuple[torch.Tensor, ...]: """Rebuild zero tensors from a cache key (shape, dtype, device tuples). @@ -361,3 +568,806 @@ class CFMGraphWrapper: static_inputs[4].detach().clone(), static_inputs[5].detach().clone(), ) + + +def _zero_padded_cnn_cache( + cnn_cache: torch.Tensor, + estimator: torch.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. + ``cnn_cache`` is indexed by block on its first axis. + """ + blocks = getattr(estimator, "blocks", None) + if pad_frames > 0 and blocks and len(blocks) == cnn_cache.shape[0]: + width = int(cnn_cache.shape[-1]) + if width > 0 and all( + hasattr(block, "conv") + and hasattr(block.conv, "block") + and len(block.conv.block) > 1 + and hasattr(block.conv.block[1], "causal_padding") + and int(block.conv.block[1].causal_padding[0]) == width + for block in blocks + ): + # _estimator_buffers packs equal-width blocks into one tensor. + # Clear their shared tail with one write instead of one per block. + cnn_cache[..., max(0, width - pad_frames) :] = 0.0 + return + + if blocks is not None: + for index, block in enumerate(blocks): + if index >= len(cnn_cache): + break + conv = getattr(block, "conv", None) + conv_block = getattr(conv, "block", None) if conv is not None else None + if conv_block is None or len(conv_block) <= 1 or not hasattr(conv_block[1], "causal_padding"): + continue + width = int(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 + + +class WholeEulerExecutionArena: + """Static staging buffers for Whole-Euler graphs, shared across graph entries. + + Time embeddings, CNN caches and speakers are allocated once per graph batch + size. The attention caches dominate the arena (timesteps x depth x CFG x + heads x frames, ~1.25 MiB per frame per request in fp32 on the shipped + DiT), so there is exactly one attention cache storage: every graph, of + every batch size and offset, takes a ``[:2B, ..., :frames]`` view of it. + A graph's output cache is ``[current(query_cap) | cache(offset)]`` and its + input cache is the ``cache`` part of that same view, so a block that + attends over its output in place (``_attend_in_place``) finds the cache + already there; upstream's block concatenates the cache it reads + (``torch.cat([k, k_cache])``) before writing it back, which is safe on + that aliasing too. Graphs replay one at a time on one stream and refill + their inputs first, so the sharing across graphs is safe. Separate + per-offset, per-batch or input/output storages each held another full copy. + Small shape-keyed buffers (x, mu, cond, mask) stay per query width. + """ + + def __init__( + self, + estimator: torch.nn.Module, + *, + n_timesteps: int = 10, + device: torch.device, + dtype: torch.dtype, + att_cache_dtype: torch.dtype = torch.float32, + time_steps: list[torch.Tensor] | None = None, + ) -> None: + self.estimator = estimator + self.n_timesteps = int(n_timesteps) + self.device = device + self.dtype = dtype + self.att_cache_dtype = att_cache_dtype + self.time_steps = list(time_steps) if time_steps is not None else [] + + blocks = estimator.blocks + self.depth = len(blocks) + block0 = blocks[0] + self.cnn_channels = int(block0.conv.in_channels + block0.conv.out_channels) + self.cnn_width = int(block0.conv.block[1].causal_padding[0]) + self.heads = int(block0.attn.num_heads) + self.att_width = int(block0.attn.head_dim * 2) + + # Static staging buffers, shared by every graph that stages the same shape. + self._buffers: dict[tuple, torch.Tensor] = {} + # The one attention cache storage; every graph takes a prefix view. + self._att: torch.Tensor | None = None + + def _buffer(self, name: str, shape: tuple[int, ...], new=torch.zeros, dtype: torch.dtype | None = None): + key = (name, shape) + buf = self._buffers.get(key) + if buf is None: + buf = new(shape, device=self.device, dtype=self.dtype if dtype is None else dtype) + self._buffers[key] = buf + return buf + + def get_time_embeddings(self, batch_size: int) -> torch.Tensor: + key = ("time_embeddings", batch_size) + t_emb = self._buffers.get(key) + if t_emb is None: + t_emb = torch.stack( + [ + self.estimator.t_embedder(self.time_steps[s].expand(2 * batch_size)).unsqueeze(1) + for s in range(self.n_timesteps) + ], + dim=0, + ) + self._buffers[key] = t_emb + return t_emb + + def _cnn_cache_shape(self, batch_size: int) -> tuple[int, ...]: + return (self.n_timesteps, self.depth, 2 * batch_size, self.cnn_channels, self.cnn_width) + + def get_cnn_cache_in(self, batch_size: int) -> torch.Tensor: + return self._buffer("cnn_in", self._cnn_cache_shape(batch_size)) + + def get_cnn_cache_out(self, batch_size: int) -> torch.Tensor: + return self._buffer("cnn_out", self._cnn_cache_shape(batch_size), new=torch.empty) + + def get_speakers_staging(self, batch_size: int, spk_dim: int) -> torch.Tensor: + return self._buffer("speakers", (2 * batch_size, spk_dim)) + + def get_x_staging(self, batch_size: int, channels: int, mel_width: int) -> torch.Tensor: + return self._buffer("x", (batch_size, channels, mel_width)) + + def get_mu_staging(self, batch_size: int, channels: int, mel_width: int) -> torch.Tensor: + return self._buffer("mu", (2 * batch_size, channels, mel_width)) + + def get_cond_staging(self, batch_size: int, channels: int, mel_width: int) -> torch.Tensor: + return self._buffer("cond", (2 * batch_size, channels, mel_width)) + + def att_cache_fits(self, batch_size: int, frames: int) -> bool: + """Whether a view of ``batch_size`` requests ending at frame ``frames`` fits the current storage.""" + storage = self._att + return storage is not None and int(storage.shape[2]) >= 2 * batch_size and int(storage.shape[4]) >= frames + + def att_cache_view( + self, batch_size: int, frames: int, *, start: int = 0, capacity: int = 0, rows: int = 0 + ) -> torch.Tensor: + """``(n_t, depth, 2B, heads, frames, width)`` view of the shared cache storage from frame ``start``. + + ``capacity`` frames and ``rows`` requests are reserved up front. A + larger request replaces the storage. Graphs captured on the old one + would keep it alive beside the new one, so + ``WholeEulerCFMGraphWrapper._capture`` flushes them first. + """ + storage = self._att + frames = int(frames) + end = int(start) + frames + if not self.att_cache_fits(batch_size, end): + old_rows = 0 if storage is None else int(storage.shape[2]) // 2 + old_frames = 0 if storage is None else int(storage.shape[4]) + storage = torch.zeros( + self.n_timesteps, + self.depth, + 2 * max(int(batch_size), int(rows), old_rows), + self.heads, + _align_up(max(end, int(capacity), old_frames, 1), _ATT_FRAME_ALIGN), + self.att_width, + device=self.device, + dtype=self.att_cache_dtype, + ) + self._att = storage + return storage[:, :, : 2 * batch_size, :, start:end, :] + + def get_mask_staging(self, batch_size: int, mel_width: int, total_len: int) -> torch.Tensor: + return self._buffer("mask", (2 * batch_size, mel_width, total_len), new=torch.ones, dtype=torch.bool) + + def get_lengths_staging(self, batch_size: int) -> torch.Tensor: + """Per-CFG-row valid query lengths read by a ragged graph body.""" + return self._buffer("lengths", (2 * batch_size,), dtype=torch.long) + + def clear(self) -> None: + self._buffers.clear() + self._att = None + + +class _WholeEulerStatics(NamedTuple): + """A Whole-Euler graph's static inputs (``lengths`` only with a ragged body).""" + + x: torch.Tensor + mu_cfg: torch.Tensor + speakers_cfg: torch.Tensor + cond_cfg: torch.Tensor + cnn_cache: torch.Tensor + att_cache: torch.Tensor + attn_mask: torch.Tensor + time_embeddings: torch.Tensor + lengths: torch.Tensor | None + + +class WholeEulerCFMGraphWrapper: + """Per-shape CUDA graph capture/replay for the entire 10-step CFM Euler ODE solver. + + Captures all 10 iterations of: + (time embedding -> in_proj -> DiT blocks -> final layer -> CFG combination -> Euler step) + into a single CUDA Graph replay. + + Replaces 10 separate step-level graph replays and 30 Python clones per chunk + with a single graph replay, reducing CPU launch overhead and eliminating host-device + synchronization bubbles during high concurrency. + + A graph must hold all 10 timesteps' attention caches in static buffers, + ten times what one step-level graph holds, so memory is the constraint: + graphs exist for the powers of two up to ``micro_batch_size``, all of them + share the one attention cache storage of ``WholeEulerExecutionArena`` + (sized for ``micro_batch_size`` requests), and larger batches run as a + sequence of micro-batch replays that copy each request's cache in and out + of that storage directly. + """ + + def __init__( + self, + estimator: torch.nn.Module, + *, + n_timesteps: int = 10, + inference_cfg_rate: float = 0.7, + att_cache_dtype: torch.dtype = torch.float32, + max_graphs: int = 32, + max_serial_batch: int | None = None, + max_graph_batch: int | None = None, + micro_batch_size: int = 4, + query_bucket_frames: int | None = None, + pad_max_rows: int | None = None, + ragged_body: Callable[..., torch.Tensor] | None = None, + ) -> None: + """``ragged_body(estimator, input, t_emb, mask, cnn, att, cnn_out, att_out, lengths)`` + replaces ``estimator.blocks_forward_chunk`` in the captured solve when + given: it takes each row's valid query length as a ``(2B,)`` tensor and + writes that row's exact CNN cache (``_blocks_forward_chunk_ragged``). + It lets one graph solve rows of different lengths, and lets a padded + query axis keep the CNN cache the unpadded solve would produce instead + of zeroing it. + """ + self.estimator = estimator + self.ragged_body = ragged_body + self.n_timesteps = int(n_timesteps) + self.inference_cfg_rate = float(inference_cfg_rate) + self.att_cache_dtype = att_cache_dtype + self.max_graphs = int(max_graphs) + self.query_bucket_frames = 0 if query_bucket_frames is None else int(query_bucket_frames) + self.max_serial_batch = 4 if max_serial_batch is None else int(max_serial_batch) + self.micro_batch_size = 4 if micro_batch_size is None else int(micro_batch_size) + if max_graph_batch is None: + if max_serial_batch is not None and max_serial_batch < self.micro_batch_size: + self.max_graph_batch = int(max_serial_batch) + else: + self.max_graph_batch = 16 + else: + self.max_graph_batch = int(max_graph_batch) + if pad_max_rows is None: + pad_max_rows = self.micro_batch_size // 4 + self.pad_max_rows = max(0, int(pad_max_rows)) + # Largest steady cache length a caller announced; sizes the arena storage. + self._att_capacity = 0 + parameter = next(estimator.parameters(), None) + if parameter is not None: + self.device = parameter.device + self.dtype = parameter.dtype + else: + self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + self.dtype = torch.float32 + + self.enabled = self.max_graphs > 0 and self.device.type == "cuda" + self._cache: dict[tuple, tuple] = {} + self._unsupported: set[tuple] = set() + self._stats = { + "calls": 0, + "hits": 0, + "captures": 0, + "flushes": 0, + "eager": 0, + } + + self.timeline, self.dt_steps = _euler_timeline(self.n_timesteps, self.device, self.dtype) + self.time_steps = [self.timeline[i] for i in range(self.n_timesteps)] + self.arena = WholeEulerExecutionArena( + estimator=self.estimator, + n_timesteps=self.n_timesteps, + device=self.device, + dtype=self.dtype, + att_cache_dtype=self.att_cache_dtype, + time_steps=self.time_steps, + ) + + def stats_snapshot(self) -> dict[str, int]: + """Bounded cumulative telemetry for the graph cache.""" + return {**self._stats, "cache_size": len(self._cache)} + + def _flush(self) -> None: + """Retire every captured graph at once.""" + torch.accelerator.synchronize(self.device) + if self._cache: + for entry in self._cache.values(): + entry[4].reset() + self._cache.clear() + self._stats["flushes"] += 1 + logger.info("Whole-Euler CFM graph cache flushed; stats=%s", self.stats_snapshot()) + self.arena.clear() + + def _disable(self, reason: str, key: tuple) -> None: + logger.warning( + "Disabling Whole-Euler CFM CUDA graphs (%s) for shape=%s; using step/eager fallback", + reason, + key, + exc_info=True, + ) + self.enabled = False + self._flush() + + def _run_euler_loop( + self, + statics: _WholeEulerStatics, + out_cnn_cache: torch.Tensor, + out_att_cache: torch.Tensor, + *, + batch_size: int, + ) -> torch.Tensor: + cur_x = statics.x + width = int(statics.mu_cfg.shape[2]) + speaker_features = statics.speakers_cfg.unsqueeze(-1).expand(-1, -1, width) + + for step in range(self.n_timesteps): + dt = self.dt_steps[step] + x_cfg = torch.cat((cur_x, cur_x), dim=0) + estimator_input = torch.cat((x_cfg, statics.mu_cfg, speaker_features, statics.cond_cfg), dim=1) + args = ( + estimator_input, + statics.time_embeddings[step], + statics.attn_mask, + statics.cnn_cache[step], + statics.att_cache[step], + out_cnn_cache[step], + out_att_cache[step], + ) + if statics.lengths is not None: + assert self.ragged_body is not None + estimate = self.ragged_body(self.estimator, *args, statics.lengths) + else: + estimate = self.estimator.blocks_forward_chunk(*args) + + cur_x = _euler_step(cur_x, estimate, dt, self.inference_cfg_rate, batch_size) + + return cur_x + + def _graph_batches(self) -> list[int]: + """Native graph batch sizes: the powers of two below ``micro_batch_size``, and it.""" + sizes = [] + size = 1 + while size < self.micro_batch_size: + sizes.append(size) + size *= 2 + sizes.append(self.micro_batch_size) + return sizes + + def _plan_groups(self, batch_size: int) -> list[tuple[int, int]] | None: + """``(graph_batch, rows)`` replays that cover ``batch_size`` requests. + + Larger batches run whole micro-batches. The remainder takes the smallest + native size that holds it if that pads at most ``pad_max_rows`` rows + (the padded rows are ignored), and otherwise the largest native size + below it, repeating on what is left. On a saturated device a small + graph's replay costs well over its share of a larger one, so a padded + row is cheaper than a further replay. Returns ``None`` when the plan + needs more replays than allowed. + """ + micro = self.micro_batch_size + sizes = self._graph_batches() + groups: list[tuple[int, int]] = [] + remainder = batch_size + while remainder >= micro: + groups.append((micro, micro)) + remainder -= micro + while remainder > 0: + up = next(size for size in sizes if size >= remainder) + if up - remainder <= self.pad_max_rows: + groups.append((up, remainder)) + break + down = max(size for size in sizes if size <= remainder) + groups.append((down, down)) + remainder -= down + max_allowed = max( + self.max_serial_batch, + (self.max_graph_batch + micro - 1) // micro + (micro - 1), + ) + if len(groups) > max_allowed: + return None + return groups + + @staticmethod + def _fill_static_inputs( + statics: _WholeEulerStatics, + *, + graph_batch: int, + start: int, + stop: int, + x: torch.Tensor, + mu_cfg: torch.Tensor, + speakers_cfg: torch.Tensor, + cond_cfg: torch.Tensor, + attn_mask: torch.Tensor, + cnn_cache: torch.Tensor | None, + att_cache: torch.Tensor | None, + att_rows: list[torch.Tensor] | None, + offset: int, + lengths: torch.Tensor | int | None = None, + ) -> None: + """Copy requests ``start:stop`` into a graph's static inputs. + + CFG inputs are stacked ``[cond x B | uncond x B]``; ``unflatten`` into + ``(2, B)`` moves both halves of a request range with one copy. Rows past + ``stop - start`` belong to a padded replay and keep whatever finite + values they held: DiT rows never mix, so their output is simply dropped. + ``lengths`` is the ragged body's per-row lengths, or one for every row. + Padded rows get length 0: every query width shares the buffer, and a + wider replay's length would index past this one's causal history. + """ + batch_size = int(x.shape[0]) + rows = stop - start + statics.x[:rows].copy_(x[start:stop]) + pairs = [ + (statics.mu_cfg, mu_cfg), + (statics.speakers_cfg, speakers_cfg), + (statics.cond_cfg, cond_cfg), + (statics.attn_mask, attn_mask), + ] + if statics.lengths is not None: + assert lengths is not None + if isinstance(lengths, int): + statics.lengths.fill_(lengths) + else: + if rows < graph_batch: + statics.lengths.zero_() + pairs.append((statics.lengths, lengths)) + for static, value in pairs: + static.unflatten(0, (2, graph_batch))[:, :rows].copy_(value.unflatten(0, (2, batch_size))[:, start:stop]) + if cnn_cache is None: + statics.cnn_cache.zero_() + else: + statics.cnn_cache.unflatten(2, (2, graph_batch))[:, :, :, :rows].copy_( + cnn_cache.unflatten(2, (2, batch_size))[:, :, :, start:stop] + ) + if offset <= 0: + return + static_att_rows = statics.att_cache.unflatten(2, (2, graph_batch)) + if att_rows is None: + assert att_cache is not None + static_att_rows[:, :, :, :rows].copy_(att_cache.unflatten(2, (2, batch_size))[:, :, :, start:stop]) + else: + for row in range(rows): + static_att_rows[:, :, :, row].copy_(att_rows[start + row]) + + def _capture( + self, + key: tuple, + *, + graph_batch: int, + channels: int, + query_cap: int, + spk_dim: int, + offset: int, + fill, + ) -> tuple | None: + arena = self.arena + # Reserve the steady offset plus two query buckets, and the largest + # graph batch, up front, so every offset on the way there + # (304 -> 368 -> 400) and every batch size share one storage. Two + # buckets because a final chunk carries the lookahead tail past one + # (50 + 6 mel frames). The wide offset-0 prompt solve fits within that; + # adding its width on top of the steady offset would hold ~300 frames + # per row that nothing uses. + capacity = max(offset + query_cap, max(offset, self._att_capacity) + 2 * self.query_bucket_frames) + rows = self.micro_batch_size + if self._cache and not arena.att_cache_fits(graph_batch, offset + query_cap): + # The captured graphs hold views of the current storage: retire + # them, or they keep it alive beside the larger replacement. + self._flush() + try: + statics = _WholeEulerStatics( + x=arena.get_x_staging(graph_batch, channels, query_cap), + mu_cfg=arena.get_mu_staging(graph_batch, channels, query_cap), + speakers_cfg=arena.get_speakers_staging(graph_batch, spk_dim), + cond_cfg=arena.get_cond_staging(graph_batch, channels, query_cap), + cnn_cache=arena.get_cnn_cache_in(graph_batch), + # Right behind the current chunk in the output view below. + att_cache=arena.att_cache_view(graph_batch, offset, start=query_cap, capacity=capacity, rows=rows), + attn_mask=arena.get_mask_staging(graph_batch, query_cap, offset + query_cap), + time_embeddings=arena.get_time_embeddings(graph_batch), + lengths=arena.get_lengths_staging(graph_batch) if self.ragged_body is not None else None, + ) + out_cnn_cache = arena.get_cnn_cache_out(graph_batch) + # The output view aliases the input view; see WholeEulerExecutionArena. + out_att_cache = arena.att_cache_view(graph_batch, offset + query_cap, capacity=capacity, rows=rows) + fill(statics) + except Exception: + logger.warning("Failed to allocate static buffers for Whole-Euler graph key: %s", key, exc_info=True) + self._unsupported.add(key) + return None + + static_x = statics.x + # The fused Euler step integrates static_x in place, so every warmup + # and the capture itself start from the same initial state. + initial_x = static_x.clone() + loop_args = (statics, out_cnn_cache, out_att_cache) + memory_before = _memory_snapshot(self.device) + try: + current_stream = torch.cuda.current_stream(self.device) + warmup_stream = torch.cuda.Stream(device=self.device) + warmup_stream.wait_stream(current_stream) + with torch.cuda.stream(warmup_stream), torch.no_grad(): + for _ in range(3): + static_x.copy_(initial_x) + self._run_euler_loop(*loop_args, batch_size=graph_batch) + current_stream.wait_stream(warmup_stream) + static_x.copy_(initial_x) + graph = CUDAGraph() + with torch.no_grad(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + static_final_x = self._run_euler_loop(*loop_args, batch_size=graph_batch) + except Exception: + self._disable("capture failed", key) + return None + + self._stats["captures"] += 1 + logger.info( + "Captured Whole-Euler CFM CUDA Graph for shape %s (cache=%d/%d, att storage rows/frames=%d/%d, stats=%s)%s", + key, + len(self._cache) + 1, + self.max_graphs, + int(arena._att.shape[2]) // 2, + int(arena._att.shape[4]), + self.stats_snapshot(), + _format_memory_delta(memory_before, _memory_snapshot(self.device)), + ) + return (statics, static_final_x, out_cnn_cache, out_att_cache, graph) + + def _entry( + self, + *, + graph_batch: int, + query_cap: int, + offset: int, + x: torch.Tensor, + spk_dim: int, + fill, + ) -> tuple | None: + # One graph per (batch, capture query width, cache offset). + # Presence of cnn/att/mask is data copied into static buffers, not a key. + key = ("whole_euler", graph_batch, query_cap, offset, str(x.dtype), str(x.device)) + if key in self._unsupported: + return None + entry = self._cache.get(key) + if entry is not None: + self._stats["hits"] += 1 + return entry + if len(self._cache) >= self.max_graphs: + self._flush() + entry = self._capture( + key, + graph_batch=graph_batch, + channels=int(x.shape[1]), + query_cap=query_cap, + spk_dim=spk_dim, + offset=offset, + fill=fill, + ) + if entry is not None: + self._cache[key] = entry + return entry + + def _group_entries(self, groups: list[tuple[int, int]], fills: list, **entry_args) -> list[tuple] | None: + """Every group's graph, acquired before any of them replays. + + A replay may update request caches in place, so no group may fall back + to eager after an earlier one ran. Capturing a later group can flush the + cache (``max_graphs``, arena growth) and retire an earlier group's + graph; the groups are then acquired once more, and fit this time. + """ + for _ in range(2): + flushes = self._stats["flushes"] + entries = [] + for (graph_batch, _), fill in zip(groups, fills, strict=True): + entry = self._entry(graph_batch=graph_batch, fill=fill, **entry_args) + if entry is None: + return None + entries.append(entry) + if self._stats["flushes"] == flushes: + return entries + return None + + def replay( + self, + *, + x: torch.Tensor, + mu_cfg: torch.Tensor, + speakers_cfg: torch.Tensor, + cond_cfg: torch.Tensor, + cnn_cache: torch.Tensor | None, + att_cache: torch.Tensor | list[torch.Tensor] | None, + attn_mask: torch.Tensor | None = None, + mel_frames: int | None = None, + pad_frames: int = 0, + valid_lengths: list[int] | None = None, + att_keep: tuple[int, int] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | list[torch.Tensor]] | None: + """Solve the chunk through captured graphs; ``None`` sends it to the step/eager path. + + ``att_cache`` is the stacked ``(n_t, depth, 2B, heads, L, 2*head_dim)`` + estimator cache, or one ``(n_t, depth, 2, heads, L, 2*head_dim)`` + tensor per request; the new cache comes back in the same form, trimmed + by ``att_keep`` (``_att_keep_ranges``) on the way out. Per request, + caches move straight into and out of the arena and no stacked copy of + the batch's cache is ever built. ``sum(att_keep)`` is also the steady + cache length the arena reserves. + + Request caches are updated in place: holding the old and new cache of + every request at once was the largest block of Stage-2 memory at high + concurrency (~7 GiB at 16 requests). A new one is allocated with room + for the steady length ``sum(att_keep)`` and returned as a view of its + first frames, so it also grows in place (``_resized_frame_view``). The + caller hands over those caches; a tensor shared by several rows (the + prompt state) is never written. + + ``valid_lengths`` (one per request; needs ``ragged_body``) solves rows + of different lengths in one graph, as ``_decode_cfm`` does eagerly: + ``mu_cfg`` is padded to the longest row and ``attn_mask`` masks each + row's padding. The new cache then always comes back per request, each + row keeping only its own current frames. + """ + self._stats["calls"] += 1 + + if not self.enabled or torch.cuda.is_current_stream_capturing() or x.device.type != "cuda": + return None + + batch_size = int(x.shape[0]) + if batch_size > self.max_graph_batch: + return None + groups = self._plan_groups(batch_size) + if groups is None: + return None + + att_rows = None if att_cache is None or isinstance(att_cache, torch.Tensor) else list(att_cache) + if att_rows is not None: + if len(att_rows) != batch_size or len({int(row.shape[4]) for row in att_rows}) != 1: + return None + offset = int(att_rows[0].shape[4]) + else: + offset = int(att_cache.shape[4]) if att_cache is not None else 0 + mel_width = int(mu_cfg.shape[2]) + att_rows_out = att_rows is not None or valid_lengths is not None + if valid_lengths is not None: + if self.ragged_body is None or len(valid_lengths) != batch_size: + return None + pad_frames = 0 + mel_frames = mel_width + elif mel_frames is None: + mel_frames = mel_width - pad_frames + if att_keep is not None: + self._att_capacity = max(self._att_capacity, sum(att_keep)) + + query_cap = _capture_query_width(mel_width, self.query_bucket_frames) + x_cap, mu_cap, cond_cap = _pad_query_for_capture(x, mu_cfg, cond_cfg, query_cap) + mask_cap = _build_capture_mask( + attn_mask=attn_mask, + batch_size=batch_size, + query_cap=query_cap, + offset=offset, + mel_width=mel_width, + mel_frames=mel_frames, + device=x.device, + ) + row_lengths = [mel_width] * batch_size if valid_lengths is None else [int(n) for n in valid_lengths] + lengths: torch.Tensor | int | None = None + if self.ragged_body is not None: + lengths = mel_width + if valid_lengths is not None: + lengths = torch.tensor((*row_lengths, *row_lengths), dtype=torch.long, device=x.device) + # Queries past a row's length take its first query's mask, as + # capture padding does, so no attention row is left empty. + invalid = torch.arange(query_cap, device=x.device).unsqueeze(0) >= lengths.unsqueeze(1) + mask_cap = torch.where(invalid.unsqueeze(2), mask_cap[:, :1, :], mask_cap) + segments_by_length = { + length: _whole_euler_att_segments(mel_width=length, query_cap=query_cap, offset=offset, keep=att_keep) + for length in set(row_lengths) + } + segments = segments_by_length[row_lengths[0]] + keep_len = sum(length for _, length in segments) + + arena = self.arena + chunk_mel = x.new_empty((batch_size, int(x.shape[1]), mel_frames)) + out_cnn = torch.empty( + (self.n_timesteps, arena.depth, 2 * batch_size, arena.cnn_channels, arena.cnn_width), + device=x.device, + dtype=self.dtype, + ) + out_cnn_rows = out_cnn.unflatten(2, (2, batch_size)) + out_att: torch.Tensor | None = None + out_att_rows = None + request_att: list[torch.Tensor] = [] + if not att_rows_out: + out_att = torch.empty( + (self.n_timesteps, arena.depth, 2 * batch_size, arena.heads, keep_len, arena.att_width), + device=x.device, + dtype=self.att_cache_dtype, + ) + out_att_rows = out_att.unflatten(2, (2, batch_size)) + # A ragged body writes each row's exact CNN cache; only the bucket + # padding (``pad_frames``) is then cleared. Otherwise every padded position is. + cnn_pad = pad_frames if self.ragged_body is not None else query_cap - mel_frames + + fills = [] + start = 0 + for graph_batch, rows in groups: + stop = start + rows + + def fill(statics, *, _graph_batch=graph_batch, _start=start, _stop=stop): + self._fill_static_inputs( + statics, + graph_batch=_graph_batch, + start=_start, + stop=_stop, + x=x_cap, + mu_cfg=mu_cap, + speakers_cfg=speakers_cfg, + cond_cfg=cond_cap, + attn_mask=mask_cap, + cnn_cache=cnn_cache, + att_cache=att_cache if att_rows is None else None, + att_rows=att_rows, + offset=offset, + lengths=lengths, + ) + + fills.append(fill) + start = stop + entries = self._group_entries( + groups, + fills, + query_cap=query_cap, + offset=offset, + x=x, + spk_dim=int(speakers_cfg.shape[1]), + ) + if entries is None: + return None + owners = Counter(row.data_ptr() for row in att_rows) if att_rows is not None else Counter() + steady_frames = sum(att_keep) if att_keep is not None else 0 + + start = 0 + for (graph_batch, rows), fill, entry in zip(groups, fills, entries, strict=True): + stop = start + rows + statics, static_final_x, out_cnn_cache, out_att_cache, graph = entry + fill(statics) + graph.replay() + + chunk_mel[start:stop].copy_(static_final_x[:rows, :, :mel_frames]) + if cnn_pad > 0: + # (n_timesteps, depth, ...): put the block axis first for the helper. + _zero_padded_cnn_cache(out_cnn_cache.transpose(0, 1), self.estimator, cnn_pad) + out_cnn_rows[:, :, :, start:stop].copy_(out_cnn_cache.unflatten(2, (2, graph_batch))[:, :, :, :rows]) + att_src = out_att_cache.unflatten(2, (2, graph_batch))[:, :, :, :rows] + if pad_frames > 0: + # The static output is scratch until the next replay, so the + # padded frames are cleared once here rather than per copy. + att_src[..., mel_frames:mel_width, :] = 0.0 + if out_att_rows is not None: + _copy_frame_segments(out_att_rows[:, :, :, start:stop], att_src, segments) + else: + for row in range(rows): + row_segments = segments_by_length[row_lengths[start + row]] + shape = ( + self.n_timesteps, + arena.depth, + 2, + arena.heads, + sum(length for _, length in row_segments), + arena.att_width, + ) + # The row's old cache was copied into the arena by ``fill``. + old = att_rows[start + row] if att_rows is not None else None + request = None + if old is not None and old.dtype == self.att_cache_dtype and owners[old.data_ptr()] == 1: + request = _resized_frame_view(old, shape[4]) + if request is None: + # Room for the steady length, so the cache grows into it in place. + room = max(shape[4], steady_frames) + request = torch.empty( + (*shape[:4], room, shape[5]), device=x.device, dtype=self.att_cache_dtype + )[..., : shape[4], :] + _copy_frame_segments(request, att_src[:, :, :, row], row_segments) + request_att.append(request) + start = stop + + return chunk_mel, out_cnn, (request_att if att_rows_out else out_att) 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 098354c44..072d51e50 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 @@ -268,14 +268,34 @@ class MiniCPMO45Code2Wav(nn.Module): raise ValueError(f"Invalid MiniCPM-o connector chunk configuration: {self._connector_config}") raw_capture_batch_sizes = extra.get("hift_graph_capture_batch_sizes") capture_batch_sizes = [1] if raw_capture_batch_sizes is None else raw_capture_batch_sizes + max_serial_batch = extra.get("max_serial_batch") + max_serial_batch = 4 if max_serial_batch is None else int(max_serial_batch) self._hift_graph_config = { "enabled": bool(extra.get("enable_hift_graph", False)), "capture_batch_sizes": capture_batch_sizes, + "max_serial_batch": max_serial_batch, } + enable_whole_euler = extra.get("enable_whole_euler") + max_graph_batch_raw = extra.get("max_graph_batch") + max_graph_batch = int(max_graph_batch_raw) if max_graph_batch_raw is not None else None + micro_batch_size_raw = extra.get("micro_batch_size") + if micro_batch_size_raw is not None: + micro_batch_size = int(micro_batch_size_raw) + else: + # The Whole-Euler arena reserves one attention cache per micro-batch + # row, so size it for the most requests this stage ever batches. + max_num_seqs = getattr(getattr(vllm_config, "scheduler_config", None), "max_num_seqs", None) + micro_batch_size = min(int(max_num_seqs), max_graph_batch or 16) if max_num_seqs else None 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)), + "capture_frames": extra.get("cfm_graph_capture_frames"), + "enable_whole_euler": enable_whole_euler is None or bool(enable_whole_euler), + "max_serial_batch": max_serial_batch, + "max_graph_batch": max_graph_batch, + "micro_batch_size": micro_batch_size, + "pad_max_rows": extra.get("whole_euler_pad_max_rows"), } self._ref_max_seconds = float(extra.get("ref_audio_max_seconds", _REF_MAX_SECONDS)) if self._ref_max_seconds <= 0: 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 4346e83b6..3314f468f 100644 --- a/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py +++ b/vllm_omni/platforms/npu/models/minicpmo_4_5_code2wav.py @@ -114,6 +114,7 @@ def _patched_estimator_step( attn_mask=None, valid_lengths=None, valid_frames=None, + time_embedding=None, ): assert _original_estimator_step is not None graph_runner = _backend_graph_runners.get(self) @@ -137,6 +138,7 @@ def _patched_estimator_step( attn_mask=attn_mask, valid_lengths=valid_lengths, valid_frames=valid_frames, + time_embedding=time_embedding, ) if (cnn_cache is None) != (att_cache is None): raise ValueError("estimator CNN and attention caches must both be present or absent") @@ -148,7 +150,8 @@ def _patched_estimator_step( # 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) + if time_embedding is None: + time_embedding = estimator.t_embedder(time).unsqueeze(1) if cnn_cache is None: return graph_runner.run( "cfm_estimator",