From 79562e97ad8e08a3f9d9480ed99cc9abee69b98a Mon Sep 17 00:00:00 2001 From: BeatSeat <137616995+BeatSeat@users.noreply.github.com> Date: Tue, 6 Oct 2026 02:36:28 +0800 Subject: [PATCH] [Perf][MiniCPM-o] Code2Wav resident Euler slot pool, fused DiT body, and exact HiFT graphs (#8443) Signed-off-by: BeatSeat Co-authored-by: amy-why-3459 --- docs/design/minicpm_o45_mrv2_performance.md | 7 +- tests/config/test_environment_variables.py | 4 +- .../minicpmo_4_5/test_code2wav_batching.py | 111 ++- .../minicpmo_4_5/test_cuda_graph_wrapper.py | 150 ++- .../models/minicpmo_4_5/test_dit_fused.py | 70 ++ .../minicpmo_4_5/test_flow_encoder_graph.py | 116 +++ .../config/environment_variable_inventory.py | 1 + .../deploy/minicpmo_4_5_duplex_h200.yaml | 3 +- vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml | 8 +- .../models/minicpmo_4_5/batched_token2wav.py | 541 +++++++++-- .../models/minicpmo_4_5/cfm_attention.py | 138 +++ .../models/minicpmo_4_5/cuda_graph_wrapper.py | 884 +++++++++++++++--- .../models/minicpmo_4_5/dit_fused.py | 284 ++++++ .../models/minicpmo_4_5/flow_encoder_graph.py | 226 +++++ .../minicpmo_4_5/minicpmo_4_5_code2wav.py | 256 ++++- .../minicpmo_4_5/minicpmo_4_5_token2wav.py | 4 +- .../models/minicpmo_4_5/ops/__init__.py | 10 + .../minicpmo_4_5/ops/fused_qkv_layer_norm.py | 173 ++++ .../ops/fused_residual_layer_norm.py | 190 ++++ .../step_audio2/step_audio2_token2wav.py | 24 + 20 files changed, 2954 insertions(+), 246 deletions(-) create mode 100644 tests/model_executor/models/minicpmo_4_5/test_dit_fused.py create mode 100644 tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py create mode 100644 vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py diff --git a/docs/design/minicpm_o45_mrv2_performance.md b/docs/design/minicpm_o45_mrv2_performance.md index 64a8c6abf..e543f7b1b 100644 --- a/docs/design/minicpm_o45_mrv2_performance.md +++ b/docs/design/minicpm_o45_mrv2_performance.md @@ -12,8 +12,11 @@ require new end-to-end performance measurements after the mainline integration. Previously reported numbers do not describe this revised codec backend. The generic MRv2 profile and the opt-in native V1 duplex H200 profile enable -TF32 matmuls within Code2Wav forward/capture only; the previous process policy -is restored afterwards. cuDNN's TF32 policy is independent. TF32 changes rounding. +ordinary TF32 for CFM DiT GEMMs within Code2Wav forward/capture only +(`torch.backends.cuda.matmul.allow_tf32` on dense QKV/MLP, Triton +`input_precision="tf32"` on tiled attention). This is not compensated TF32x3. +The previous process matmul policy is restored afterwards. cuDNN's TF32 policy +is independent. HiFT stays IEEE FP32. TF32 changes rounding. Duplex keeps the mainline V1 session path. Turn results do not establish duplex performance or interruption correctness. diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index 58e2a31f1..a32c1cfec 100644 --- a/tests/config/test_environment_variables.py +++ b/tests/config/test_environment_variables.py @@ -167,7 +167,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): EnvironmentVariableCategory.PUBLIC_OMNI: 32, EnvironmentVariableCategory.INHERITED_VLLM: 20, EnvironmentVariableCategory.PLATFORM_EXTERNAL: 29, - EnvironmentVariableCategory.MODEL_SPECIFIC: 95, + EnvironmentVariableCategory.MODEL_SPECIFIC: 96, EnvironmentVariableCategory.BENCHMARK_TRANSITIONAL: 21, EnvironmentVariableCategory.INTERNAL: 7, } @@ -178,7 +178,7 @@ def test_inventory_matches_reviewed_snapshot_counts(): if item.category is EnvironmentVariableCategory.MODEL_SPECIFIC ) assert {disposition: disposition_counts[disposition] for disposition in ModelEnvironmentVariableDisposition} == { - ModelEnvironmentVariableDisposition.PROMOTE: 66, + ModelEnvironmentVariableDisposition.PROMOTE: 67, ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5, ModelEnvironmentVariableDisposition.EXTERNAL: 0, ModelEnvironmentVariableDisposition.INTERNALIZE: 17, 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 e97eb8e5a..02f5d3c80 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 @@ -173,7 +173,13 @@ 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)] + ("extra", "max_num_seqs", "micro"), + [ + ({}, 6, 6), + ({}, 64, 16), + ({"micro_batch_size": 2}, 6, 2), + ({"max_graph_batch": 8}, 64, 8), + ], ) def test_whole_euler_micro_batch_defaults_to_stage_concurrency(extra, max_num_seqs, micro): config = _config() @@ -183,6 +189,8 @@ def test_whole_euler_micro_batch_defaults_to_stage_concurrency(extra, max_num_se model = MiniCPMO45Code2Wav(vllm_config=config) assert model._cfm_graph_config["micro_batch_size"] == micro + if "max_graph_batch" in extra: + assert model._cfm_graph_config["max_graph_batch"] == extra["max_graph_batch"] def _model(initial: int = 0, minimum: int = 1, *, runtime_prompt_cache_size: int = 4, setup_cache_size: int = 1): @@ -1701,7 +1709,17 @@ def _padded_decode_adapter(bucket_frames: int) -> tuple[BatchedToken2Wav, int, i Returns the adapter, the valid width, and the padding width. Every call seeds the same way, so two adapters built here agree on weights and noise. """ - from cosyvoice2.flow.decoder_dit import DiT + for name in ("cosyvoice2.flow.decoder_dit", "stepaudio2.cosyvoice2.flow.decoder_dit"): + try: + import importlib + + decoder_dit = importlib.import_module(name) + break + except ImportError: + pass + else: + decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit") + DiT = decoder_dit.DiT torch.manual_seed(23) estimator = DiT( @@ -1809,3 +1827,92 @@ def test_padding_does_not_change_the_valid_frames(): assert int(padded_x.shape[2]) == int(exact_x.shape[2]) == mel_frames assert torch.allclose(padded_x, exact_x, atol=1e-5) + + +def test_decode_ragged_batch_mixed_onset_and_continuation(): + """Ragged batch merging fresh onset and continuation states does not crash stacking conformer caches.""" + adapter = BatchedToken2Wav(_FakeToken2Wav()) + _enable_fake_ragged_kernel(adapter) + prompt = adapter.prepare_prompt("shared_p", "/fake/prompt.wav") + onset_states = adapter.setup_batch(prompt, 1) + state_onset = onset_states[0] + + # Run one decode step on another state so its conformer/estimator caches advance. + cont_states_init = adapter.setup_batch(prompt, 1) + _, cont_states = adapter.decode_batch( + torch.tensor([[10, 11]]), + prompt, + cont_states_init, + last_chunk=False, + ) + state_continuation = cont_states[0] + + # Verify conformer caches have mismatched shapes between onset and continuation states. + assert ( + state_onset.flow_cache["conformer_att_cache"].shape + != state_continuation.flow_cache["conformer_att_cache"].shape + ) + + # Ragged decode with both onset and continuation states in a single batch. + audios, next_states = adapter.decode_ragged_batch( + [torch.tensor([20, 21]), torch.tensor([30, 31])], + prompt, + [state_onset, state_continuation], + last_chunks=[False, False], + ) + assert len(audios) == 2 + assert len(next_states) == 2 + + +def test_row_offset_fallback_reaches_eager_when_replay_declines(): + """Declined Whole-Euler replay must eager-solve, not recurse through per-offset grouping.""" + adapter = BatchedToken2Wav(_FakeToken2Wav(), cfm_graph_config={"row_offset_merge": True}) + _enable_fake_ragged_kernel(adapter) + mu, speakers, cond = torch.ones((2, 1, 2)), torch.ones((2, 1)), torch.zeros((2, 1, 2)) + _, _, stacked = adapter._decode_cfm(mu, speakers, cond, cnn_cache=None, att_cache=None) + equal_rows = [ + torch.cat((stacked[:, :, row : row + 1], stacked[:, :, 2 + row : 3 + row]), dim=2) for row in range(2) + ] + mixed_rows = [equal_rows[0], torch.nn.functional.pad(equal_rows[1], (0, 0, 0, 3))] + adapter._whole_euler_graph_wrapper = SimpleNamespace(enabled=True, ragged_body=object(), replay=lambda **_: None) + calls = {"n": 0} + original = adapter._decode_cfm + + def counting(*args, **kwargs): + calls["n"] += 1 + assert calls["n"] < 8, "row-offset fallback recursed instead of eager-solving" + return original(*args, **kwargs) + + adapter._decode_cfm = counting + previous_limit = sys.getrecursionlimit() + sys.setrecursionlimit(64) + try: + equal_x, _, equal_att = adapter._decode_cfm( + mu, speakers, cond, cnn_cache=None, att_cache=equal_rows, valid_lengths=[2, 2] + ) + mixed_x, _, mixed_att = adapter._decode_cfm( + mu, speakers, cond, cnn_cache=None, att_cache=mixed_rows, valid_lengths=[2, 2] + ) + finally: + sys.setrecursionlimit(previous_limit) + assert equal_x.shape[0] == mixed_x.shape[0] == 2 + assert len(equal_att) == len(mixed_att) == 2 + + +def test_tf32_mode_reads_shipped_yaml_key(monkeypatch): + from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import _tf32_mode + + monkeypatch.delenv("MINICPMO_CODE2WAV_TF32", raising=False) + assert _tf32_mode({}) == "tf32" + assert _tf32_mode({"token2wav_allow_tf32": False}) == "off" + assert _tf32_mode({"code2wav_allow_tf32": False}) == "off" + assert _tf32_mode({"token2wav_allow_tf32": True}) == "tf32" + assert _tf32_mode({"code2wav_allow_tf32": True}) == "tf32" + monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", "0") + assert _tf32_mode({}) == "off" + monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", "off") + assert _tf32_mode({}) == "off" + monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", "tf32") + assert _tf32_mode({}) == "tf32" + monkeypatch.setenv("MINICPMO_CODE2WAV_TF32", "tf32x3") + assert _tf32_mode({}) == "tf32" 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 7519b3b5f..e83cb3819 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 @@ -103,6 +103,7 @@ def _fake_wrapper(monkeypatch: pytest.MonkeyPatch) -> HiFTGraphWrapper: monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) wrapper = object.__new__(HiFTGraphWrapper) wrapper.capture_batch_sizes = [1] + wrapper._legit_shapes = {(7, 0), (9, 0)} wrapper.graph = {} wrapper.static_speech_inputs = {} wrapper.static_cache_source_inputs = {} @@ -880,7 +881,7 @@ def test_whole_euler_graph_with_padding_matches_eager( @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") -def test_whole_euler_cache_flushes_whole_generation( +def test_whole_euler_graph_boundary_enforces_budget_and_falls_back_to_eager( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(current_platform, "get_global_graph_pool", lambda: torch.cuda.graph_pool_handle()) @@ -903,7 +904,8 @@ def test_whole_euler_cache_flushes_whole_generation( att_cache=None, ) - _call(10) + res10 = _call(10) + assert res10 is not None assert len(wrapper._cache) == 1 assert wrapper._stats["captures"] == 1 @@ -912,46 +914,110 @@ def test_whole_euler_cache_flushes_whole_generation( assert wrapper._stats["hits"] == 1 # Second distinct shape fills cache to max_graphs=2 - _call(12) + res12 = _call(12) + assert res12 is not None 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 + # A third distinct shape exceeds max_graphs=2: falls back to eager (replay returns None) without expanding budget + res14 = _call(14) + assert res14 is None + assert wrapper._stats["flushes"] == 0 + assert len(wrapper._cache) == 2 + assert wrapper.max_graphs == 2 + assert wrapper._stats["captures"] == 2 - # 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(), + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_whole_euler_precapture_enforces_budget_boundary( + 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() + # Configure max_graphs=1 with query_bucket_frames=16 + wrapper = WholeEulerCFMGraphWrapper( + estimator=estimator, + n_timesteps=10, + max_graphs=1, + max_graph_batch=2, + micro_batch_size=2, + query_bucket_frames=16, + ) + + count = wrapper.precapture( + offsets=[0, 16, 32], + steady=32, + channels=4, + spk_dim=4, + ) + # Must capture at most 1 graph and strictly respect max_graphs=1 + assert count <= 1 + assert wrapper.max_graphs == 1 + assert len(wrapper._cache) <= 1 + + # An uncached shape falls back to eager (None) because max_graphs=1 is already exhausted + x_uncached = torch.randn(1, 4, 10, device="cuda") + mu_uncached = torch.randn(2, 4, 10, device="cuda") + spk_uncached = torch.randn(2, 4, device="cuda") + cond_uncached = torch.randn(2, 4, 10, device="cuda") + res_uncached = wrapper.replay( + x=x_uncached, + mu_cfg=mu_uncached, + speakers_cfg=spk_uncached, + cond_cfg=cond_uncached, 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, + assert res_uncached is None # Eager fallback + assert wrapper.max_graphs == 1 + assert len(wrapper._cache) == 1 + wrapper._flush() + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") +def test_whole_euler_slot_entry_enforces_budget_boundary(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, + att_slots=2, + ragged_body=Mock(), + modulation_fn=lambda t: t, + query_bucket_frames=16, ) - torch.testing.assert_close(res10[0], eager_x10, rtol=1e-4, atol=1e-5) + # `_record` returns (statics, x, cnn, att, graph); `_retire` calls entry[4].reset(). + monkeypatch.setattr(wrapper, "_record", lambda key, *args, **kwargs: (Mock(), Mock(), Mock(), Mock(), Mock())) + pool = wrapper._ensure_slot_pool((0, 4)) + assert pool is not None + + fill = wrapper._precapture_fill + # Capture 1 slot graph + entry1 = wrapper._slot_entry(graph_batch=1, query_cap=10, channels=4, spk_dim=4, fill=fill) + assert entry1 is not None + assert len(wrapper._slot_graphs) == 1 + assert wrapper.stats_snapshot()["cache_size"] == 1 + + # Capture 1 arena graph + x = torch.empty((1, 4, 1), device="cuda", dtype=torch.float32) + entry2 = wrapper._entry(graph_batch=1, query_cap=12, offset=0, x=x, spk_dim=4, fill=fill) + assert entry2 is not None + assert len(wrapper._cache) == 1 + assert wrapper.stats_snapshot()["cache_size"] == 2 + + # A 3rd graph (slot or arena) exceeds max_graphs=2: falls back to eager (returns None) + entry3_slot = wrapper._slot_entry(graph_batch=1, query_cap=14, channels=4, spk_dim=4, fill=fill) + assert entry3_slot is None + + entry3_arena = wrapper._entry(graph_batch=1, query_cap=16, offset=0, x=x, spk_dim=4, fill=fill) + assert entry3_arena is None + + assert wrapper.stats_snapshot()["cache_size"] == 2 wrapper._flush() @@ -2128,7 +2194,16 @@ def test_whole_euler_disabled_via_serving_config() -> None: 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") + for name in ("cosyvoice2.flow.decoder_dit", "stepaudio2.cosyvoice2.flow.decoder_dit"): + try: + import importlib + + decoder_dit = importlib.import_module(name) + break + except ImportError: + pass + else: + decoder_dit = pytest.importorskip("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(): @@ -2315,7 +2390,7 @@ def test_whole_euler_request_caches_grow_and_update_in_place(monkeypatch: pytest @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: +def test_whole_euler_groups_are_acquired_again_after_a_flush() -> 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] = [] @@ -2326,13 +2401,12 @@ def test_whole_euler_groups_are_acquired_again_after_a_flush(monkeypatch: pytest 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 wrapper._group_entries(entry, 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 + assert wrapper._group_entries(entry, groups, fills) is None @pytest.mark.parametrize("fail", [False, True]) diff --git a/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py b/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py new file mode 100644 index 000000000..e24da281e --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_dit_fused.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused Code2Wav DiT body matches the ragged reference.""" + +import pytest +import torch +import torch.nn as nn + +from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import BatchedToken2Wav +from vllm_omni.model_executor.models.minicpmo_4_5.dit_fused import blocks_forward_chunk_fused, supports_fused_body + +pytestmark = [pytest.mark.core_model, pytest.mark.cuda] + +_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else []) + + +def _upstream_dit(device: str) -> nn.Module: + for name in ("cosyvoice2.flow.decoder_dit", "stepaudio2.cosyvoice2.flow.decoder_dit"): + try: + import importlib + + decoder_dit = importlib.import_module(name) + break + except ImportError: + pass + else: + decoder_dit = pytest.importorskip("cosyvoice2.flow.decoder_dit") + torch.manual_seed(0) + estimator = decoder_dit.DiT(in_channels=320, out_channels=80, depth=2, num_heads=4, head_dim=16, hidden_size=64) + with torch.no_grad(): + for parameter in estimator.parameters(): + parameter.normal_(0.0, 0.08) + return estimator.eval().to(device) + + +def _run(body, estimator, *, rows: int, frames: int, cached: int, lengths): + torch.manual_seed(1) + depth = len(estimator.blocks) + channels = estimator.blocks[0].conv.block[1].in_channels + attn = estimator.blocks[0].attn + device = next(estimator.parameters()).device + estimator_input = torch.randn(rows, 320, frames, device=device) + time_embedding = estimator.t_embedder(torch.rand((), device=device).expand(rows)).unsqueeze(1) + mask = torch.rand(rows, frames, cached + frames, device=device) > 0.3 + mask[..., 0] = True + cnn = [torch.randn(rows, 2 * channels, 2, device=device) for _ in range(depth)] + att = [torch.randn(rows, attn.num_heads, cached, 2 * attn.head_dim, device=device) for _ in range(depth)] + cnn_out = torch.zeros((depth, rows, 2 * channels, 2), device=device) + att_out = torch.zeros((depth, rows, attn.num_heads, cached + frames, 2 * attn.head_dim), device=device) + with torch.no_grad(): + out = body(estimator, estimator_input, time_embedding, mask, cnn, att, cnn_out, att_out, lengths) + return out, cnn_out, att_out + + +@pytest.mark.parametrize("device", _DEVICES) +def test_fused_body_matches_ragged_body(device: str) -> None: + estimator = _upstream_dit(device) + assert supports_fused_body(estimator) + expected = _run( + BatchedToken2Wav._blocks_forward_chunk_ragged, estimator, rows=4, frames=6, cached=5, lengths=[6, 4] + ) + actual = _run(blocks_forward_chunk_fused, estimator, rows=4, frames=6, cached=5, lengths=[6, 4]) + for want, got in zip(expected, actual, strict=True): + torch.testing.assert_close(got, want, rtol=1e-4, atol=2e-4) + + +def test_unsupported_estimator_is_refused() -> None: + estimator = _upstream_dit("cpu") + estimator.blocks[0].attn.q_norm = nn.Identity() + assert not supports_fused_body(estimator) diff --git a/tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py b/tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py new file mode 100644 index 000000000..525df6c0b --- /dev/null +++ b/tests/model_executor/models/minicpmo_4_5/test_flow_encoder_graph.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exact-shape flow-encoder graphs (``cfm_encoder_cuda_graph``): keys, eager fallbacks, bit-exact replays.""" + +import functools +from types import SimpleNamespace + +import pytest +import torch + +import vllm_omni.model_executor.models.minicpmo_4_5.flow_encoder_graph as graph_module +from vllm_omni.model_executor.models.minicpmo_4_5.flow_encoder_graph import ( + FlowEncoderGraphs, + reachable_conformer_cache_frames, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +_DEPTH, _HEADS, _WIDTH, _HIDDEN = 2, 2, 4, 8 +_CNN = (3, 2) +_LOOKAHEAD = 1 + + +def _encode(tokens, *, cnn_cache, att_cache, weight): + """Row-independent stand-in for embedding + ``forward_chunk`` + projection: 1-token lookahead, 2x upsample.""" + rows = int(tokens.shape[0]) + x = (tokens.to(weight.dtype).unsqueeze(-1) @ weight[:1])[:, :-_LOOKAHEAD] + x = x.unsqueeze(2).expand(-1, -1, 2, -1).reshape(rows, -1, _HIDDEN) + x = x + cnn_cache.sum(dim=(1, 2))[:, None, None] + att_cache.mean(dim=(0, 2, 3, 4))[:, None, None] + hidden = x @ weight + new_cnn = hidden[:, -_CNN[1] :, : _CNN[0]].transpose(1, 2).contiguous() + added = hidden[:, :, :_WIDTH][None, :, None].expand(_DEPTH, rows, _HEADS, -1, _WIDTH) + return hidden, new_cnn, torch.cat((att_cache, added), dim=3) + + +@pytest.fixture +def fake_capture(monkeypatch: pytest.MonkeyPatch) -> None: + # A "graph" whose replay runs the recorded body on the static views. + record = lambda self, views: SimpleNamespace(replay=functools.partial(self._body, views)) # noqa: E731 + monkeypatch.setattr(FlowEncoderGraphs, "_record", record) + monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: False) + + +def _graphs(device: str, rows=(1, 2, 3), frames=(8, 18)): + weight = torch.randn(_HIDDEN, _HIDDEN, generator=torch.Generator().manual_seed(0)).to(device) + encode = functools.partial(_encode, weight=weight) + graphs = FlowEncoderGraphs(encode, rows=rows, token_widths=[6], lookahead=_LOOKAHEAD, upsample=2) + keys = [(r, 6, f) for r in rows for f in frames] + layout = {"cnn_shape": _CNN, "att_layout": (_DEPTH, _HEADS, _WIDTH), "hidden_dim": _HIDDEN} + assert graphs.capture(keys, **layout, dtype=torch.float32, device=torch.device(device)) == len(keys) + return graphs, encode + + +def _inputs(rows: int, width: int, frames: int, device: str, seed: int = 1): + generator = torch.Generator().manual_seed(seed) + tokens = torch.randint(0, 50, (rows, width), generator=generator) + cnn = [torch.randn(1, *_CNN, generator=generator) for _ in range(rows)] + att = [torch.randn(_DEPTH, 1, _HEADS, frames, _WIDTH, generator=generator) for _ in range(rows)] + return tokens.to(device), [c.to(device) for c in cnn], [a.to(device) for a in att] + + +@pytest.mark.parametrize( + ("growths", "expected"), + [([50], [300, 350, 400]), ([150], [300, 400]), ([30], [300, 330, 360, 390, 400]), ([], [300])], +) +def test_reachable_cache_frames_follow_the_streaming_trim(growths, expected) -> None: + # 6 s default voice (300 frames); 28-token duplex units grow the cache by 50 frames. + assert reachable_conformer_cache_frames(start=300, prompt_len=300, suffix=100, growths=growths) == expected + + +def test_keys_cover_every_row_count_width_and_cache_length() -> None: + graphs = FlowEncoderGraphs(lambda *a, **k: a, rows=[2, 1, 2], token_widths=[3, 28], lookahead=3, upsample=2) + assert graphs.token_widths == (28,) # a width with no new frames is dropped + assert graphs.keys_for(start=300, prompt_len=300, suffix=100) == [ + (rows, 28, frames) for rows in (1, 2) for frames in (300, 350, 400) + ] + + +def _check_replay(graphs, encode, rows: int, frames: int, device: str) -> None: + tokens, cnn, att = _inputs(rows, 6, frames, device, seed=rows * 100 + frames) + graphed = graphs.run(tokens, cnn, att) + expected = encode(tokens, cnn_cache=torch.cat(cnn, dim=0), att_cache=torch.cat(att, dim=1)) + assert graphed is not None + for actual, reference in zip(graphed, expected, strict=True): + assert torch.equal(actual, reference) + + +@pytest.mark.parametrize(("rows", "frames"), [(1, 8), (3, 18), (3, 8)]) +def test_replay_matches_the_eager_encoder_exactly(fake_capture, rows, frames) -> None: + graphs, encode = _graphs("cpu") + _check_replay(graphs, encode, rows, frames, "cpu") + assert graphs.replays == 1 + + +def test_other_shapes_and_layouts_run_eager(fake_capture, monkeypatch: pytest.MonkeyPatch) -> None: + graphs, _ = _graphs("cpu") + for rows, width, frames in ((4, 6, 8), (1, 5, 8), (1, 6, 10)): # never padded, never another shape + assert graphs.run(*_inputs(rows, width, frames, "cpu")) is None + tokens, cnn, att = _inputs(2, 6, 8, "cpu") + assert graphs.run(tokens, cnn, [att[0], att[1].double()]) is None + assert graphs.run(tokens.int(), cnn, att) is None + assert graphs.run(tokens[:1], cnn, att) is None + with monkeypatch.context() as patch: + patch.setattr(graph_module, "_precision_state", lambda: ("other precision flags",)) + assert graphs.run(tokens, cnn, att) is None + assert graphs.run(tokens, cnn, att) is not None + assert graphs.replays == 1 + + +@pytest.mark.cuda +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for graph capture") +def test_captured_graphs_match_eager_bit_exactly_on_cuda() -> None: + graphs, encode = _graphs("cuda") + for rows, frames in ((1, 8), (3, 18), (1, 18), (3, 8)): + _check_replay(graphs, encode, rows, frames, "cuda") + assert graphs.replays == 4 diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index 965e8d340..0e221317f 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/vllm_omni/config/environment_variable_inventory.py @@ -172,6 +172,7 @@ _MODEL_PROMOTE = ( "MIMO_AUDIO_TOKENIZER_DEVICE", "MIMO_AUDIO_TOKENIZER_PATH", "MING_CFM_CUDAGRAPH", + "MINICPMO_CODE2WAV_TF32", "MINICPMO_TOKEN2WAV_TRT", "MINICPMO_TOKEN2WAV_TRT_CACHE", "MINICPMO_TOKEN2WAV_TRT_DTYPE", diff --git a/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml b/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml index c743ee260..dbabc9172 100644 --- a/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_duplex_h200.yaml @@ -1,7 +1,8 @@ # Opt-in single-H200 native full-duplex codec profile (V1). # Inherits the native session/stage capacities, sampling and chunk sizes. # Uses the mainline Whole-Euler Flow and shared prompt-state implementation. -# TF32 changes FP32 GEMM rounding; validate audio quality for your workload. +# Ordinary TF32 for CFM DiT GEMMs (cuBLAS allow_tf32 + Triton tf32 dots). +# HiFT stays IEEE FP32. Validate audio quality for your workload. base_config: minicpmo_4_5.yaml connectors: diff --git a/vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml b/vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml index a62dcf989..d0460108d 100644 --- a/vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml +++ b/vllm_omni/deploy/minicpmo_4_5_turn_mrv2.yaml @@ -14,9 +14,11 @@ connectors: connector_of_shared_memory: extra: hift_graph_capture_batch_sizes: [1, 2, 4, 8] - # TF32 tensor cores for the vocoder's fp32 flow GEMMs: ~25% less GPU - # time per chunk (batch 1-8). A numerics change: the HiFT input mel moves - # by ~3e-4 relative L2; dev-short WER/UTMOS unchanged (64 prompts). + # Ordinary TF32 for CFM DiT GEMMs (~25% less GPU time per chunk at + # batch 1-8): cuBLAS allow_tf32 on dense QKV/MLP and Triton + # input_precision=tf32 on tiled attention. HiFT stays IEEE FP32. + # Numerics: HiFT input mel moves by ~3e-4 relative L2; dev-short + # WER/UTMOS unchanged. token2wav_allow_tf32: true stages: 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 59998217d..fc5ad1edf 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 @@ -4,6 +4,8 @@ from __future__ import annotations +import functools +import math from collections import OrderedDict from collections.abc import Mapping from contextlib import nullcontext @@ -15,16 +17,32 @@ import torch.nn as nn import torch.nn.functional as F from vllm.logger import init_logger +from vllm_omni.utils.device_copy import index_to_device + from .cuda_graph_wrapper import ( CFMGraphWrapper, HiFTGraphWrapper, + ResidentAttCache, WholeEulerCFMGraphWrapper, _att_keep_ranges, _copy_frame_segments, _euler_step, _euler_timeline, + _format_memory_delta, + _materialize_att_rows, + _memory_snapshot, _zero_padded_cnn_cache, + empty_hift_outputs, ) +from .dit_fused import ( + blocks_forward_chunk_fused, + causal_cache_index, + dit_modulation, + gather_causal_cache, + supports_fused_body, + tiled_attention_supported, +) +from .flow_encoder_graph import FlowEncoderGraphs logger = init_logger(__name__) @@ -37,10 +55,20 @@ _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 +_TIMESTEP_MAX_PERIOD = 10000 + + +def _timestep_frequencies(dim: int) -> torch.Tensor: + """Host frequency table of ``TimestepEmbedder.timestep_embedding``.""" + half = dim // 2 + return torch.exp(-math.log(_TIMESTEP_MAX_PERIOD) * torch.arange(start=0, end=half) / half) def _trim_streaming_cache(cache: torch.Tensor, prompt_len: int) -> torch.Tensor: """The streaming trim of a cache whose frame axis is ``-2``.""" + if isinstance(cache, ResidentAttCache): + # Its slot-pool solve already trimmed it (``WholeEulerCFMGraphWrapper._replay_slots``). + return cache ranges = _att_keep_ranges(int(cache.shape[-2]), (prompt_len, _CACHE_TRIM_SUFFIX)) if len(ranges) == 1: return cache @@ -207,6 +235,21 @@ def state_shape_signature(state: BatchedToken2WavState) -> tuple[Any, ...]: return flow, hift +# Cache frame axes a row-offset merge keeps per row: the encoder runs once per conformer cache +# length (``decode_ragged_batch``), and a slot-pool solve attends each estimator cache in its slot. +_ROW_OFFSET_FRAME_AXES = {"conformer_att_cache": 3, "estimator_att_cache": 4} + + +def row_offset_signature(state: BatchedToken2WavState) -> tuple[Any, ...]: + """``state_shape_signature`` without the cache lengths a row-offset merge keeps per row.""" + flow, hift = state_shape_signature(state) + masked = [] + for name, (shape, dtype, device) in flow: + axis = _ROW_OFFSET_FRAME_AXES.get(name) + masked.append((name, (shape if axis is None else (*shape[:axis], -1, *shape[axis + 1 :]), dtype, device))) + return tuple(masked), hift + + @dataclass(frozen=True) class PromptFeatures: cache_key: tuple[str, str] @@ -240,6 +283,23 @@ def _undecorate_dynamo(module: nn.Module, method: str) -> None: logger.info("Bypassed TorchDynamo wrapper on %s.%s", type(module).__name__, method) +def _cfm_matmul_tf32(fn): + """Run the flow CFM (and capture its graphs) with ordinary TF32 GEMMs when ``cfm_tf32``; HiFT stays fp32.""" + + @functools.wraps(fn) + def wrapper(self, *args, **kwargs): + if not getattr(self, "_cfm_tf32", False): + return fn(self, *args, **kwargs) + previous = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = True + try: + return fn(self, *args, **kwargs) + finally: + torch.backends.cuda.matmul.allow_tf32 = previous + + return wrapper + + class BatchedToken2Wav(nn.Module): """Drive Token2wav's modules with dynamically-sized, request-owned caches. @@ -258,11 +318,18 @@ class BatchedToken2Wav(nn.Module): cfm_graph_config: Mapping[str, Any] | None = None, bfloat16_attention_cache: bool = False, setup_cache_size: int = 1, + cfm_tf32: bool = False, + encoder_graph_config: Mapping[str, Any] | None = None, ): super().__init__() if setup_cache_size < 0: raise ValueError("setup_cache_size must be >= 0") self._token2wav = token2wav + self._cfm_tf32 = bool(cfm_tf32) + cfm_graph_cfg = dict(cfm_graph_config or {}) + # Ragged solves of rows whose caches differ in length (a stream's second chunk next to steady + # ones): each row keeps its own cache offset, in one slot-pool Whole-Euler replay when it can. + self._row_offset_merge = bool(cfm_graph_cfg.get("row_offset_merge", False)) # Optional TrtDiTStepper (step_audio2_dit_trt): replaces only the # per-timestep DiT estimator call; encoder and HiFT stay on torch. self._trt_stepper = trt_stepper @@ -274,6 +341,7 @@ class BatchedToken2Wav(nn.Module): encoder = getattr(self.flow, "encoder", None) if encoder is not None: _undecorate_dynamo(encoder, "forward_chunk") + self._encoder_graphs = self._build_encoder_graphs(encoder_graph_config or {}, connector_config or {}) hift_parameter = next(self.hift.parameters(), None) if hift_parameter is not None and hift_parameter.device.type == "cuda": # Prime the CUDA state used by HiFT during backend construction. @@ -336,13 +404,16 @@ class BatchedToken2Wav(nn.Module): 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") + # Captured by the first forward (``precapture_hift``): graphs + # captured while vLLM loads the weights hold far more memory. 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)): + # Whether the ragged DiT body is the fused one (``dit_fused.py``). + self._ragged_fused_body = False + # On NPU the platform graph runner captures instead of the CUDA wrappers, so bucketing and + # padding key off the requested flag. + self._cfm_graph_enabled = bool(cfm_graph_cfg.get("enabled", False)) + if self._cfm_graph_enabled: flow_parameter = next(self.flow.parameters(), None) if flow_parameter is not None and flow_parameter.device.type == "cuda": estimator = self.flow.decoder.estimator @@ -352,7 +423,26 @@ class BatchedToken2Wav(nn.Module): 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: + enable_whole_euler = bool(cfm_graph_cfg.get("enable_whole_euler", True)) + if bool(cfm_graph_cfg.get("fused_body", False)): + # A reduced-precision attention cache keeps the ragged body, + # which attends in the activation dtype through upstream. + self._ragged_fused_body = ( + supports_fused_body(estimator) and self._estimator_att_cache_dtype == torch.float32 + ) + if not self._ragged_fused_body: + logger.warning( + "cfm_fused_body requested but unsupported here (DiT blocks or a non-fp32 attention cache)" + ) + # Resident request caches need the fused body and its tiled + # attention, which reads keys/values out of the pool in place. + slot_pool = bool(cfm_graph_cfg.get("slot_pool", False)) + if slot_pool and not ( + self._ragged_fused_body and tiled_attention_supported(flow_parameter.get_device()) + ): + logger.warning("cfm_slot_pool requested but needs cfm_fused_body on NVIDIA SM80+; using the arena") + slot_pool = False + if enable_whole_euler 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 @@ -361,7 +451,12 @@ class BatchedToken2Wav(nn.Module): 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) + if isinstance(capture_frames, (list, tuple)): + # Several widths (e.g. a 25-token duplex unit and a + # 75-token turn chunk) each get their own graphs. + query_bucket_frames: int | tuple[int, ...] = tuple(int(w) for w in capture_frames) + else: + 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, @@ -373,24 +468,30 @@ class BatchedToken2Wav(nn.Module): 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 - ), + offset_bucket_frames=int(cfm_graph_cfg.get("offset_bucket_frames", 0)), + # Resident request caches, a slot per stream plus a few for streams + # that end while new ones start (``AttSlotPool``). + att_slots=(micro_batch_size + max(2, micro_batch_size // 8)) if slot_pool else 0, + # ``_ragged_body`` picks the configured body on every call, so a + # replaced body is the one a capture records. + ragged_body=self._ragged_body if _supports_ragged_kernel(estimator) else None, + # The fused body takes each timestep's modulation, computed once outside the graphs. + modulation_fn=functools.partial(dit_modulation, estimator) if self._ragged_fused_body else None, + row_offsets=self._row_offset_merge, ) 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)", + "micro_batch_size=%d, query_bucket_frames=%s, offset_bucket_frames=%d, fused_body=%s)", max_graphs, max_serial_batch, str(max_graph_batch), micro_batch_size, query_bucket_frames, + self._whole_euler_graph_wrapper.offset_bucket_frames, + self._ragged_fused_body, ) - elif self._trt_stepper is not None and bool(cfm_graph_cfg.get("enable_whole_euler", True)): + elif enable_whole_euler: logger.info("Whole-Euler CFM CUDA Graph disabled because TensorRT stepper is configured") self._cfm_graph_wrapper = CFMGraphWrapper( graph_fn=estimator.blocks_forward_chunk, @@ -407,7 +508,12 @@ class BatchedToken2Wav(nn.Module): # 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 or self._whole_euler_graph_wrapper is not None) + if ( + self._cfm_graph_wrapper is not None + or self._whole_euler_graph_wrapper is not None + # NPU has no wrappers; the platform runner is keyed by this flag. + or self._cfm_graph_enabled + ) else 0 ) if self._cfm_graph_bucket_frames > 1: @@ -415,6 +521,8 @@ class BatchedToken2Wav(nn.Module): "CFM CUDA Graph bucketing enabled (bucket_frames=%d)", self._cfm_graph_bucket_frames, ) + # Device copies of the DiT timestep-embedding frequency table. + self._timestep_freqs: dict[tuple[Any, ...], torch.Tensor] = {} self._prompt_features: dict[tuple[str, str], PromptFeatures] = {} self._setup_cache_size = setup_cache_size self._setup_cache: OrderedDict[tuple[tuple[str, str], int], BatchedToken2WavState] = OrderedDict() @@ -424,6 +532,9 @@ class BatchedToken2Wav(nn.Module): mel: torch.Tensor, source_cache: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: + if int(mel.shape[2]) == 0: + # No frames, no samples: HiFT itself cannot run on an empty mel. + return empty_hift_outputs(mel) if self.hift_graph_wrapper is None: return self.hift.inference(mel, source_cache) return self.hift_graph_wrapper.replay(mel, source_cache) @@ -467,6 +578,17 @@ class BatchedToken2Wav(nn.Module): features.mels.expand(batch_size, -1, -1), ) + def _project_speakers(self, speakers: torch.Tensor) -> torch.Tensor: + """``spk_embed_affine_layer(normalize(speakers))`` with M=1, independent of the batch size. + + Rows expand one prompt speaker, so one row is projected: M=1 and M>=2 + GEMMs round differently (~6e-8), moving a row's CFM mel by ~1e-3. + """ + if speakers.shape[0] > 1 and speakers.stride(0) == 0: + projected = self.flow.spk_embed_affine_layer(F.normalize(speakers[:1], dim=1)) + return projected.expand(speakers.shape[0], -1).contiguous() + return self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + def _autocast(self, device: torch.device): if device.type != "cuda": return nullcontext() @@ -538,6 +660,96 @@ class BatchedToken2Wav(nn.Module): ) return self.flow.encoder_proj(hidden), new_cnn, new_att + def _flow_on_cuda(self) -> bool: + flow_parameter = next(self.flow.parameters(), None) + return flow_parameter is not None and flow_parameter.device.type == "cuda" + + def _build_encoder_graphs(self, config: Mapping[str, Any], connector: Mapping) -> FlowEncoderGraphs | None: + """``cfm_encoder_cuda_graph`` (default off): exact-shape graphs of the flow encoder's continuation chunk.""" + if not config.get("enabled", False): + return None + if not self._flow_on_cuda(): + logger.info("Flow encoder CUDA Graph is disabled off CUDA") + return None + encoder = self.flow.encoder + embed = getattr(encoder, "embed", None) + lookahead = self._pre_lookahead_len() + if ( + lookahead is None + or getattr(encoder, "up_layer", None) is None + or not isinstance(getattr(getattr(embed, "pos_enc", None), "pe", None), torch.Tensor) + # ``_ensure_relpos_pe`` would grow such a table between calls. + or callable(getattr(embed, "extend_pe", None)) + ): + logger.warning("Flow encoder CUDA Graph needs the CosyVoice2 upsample-conformer encoder; staying eager") + return None + # Every ``position_encoding`` slice a graph reads lives in these tables. + tables = [ + module + for module in (embed, getattr(encoder, "up_embed", None)) + if isinstance(getattr(getattr(module, "pos_enc", None), "pe", None), torch.Tensor) + ] + # Default width: the steady duplex chunk, the left context plus one unit of new frames. + unit = int(connector.get("initial_codec_chunk_frames", 0) or 0) or int(connector.get("codec_chunk_frames", 25)) + graphs = FlowEncoderGraphs( + # With no ``extend_pe`` (checked above) this runs exactly the eager continuation's ops. + functools.partial(self._encode_chunk, last_chunk=False), + rows=config.get("rows") or [1], + token_widths=config.get("token_widths") or [int(connector.get("codec_left_context_frames", 3)) + unit], + lookahead=lookahead, + upsample=self._upsample_stride(), + held_tensors=lambda: tuple(module.pos_enc.pe for module in tables), + ) + logger.info( + "Flow encoder CUDA Graph enabled for continuation chunks (rows %s, tokens %s)", + list(graphs.rows), + list(graphs.token_widths), + ) + return graphs + + def _graph_encode( + self, tokens: torch.Tensor, states: list[BatchedToken2WavState], copy: bool = False + ) -> tuple[torch.Tensor, ...] | None: + """A continuation chunk's encoder outputs from its graph: views rewritten by the next replay, or copies.""" + if self._encoder_graphs is None: + return None + flows = [state.flow_cache for state in states] + outputs = self._encoder_graphs.run( + tokens, [flow["conformer_cnn_cache"] for flow in flows], [flow["conformer_att_cache"] for flow in flows] + ) + return tuple(value.clone() for value in outputs) if outputs and copy else outputs + + @torch.inference_mode() + def precapture_flow_encoder(self, features: PromptFeatures) -> int: + """Capture the encoder graphs of every continuation chunk this prompt's streams run; returns the count.""" + graphs = self._encoder_graphs + if graphs is None or graphs.graphs: + return 0 + (state,) = self.setup_batch(features, 1) + cnn, att = state.flow_cache["conformer_cnn_cache"], state.flow_cache["conformer_att_cache"] + prompt_len = int(features.mels.shape[1]) + keys = graphs.keys_for(start=int(att.shape[3]), prompt_len=prompt_len, suffix=_CACHE_TRIM_SUFFIX) + torch.accelerator.synchronize(att.device) + memory_before = _memory_snapshot(att.device) + with self._autocast(att.device): + captured = graphs.capture( + keys, + cnn_shape=(int(cnn.shape[1]), int(cnn.shape[2])), + att_layout=(int(att.shape[0]), int(att.shape[2]), int(att.shape[4])), + hidden_dim=int(self.flow.encoder_proj.out_features), + dtype=att.dtype, + device=att.device, + ) + torch.accelerator.synchronize(att.device) + logger.info( + "Captured %d flow encoder CUDA Graphs (cache frames %s): shared buffers %.1f MiB%s", + captured, + sorted({key[2] for key in keys}), + graphs.storage_bytes() / 2**20, + _format_memory_delta(memory_before, _memory_snapshot(att.device)), + ) + return captured + @staticmethod def _estimator_buffers( estimator: nn.Module, @@ -558,6 +770,28 @@ class BatchedToken2Wav(nn.Module): att = x.new_empty((depth, batch_size, heads, old_att_len + chunk_size, att_width)) return cnn, att + def _time_embedding(self, estimator: nn.Module, time: torch.Tensor) -> torch.Tensor: + """``estimator.t_embedder(time)`` without a host sync per call (bitwise identical). + + Upstream rebuilds its frequency table on the host and moves it with a + pageable ``.to(t)`` once per CFM step, blocking the host until the GPU + drains; the table is built the same way once and kept on the device. + """ + embedder = estimator.t_embedder + dim, scale, mlp = (getattr(embedder, name, None) for name in ("frequency_embedding_size", "scale", "mlp")) + if not isinstance(dim, int) or scale is None or mlp is None: + return embedder(time) + # The host table takes the default dtype before ``.to(t)``, as upstream. + key = (time.device, time.dtype, torch.get_default_dtype(), dim) + freqs = self._timestep_freqs.get(key) + if freqs is None: + freqs = self._timestep_freqs[key] = _timestep_frequencies(dim).to(time) + args = (time * scale)[:, None] * freqs[None] + embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + if dim % 2: + embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1) + return mlp(embedding) + def _estimator_step( self, estimator: nn.Module, @@ -585,8 +819,7 @@ class BatchedToken2Wav(nn.Module): att_cache=att_cache, ) return out.to(mu.dtype), new_cnn, new_att - if time_embedding is None: - time_embedding = estimator.t_embedder(time).unsqueeze(1) + time_embedding = self._time_embedding(estimator, 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: @@ -621,7 +854,7 @@ class BatchedToken2Wav(nn.Module): if valid_lengths is not None: if not hasattr(estimator, "in_proj"): raise RuntimeError('MiniCPMO45Code2WavBatchError {"reason":"ragged_kernel_unavailable"}') - result = self._blocks_forward_chunk_ragged( + result = self._ragged_body( estimator, estimator_input, time_embedding, @@ -644,15 +877,21 @@ class BatchedToken2Wav(nn.Module): ) return result, cnn_out, att_out - @staticmethod - def _gather_causal_cache( - history: torch.Tensor, - valid_lengths: torch.Tensor, - width: int, + def _ragged_body( + self, + *args: Any, + modulation: torch.Tensor | None = None, + slots: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> 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(1, indices[:, :, None].expand(-1, -1, int(history.shape[2]))).transpose(1, 2) + """The configured ragged DiT body, looked up per call so a capture records a replaced one. + + ``modulation`` (fused body only) comes from the Whole-Euler arena. + """ + estimator_input = args[1] + if slots is not None or (self._ragged_fused_body and estimator_input.dtype == torch.float32): + # Slot-pool graphs are only built over the fused body (``slot_pool`` above). + return blocks_forward_chunk_fused(*args, modulation=modulation, slots=slots) + return self._blocks_forward_chunk_ragged(*args) @staticmethod def _causal_conv_frames_major(conv: nn.Conv1d, history: torch.Tensor) -> torch.Tensor: @@ -765,6 +1004,8 @@ class BatchedToken2Wav(nn.Module): device=estimator_input.device, dtype=torch.long, ) + # Causal cache gather index per conv width, shared by the blocks. + cache_index: dict[int, torch.Tensor] = {} x = estimator.in_proj(estimator_input.transpose(1, 2)) for block_index, block in enumerate(estimator.blocks): ( @@ -803,11 +1044,13 @@ class BatchedToken2Wav(nn.Module): ) ) width = int(old_cnn.shape[2]) + if width not in cache_index: + cache_index[width] = causal_cache_index(lengths, width) history = torch.cat((old_cnn[:, : conv.in_channels].transpose(1, 2), conv_input), dim=1) - new_cnn1 = BatchedToken2Wav._gather_causal_cache(history, lengths, width) + new_cnn1 = gather_causal_cache(history, cache_index[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) + new_cnn2 = gather_causal_cache(history, cache_index[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) @@ -829,6 +1072,7 @@ class BatchedToken2Wav(nn.Module): self._timeline_cache[key] = cached return cached + @_cfm_matmul_tf32 def _decode_cfm( self, mu: torch.Tensor, @@ -859,7 +1103,18 @@ class BatchedToken2Wav(nn.Module): offset = int(att_rows[0].shape[4]) else: offset = int(att_cache.shape[4]) if att_cache is not None else 0 + # ``row_offset_merge``: ragged rows whose caches differ in length each + # start from the noise at their own offset and attend their own cache. + row_offsets = None + if att_rows is not None and valid_lengths is not None: + lengths = [int(row.shape[4]) for row in att_rows] + if len(set(lengths)) > 1 or self._row_offset_merge: + row_offsets = lengths + if len(set(row_offsets)) == 1 and not self._row_offset_merge: + row_offsets = None mel_frames = int(mu.shape[2]) + # Padding only pays off while replay is active: Whole-Euler, the CFM wrapper (still enabled, + # no TRT stepper), or on NPU the platform runner keyed by `_cfm_graph_enabled`. graphs_active = valid_lengths is None and ( self._whole_euler_active() or ( @@ -867,6 +1122,7 @@ class BatchedToken2Wav(nn.Module): and getattr(self._cfm_graph_wrapper, "enabled", True) and self._trt_stepper is None ) + or (self._cfm_graph_wrapper is None and self._cfm_graph_enabled) ) pad_frames = _cfm_pad_frames( mel_frames=mel_frames, @@ -881,14 +1137,18 @@ class BatchedToken2Wav(nn.Module): # carry less of a step change into the attention and CNN caches. mu = torch.nn.functional.pad(mu, (0, pad_frames), mode="replicate") cond = torch.nn.functional.pad(cond, (0, pad_frames), mode="replicate") - end = offset + int(mu.shape[2]) + end = (offset if row_offsets is None else max(row_offsets)) + int(mu.shape[2]) if end > int(decoder.rand_noise.shape[2]): raise RuntimeError( "MiniCPMO45Code2WavBatchError " f'{{"reason":"noise_capacity","required":{end},' f'"available":{int(decoder.rand_noise.shape[2])}}}' ) - x = decoder.rand_noise[:, :, offset:end].expand(batch_size, -1, -1).clone() + if row_offsets is None: + x = decoder.rand_noise[:, :, offset:end].expand(batch_size, -1, -1).clone() + else: + width = int(mu.shape[2]) + x = torch.cat([decoder.rand_noise[:, :, start : start + width] for start in row_offsets], dim=0) if pad_frames: # The padded columns would otherwise carry real noise values. They # are excluded from this chunk's attention by ``attn_mask`` below @@ -904,18 +1164,21 @@ class BatchedToken2Wav(nn.Module): if valid_lengths is not None: if len(valid_lengths) != batch_size: raise ValueError(f"valid length count {len(valid_lengths)} != batch {batch_size}") - cfg_lengths = torch.tensor( - (*valid_lengths, *valid_lengths), - device=mu.device, - ) + cfg_lengths = index_to_device((*valid_lengths, *valid_lengths), mu.device) positions = torch.arange(int(mu.shape[2]), device=mu.device) valid_queries = positions.unsqueeze(0) < cfg_lengths.unsqueeze(1) current_keys = valid_queries.unsqueeze(1).expand(-1, int(mu.shape[2]), -1) - old_keys = torch.ones( - (2 * batch_size, int(mu.shape[2]), offset), - dtype=torch.bool, - device=mu.device, - ) + if row_offsets is None: + old_keys = torch.ones( + (2 * batch_size, int(mu.shape[2]), offset), + dtype=torch.bool, + device=mu.device, + ) + else: + # Cache columns span the longest cache; a shorter one's tail stays masked. + cfg_offsets = index_to_device((*row_offsets, *row_offsets), mu.device) + cached = torch.arange(max(row_offsets), device=mu.device).unsqueeze(0) < cfg_offsets.unsqueeze(1) + old_keys = cached.unsqueeze(1).expand(-1, int(mu.shape[2]), -1) attn_mask = valid_queries.unsqueeze(2) & torch.cat((current_keys, old_keys), dim=2) elif pad_frames: # Mask the padded keys instead of only zeroing their content: a @@ -948,15 +1211,26 @@ class BatchedToken2Wav(nn.Module): ) if whole_euler_result is not None: return whole_euler_result + if row_offsets is not None and len(set(row_offsets)) > 1: + # Mixed cache lengths only. A same-length group (including every + # ``_decode_cfm_per_offset`` subgroup, and equal-length rows under + # ``row_offset_merge``) must fall through to eager: otherwise a + # declined Whole-Euler replay re-enters this fallback forever. + return self._decode_cfm_per_offset( + mu, + speakers, + cond, + cnn_cache=cnn_cache, + att_rows=att_rows, + valid_lengths=valid_lengths, + att_keep=att_keep, + ) 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 - 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 @@ -974,7 +1248,6 @@ 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) @@ -1032,6 +1305,48 @@ class BatchedToken2Wav(nn.Module): assert next_att_cache is not None return x, torch.stack(next_cnn), next_att_cache + def _decode_cfm_per_offset( + self, + mu: torch.Tensor, + speakers: torch.Tensor, + cond: torch.Tensor, + *, + cnn_cache: torch.Tensor | None, + att_rows: list, + valid_lengths: list[int], + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor, list]: + """``_decode_cfm`` once per cache length, in row order: a row-offset merge no single solve took.""" + batch_size = len(att_rows) + groups: dict[int, list[int]] = {} + for row, cache in enumerate(att_rows): + groups.setdefault(int(cache.shape[4]), []).append(row) + x: torch.Tensor | None = None + cnn: torch.Tensor | None = None + att: list = [None] * batch_size + for rows in groups.values(): + index = index_to_device(rows, mu.device) + # CFG rows are ``[cond x B | uncond x B]``. + cfg_index = index_to_device([*rows, *(batch_size + row for row in rows)], mu.device) + group_x, group_cnn, group_att = self._decode_cfm( + mu.index_select(0, index), + speakers.index_select(0, index), + cond.index_select(0, index), + cnn_cache=None if cnn_cache is None else cnn_cache.index_select(2, cfg_index), + att_cache=[att_rows[row] for row in rows], + valid_lengths=[valid_lengths[row] for row in rows], + **kwargs, + ) + if x is None or cnn is None: + x = group_x.new_empty((batch_size, *group_x.shape[1:])) + cnn = group_cnn.new_empty((*group_cnn.shape[:2], 2 * batch_size, *group_cnn.shape[3:])) + x.index_copy_(0, index, group_x) + cnn.index_copy_(2, cfg_index, group_cnn) + for row, cache in zip(rows, group_att, strict=True): + att[row] = cache + assert x is not None and cnn is not None + return x, cnn, att + def _split_flow_cache( self, cache: dict[str, Any], @@ -1085,6 +1400,7 @@ class BatchedToken2Wav(nn.Module): @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]``.""" + rows = _materialize_att_rows(rows) return torch.cat((*[row[:, :, 0:1] for row in rows], *[row[:, :, 1:2] for row in rows]), dim=2) def _stack_flow_cache( @@ -1092,19 +1408,26 @@ class BatchedToken2Wav(nn.Module): states: list[BatchedToken2WavState], *, include_estimator_att: bool = True, + include_conformer: 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``). + skips stacking it (``include_estimator_att=False``). A caller that + already ran the encoder skips the conformer caches + (``include_conformer=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] 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), + "conformer_cnn_cache": ( + torch.cat([flow["conformer_cnn_cache"] for flow in flows], dim=0) if include_conformer else None + ), + "conformer_att_cache": ( + torch.cat([flow["conformer_att_cache"] for flow in flows], dim=1) if include_conformer else None + ), "estimator_cnn_cache": torch.cat((*conditional_cnn, *unconditional_cnn), dim=2), "estimator_att_cache": ( self._stack_estimator_att([flow["estimator_att_cache"] for flow in flows]) @@ -1121,6 +1444,69 @@ class BatchedToken2Wav(nn.Module): """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 _ragged_merge_fits(self, states: list[BatchedToken2WavState], token_counts: list[int]) -> bool: + """Ragged Whole-Euler is on, the rows share a CNN cache shape, and each row's tokens encode at once.""" + return ( + 2 <= len(states) == len(token_counts) + and self._whole_euler_ragged_active() + and len({tuple(state.flow_cache["estimator_cnn_cache"].shape) for state in states}) == 1 + and all(count <= self._max_encode_token_frames([s]) for count, s in zip(token_counts, states, strict=True)) + ) + + def can_merge_state_shapes(self, states: list[BatchedToken2WavState], token_counts: list[int]) -> bool: + """Whether mixed onset/continuation rows fit ragged Whole-Euler.""" + caches = [state.flow_cache.get("estimator_att_cache") for state in states] + return ( + all(isinstance(cache, torch.Tensor) for cache in caches) + and len({(tuple(cache.shape[:4]), cache.dtype, cache.device) for cache in caches}) == 1 + and self._ragged_merge_fits(states, token_counts) + ) + + def can_merge_row_offsets(self, states: list[BatchedToken2WavState], token_counts: list[int]) -> bool: + """Whether continuation rows whose caches differ only in length fit one slot-pool ragged solve. + + ``row_offset_merge``: every estimator cache must be resident in the + Whole-Euler slot pool, where each row attends its own frames; + ``decode_ragged_batch`` runs the encoder once per conformer cache length. + """ + if not (self._row_offset_merge and self._ragged_merge_fits(states, token_counts)): + return False + pool = self._whole_euler_graph_wrapper.slot_pool + caches = [state.flow_cache.get("estimator_att_cache") for state in states] + return pool is not None and all(isinstance(cache, ResidentAttCache) and cache.pool is pool for cache in caches) + + @torch.inference_mode() + def precapture_hift(self) -> int: + """Capture every configured HiFT graph (prompt-independent) once, before serving; returns the count.""" + wrapper = self.hift_graph_wrapper + if wrapper is None or wrapper.graph: + return 0 + with _autocast_disabled(wrapper.device): + wrapper.capture() + return len(wrapper.graph) + + @torch.inference_mode() + @_cfm_matmul_tf32 + def precapture_whole_euler(self, features: PromptFeatures) -> int: + """Capture the Whole-Euler graphs of every cache length this prompt's streams pass through. + + Runs in ``_decode_cfm``'s precision (``_cfm_matmul_tf32``, ``_autocast``): + a graph replays the kernels it was captured with. Returns the count. + """ + if not self._whole_euler_active(): + return 0 + wrapper = self._whole_euler_graph_wrapper + prompt_len = int(features.mels.shape[1]) + steady = prompt_len + _CACHE_TRIM_SUFFIX + with self._autocast(wrapper.device): + return wrapper.precapture( + offsets=range(prompt_len, steady + 1), + steady=steady, + channels=int(features.mels.shape[2]), + spk_dim=int(self.flow.spk_embed_affine_layer.out_features), + keep=(prompt_len, _CACHE_TRIM_SUFFIX), + ) + def _create_initial_states( self, features: PromptFeatures, @@ -1139,7 +1525,7 @@ class BatchedToken2Wav(nn.Module): cnn_cache=None, att_cache=None, ) - projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + projected_speakers = self._project_speakers(speakers) # 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( @@ -1301,18 +1687,23 @@ class BatchedToken2Wav(nn.Module): # 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) + # A captured encoder graph stacks the conformer caches into its own inputs. + graphed = None if last_chunk or flush_encoder else self._graph_encode(tokens, states) + flow_cache = self._stack_flow_cache( + states, include_estimator_att=not per_request_att, include_conformer=graphed is None + ) 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( + # A graph's shared results are read below and copied per row by _split_flow_cache. + hidden, conformer_cnn, conformer_att = graphed or self._encode_chunk( tokens, last_chunk=last_chunk or flush_encoder, cnn_cache=flow_cache["conformer_cnn_cache"], att_cache=flow_cache["conformer_att_cache"], ) - projected_speakers = self.flow.spk_embed_affine_layer(F.normalize(speakers, dim=1)) + projected_speakers = self._project_speakers(speakers) cond = torch.zeros_like(hidden).transpose(1, 2).contiguous() chunk_mel, estimator_cnn, estimator_att = self._decode_cfm( hidden.transpose(1, 2).contiguous(), @@ -1432,20 +1823,30 @@ class BatchedToken2Wav(nn.Module): audios[row] = group_audio[group_row] next_states[row] = group_states[group_row] return self._require_complete_ragged_outputs(audios, next_states) - encoder_groups: dict[tuple[int, bool], list[int]] = {} - for row, (row_tokens, last_chunk) in enumerate(zip(tokens, last_chunks, strict=True)): - encoder_groups.setdefault((int(row_tokens.numel()), last_chunk), []).append(row) + encoder_groups: dict[tuple[Any, ...], list[int]] = {} + for row, (row_tokens, state, last_chunk) in enumerate(zip(tokens, states, last_chunks, strict=True)): + # Onset and continuation rows may differ in conformer cache length: encode them apart; + # the estimator caches stay per-request rows in the ragged Whole-Euler solve. + key = ( + tuple(state.flow_cache["conformer_cnn_cache"].shape), + tuple(state.flow_cache["conformer_att_cache"].shape), + int(row_tokens.numel()), + last_chunk, + ) + encoder_groups.setdefault(key, []).append(row) hidden_rows: list[torch.Tensor | None] = [None] * batch_size conformer_cnn_rows: list[torch.Tensor | None] = [None] * batch_size conformer_att_rows: list[torch.Tensor | None] = [None] * batch_size - for (_, last_chunk), rows in encoder_groups.items(): + for (*_, last_chunk), rows in encoder_groups.items(): group_states = [states[row] for row in rows] - # 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 several groups the next replay rewrites the shared graph results, so they are copied. + graphed = None if last_chunk else self._graph_encode(group_tokens, group_states, len(encoder_groups) > 1) + # The encoder reads only the conformer caches. + group_cache = {} if graphed else self._stack_flow_cache(group_states, include_estimator_att=False) with self._autocast(group_tokens.device): - hidden, conformer_cnn, conformer_att = self._encode_chunk( + hidden, conformer_cnn, conformer_att = graphed or self._encode_chunk( group_tokens, last_chunk=last_chunk, cnn_cache=group_cache["conformer_cnn_cache"], @@ -1470,11 +1871,16 @@ class BatchedToken2Wav(nn.Module): # 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) + # Row-offset merges hold caches of different lengths, which never stack: the solve takes them + # per row, and the encoder groups above already read the conformer caches. + per_row_att = per_request_att or ( + len({int(state.flow_cache["estimator_att_cache"].shape[4]) for state in states}) > 1 + ) + flow_cache = self._stack_flow_cache(states, include_estimator_att=not per_row_att, include_conformer=False) 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)) + projected_speakers = self._project_speakers(speakers) cond = torch.zeros_like(padded_hidden).transpose(1, 2).contiguous() chunk_mel, estimator_cnn, estimator_att = self._decode_cfm( padded_hidden.transpose(1, 2).contiguous(), @@ -1483,7 +1889,7 @@ class BatchedToken2Wav(nn.Module): cnn_cache=flow_cache["estimator_cnn_cache"], att_cache=( [state.flow_cache["estimator_att_cache"] for state in states] - if per_request_att + if per_row_att else flow_cache["estimator_att_cache"] ), valid_lengths=hidden_lengths, @@ -1516,10 +1922,17 @@ class BatchedToken2Wav(nn.Module): audios: list[torch.Tensor | None] = [None] * batch_size next_states: list[BatchedToken2WavState | None] = [None] * batch_size - vocoder_groups: dict[tuple[int, bool], list[int]] = {} + vocoder_groups: dict[tuple[Any, ...], list[int]] = {} for row, last_chunk in enumerate(last_chunks): - vocoder_groups.setdefault((hidden_lengths[row], last_chunk), []).append(row) - for (hidden_length, last_chunk), rows in vocoder_groups.items(): + key = ( + hidden_lengths[row], + last_chunk, + tuple(states[row].hift_cache["mel"].shape), + tuple(states[row].hift_cache["source"].shape), + tuple(states[row].hift_cache["speech"].shape), + ) + vocoder_groups.setdefault(key, []).append(row) + for (hidden_length, last_chunk, *_), rows in vocoder_groups.items(): old_mel = torch.cat([states[row].hift_cache["mel"] for row in rows], dim=0) old_source = torch.cat([states[row].hift_cache["source"] for row in rows], dim=0) old_speech = torch.cat([states[row].hift_cache["speech"] for row in rows], dim=0) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py b/vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py new file mode 100644 index 000000000..3ab3f7f96 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/cfm_attention.py @@ -0,0 +1,138 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Tiled CFM attention: float32 tensors, ordinary TF32 Tensor Core dots.""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _attention( + q, + k, + v, + mask, + kv_rows, + out, + qb: tl.constexpr, + qh: tl.constexpr, + qt: tl.constexpr, + qd: tl.constexpr, + kb: tl.constexpr, + kh: tl.constexpr, + kt: tl.constexpr, + kd: tl.constexpr, + vb: tl.constexpr, + vh: tl.constexpr, + vt: tl.constexpr, + vd: tl.constexpr, + mb: tl.constexpr, + mt: tl.constexpr, + mk: tl.constexpr, + heads: tl.constexpr, + nq: tl.constexpr, + nk: tl.constexpr, + dim: tl.constexpr, + has_mask: tl.constexpr, + has_kv_rows: tl.constexpr, + scale: tl.constexpr, + block_m: tl.constexpr, + block_n: tl.constexpr, + block_d: tl.constexpr, +): + b, h, tile = tl.program_id(0), tl.program_id(1), tl.program_id(2) + # Keys and values of a pooled cache live in row ``kv_rows[b]`` of the pool. + kv_b = b + if has_kv_rows: + kv_b = tl.load(kv_rows + b).to(tl.int64) + rows = tile * block_m + tl.arange(0, block_m) + ds = tl.arange(0, block_d) + cols = tl.arange(0, block_n) + queries = tl.load( + q + b * qb + h * qh + rows[:, None] * qt + ds[None, :] * qd, (rows[:, None] < nq) & (ds[None, :] < dim), 0 + ) + maximum = tl.full((block_m,), -float("inf"), tl.float32) + denominator = tl.full((block_m,), 0, tl.float32) + accumulator = tl.full((block_m, block_d), 0, tl.float32) + for start in range(tl.cdiv(nk, block_n)): + keys = start * block_n + cols + kval = tl.load( + k + kv_b * kb + h * kh + keys[None, :] * kt + ds[:, None] * kd, + (keys[None, :] < nk) & (ds[:, None] < dim), + 0, + ) + score = tl.dot(queries, kval, input_precision="tf32") * scale + valid = (rows[:, None] < nq) & (keys[None, :] < nk) + if has_mask: + keep = tl.load(mask + b * mb + rows[:, None] * mt + keys[None, :] * mk, valid, 0) + valid = valid & (keep != 0) + score = tl.where(valid, score, -float("inf")) + new_max = tl.maximum(maximum, tl.max(score, 1)) + # Fully masked rows produce zero, matching PyTorch SDPA. + # Keep the running maximum at -inf until a valid tile arrives. A + # permanent zero would underflow a later tile with very negative logits. + safe_max = tl.where(new_max == -float("inf"), 0.0, new_max) + correction = tl.exp(maximum - safe_max) + prob = tl.exp(score - safe_max[:, None]) + denominator = denominator * correction + tl.sum(prob, 1) + values = tl.load( + v + kv_b * vb + h * vh + keys[:, None] * vt + ds[None, :] * vd, + (keys[:, None] < nk) & (ds[None, :] < dim), + 0, + ) + accumulator = accumulator * correction[:, None] + tl.dot(prob, values, input_precision="tf32") + maximum = new_max + result = accumulator / tl.where(denominator > 0, denominator, 1.0)[:, None] + tl.store( + out + ((b * nq + rows[:, None]) * heads + h) * dim + ds[None, :], + result, + (rows[:, None] < nq) & (ds[None, :] < dim), + ) + + +@torch.library.custom_op("vllm_omni::cfm_tiled_attention", mutates_args=()) +def cfm_attention( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + mask: torch.Tensor | None = None, + kv_rows: torch.Tensor | None = None, +) -> torch.Tensor: + """Tiled attention with TF32 dots; optional ``kv_rows`` indexes a pooled K/V cache.""" + b, heads, nq, dim = q.shape + nk = k.shape[2] + if q.dtype != torch.float32 or k.dtype != torch.float32 or v.dtype != torch.float32: + raise ValueError("CFM tiled attention requires float32 Q/K/V") + num_stages = 2 + out = torch.empty((b, nq, heads, dim), device=q.device, dtype=q.dtype) + _attention[(b, heads, triton.cdiv(nq, 32))]( + q, + k, + v, + mask, + kv_rows, + out, + *q.stride(), + *k.stride(), + *v.stride(), + *(mask.stride() if mask is not None else (0, 0, 0)), + heads, + nq, + nk, + dim, + mask is not None, + kv_rows is not None, + dim**-0.5, + 32, + 64, + triton.next_power_of_2(dim), + num_warps=4, + num_stages=num_stages, + ) + return out.transpose(1, 2) + + +@cfm_attention.register_fake +def _cfm_attention_fake(q, k, v, mask=None, kv_rows=None): + b, heads, nq, dim = q.shape + return q.new_empty((b, nq, heads, dim)).transpose(1, 2) 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 a37218fe5..95ba9de16 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,16 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +import functools +import itertools +import time +import weakref from collections import Counter -from collections.abc import Callable +from collections.abc import Callable, Iterable from typing import NamedTuple +import numpy as np 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 +from vllm_omni.utils.device_copy import index_to_device, to_device_nonblocking + logger = init_logger(__name__) @@ -34,12 +41,28 @@ def _euler_timeline(n_timesteps: int, device: torch.device, dtype: torch.dtype) return timeline, (timeline[1:] - timeline[:-1]).tolist() +def codec_frame_range(value, *, name: str) -> range: + """An inclusive ``[first, last]`` codec-frame range from the connector extra (empty: off).""" + if value is None or (isinstance(value, (list, tuple)) and len(value) == 0): + return range(0) + if not isinstance(value, (list, tuple)) or len(value) != 2 or not 1 <= int(value[0]) <= int(value[1]): + raise ValueError(f"MiniCPM-o {name} must be [first, last] codec frames with 1 <= first <= last, got {value!r}") + return range(int(value[0]), int(value[1]) + 1) + + +def empty_hift_outputs(speech_feat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """HiFT ``(speech, source)`` for a zero-width mel, without launching the vocoder.""" + batch_size = int(speech_feat.shape[0]) + return speech_feat.new_empty((batch_size, 0)), speech_feat.new_empty((batch_size, 1, 0)) + + class HiFTGraphWrapper: 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 self.codec_chunk_frames = connector_config["codec_chunk_frames"] + self.initial_codec_chunk_frames = int(connector_config.get("initial_codec_chunk_frames", 0)) self.codec_left_context_frames = connector_config["codec_left_context_frames"] lookahead_layer = getattr(token2wav.flow.encoder, "pre_lookahead_layer", None) pre_lookahead_len = getattr(lookahead_layer, "pre_lookahead_len", None) @@ -48,7 +71,24 @@ class HiFTGraphWrapper: self.source_cache_len = int(token2wav.source_cache_len) self.mel_frames = int(token2wav.hift.conv_pre.in_channels) self.flow_upsample_rate = int(getattr(token2wav.flow, "token_mel_ratio", 2)) + self.extra_codec_chunk_frames = [int(c) for c in connector_config.get("hift_graph_codec_chunk_frames") or ()] self.capture_bucket_size, self.capture_source_cache_len = self.derive_capture_bucket_size() + # The only (mel_frames, cache_source_len) shapes ever lazily captured: a duplex first chunk's + # arbitrary width is seen once, and capturing it would stall every stream for one replay. + self._legit_shapes = set(zip(self.capture_bucket_size, self.capture_source_cache_len, strict=True)) + # Exact-shape graphs (default off): vocoder shapes outside the buckets (a short first chunk, a merged + # backlog continuation) captured with the buckets at ``hift_graph_exact_batch_sizes`` and replayed + # unpadded, so they run eager HiFT's kernels; other batch sizes of these shapes stay eager. + self.first_chunk_frames, self.continuation_frames = ( + codec_frame_range(connector_config.get(name), name=name) + for name in ("hift_graph_first_chunk_frames", "hift_graph_continuation_frames") + ) + self.exact_batch_sizes = sorted( + {int(b) for b in connector_config.get("hift_graph_exact_batch_sizes") or (1,) if int(b) > 0} + ) + self.exact_shapes = self.derive_exact_shapes() + # Ordered (capture order) and hashable (replay lookups). + self._exact_keys = dict.fromkeys((b, *shape) for b in self.exact_batch_sizes for shape in self.exact_shapes) self.capture_batch_sizes = capture_batch_sizes self.graph: dict[tuple[int, int, int], torch.cuda.CUDAGraph] = {} self.static_speech_inputs: dict[tuple[int, int, int], torch.Tensor] = {} @@ -59,19 +99,37 @@ class HiFTGraphWrapper: parameter = next(token2wav.hift.parameters()) self.device = parameter.device self.dtype = parameter.dtype - self.max_lazy_graphs = 8 + self.max_lazy_graphs = int(connector_config.get("hift_max_lazy_graphs", 8)) + if self.max_lazy_graphs < 0: + raise ValueError("MiniCPM-o hift_max_lazy_graphs must be >= 0") 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 = ( - self.codec_chunk_frames + self.codec_left_context_frames - self.pre_lookahead_len - ) * self.flow_upsample_rate + def _chunk_mel_frames(self, codec_frames: int) -> int: + return (int(codec_frames) + self.codec_left_context_frames - self.pre_lookahead_len) * self.flow_upsample_rate - return [chunk_mel_frames, chunk_mel_frames + self.mel_cache_len], [ - 0, - self.source_cache_len, + def derive_capture_bucket_size(self): + chunk_mel_frames = self._chunk_mel_frames(self.codec_chunk_frames) + shapes = [(chunk_mel_frames, 0), (chunk_mel_frames + self.mel_cache_len, self.source_cache_len)] + if self.initial_codec_chunk_frames and self.initial_codec_chunk_frames != self.codec_chunk_frames: + first_frames = self._chunk_mel_frames(self.initial_codec_chunk_frames) + if first_frames <= self.mel_cache_len: + raise ValueError("MiniCPM-o initial codec chunk must emit audio beyond the HiFT mel cache") + shapes.append((first_frames, 0)) + for codec_frames in self.extra_codec_chunk_frames: + mel_frames = self._chunk_mel_frames(codec_frames) + if mel_frames <= self.mel_cache_len: + raise ValueError("MiniCPM-o hift_graph_codec_chunk_frames must emit audio beyond the HiFT mel cache") + shapes += [(mel_frames, 0), (mel_frames + self.mel_cache_len, self.source_cache_len)] + shapes = list(dict.fromkeys(shapes)) + return [shape[0] for shape in shapes], [shape[1] for shape in shapes] + + def derive_exact_shapes(self) -> list[tuple[int, int]]: + """``(mel_frames, cache_source_len)`` of the exact-shape graphs, outside the buckets.""" + shapes = [(self._chunk_mel_frames(f), 0) for f in self.first_chunk_frames] + [ + (self._chunk_mel_frames(f) + self.mel_cache_len, self.source_cache_len) for f in self.continuation_frames ] + return list(dict.fromkeys(s for s in shapes if s[0] > 0 and s not in self._legit_shapes)) def capture(self): for batch_size in self.capture_batch_sizes: @@ -81,6 +139,23 @@ class HiFTGraphWrapper: strict=True, ): self._capture(batch_size, mel_frames, source_cache_len) + self.capture_exact() + + def capture_exact(self) -> int: + """Capture the exact-shape graphs (``derive_exact_shapes``); returns how many are new.""" + before, started, memory_before = len(self.graph), time.perf_counter(), _memory_snapshot(self.device) + for key in self._exact_keys: + self._capture(*key) + # Builds the ISTFT window state before the first replay needs it. + self.finalize_fn(self.static_magnitude_outputs[key][:1], self.static_phase_outputs[key][:1]) + if self._exact_keys: + logger.info( + "Captured %d exact-shape HiFT CUDA Graphs in %.1f s%s", + len(self.graph) - before, + time.perf_counter() - started, + _format_memory_delta(memory_before, _memory_snapshot(self.device)), + ) + return len(self.graph) - before def _capture( self, @@ -124,6 +199,8 @@ class HiFTGraphWrapper: logger.info("Captured HiFT CUDA Graph for shape %s", key) def replay(self, speech_feat, cache_source): + if int(speech_feat.shape[2]) == 0: + return empty_hift_outputs(speech_feat) if torch.cuda.is_current_stream_capturing(): logger.info("Falling back to eager HiFT inference during an active stream capture") return self.decode_fn(speech_feat, cache_source) @@ -131,6 +208,17 @@ class HiFTGraphWrapper: batch_size = speech_feat.shape[0] num_frames = speech_feat.shape[2] cache_source_len = cache_source.shape[2] + + exact_key = (batch_size, num_frames, cache_source_len) + if exact_key in getattr(self, "_exact_keys", ()) and exact_key in self.graph: + # An exact-shape graph: same shape as the eager call, no padding. + return self._replay_key(exact_key, speech_feat, cache_source) + + if (num_frames, cache_source_len) not in self._legit_shapes: + # Never lazily captured (see ``_legit_shapes``): one eager call over the whole batch. + logger.info("Eager HiFT for shape %s outside the capture buckets", exact_key) + return self.decode_fn(speech_feat, cache_source) + target_b = next((b for b in sorted(self.capture_batch_sizes) if b >= batch_size), None) if target_b is None: @@ -155,6 +243,11 @@ class HiFTGraphWrapper: self._capture(*key) self.lazy_graph_count += 1 + return self._replay_key(key, speech_feat, cache_source) + + def _replay_key(self, key, speech_feat, cache_source): + """Replay the captured graph ``key`` on a batch of at most ``key[0]`` rows.""" + batch_size = speech_feat.shape[0] static_speech_inputs = self.static_speech_inputs[key].zero_() static_speech_inputs[:batch_size].copy_(speech_feat) static_cache_sources = self.static_cache_source_inputs[key].zero_() @@ -194,6 +287,25 @@ def _memory_snapshot(device: torch.device) -> tuple[int, int] | None: return None +def _memory_peak(device: torch.device) -> tuple[int, int] | None: + """Return allocator peak (allocated, reserved) bytes when available. + + A before/after delta misses transient capture workspaces. Peak counters + are the useful signal for deciding whether another graph bucket is safe; + they are sampled without synchronizing and are therefore cheap enough for + capture telemetry. + """ + if device.type != "cuda": + return None + try: + return ( + int(torch.accelerator.max_memory_allocated(device)), + int(torch.accelerator.max_memory_reserved(device)), + ) + except Exception: + return None + + def _format_memory_delta(before: tuple[int, int] | None, after: tuple[int, int] | None) -> str: if before is None or after is None: return "" @@ -206,6 +318,7 @@ 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 +_minus_ones = functools.partial(torch.full, fill_value=-1) def _align_up(n: int, bucket: int) -> int: @@ -214,7 +327,7 @@ def _align_up(n: int, bucket: int) -> int: return ((n + bucket - 1) // bucket) * bucket -def _capture_query_width(mel_width: int, bucket: int) -> int: +def _capture_query_width(mel_width: int, bucket: int | tuple[int, ...]) -> int: """Snap the query axis onto the Whole-Euler capture grid. ``_cfm_pad_frames`` already aligns onto ``bucket_frames=16``, which is why @@ -224,7 +337,12 @@ def _capture_query_width(mel_width: int, bucket: int) -> int: 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. + + Ascending widths (``query_widths``) snap to the narrowest one holding the chunk (a duplex + unit and a turn chunk keep their own width), past the widest aligns up on it. """ + if isinstance(bucket, tuple): + bucket = next((width for width in bucket if mel_width <= width), bucket[-1]) if bucket else 0 if bucket <= 1: return int(mel_width) if mel_width <= bucket: @@ -232,6 +350,21 @@ def _capture_query_width(mel_width: int, bucket: int) -> int: return _align_up(int(mel_width), int(bucket)) +def _capture_offset(offset: int, bucket: int, steady: int) -> int: + """Snap an estimator-cache length onto the Whole-Euler offset grid (one graph per first-chunk size otherwise). + + Anchored at the steady ``prompt + 100`` length so the steady solve is never + padded; shorter caches round up in ``bucket``-frame steps, the padded frames + masked out of attention (``_build_capture_mask``), so no row's result changes. + """ + offset = int(offset) + if bucket <= 1 or offset <= 0: + return offset + if 0 < steady and offset <= steady: + return steady - ((steady - offset) // bucket) * bucket + return _align_up(offset, bucket) + + def _pad_query_for_capture( x: torch.Tensor, mu_cfg: torch.Tensor, @@ -262,6 +395,7 @@ def _build_capture_mask( mel_width: int, mel_frames: int, device: torch.device, + offset_cap: int | None = None, ) -> torch.Tensor: """Always-on attention mask in the capture layout, so presence-of-mask is not a graph key. @@ -279,26 +413,37 @@ def _build_capture_mask( query_cap = int(query_cap) offset = int(offset) mel_width = int(mel_width) + # ``offset_cap`` columns of cache follow the query block; the ones past the + # request's own ``offset`` are ``_capture_offset`` padding and stay False. + offset_cap = offset if offset_cap is None else int(offset_cap) mask = torch.zeros( 2 * int(batch_size), query_cap, - query_cap + offset, + query_cap + offset_cap, 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] + mask[:, :mel_width, query_cap : query_cap + offset] = 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 + mask[:, :mel_width, query_cap : query_cap + offset] = True if query_cap > mel_width: mask[:, mel_width:] = mask[:, :1] return mask +def _pad_invalid_queries(mask: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor: + """Queries past each CFG row's ``lengths`` take its first query's keys: dropped, but never empty.""" + invalid = torch.arange(int(mask.shape[1]), device=mask.device).unsqueeze(0) >= lengths.unsqueeze(1) + return torch.where(invalid.unsqueeze(2), mask[:, :1], 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. @@ -433,12 +578,20 @@ class CFMGraphWrapper: "captures": 0, "flushes": 0, "eager": 0, + "peak_allocated": 0, + "peak_reserved": 0, } def stats_snapshot(self) -> dict[str, int]: """Bounded cumulative telemetry for the graph cache.""" return {**self._stats, "cache_size": len(self._cache)} + def _record_peak(self) -> None: + peak = _memory_peak(self.device) + if peak is not None: + self._stats["peak_allocated"] = max(self._stats["peak_allocated"], peak[0]) + self._stats["peak_reserved"] = max(self._stats["peak_reserved"], peak[1]) + def _call_graph_fn(self, args: tuple[torch.Tensor, ...]) -> torch.Tensor: return self.graph_fn(args[0], args[1], args[6], args[2], args[3], args[4], args[5]) @@ -517,6 +670,7 @@ class CFMGraphWrapper: return None self._stats["captures"] += 1 + self._record_peak() logger.info( "Captured CFM CUDA Graph for shape %s (cache=%d/%d, stats=%s)%s", key, @@ -665,6 +819,8 @@ class WholeEulerExecutionArena: self._buffers: dict[tuple, torch.Tensor] = {} # The one attention cache storage; every graph takes a prefix view. self._att: torch.Tensor | None = None + # Per-timestep modulation table for a fused body (``get_modulation``). + self._modulation: torch.Tensor | None = None def _buffer(self, name: str, shape: tuple[int, ...], new=torch.zeros, dtype: torch.dtype | None = None): key = (name, shape) @@ -688,6 +844,14 @@ class WholeEulerExecutionArena: self._buffers[key] = t_emb return t_emb + def get_modulation(self, modulation_fn: Callable[[torch.Tensor], torch.Tensor]) -> torch.Tensor: + """``modulation_fn`` of each timestep's one-row time embedding, stacked; computed once for every graph.""" + if self._modulation is None: + embed = self.estimator.t_embedder + steps = [modulation_fn(embed(self.time_steps[s].expand(1)).unsqueeze(1)) for s in range(self.n_timesteps)] + self._modulation = torch.stack(steps, dim=0) + return self._modulation + 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) @@ -750,9 +914,97 @@ class WholeEulerExecutionArena: """Per-CFG-row valid query lengths read by a ragged graph body.""" return self._buffer("lengths", (2 * batch_size,), dtype=torch.long) + def att_bytes(self) -> int: + """Device bytes of the shared attention cache storage.""" + storage = self._att + return 0 if storage is None else storage.numel() * storage.element_size() + def clear(self) -> None: self._buffers.clear() self._att = None + self._modulation = None + + +class AttSlotPool: + """Resident estimator attention caches, one slot per request, for Whole-Euler graphs.""" + + def __init__(self, arena: WholeEulerExecutionArena, *, slots: int, frames: int, suffix: int) -> None: + """Slots shaped like ``arena``'s attention cache rows.""" + self.n_timesteps, self.depth = arena.n_timesteps, arena.depth + self.heads, self.width = arena.heads, arena.att_width + self.slots, self.frames, self.suffix = int(slots), int(frames), int(suffix) + self.device, self.dtype = arena.device, arena.att_cache_dtype + shape = (self.n_timesteps, self.depth, self.slots, 2, self.heads, self.frames, self.width) + self.storage = torch.zeros(shape, device=self.device, dtype=self.dtype) + # Lowest slot first, so a lightly loaded pool keeps touching the same pages. + self._free = list(range(self.slots - 1, -1, -1)) + + def rows(self) -> torch.Tensor: + """``(n_t, depth, 2 * slots, heads, frames, width)``: one row per request CFG half.""" + return self.storage.flatten(2, 3) + + @property + def free_slots(self) -> int: + return len(self._free) + + def bytes(self) -> int: + return self.storage.numel() * self.storage.element_size() + + def fits(self, frames: int) -> bool: + """Whether a logical cache of ``frames`` frames (after a chunk) fits one slot.""" + return self.suffix <= frames <= self.frames + + def adopt(self, cache: torch.Tensor) -> "ResidentAttCache | None": + """A slot holding the logical cache; ``None`` when the pool is full.""" + frames = int(cache.shape[4]) + if not self._free or not self.fits(frames): + return None + slot = self._free.pop() + streaming = frames - self.suffix + _copy_frame_segments(self.storage[:, :, slot], cache, [(streaming, self.suffix), (0, streaming)]) + return ResidentAttCache(self, slot, list(range(self.suffix, self.suffix + streaming))) + + +class ResidentAttCache: + """A request's estimator attention cache resident in an ``AttSlotPool`` slot. + + Stands in for the logical ``(n_t, depth, 2, heads, L, width)`` tensor. The slot returns + to the pool when the handle is collected; work queued on it before then runs first (one stream). + """ + + __slots__ = ("pool", "slot", "layout", "__weakref__") + + def __init__(self, pool: AttSlotPool, slot: int, layout: list[int]) -> None: + self.pool = pool + self.slot = slot + self.layout = layout + weakref.finalize(self, pool._free.append, slot) + + @property + def shape(self) -> torch.Size: + pool = self.pool + return torch.Size((pool.n_timesteps, pool.depth, 2, pool.heads, len(self.layout) + pool.suffix, pool.width)) + + @property + def dtype(self) -> torch.dtype: + return self.pool.dtype + + @property + def device(self) -> torch.device: + return self.pool.device + + def physical_frames(self) -> list[int]: + """The physical frame of every logical frame, in logical order.""" + return [*self.layout, *range(self.pool.suffix)] + + def materialize(self) -> torch.Tensor: + """The logical cache as a tensor of its own.""" + index = index_to_device(self.physical_frames(), self.pool.device) + return self.pool.storage[:, :, self.slot].index_select(4, index) + + +def _materialize_att_rows(rows: list) -> list[torch.Tensor]: + return [row.materialize() if isinstance(row, ResidentAttCache) else row for row in rows] class _WholeEulerStatics(NamedTuple): @@ -767,6 +1019,11 @@ class _WholeEulerStatics(NamedTuple): attn_mask: torch.Tensor time_embeddings: torch.Tensor lengths: torch.Tensor | None + # (n_timesteps, ...) per-timestep modulation for a body that takes it. + modulation: torch.Tensor | None = None + # Slot-pool graphs: each row's pool row, and the pool frame each query writes (-1: none). + slot_rows: torch.Tensor | None = None + slot_positions: torch.Tensor | None = None class WholeEulerCFMGraphWrapper: @@ -800,9 +1057,13 @@ class WholeEulerCFMGraphWrapper: max_serial_batch: int | None = None, max_graph_batch: int | None = None, micro_batch_size: int = 4, - query_bucket_frames: int | None = None, + query_bucket_frames: int | tuple[int, ...] | None = None, pad_max_rows: int | None = None, ragged_body: Callable[..., torch.Tensor] | None = None, + offset_bucket_frames: int = 0, + modulation_fn: Callable[[torch.Tensor], torch.Tensor] | None = None, + att_slots: int = 0, + row_offsets: bool = False, ) -> 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 @@ -811,16 +1072,32 @@ class WholeEulerCFMGraphWrapper: 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. + + ``query_bucket_frames`` is one capture width or several + (``_capture_query_width``); ``offset_bucket_frames`` > 1 snaps cache + lengths onto a grid (``_capture_offset``); ``modulation_fn`` gives a + fused body its per-timestep table. ``att_slots`` > 0 keeps request + caches resident in an ``AttSlotPool``: streaming chunks replay graphs + keyed by batch and width alone, with no cache copies, and with + ``row_offsets`` rows of different cache lengths share one replay. """ self.estimator = estimator self.ragged_body = ragged_body + self.modulation_fn = modulation_fn if ragged_body is not None else None 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) + # The capture widths, narrowest first (``_capture_query_width``); the widest sizes the arena. + if isinstance(query_bucket_frames, (tuple, list)): + self.query_widths = tuple(sorted({int(w) for w in query_bucket_frames if int(w) > 1})) + self.query_bucket_frames = self.query_widths[-1] if self.query_widths else 0 + else: + self.query_bucket_frames = int(query_bucket_frames or 0) + self.query_widths = (self.query_bucket_frames,) if self.query_bucket_frames > 1 else () + self.offset_bucket_frames = int(offset_bucket_frames) self.micro_batch_size = 4 if micro_batch_size is None else int(micro_batch_size) + # ``max_serial_batch`` only caps the default ``max_graph_batch``. 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) @@ -850,30 +1127,52 @@ class WholeEulerCFMGraphWrapper: "captures": 0, "flushes": 0, "eager": 0, + "slot_replays": 0, + "peak_allocated": 0, + "peak_reserved": 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, - ) + arena_args = { + "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, + } + self.arena = WholeEulerExecutionArena(**arena_args) + # Resident request caches (``AttSlotPool``), sized once the prompt + # length is known; their graphs and staging buffers are kept apart + # from the arena's so an arena flush never retires them. + self.att_slots = int(att_slots) if ragged_body is not None and modulation_fn is not None else 0 + self.slot_pool: AttSlotPool | None = None + self.row_offsets = bool(row_offsets) and self.att_slots > 0 + self._slot_arena = WholeEulerExecutionArena(**arena_args) + self._slot_graphs: dict[tuple, tuple] = {} def stats_snapshot(self) -> dict[str, int]: """Bounded cumulative telemetry for the graph cache.""" - return {**self._stats, "cache_size": len(self._cache)} + return {**self._stats, "cache_size": len(self._cache) + len(self._slot_graphs)} + + def _record_peak(self) -> None: + peak = _memory_peak(self.device) + if peak is not None: + self._stats["peak_allocated"] = max(self._stats["peak_allocated"], peak[0]) + self._stats["peak_reserved"] = max(self._stats["peak_reserved"], peak[1]) + + @staticmethod + def _retire(graphs: dict[tuple, tuple]) -> None: + for entry in graphs.values(): + entry[4].reset() + graphs.clear() 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._retire(self._cache) self._stats["flushes"] += 1 logger.info("Whole-Euler CFM graph cache flushed; stats=%s", self.stats_snapshot()) self.arena.clear() @@ -887,6 +1186,7 @@ class WholeEulerCFMGraphWrapper: ) self.enabled = False self._flush() + self._retire(self._slot_graphs) def _run_euler_loop( self, @@ -898,24 +1198,29 @@ class WholeEulerCFMGraphWrapper: ) -> torch.Tensor: cur_x = statics.x width = int(statics.mu_cfg.shape[2]) - speaker_features = statics.speakers_cfg.unsqueeze(-1).expand(-1, -1, width) - + # A slot-pool graph attends to the pool in place: no cache goes in. + slotted = statics.slot_rows is not None + no_cache = [None] * len(self.estimator.blocks) for step in range(self.n_timesteps): dt = self.dt_steps[step] x_cfg = torch.cat((cur_x, cur_x), dim=0) + speaker_features = statics.speakers_cfg.unsqueeze(-1).expand(-1, -1, width) 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], + no_cache if slotted else 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) + kwargs = {} if statics.modulation is None else {"modulation": statics.modulation[step]} + if slotted: + kwargs["slots"] = (statics.slot_rows, statics.slot_positions) + estimate = self.ragged_body(self.estimator, *args, statics.lengths, **kwargs) else: estimate = self.estimator.blocks_forward_chunk(*args) @@ -933,7 +1238,7 @@ class WholeEulerCFMGraphWrapper: sizes.append(self.micro_batch_size) return sizes - def _plan_groups(self, batch_size: int) -> list[tuple[int, int]] | None: + def _plan_groups(self, batch_size: int) -> list[tuple[int, int]]: """``(graph_batch, rows)`` replays that cover ``batch_size`` requests. Larger batches run whole micro-batches. The remainder takes the smallest @@ -941,16 +1246,15 @@ class WholeEulerCFMGraphWrapper: (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. + row is cheaper than a further replay. """ - micro = self.micro_batch_size sizes = self._graph_batches() + largest = max(sizes) groups: list[tuple[int, int]] = [] remainder = batch_size - while remainder >= micro: - groups.append((micro, micro)) - remainder -= micro + while remainder >= largest: + groups.append((largest, largest)) + remainder -= largest while remainder > 0: up = next(size for size in sizes if size >= remainder) if up - remainder <= self.pad_max_rows: @@ -959,12 +1263,6 @@ class WholeEulerCFMGraphWrapper: 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 @@ -1022,13 +1320,57 @@ class WholeEulerCFMGraphWrapper: ) if offset <= 0: return - static_att_rows = statics.att_cache.unflatten(2, (2, graph_batch)) + # A graph on the ``_capture_offset`` grid holds at least ``offset`` + # cache frames; the rest are masked out and keep the finite values the + # zero-initialized storage last held, like padded rows do. + static_att_rows = statics.att_cache.unflatten(2, (2, graph_batch))[..., :offset, :] 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]) + for row, source in enumerate(_materialize_att_rows(att_rows[start:stop])): + static_att_rows[:, :, :, row].copy_(source[..., :offset, :]) + + def _group_fills(self, groups: list[tuple[int, int]], after=None, **inputs) -> list[Callable]: + """Per replay group, a ``fill(statics)`` that copies the group's requests in, then runs ``after``.""" + fills = [] + start = 0 + for graph_batch, rows in groups: + + def fill(statics, *, _graph_batch=graph_batch, _start=start, _stop=start + rows): + self._fill_static_inputs(statics, graph_batch=_graph_batch, start=_start, stop=_stop, **inputs) + if after is not None: + after(statics, _graph_batch, _start, _stop) + + fills.append(fill) + start += rows + return fills + + def _statics( + self, + arena: WholeEulerExecutionArena, + *, + graph_batch: int, + channels: int, + query_cap: int, + spk_dim: int, + att_cache: torch.Tensor, + key_frames: int, + **slots: torch.Tensor, + ) -> _WholeEulerStatics: + return _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), + att_cache=att_cache, + attn_mask=arena.get_mask_staging(graph_batch, query_cap, key_frames), + time_embeddings=arena.get_time_embeddings(graph_batch), + lengths=arena.get_lengths_staging(graph_batch) if self.ragged_body is not None else None, + modulation=arena.get_modulation(self.modulation_fn) if self.modulation_fn is not None else None, + **slots, + ) def _capture( self, @@ -1051,22 +1393,24 @@ class WholeEulerCFMGraphWrapper: # 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): + if self.att_slots: + # Streaming chunks run from the slot pool; the arena only holds + # what reaches it (the prompt solve or fallback), not a steady batch. + capacity, rows = max(capacity, offset + query_cap), max(rows, graph_batch) + if self._cache and not arena.att_cache_fits(max(graph_batch, rows), 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), + statics = self._statics( + arena, + graph_batch=graph_batch, + channels=channels, + query_cap=query_cap, + spk_dim=spk_dim, # 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, + key_frames=offset + query_cap, ) out_cnn_cache = arena.get_cnn_cache_out(graph_batch) # The output view aliases the input view; see WholeEulerExecutionArena. @@ -1076,7 +1420,27 @@ class WholeEulerCFMGraphWrapper: logger.warning("Failed to allocate static buffers for Whole-Euler graph key: %s", key, exc_info=True) self._unsupported.add(key) return None + storage = arena._att + return self._record( + key, + statics, + out_cnn_cache, + out_att_cache, + graph_batch=graph_batch, + where=f"att storage rows/frames={int(storage.shape[2]) // 2}/{int(storage.shape[4])}", + ) + def _record( + self, + key: tuple, + statics: _WholeEulerStatics, + out_cnn_cache: torch.Tensor, + out_att_cache: torch.Tensor, + *, + graph_batch: int, + where: str, + ) -> tuple | None: + """Warm up and capture one graph over filled ``statics``; ``None`` (and disabled) on failure.""" 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. @@ -1101,13 +1465,13 @@ class WholeEulerCFMGraphWrapper: return None self._stats["captures"] += 1 + self._record_peak() logger.info( - "Captured Whole-Euler CFM CUDA Graph for shape %s (cache=%d/%d, att storage rows/frames=%d/%d, stats=%s)%s", + "Captured Whole-Euler CFM CUDA Graph for shape %s (cache=%d/%d, %s, stats=%s)%s", key, - len(self._cache) + 1, + len(self._cache) + len(self._slot_graphs) + 1, self.max_graphs, - int(arena._att.shape[2]) // 2, - int(arena._att.shape[4]), + where, self.stats_snapshot(), _format_memory_delta(memory_before, _memory_snapshot(self.device)), ) @@ -1132,8 +1496,9 @@ class WholeEulerCFMGraphWrapper: if entry is not None: self._stats["hits"] += 1 return entry - if len(self._cache) >= self.max_graphs: - self._flush() + if len(self._cache) + len(self._slot_graphs) >= self.max_graphs: + self._unsupported.add(key) + return None entry = self._capture( key, graph_batch=graph_batch, @@ -1147,7 +1512,9 @@ class WholeEulerCFMGraphWrapper: self._cache[key] = entry return entry - def _group_entries(self, groups: list[tuple[int, int]], fills: list, **entry_args) -> list[tuple] | None: + def _group_entries( + self, entry_fn: Callable[..., tuple | None], 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 @@ -1159,7 +1526,7 @@ class WholeEulerCFMGraphWrapper: 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) + entry = entry_fn(graph_batch=graph_batch, fill=fill, **entry_args) if entry is None: return None entries.append(entry) @@ -1167,6 +1534,297 @@ class WholeEulerCFMGraphWrapper: return entries return None + def _ensure_slot_pool(self, keep: tuple[int, int]) -> AttSlotPool | None: + """The slot pool for streaming trim ``keep``: a steady cache plus the widest chunk per slot.""" + if not self.att_slots or not self.enabled: + return None + prefix, suffix = int(keep[0]), int(keep[1]) + pool = self.slot_pool + if pool is None: + if prefix < 0 or suffix <= 0 or self.query_bucket_frames <= 1: + return None + max_query = max(self.query_widths) if self.query_widths else self.query_bucket_frames + max_query = max(max_query, 200) + frames = _align_up(prefix + suffix + max_query, _ATT_FRAME_ALIGN) + try: + pool = AttSlotPool(self.arena, slots=self.att_slots, frames=frames, suffix=suffix) + except torch.OutOfMemoryError: + logger.warning("Whole-Euler slot pool (%d x %d frames) does not fit", self.att_slots, frames) + self.att_slots = 0 + return None + self.slot_pool = pool + gib = pool.bytes() / 2**30 + logger.info("Whole-Euler slot pool: %d x %d frames, %.2f GiB", pool.slots, pool.frames, gib) + return pool if pool.suffix == suffix else None + + def _slot_entry(self, *, graph_batch: int, query_cap: int, channels: int, spk_dim: int, fill) -> tuple | None: + """The slot-pool graph of ``graph_batch`` rows at width ``query_cap``, captured on first use.""" + key = ("whole_euler_slots", graph_batch, query_cap) + entry = self._slot_graphs.get(key) + if entry is not None: + self._stats["hits"] += 1 + return entry + if key in self._unsupported: + return None + if len(self._cache) + len(self._slot_graphs) >= self.max_graphs: + self._unsupported.add(key) + return None + pool = self.slot_pool + assert pool is not None + arena = self._slot_arena + try: + statics = self._statics( + arena, + graph_batch=graph_batch, + channels=channels, + query_cap=query_cap, + spk_dim=spk_dim, + att_cache=pool.rows(), + key_frames=pool.frames, + slot_rows=arena._buffer("slot_rows", (2 * graph_batch,), dtype=torch.long), + # Nothing is written until a replay says where. + slot_positions=arena._buffer( + "slot_positions", (2 * graph_batch, query_cap), new=_minus_ones, dtype=torch.int32 + ), + ) + out_cnn_cache = arena.get_cnn_cache_out(graph_batch) + 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 + where = f"slot pool {pool.slots} x {pool.frames} frames" + entry = self._record(key, statics, out_cnn_cache, pool.rows(), graph_batch=graph_batch, where=where) + if entry is not None: + self._slot_graphs[key] = entry + return entry + + def _replay_slots( + self, + *, + x_cap: torch.Tensor, + mu_cap: torch.Tensor, + cond_cap: torch.Tensor, + speakers_cfg: torch.Tensor, + cnn_cache: torch.Tensor | None, + att_rows: list, + attn_mask: torch.Tensor | None, + mel_width: int, + query_cap: int, + row_lengths: list[int], + lengths: torch.Tensor | int, + att_keep: tuple[int, int], + groups: list[tuple[int, int]], + ) -> tuple[torch.Tensor, torch.Tensor, list[ResidentAttCache]] | None: + """``replay`` on resident caches (``AttSlotPool``); ``None`` leaves the batch to the arena. + + Each row's new frames take free frames of its slot, its logical + ``[current | cache]`` mask is scattered onto the slot, the graph writes + and attends in place, and the trim drops frames from the row's layout. + With ``row_offsets`` the cache columns span the longest row's cache. + """ + pool = self._ensure_slot_pool(att_keep) + if pool is None: + return None + if attn_mask is None and any(length != mel_width for length in row_lengths): + # Without a mask the arena path lets every row attend all current + # frames, a row's padding included; a slot never holds padding. + return None + batch_size = len(att_rows) + offsets = [int(row.shape[4]) for row in att_rows] + # Everything that can refuse the pool is checked before a slot changes. + if min(offsets) < pool.suffix or any( + offset + length > pool.frames for offset, length in zip(offsets, row_lengths, strict=True) + ): + return None + seen: set[int] = set() + adopt = [] + for row in att_rows: + resident = isinstance(row, ResidentAttCache) and row.pool is pool and id(row) not in seen + if resident: + seen.add(id(row)) + elif row.dtype != pool.dtype: + return None + adopt.append(not resident) + if sum(adopt) > pool.free_slots: + return None + handles: list[ResidentAttCache] = [] + for row, needs in zip(att_rows, adopt, strict=True): + handle = pool.adopt(row.materialize() if isinstance(row, ResidentAttCache) else row) if needs else row + assert handle is not None + handles.append(handle) + + # Per CFG row, in one host buffer (one copy to the device): its pool + # row, the slot frame each query writes (-1: none), and the slot frame + # of each logical key (a row's current padding maps past the slot and + # is dropped). + rows_cfg, keys = 2 * batch_size, mel_width + max(offsets) + host = np.full((rows_cfg, 1 + query_cap + keys), -1, dtype=np.int64) + host[:, 1 + query_cap :] = pool.frames + new_frames: list[np.ndarray] = [] + for row, (handle, length) in enumerate(zip(handles, row_lengths, strict=True)): + free = np.ones(pool.frames, dtype=bool) + free[: pool.suffix] = False + free[handle.layout] = False + frames = np.flatnonzero(free)[:length] + new_frames.append(frames) + halves = [row, row + batch_size] + host[halves, 0] = [2 * handle.slot, 2 * handle.slot + 1] + host[halves, 1 : 1 + length] = frames + host[halves, 1 + query_cap : 1 + query_cap + length] = frames + cached = handle.physical_frames() + host[halves, 1 + query_cap + mel_width : 1 + query_cap + mel_width + len(cached)] = cached + device = x_cap.device + index = to_device_nonblocking(torch.from_numpy(host), device) + slot_rows, positions, columns = index.split((1, query_cap, keys), dim=1) + columns = columns.unsqueeze(1).expand(-1, mel_width, -1) + slot_mask = torch.zeros((rows_cfg, query_cap, pool.frames + 1), dtype=torch.bool, device=device) + if attn_mask is None: + slot_mask[:, :mel_width].scatter_(2, columns, True) + else: + slot_mask[:, :mel_width].scatter_(2, columns, attn_mask[:, :mel_width, :keys]) + cfg_lengths = lengths if isinstance(lengths, torch.Tensor) else index_to_device(row_lengths * 2, device) + slot_mask = _pad_invalid_queries(slot_mask[..., : pool.frames], cfg_lengths) + request_rows = slot_rows.squeeze(1).unflatten(0, (2, batch_size)) + request_positions = positions.unflatten(0, (2, batch_size)) + + def copy_slots(statics: _WholeEulerStatics, graph_batch: int, start: int, stop: int) -> None: + count = stop - start + static_rows = statics.slot_rows.unflatten(0, (2, graph_batch)) + static_positions = statics.slot_positions.unflatten(0, (2, graph_batch)) + static_rows[:, :count].copy_(request_rows[:, start:stop]) + static_positions[:, :count].copy_(request_positions[:, start:stop]) + if count < graph_batch: + # Padded rows only read (the first row's slot) and never write. + static_rows[:, count:].copy_(request_rows[:, start : start + 1].expand(-1, graph_batch - count)) + static_positions[:, count:].fill_(-1) + + fills = self._group_fills( + groups, + copy_slots, + x=x_cap, + mu_cfg=mu_cap, + speakers_cfg=speakers_cfg, + cond_cfg=cond_cap, + attn_mask=slot_mask, + cnn_cache=cnn_cache, + att_cache=None, + att_rows=None, + offset=0, + lengths=lengths, + ) + entries = self._group_entries( + self._slot_entry, + groups, + fills, + query_cap=query_cap, + channels=int(x_cap.shape[1]), + spk_dim=int(speakers_cfg.shape[1]), + ) + if entries is None: + return None + + chunk_mel = x_cap.new_empty((batch_size, int(x_cap.shape[1]), mel_width)) + out_cnn = torch.empty(self._slot_arena._cnn_cache_shape(batch_size), device=device, dtype=self.dtype) + out_cnn_rows = out_cnn.unflatten(2, (2, batch_size)) + 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, _, graph = entry + fill(statics) + graph.replay() + chunk_mel[start:stop].copy_(static_final_x[:rows, :, :mel_width]) + out_cnn_rows[:, :, :, start:stop].copy_(out_cnn_cache.unflatten(2, (2, graph_batch))[:, :, :, :rows]) + start = stop + + for handle, frames in zip(handles, new_frames, strict=True): + layout = [*frames.tolist(), *handle.layout] + # The streaming frames the trim keeps are its first range. + _, kept = _att_keep_ranges(len(layout) + pool.suffix, att_keep)[0] + handle.layout = layout[:kept] + self._stats["slot_replays"] += 1 + return chunk_mel, out_cnn, handles + + def _precapture_slots(self, *, channels: int, spk_dim: int) -> int: + """Capture every slot-pool graph (batch grid x capture widths); the cache offset is not a key.""" + before = self._stats["captures"] + shape = {"channels": channels, "spk_dim": spk_dim, "fill": self._precapture_fill} + keys = list(itertools.product(sorted(self._graph_batches(), reverse=True), self.query_widths)) + room = self.max_graphs - (len(self._cache) + len(self._slot_graphs)) + if len(keys) > room: + logger.warning( + "Whole-Euler slot precapture requested %d graphs; max_graphs budget is %d (room=%d). " + "Excess shapes will execute eagerly", + len(keys), + self.max_graphs, + room, + ) + keys = keys[: max(0, room)] + for graph_batch, query_cap in keys: + if self._slot_entry(graph_batch=graph_batch, query_cap=query_cap, **shape) is None: + break + return self._stats["captures"] - before + + @staticmethod + def _precapture_fill(statics: _WholeEulerStatics) -> None: + # The zero-initialized staging buffers are finite; a ragged body also + # needs lengths within its query width for the causal cache gather. + if statics.lengths is not None: + statics.lengths.fill_(int(statics.x.shape[2])) + + def precapture( + self, + *, + offsets: Iterable[int], + steady: int, + channels: int, + spk_dim: int, + max_graph_bytes: int = 1 << 30, + keep: tuple[int, int] | None = None, + ) -> int: + """Capture Whole-Euler graphs for one prompt before serving. Returns the capture count.""" + if not self.enabled: + return 0 + if self.att_slots and keep is not None and self._ensure_slot_pool(keep) is not None: + return self._precapture_slots(channels=int(channels), spk_dim=int(spk_dim)) + self._att_capacity = max(self._att_capacity, int(steady)) + grid = sorted({_capture_offset(int(o), self.offset_bucket_frames, int(steady)) for o in offsets}, reverse=True) + widths = self.query_widths + if not widths: + return 0 + keys = [(b, w, o) for b in sorted(self._graph_batches(), reverse=True) for w in widths for o in grid] + room = self.max_graphs - (len(self._cache) + len(self._slot_graphs)) + if len(keys) > room: + logger.warning( + "Whole-Euler precapture requested %d graphs; max_graphs budget is %d (room=%d). " + "Excess shapes will execute eagerly", + len(keys), + self.max_graphs, + room, + ) + keys = keys[: max(0, room)] + x = torch.empty((1, int(channels), 1), device=self.device, dtype=self.dtype) + entry_fn = functools.partial(self._entry, x=x, spk_dim=int(spk_dim), fill=self._precapture_fill) + before = self._stats["captures"] + for graph_batch, query_cap, offset in keys: + memory_before, arena_before = _memory_snapshot(self.device), self.arena.att_bytes() + entry = entry_fn(graph_batch=graph_batch, query_cap=query_cap, offset=offset) + if entry is None: + break + memory_after = _memory_snapshot(self.device) + if memory_before is not None and memory_after is not None: + # What the graph itself holds, apart from the shared arena. + held = memory_after[1] - memory_before[1] - (self.arena.att_bytes() - arena_before) + if held > max_graph_bytes: + # A capture failure disables every graph (see _capture): stop before one runs the device dry. + logger.warning( + "Whole-Euler precapture stopped at %s: one graph reserved %.2f GiB; the rest capture lazily", + (graph_batch, query_cap, offset), + held / 2**30, + ) + break + return self._stats["captures"] - before + def replay( self, *, @@ -1215,12 +1873,12 @@ class WholeEulerCFMGraphWrapper: 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) + # Per-request caches of different lengths only replay from the slot pool (``row_offsets``). + mixed_offsets = att_rows is not None and len({int(row.shape[4]) for row in att_rows}) != 1 if att_rows is not None: - if len(att_rows) != batch_size or len({int(row.shape[4]) for row in att_rows}) != 1: + if len(att_rows) != batch_size or (mixed_offsets and not self.row_offsets): return None offset = int(att_rows[0].shape[4]) else: @@ -1237,8 +1895,38 @@ class WholeEulerCFMGraphWrapper: 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) + query_cap = _capture_query_width(mel_width, self.query_widths) x_cap, mu_cap, cond_cap = _pad_query_for_capture(x, mu_cfg, cond_cfg, query_cap) + 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) + if self.att_slots and att_rows is not None and pad_frames == 0 and att_keep is not None: + slotted = self._replay_slots( + x_cap=x_cap, + mu_cap=mu_cap, + cond_cap=cond_cap, + speakers_cfg=speakers_cfg, + cnn_cache=cnn_cache, + att_rows=att_rows, + attn_mask=attn_mask, + mel_width=mel_width, + query_cap=query_cap, + row_lengths=row_lengths, + lengths=lengths, + att_keep=att_keep, + groups=groups, + ) + if slotted is not None: + return slotted + if mixed_offsets: + return None + if att_rows is not None: + # A refused slot replay can leave resident rows; the arena path reads tensors. + att_rows = _materialize_att_rows(att_rows) + offset_cap = _capture_offset(offset, self.offset_bucket_frames, sum(att_keep) if att_keep is not None else 0) mask_cap = _build_capture_mask( attn_mask=attn_mask, batch_size=batch_size, @@ -1247,17 +1935,12 @@ class WholeEulerCFMGraphWrapper: mel_width=mel_width, mel_frames=mel_frames, device=x.device, + offset_cap=offset_cap, ) - 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) + if isinstance(lengths, torch.Tensor): + # Queries past a row's length take its first query's mask, as + # capture padding does. + mask_cap = _pad_invalid_queries(mask_cap, lengths) 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) @@ -1287,36 +1970,25 @@ class WholeEulerCFMGraphWrapper: # 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 + fills = self._group_fills( + groups, + 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, + ) entries = self._group_entries( + self._entry, groups, fills, query_cap=query_cap, - offset=offset, + offset=offset_cap, x=x, spk_dim=int(speakers_cfg.shape[1]), ) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py b/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py new file mode 100644 index 000000000..7d5390844 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/dit_fused.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused Code2Wav CFM DiT block forward (~15 kernels vs ~47 eager ops per block).""" + +from __future__ import annotations + +import functools +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .ops import qkv_head_layer_norm, residual_layer_norm + +# adaLN chunk order of ``DiTBlock.adaLN_modulation``. +_SHIFT_MSA, _SCALE_MSA, _GATE_MSA, _SHIFT_MLP, _SCALE_MLP, _GATE_MLP, _SHIFT_CONV, _SCALE_CONV, _GATE_CONV = range(9) + + +def supports_fused_body(estimator: nn.Module) -> bool: + """Whether ``blocks_forward_chunk_fused`` computes this estimator's blocks.""" + blocks = getattr(estimator, "blocks", None) + final_layer = getattr(estimator, "final_layer", None) + if not blocks or not hasattr(estimator, "in_proj") or final_layer is None: + return False + if not isinstance(getattr(final_layer, "norm_final", None), nn.LayerNorm) or not hasattr(final_layer, "linear"): + return False + for block in blocks: + attn = getattr(block, "attn", None) + conv_block = getattr(getattr(block, "conv", None), "block", None) + mlp = getattr(block, "mlp", None) + if attn is None or conv_block is None or mlp is None: + return False + if not all(hasattr(attn, name) for name in ("to_q", "to_k", "to_v", "proj", "num_heads", "head_dim")): + return False + if not (isinstance(attn.q_norm, nn.LayerNorm) and isinstance(attn.k_norm, nn.LayerNorm)): + return False + if not all(isinstance(getattr(block, name, None), nn.LayerNorm) for name in ("norm1", "norm2", "norm3")): + return False + if len(conv_block) < 7 or not isinstance(conv_block[3], nn.LayerNorm) or not isinstance(conv_block[4], nn.Mish): + return False + for conv in (conv_block[1], conv_block[6]): + if not isinstance(conv, nn.Conv1d) or conv.stride != (1,) or conv.dilation != (1,) or conv.groups != 1: + return False + if conv.padding != (0,) or conv.kernel_size[0] < 2: + return False + if not (hasattr(mlp, "fc1") and hasattr(mlp, "fc2") and isinstance(mlp.act, nn.GELU)): + return False + if not isinstance(getattr(mlp, "norm", nn.Identity()), nn.Identity): + return False + return True + + +def dit_modulation(estimator: nn.Module, time_embedding: torch.Tensor) -> torch.Tensor: + """adaLN table ``(depth + 1, 9, C)`` for one timestep; scales stored as ``1 + scale``.""" + embedding = time_embedding.reshape(1, 1, -1) + table = torch.stack([block.adaLN_modulation(embedding).reshape(9, -1) for block in estimator.blocks], dim=0) + table[:, (_SCALE_MSA, _SCALE_MLP, _SCALE_CONV)] += 1.0 + final_shift, final_scale = estimator.final_layer.adaLN_modulation(embedding).reshape(2, -1) + final = table.new_zeros((1, 9, int(table.shape[2]))) + # Where a block keeps its MSA shift/scale, so the last block's closing LayerNorm reads it like a next block. + final[0, _SHIFT_MSA] = final_shift + final[0, _SCALE_MSA] = final_scale + 1.0 + return torch.cat((table, final), dim=0).contiguous() + + +def _cached(module: nn.Module, key: str, sources: tuple[torch.Tensor, ...], build) -> Any: + version = tuple((t.data_ptr(), t._version, t.dtype) for t in sources) + cached = module.__dict__.get(key) + if cached is None or cached[0] != version: + with torch.no_grad(): + cached = (version, build()) + module.__dict__[key] = cached + return cached[1] + + +def packed_qkv(attn: nn.Module) -> tuple[torch.Tensor, torch.Tensor | None]: + """``[to_q; to_k; to_v]`` weight and bias as one projection, rebuilt only when a source changes.""" + sources = (attn.to_q.weight, attn.to_k.weight, attn.to_v.weight) + + def build(): + weight = torch.cat(sources, dim=0).contiguous() + bias = None + if attn.to_q.bias is not None: + bias = torch.cat((attn.to_q.bias, attn.to_k.bias, attn.to_v.bias), dim=0).contiguous() + return weight, bias + + return _cached(attn, "_packed_qkv", sources, build) + + +def _conv_taps_weight(conv: nn.Conv1d) -> torch.Tensor: + """``(C_out, C_in, K)`` -> ``(K * C_out, C_in)``: tap ``k`` is output rows ``[k*C_out, (k+1)*C_out)``.""" + return _cached( + conv, + "_fused_taps_weight", + (conv.weight,), + lambda: conv.weight.detach().permute(2, 0, 1).reshape(-1, conv.in_channels).contiguous(), + ) + + +def _conv_history( + old_cache: torch.Tensor | None, rows: int, width: int, frames: int, channels: int, like: torch.Tensor +) -> torch.Tensor: + """``(N, width + frames, C)``: the cached frames, then room for the current ones.""" + history = like.new_empty((rows, width + frames, channels)) + if old_cache is None: + history[:, :width].zero_() + else: + history[:, :width].copy_(old_cache.transpose(1, 2)) + return history + + +def causal_cache_index(lengths: torch.Tensor, width: int) -> torch.Tensor: + """``(N, width)`` history frames that end each row's valid frames, for ``gather_causal_cache``.""" + return lengths[:, None] + torch.arange(width, device=lengths.device)[None, :] + + +def gather_causal_cache(history: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + """The ``index`` frames of ``history`` (N, width + T, C), as the (N, C, width) causal cache.""" + return history.gather(1, index[:, :, None].expand(-1, -1, int(history.shape[2]))).transpose(1, 2) + + +def _mlp(mlp: nn.Module, hidden: torch.Tensor) -> torch.Tensor: + rows = hidden.reshape(-1, int(hidden.shape[-1])) + if mlp.act.approximate == "tanh" and rows.is_cuda and mlp.fc1.bias is not None: + # cuBLASLt applies bias + tanh-GELU in the GEMM epilogue. + inner = torch._addmm_activation(mlp.fc1.bias, rows, mlp.fc1.weight.t(), use_gelu=True) + else: + inner = mlp.act(F.linear(rows, mlp.fc1.weight, mlp.fc1.bias)) + return F.linear(inner, mlp.fc2.weight, mlp.fc2.bias).view(*hidden.shape[:-1], -1) + + +@functools.cache +def tiled_attention_supported(device_index: int) -> bool: + """``cfm_attention`` TF32 dots need NVIDIA tensor cores (SM80+).""" + return torch.version.hip is None and torch.cuda.get_device_capability(device_index) >= (8, 0) + + +def _attention( + attn: nn.Module, + hidden: torch.Tensor, + att_cache: torch.Tensor | None, + attn_mask: torch.Tensor | None, + kv: torch.Tensor, + slots: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """``attn.forward_chunk`` over ``hidden``, writing new keys/values into ``kv``.""" + batch, frames, _ = hidden.shape + weight, bias = packed_qkv(attn) + qkv = F.linear(hidden.reshape(batch * frames, -1), weight, bias).view(batch, frames, -1) + rows, positions = slots if slots is not None else (None, None) + if slots is None: + # ``att_cache`` behind the current frames: left where it is when it + # already sits there (the Whole-Euler arena's aliasing), else copied once. + 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) + q = qkv_head_layer_norm( + qkv, + kv, + num_heads=int(attn.num_heads), + head_dim=int(attn.head_dim), + q_norm=attn.q_norm, + k_norm=attn.k_norm, + rows=rows, + positions=positions, + ) + head_dim = int(q.shape[-1]) + keys, values = kv[..., :head_dim], kv[..., head_dim:] + if rows is not None or ( + kv.is_cuda and kv.dtype == q.dtype == torch.float32 and tiled_attention_supported(kv.get_device()) + ): + from .cfm_attention import cfm_attention + + out = cfm_attention(q, keys, values, attn_mask, rows) + else: + out = F.scaled_dot_product_attention( + q, keys, values, attn_mask=None if attn_mask is None else attn_mask.unsqueeze(1) + ) + out = out.transpose(1, 2).reshape(batch * frames, -1) + return F.linear(out, attn.proj.weight, attn.proj.bias).view(batch, frames, -1) + + +def blocks_forward_chunk_fused( + estimator: nn.Module, + estimator_input: torch.Tensor, + time_embedding: torch.Tensor, + attn_mask: torch.Tensor | None, + cnn_cache: Any, + att_cache: Any, + cnn_cache_buffer: torch.Tensor, + att_cache_buffer: torch.Tensor, + valid_lengths: list[int] | torch.Tensor, + modulation: torch.Tensor | None = None, + slots: tuple[torch.Tensor, torch.Tensor] | None = None, +) -> torch.Tensor: + """``BatchedToken2Wav._blocks_forward_chunk_ragged`` with fused norms/GEMM/attention.""" + if isinstance(valid_lengths, torch.Tensor): + lengths = valid_lengths + else: + lengths = torch.tensor((*valid_lengths, *valid_lengths), device=estimator_input.device, dtype=torch.long) + if modulation is None: + modulation = dit_modulation(estimator, time_embedding[:1]) + blocks = estimator.blocks + depth = len(blocks) + # Causal cache gather index per conv width, shared by the blocks. + cache_index: dict[int, torch.Tensor] = {} + + x = estimator.in_proj(estimator_input.transpose(1, 2)).contiguous() + rows, frames, _ = x.shape + first = modulation[0] + hidden = residual_layer_norm(x, weight=first[_SCALE_MSA], bias=first[_SHIFT_MSA], eps=blocks[0].norm1.eps) + for block_index, block in enumerate(blocks): + mod = modulation[block_index] + x_att = _attention(block.attn, hidden, att_cache[block_index], attn_mask, att_cache_buffer[block_index], slots) + + # CausalConvBlock: [conv1, LayerNorm, Mish, conv2] on (N, width + T, C) histories. + conv_layers = block.conv.block + conv1, conv_norm, conv2 = conv_layers[1], conv_layers[3], conv_layers[6] + width = int(conv1.kernel_size[0]) - 1 + old_cnn = cnn_cache[block_index] + split = conv1.in_channels + old_cnn1, old_cnn2 = (None, None) if old_cnn is None else (old_cnn[:, :split], old_cnn[:, split:]) + history = _conv_history(old_cnn1, rows, width, frames, split, x) + # x += gate_msa * attention; the conv block's modulated LayerNorm lands in the history. + residual_layer_norm( + x, + x_att, + gate=mod[_GATE_MSA], + weight=mod[_SCALE_CONV], + bias=mod[_SHIFT_CONV], + eps=block.norm3.eps, + residual_out=x, + out=history[:, width:], + ) + taps1 = F.linear(history.view(-1, conv1.in_channels), _conv_taps_weight(conv1)) + history2 = _conv_history(old_cnn2, rows, width, frames, conv2.in_channels, x) + residual_layer_norm( + None, + taps1.view(rows, width + frames, -1), + taps=width + 1, + y_bias=conv1.bias, + weight=conv_norm.weight, + bias=conv_norm.bias, + eps=conv_norm.eps, + activation="mish", + out=history2[:, width:], + ) + taps2 = F.linear(history2.view(-1, conv2.in_channels), _conv_taps_weight(conv2)) + # x += gate_conv * conv2; the MLP's modulated LayerNorm. + mlp_in = residual_layer_norm( + x, + taps2.view(rows, width + frames, -1), + taps=width + 1, + y_bias=conv2.bias, + gate=mod[_GATE_CONV], + weight=mod[_SCALE_MLP], + bias=mod[_SHIFT_MLP], + eps=block.norm2.eps, + residual_out=x, + ) + mlp_out = _mlp(block.mlp, mlp_in) + # x += gate_mlp * MLP, then the next block's (or the final layer's) modulated LayerNorm: + # table row ``depth`` holds the final layer's shift and 1 + scale where a block keeps its MSA pair. + last = block_index + 1 == depth + hidden = residual_layer_norm( + x, + mlp_out, + gate=mod[_GATE_MLP], + weight=modulation[block_index + 1][_SCALE_MSA], + bias=modulation[block_index + 1][_SHIFT_MSA], + eps=estimator.final_layer.norm_final.eps if last else blocks[block_index + 1].norm1.eps, + residual_out=None if last else x, + ) + + if width not in cache_index: + cache_index[width] = causal_cache_index(lengths, width) + cnn_out = cnn_cache_buffer[block_index] + cnn_out[:, :split].copy_(gather_causal_cache(history, cache_index[width])) + cnn_out[:, split:].copy_(gather_causal_cache(history2, cache_index[width])) + + return estimator.final_layer.linear(hidden).transpose(1, 2) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py b/vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py new file mode 100644 index 000000000..8ddd39a8e --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/flow_encoder_graph.py @@ -0,0 +1,226 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Exact-shape CUDA graphs of the Code2Wav flow encoder's continuation chunk. + +The CosyVoice2 upsample-conformer encoder (embedding, ``forward_chunk``, +``encoder_proj``) is ~320 small kernels whose launches, not their GPU time, set +a serving step's latency; a graph replays them with one launch. One graph per +exact ``(rows, tokens, conformer cache frames)``, replayed only for calls of +exactly that shape, so it runs the eager kernels and changes no output. Nothing +is padded: rounding 3 rows up to a 4-row graph changes the M of every GEMM, and +with it cuBLAS's kernel and FP32 rounding, which the TF32 CFM and HiFT amplified +to 0.81 dB. The streaming trim caps the cache at ``prompt + 100`` frames, so one +prompt's streams only pass through the lengths ``reachable_conformer_cache_frames`` +lists. All graphs take prefix views of one storage per static tensor and share +one private pool; a replay's results stay valid until the next replay. +""" + +from __future__ import annotations + +import math +from collections.abc import Callable, Iterable, Sequence +from typing import Any + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def reachable_conformer_cache_frames(*, start: int, prompt_len: int, suffix: int, growths: Iterable[int]) -> list[int]: + """Conformer attention-cache lengths a stream passes through from ``start``, for chunks of any mix of ``growths``. + + Past ``prompt_len + suffix`` frames the streaming trim keeps that many (``_trim_streaming_cache``). + """ + steady = int(prompt_len) + int(suffix) + steps = sorted({int(g) for g in growths if int(g) > 0}) + seen, frontier = {int(start)}, [int(start)] + while frontier: + frames = frontier.pop() + for following in (min(frames + growth, steady) for growth in steps): + if following not in seen: + seen.add(following) + frontier.append(following) + return sorted(seen) + + +def _precision_state() -> tuple[Any, ...]: + """The matmul/conv precision flags a captured graph's kernels were chosen under.""" + return ( + bool(torch.backends.cuda.matmul.allow_tf32), + bool(torch.backends.cudnn.allow_tf32), + torch.get_float32_matmul_precision(), + ) + + +class FlowEncoderGraphs: + """Exact-shape graphs of ``encode_fn(tokens, cnn_cache=, att_cache=)``, the eager continuation chunk. + + Caches use the encoder layouts: ``cnn_cache`` ``(rows, channels, width)``, + ``att_cache`` ``(depth, rows, heads, frames, 2 * head_dim)``; a chunk of + ``tokens`` tokens emits ``upsample * (tokens - lookahead)`` frames. + ``held_tensors`` returns tensors the graphs read in place (the RelPos + tables): kept alive, and a replaced one disables the graphs. + """ + + def __init__( + self, + encode_fn: Callable[..., tuple[torch.Tensor, torch.Tensor, torch.Tensor]], + *, + rows: Iterable[int], + token_widths: Iterable[int], + lookahead: int, + upsample: int, + held_tensors: Callable[[], Sequence[torch.Tensor]] = tuple, + ) -> None: + self.encode_fn = encode_fn + self.rows = tuple(sorted({int(r) for r in rows if int(r) > 0})) + self.lookahead, self.upsample = int(lookahead), int(upsample) + self.token_widths = tuple(sorted({int(w) for w in token_widths if int(w) > self.lookahead})) + self._held_tensors = held_tensors + self._held: tuple[torch.Tensor, ...] = () + self._precision: tuple[Any, ...] | None = None + # (rows, tokens, cache frames) -> (graph, static views by name) + self.graphs: dict[tuple[int, int, int], tuple[Any, dict[str, torch.Tensor]]] = {} + self.replays = 0 + self.enabled = True + self._storage: dict[str, torch.Tensor] = {} + self._pool: Any = None + + def output_frames(self, tokens: int) -> int: + return self.upsample * (int(tokens) - self.lookahead) + + def keys_for(self, *, start: int, prompt_len: int, suffix: int) -> list[tuple[int, int, int]]: + """``(rows, tokens, cache frames)`` of every continuation chunk a prompt's streams can produce.""" + growths = [self.output_frames(width) for width in self.token_widths] + frames = reachable_conformer_cache_frames(start=start, prompt_len=prompt_len, suffix=suffix, growths=growths) + return [(rows, width, length) for rows in self.rows for width in self.token_widths for length in frames] + + def _body(self, views: dict[str, torch.Tensor]) -> None: + outputs = self.encode_fn(views["tokens"], cnn_cache=views["cnn"], att_cache=views["att"]) + for name, value in zip(("hidden", "new_cnn", "new_att"), outputs, strict=True): + target = views[name] + if value.shape != target.shape or value.dtype != target.dtype: + raise RuntimeError( + f"flow encoder {name} output {tuple(value.shape)}/{value.dtype} does not match " + f"the captured layout {tuple(target.shape)}/{target.dtype}" + ) + target.copy_(value) + + def _record(self, views: dict[str, torch.Tensor]) -> Any: + """Warm up on a side stream, then capture ``_body`` into the shared private pool.""" + device = views["att"].device + current = torch.cuda.current_stream(device) + side = torch.cuda.Stream(device=device) + side.wait_stream(current) + with torch.cuda.stream(side): + for _ in range(2): + self._body(views) + current.wait_stream(side) + torch.accelerator.synchronize(device) + if self._pool is None: + self._pool = torch.cuda.graph_pool_handle() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=self._pool): + self._body(views) + return graph + + def capture( + self, + keys: Iterable[tuple[int, int, int]], + *, + cnn_shape: tuple[int, int], + att_layout: tuple[int, int, int], + hidden_dim: int, + dtype: torch.dtype, + device: torch.device, + ) -> int: + """Capture ``keys``, all or none; ``cnn_shape`` ``(channels, width)``, ``att_layout`` ``(depth, heads, 2D)``.""" + if self.graphs: + raise RuntimeError("flow encoder graphs are already captured") + depth, heads, att_width = att_layout + + def shapes(key: tuple[int, int, int]) -> dict[str, tuple[int, ...]]: + rows, width, frames = key + out = self.output_frames(width) + return { + "tokens": (rows, width), + "cnn": (rows, *cnn_shape), + "att": (depth, rows, heads, frames, att_width), + "hidden": (rows, out, hidden_dim), + "new_cnn": (rows, *cnn_shape), + "new_att": (depth, rows, heads, frames + out, att_width), + } + + # Largest first, so the private pool grows once and later graphs reuse it. + keys = sorted( + {(int(r), int(w), int(f)) for r, w, f in keys}, + key=lambda key: (key[0] * (key[2] + self.output_frames(key[1])), key[0] * key[1]), + reverse=True, + ) + if not keys: + return 0 + layouts = [shapes(key) for key in keys] + # Ordinary tensors, not inference tensors: replays write them in and outside inference mode. + with torch.inference_mode(False): + for name in layouts[0]: + numel = max(math.prod(layout[name]) for layout in layouts) + self._storage[name] = torch.zeros(numel, dtype=torch.long if name == "tokens" else dtype, device=device) + self._precision = _precision_state() + try: + for key, layout in zip(keys, layouts, strict=True): + with torch.inference_mode(False): + views = {name: self._storage[name][: math.prod(s)].view(s) for name, s in layout.items()} + self.graphs[key] = (self._record(views), views) + except Exception: + # A partial set would replay some shapes and not others. + self.graphs.clear() + self._storage.clear() + self._pool = None + self.enabled = False + raise + self._held = tuple(self._held_tensors()) + return len(self.graphs) + + def storage_bytes(self) -> int: + """Bytes of the shared static inputs and results (not the private pool).""" + return sum(t.numel() * t.element_size() for t in self._storage.values()) + + def run( + self, + tokens: torch.Tensor, + cnn_rows: Sequence[torch.Tensor], + att_rows: Sequence[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + """The continuation chunk of these streams (one cache row each) from its graph, or ``None`` (run eager). + + Returns the shared result views: copy what must outlive the next replay. + """ + rows = len(att_rows) + if not (self.enabled and rows and tokens.dim() == 2 and len(tokens) == len(cnn_rows) == rows): + return None + key = (rows, int(tokens.shape[1]), int(att_rows[0].shape[3]) if att_rows[0].dim() == 5 else -1) + if (entry := self.graphs.get(key)) is None: + return None + graph, views = entry + static_att, static_cnn = views["att"], views["cnn"] + att_shape, cnn_shape = (static_att.shape[0], 1, *static_att.shape[2:]), (1, *static_cnn.shape[1:]) + if (tokens.dtype, tokens.device) != (views["tokens"].dtype, views["tokens"].device) or not all( + (att.shape, att.dtype, att.device, cnn.shape, cnn.dtype, cnn.device) + == (att_shape, static_att.dtype, static_att.device, cnn_shape, static_cnn.dtype, static_cnn.device) + for att, cnn in zip(att_rows, cnn_rows, strict=True) + ): + return None + if _precision_state() != self._precision or torch.cuda.is_current_stream_capturing(): + return None + if [id(t) for t in self._held_tensors()] != [id(t) for t in self._held]: # ``_held`` keeps them alive + self.enabled = False + logger.warning("Flow encoder RelPos tables were replaced after capture; encoder graphs disabled") + return None + with torch.no_grad(): # a replay is not differentiable; serving runs in inference mode anyway + torch.cat(list(cnn_rows), dim=0, out=static_cnn) + torch.cat(list(att_rows), dim=1, out=static_att) + views["tokens"].copy_(tokens) + graph.replay() + self.replays += 1 + return views["hidden"], views["new_cnn"], views["new_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 3f93f0f30..f5e374446 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 @@ -7,10 +7,11 @@ from __future__ import annotations import json import os import tempfile +import time from collections import OrderedDict -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence from dataclasses import dataclass -from functools import lru_cache +from functools import lru_cache, partial from hashlib import sha256 from pathlib import Path from typing import Any @@ -27,8 +28,10 @@ from vllm_omni.utils.device_copy import to_device_nonblocking from .batched_token2wav import ( BatchedToken2Wav, BatchedToken2WavState, + row_offset_signature, state_shape_signature, ) +from .cuda_graph_wrapper import ResidentAttCache, _format_memory_delta, _memory_snapshot logger = init_logger(__name__) @@ -47,6 +50,29 @@ def _resolve_model_dir(model_ref: str, revision: str | None = None) -> str: return hf_api().snapshot_download(model_ref, revision=revision, allow_patterns=["assets/*"]) +def _tf32_mode(extra: Mapping[str, Any]) -> str: + """Ordinary TF32 for CFM DiT GEMMs (``"off"`` / ``"tf32"``). + + Enables ``allow_tf32`` on dense QKV/MLP and Triton ``input_precision="tf32"`` + dots. Not compensated TF32x3. ``token2wav_allow_tf32`` is the shipped YAML + key; ``code2wav_allow_tf32`` and legacy env ``tf32x3`` / ``3xtf32`` alias it. + Defaults to ``"tf32"`` unless explicitly disabled. + """ + raw = os.environ.get("MINICPMO_CODE2WAV_TF32") + value = raw if raw not in (None, "") else extra.get("token2wav_allow_tf32", extra.get("code2wav_allow_tf32", True)) + if isinstance(value, str): + return { + "tf32": "tf32", + "tf32x3": "tf32", + "3xtf32": "tf32", + "flow": "tf32", + "1": "tf32", + "true": "tf32", + "yes": "tf32", + }.get(value.strip().lower(), "off") + return "tf32" if value else "off" + + def _batch_error(reason: str, **details: Any) -> RuntimeError: payload = {"reason": reason, **details} return RuntimeError(f"MiniCPMO45Code2WavBatchError {json.dumps(payload, sort_keys=True)}") @@ -225,6 +251,12 @@ class _WorkItem: has_payload: bool = True +def _slot_resident(item: _WorkItem) -> bool: + if item.previous is None: + return False + return isinstance(item.previous.token2wav.flow_cache.get("estimator_att_cache"), ResidentAttCache) + + class MiniCPMO45Code2Wav(nn.Module): """LLM_GENERATION model with request-owned state and compatible batching.""" @@ -265,12 +297,22 @@ class MiniCPMO45Code2Wav(nn.Module): raise ValueError("MiniCPM-o Code2Wav setup cache capacity must be >= 0") self._connector_config = { "codec_chunk_frames": int(extra.get("codec_chunk_frames", 25)), + "initial_codec_chunk_frames": int(extra.get("initial_codec_chunk_frames", 0)), "codec_left_context_frames": int(extra.get("codec_left_context_frames", 3)), + "hift_max_lazy_graphs": int(extra.get("hift_max_lazy_graphs", 8)), + "hift_graph_codec_chunk_frames": list(extra.get("hift_graph_codec_chunk_frames") or ()), + "hift_graph_first_chunk_frames": list(extra.get("hift_graph_first_chunk_frames") or ()), + "hift_graph_continuation_frames": list(extra.get("hift_graph_continuation_frames") or ()), + "hift_graph_exact_batch_sizes": list(extra.get("hift_graph_exact_batch_sizes") or (1,)), } if self._connector_config["codec_chunk_frames"] <= 0 or self._connector_config["codec_left_context_frames"] < 0: raise ValueError(f"Invalid MiniCPM-o connector chunk configuration: {self._connector_config}") + max_num_seqs = getattr(getattr(vllm_config, "scheduler_config", None), "max_num_seqs", None) 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 + # Default: powers of two below the scheduler's batch (capped at 32) and the cap itself. + batch_cap = min(int(max_num_seqs), 32) if max_num_seqs else 32 + pow2 = [1 << i for i in range(batch_cap.bit_length()) if 1 << i < batch_cap] + [batch_cap] + capture_batch_sizes = pow2 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 = { @@ -294,11 +336,26 @@ class MiniCPMO45Code2Wav(nn.Module): "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"), + "offset_bucket_frames": int(extra.get("cfm_graph_offset_bucket_frames", 50)), "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"), + "fused_body": bool(extra.get("cfm_fused_body", False)), + "slot_pool": bool(extra.get("cfm_slot_pool", False)), + "row_offset_merge": extra.get("cfm_row_offset_merge", False) is True, + } + # Exact-shape flow-encoder graphs (``FlowEncoderGraphs``, default off). + # ``cfm_encoder_graph_rows`` is the largest row count or a list of counts. + rows = extra.get("cfm_encoder_graph_rows", 8) + rows = sorted({int(r) for r in rows}) if isinstance(rows, (list, tuple)) else list(range(1, int(rows) + 1)) + if any(r < 1 for r in rows): + raise ValueError("MiniCPM-o cfm_encoder_graph_rows must be positive") + self._encoder_graph_config = { + "enabled": extra.get("cfm_encoder_cuda_graph", False) is True, + "rows": [r for r in rows if not max_num_seqs or r <= int(max_num_seqs)], + "token_widths": [int(w) for w in extra.get("cfm_encoder_graph_token_widths") or ()], } self._ref_max_seconds = float(extra.get("ref_audio_max_seconds", _REF_MAX_SECONDS)) if self._ref_max_seconds <= 0: @@ -314,6 +371,10 @@ class MiniCPMO45Code2Wav(nn.Module): self._default_prompt_id = str(extra.get("prompt_cache_id", "HT_ref_audio")) self._prompt_wav_override = extra.get("prompt_wav") self._default_prompt_normalized: tuple[str, str] | None = None + self._onset_merge = bool(extra.get("cfm_onset_merge", False)) + self._cross_turn_buckets = extra.get("cfm_cross_turn_buckets", False) is True + self._row_offset_merge = self._cfm_graph_config["row_offset_merge"] + self._precapture_empty_cache = extra.get("cfm_precapture_empty_cache", False) is True @property def _default_prompt_wav(self) -> str: @@ -658,12 +719,17 @@ class MiniCPMO45Code2Wav(nn.Module): ) @staticmethod - def _bucket_key(item: _WorkItem) -> tuple[Any, ...]: + def _bucket_key(item: _WorkItem, *, cross_turn: bool = False, row_offsets: bool = False) -> tuple[Any, ...]: cache_signature: Any if item.previous is None: cache_signature = ("uninitialized",) + elif row_offsets and _slot_resident(item): + # A slot-pool row keeps its cache lengths per row (``cfm_row_offset_merge``). + cache_signature = ("row_offsets", row_offset_signature(item.previous.token2wav)) else: cache_signature = state_shape_signature(item.previous.token2wav) + if cross_turn: + return (item.prompt_cache_id, item.prompt_wav, cache_signature) return ( item.prompt_cache_id, item.prompt_wav, @@ -671,6 +737,52 @@ class MiniCPMO45Code2Wav(nn.Module): item.cache_epoch, ) + def _signature_groups( + self, bucket: list[_WorkItem], states: list[BatchedToken2WavState], *, row_offsets: bool + ) -> list[list[int]]: + """A bucket's rows grouped by cache signature, for a bucket that cannot decode as one. + + With ``row_offsets``, slot-pool groups differing only in cache length stay one if the backend can solve it. + """ + groups: dict[tuple[Any, ...], list[int]] = {} + for row, state in enumerate(states): + groups.setdefault(state_shape_signature(state), []).append(row) + if not row_offsets: + return list(groups.values()) + clusters: dict[tuple[Any, ...], list[list[int]]] = {} + for signature, rows in groups.items(): + if _slot_resident(bucket[rows[0]]): + signature = ("row_offsets", row_offset_signature(states[rows[0]])) + clusters.setdefault(signature, []).append(rows) + result: list[list[int]] = [] + for parts in clusters.values(): + rows = sorted(row for part in parts for row in part) + if len(parts) > 1 and self.backend.can_merge_row_offsets( + [states[row] for row in rows], [int(bucket[row].tokens.numel()) for row in rows] + ): + result.append(rows) + else: + result.extend(parts) + return result + + def _decode_rows( + self, items: list[_WorkItem], states: list[BatchedToken2WavState], features: Any, *, ragged: bool + ) -> tuple[list[Any], list[Any]]: + """Decode ``items`` together: ragged when asked to, or when their token counts or last-chunk flags differ.""" + tokens = [item.tokens for item in items] + last_chunks = [item.last_chunk for item in items] + if ragged or len({int(t.numel()) for t in tokens}) > 1 or len(set(last_chunks)) > 1: + return self.backend.decode_ragged_batch(tokens, features, states, last_chunks=last_chunks) + return self.backend.decode_batch(torch.stack(tokens, dim=0), features, states, last_chunk=last_chunks[0]) + + @staticmethod + def _onset_group_key(key: tuple[Any, ...], *, cross_turn: bool = False) -> tuple[Any, ...]: + """The onset-merge group of a bucket key: its prompt, plus its cache epoch unless rows merge across turns.""" + if cross_turn: + return key[:2] + prompt_cache_id, prompt_wav, _, cache_epoch = key + return (prompt_cache_id, prompt_wav, cache_epoch) + def _iter_limited_batches( self, buckets: Iterable[list[_WorkItem]], @@ -738,7 +850,8 @@ class MiniCPMO45Code2Wav(nn.Module): # policy after eager execution/capture; cuDNN's policy is independent. previous_tf32 = torch.backends.cuda.matmul.allow_tf32 try: - if self._extra_config().get("token2wav_allow_tf32", False): + extra = self._extra_config() + if _tf32_mode(extra) != "off": torch.backends.cuda.matmul.allow_tf32 = True return self._forward_impl( input_ids, @@ -762,6 +875,12 @@ class MiniCPMO45Code2Wav(nn.Module): **kwargs: Any, ) -> OmniOutput: del positions, intermediate_tensors, inputs_embeds + if getattr(self, "_precapture_pending", False) and self.backend is not None: + try: + self._precapture_default_prompt() + finally: + if getattr(self, "_precapture_empty_cache", False): + self._release_precapture_cache() ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.empty(0, dtype=torch.long) segments = self._split_segments(ids, kwargs.get("seq_token_counts")) empty = torch.empty(0, dtype=torch.float32, device=ids.device) @@ -845,9 +964,11 @@ class MiniCPMO45Code2Wav(nn.Module): self._trim_runtime_prompts() raise _batch_error("empty_nonfinal_chunk", request_ids=invalid_empty) + cross_turn = getattr(self, "_cross_turn_buckets", False) is True + row_offsets = getattr(self, "_row_offset_merge", False) is True buckets: dict[tuple[Any, ...], list[_WorkItem]] = {} for item in compute_items: - buckets.setdefault(self._bucket_key(item), []).append(item) + buckets.setdefault(self._bucket_key(item, cross_turn=cross_turn, row_offsets=row_offsets), []).append(item) undersized = [ { "size": len(bucket), @@ -921,34 +1042,57 @@ class MiniCPMO45Code2Wav(nn.Module): prompt_wav=item.prompt_wav, token2wav=state, ) - for bucket in self._iter_decode_batches(buckets.values()): + decode_buckets: Iterable[list[_WorkItem]] + if self._onset_merge and self.backend is not None and self.backend._whole_euler_ragged_active(): + # Offer a fresh onset and a continuation of one prompt/epoch, ready together, to ragged + # Whole-Euler; other exact-shape buckets stay as they are. + grouped: dict[tuple[Any, ...], list[list[_WorkItem]]] = {} + for key, bucket in buckets.items(): + grouped.setdefault(self._onset_group_key(key, cross_turn=cross_turn), []).append(bucket) + merged: list[list[_WorkItem]] = [] + for entries in grouped.values(): + combined = [item for bucket in entries for item in bucket] + mixed = len({item.previous is None for item in combined}) > 1 + merged.extend([combined] if mixed else entries) + decode_buckets = merged + else: + decode_buckets = buckets.values() + for bucket in self._iter_decode_batches(decode_buckets): batch_size = len(bucket) try: features = self.backend.prepare_prompt( bucket[0].prompt_cache_id, bucket[0].prompt_wav, ) - if bucket[0].previous is None: - states = self.backend.setup_batch(features, batch_size) + fresh = [item for item in bucket if item.previous is None] + setup = iter(self.backend.setup_batch(features, len(fresh)) if fresh else []) + states = [next(setup) if item.previous is None else item.previous.token2wav for item in bucket] + onset_mix = self._onset_merge and 0 < len(fresh) < batch_size + onset_ready = onset_mix and self.backend.can_merge_state_shapes( + states, [int(item.tokens.numel()) for item in bucket] + ) + row_groups: list[list[int]] | None = None + offset_mix = False + if onset_mix and not onset_ready: + # Cache layouts ragged Whole-Euler cannot take split back into signature buckets. + row_groups = self._signature_groups(bucket, states, row_offsets=row_offsets) + elif row_offsets and not onset_ready and batch_size > 1: + groups = self._signature_groups(bucket, states, row_offsets=True) + if len(groups) > 1: + row_groups = groups + else: + # One decode, ragged when it holds rows of different cache lengths. + offset_mix = len({state_shape_signature(state) for state in states}) > 1 + if row_groups is None: + audios, next_states = self._decode_rows(bucket, states, features, ragged=onset_ready or offset_mix) else: - states = [item.previous.token2wav for item in bucket if item.previous is not None] - token_lengths = {int(item.tokens.numel()) for item in bucket} - last_chunk_values = {item.last_chunk for item in bucket} - if len(token_lengths) > 1 or len(last_chunk_values) > 1: - audios, next_states = self.backend.decode_ragged_batch( - [item.tokens for item in bucket], - features, - states, - last_chunks=[item.last_chunk for item in bucket], - ) - else: - tokens = torch.stack([item.tokens for item in bucket], dim=0) - audios, next_states = self.backend.decode_batch( - tokens, - features, - states, - last_chunk=bucket[0].last_chunk, - ) + audios, next_states = [None] * batch_size, [None] * batch_size + for rows in row_groups: + group_states = [states[row] for row in rows] + mixed = row_offsets and len({state_shape_signature(state) for state in group_states}) > 1 + decoded = self._decode_rows([bucket[row] for row in rows], group_states, features, ragged=mixed) + for row, audio, next_state in zip(rows, *decoded, strict=True): + audios[row], next_states[row] = audio, next_state except Exception as exc: self._trim_runtime_prompts() if isinstance(exc, RuntimeError) and str(exc).startswith("MiniCPMO45Code2WavBatchError "): @@ -1055,6 +1199,12 @@ class MiniCPMO45Code2Wav(nn.Module): token2wav_path = Path(self.model_path) / "assets" / "token2wav" if not token2wav_path.is_dir(): raise FileNotFoundError(f"MiniCPM-o Code2Wav assets not found: {token2wav_path}") + # Token2wav runs in fp32, so without TF32 every flow-DiT GEMM runs on + # SIMT cores. Ordinary TF32 uses SM80+ Tensor Cores for dense GEMMs and + # Triton CFM attention dots. HiFT stays IEEE FP32. + tf32_mode = _tf32_mode(extra) if current_omni_platform.is_cuda() else "off" + if tf32_mode != "off": + logger.info("MiniCPM-o Code2Wav: ordinary TF32 enabled for CFM DiT GEMMs") use_float16 = bool(extra.get("token2wav_float16", False)) previous_dtype = torch.get_default_dtype() try: @@ -1066,6 +1216,9 @@ class MiniCPMO45Code2Wav(nn.Module): str(token2wav_path), float16=use_float16, n_timesteps=int(extra.get("token2wav_n_timesteps", 10)), + # The batched backend builds its estimator caches itself and + # never runs the upstream chunk path that owns these buffers. + drop_upstream_chunk_att_buffers=extra.get("cfm_drop_upstream_att_buffer", False) is True, ) finally: torch.set_default_dtype(previous_dtype) @@ -1099,4 +1252,53 @@ class MiniCPMO45Code2Wav(nn.Module): cfm_graph_config=self._cfm_graph_config, bfloat16_attention_cache=bool(extra.get("code2wav_bfloat16_attention_cache", False)), setup_cache_size=self._setup_cache_size, + cfm_tf32=tf32_mode != "off", + encoder_graph_config=self._encoder_graph_config, ) + # Captured by the first forward (the warmup run): under vLLM's weight load a graph held GiBs. + self._precapture_pending = bool(extra.get("cfm_graph_precapture", True)) + + def _release_precapture_cache(self) -> None: + """Return startup's unused cached allocator blocks to the driver (``cfm_precapture_empty_cache``). + + Live tensors, the slot pool and graph pools stay, so no later solve changes. + """ + if torch.cuda.is_available() and torch.cuda.is_initialized(): + device = torch.device("cuda", torch.accelerator.current_device_index()) + torch.accelerator.synchronize(device) + before = _memory_snapshot(device) + torch.accelerator.empty_cache() + delta = _format_memory_delta(before, _memory_snapshot(device)) + logger.info("MiniCPM-o Code2Wav: emptied the allocator cache after precapture%s", delta) + + def _precapture_default_prompt(self) -> None: + """Capture the HiFT, default-voice Whole-Euler and flow-encoder graphs before serving. + + Otherwise each graph is captured by the first chunk that needs it, + stalling every stream for 1-2 s per graph. + """ + self._precapture_pending = False + backend = self.backend + assert backend is not None + started = time.perf_counter() + + def precapture(name: str, capture: Callable[[], int]) -> bool: + try: + captured = capture() + except Exception: + logger.warning("MiniCPM-o Code2Wav: %s precapture failed; graphs capture lazily", name, exc_info=True) + return False + if captured: + elapsed = time.perf_counter() - started + logger.info("MiniCPM-o Code2Wav: precaptured %d %s CUDA Graphs in %.1f s", captured, name, elapsed) + return True + + precapture("HiFT", backend.precapture_hift) + prompt_wav, prompt_cache_id = self._normalized_default_prompt() + try: + features = backend.prepare_prompt(prompt_cache_id, prompt_wav) + except Exception: + logger.warning("MiniCPM-o Code2Wav: default prompt failed; graphs capture lazily", exc_info=True) + return + if precapture("Whole-Euler", partial(backend.precapture_whole_euler, features)): + precapture("flow encoder", partial(backend.precapture_flow_encoder, features)) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py index 9451c2560..ce531e61e 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_token2wav.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """MiniCPM-o 4.5 Token2wav adapter over in-tree ``StepAudio2Token2WavCore``. ``minicpmo_4_5_omni_tts`` historically depended on the external @@ -64,6 +64,7 @@ class MiniCPMO45Token2wav: float16: bool = False, n_timesteps: int = 10, device: str | torch.device | None = None, + drop_upstream_chunk_att_buffers: bool = False, ): self.float16 = float16 self.n_timesteps = n_timesteps @@ -73,6 +74,7 @@ class MiniCPMO45Token2wav: float16=float16, device=self.device, n_timesteps=n_timesteps, + drop_upstream_chunk_att_buffers=drop_upstream_chunk_att_buffers, ) # Eager-load so construction failures surface at init time (same as # the external Token2wav package), not on the first request. diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py b/vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py new file mode 100644 index 000000000..0593dfeb8 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/ops/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from .fused_qkv_layer_norm import qkv_head_layer_norm +from .fused_residual_layer_norm import residual_layer_norm + +__all__ = [ + "qkv_head_layer_norm", + "residual_layer_norm", +] diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py b/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py new file mode 100644 index 000000000..e7a21c497 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_qkv_layer_norm.py @@ -0,0 +1,173 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# ruff: noqa: N803 + +"""Fused packed-QKV split + per-head LayerNorm of q/k into the CFM KV cache.""" + +import torch +from vllm.triton_utils import tl, triton + +from .fused_residual_layer_norm import _use_triton + + +@triton.jit +def _head_layer_norm(x, mask, dims, weight_ptr, bias_ptr, eps, HEAD_DIM: tl.constexpr, HAS_AFFINE: tl.constexpr): + """LayerNorm over the last axis of ``(heads, dims)`` in fp32, with an optional per-dim affine.""" + x = x.to(tl.float32) + mean = tl.sum(x, axis=1) / HEAD_DIM + centered = tl.where(mask, x - mean[:, None], 0.0) + variance = tl.sum(centered * centered, axis=1) / HEAD_DIM + x = centered * tl.rsqrt(variance + eps)[:, None] + if HAS_AFFINE: + x = x * tl.load(weight_ptr + dims, mask=dims < HEAD_DIM, other=1.0).to(tl.float32)[None, :] + x = x + tl.load(bias_ptr + dims, mask=dims < HEAD_DIM, other=0.0).to(tl.float32)[None, :] + return x + + +@triton.jit +def _qkv_head_layer_norm_kernel( + q_out_ptr, + kv_ptr, + qkv_ptr, + q_weight_ptr, + q_bias_ptr, + k_weight_ptr, + k_bias_ptr, + rows_ptr, + positions_ptr, + frames, + stride_qkv_n, + stride_qkv_t, + stride_qn, + stride_qh, + stride_qt, + stride_kvn, + stride_kvh, + stride_kvt, + q_eps, + k_eps, + HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_D: tl.constexpr, + HAS_Q_AFFINE: tl.constexpr, + HAS_K_AFFINE: tl.constexpr, + HAS_ROWS: tl.constexpr, + HAS_POSITIONS: tl.constexpr, +): + row = tl.program_id(0) + n = row // frames + t = row % frames + heads = tl.arange(0, BLOCK_H) + dims = tl.arange(0, BLOCK_D) + mask = (heads[:, None] < HEADS) & (dims[None, :] < HEAD_DIM) + inner = HEADS * HEAD_DIM + src = qkv_ptr + n * stride_qkv_n + t * stride_qkv_t + heads[:, None] * HEAD_DIM + dims[None, :] + + q = tl.load(src, mask=mask, other=0.0) + q = _head_layer_norm(q, mask, dims, q_weight_ptr, q_bias_ptr, q_eps, HEAD_DIM, HAS_Q_AFFINE) + tl.store(q_out_ptr + n * stride_qn + heads[:, None] * stride_qh + t * stride_qt + dims[None, :], q, mask=mask) + + k = tl.load(src + inner, mask=mask, other=0.0) + k = _head_layer_norm(k, mask, dims, k_weight_ptr, k_bias_ptr, k_eps, HEAD_DIM, HAS_K_AFFINE) + kv_row = n + if HAS_ROWS: + kv_row = tl.load(rows_ptr + n).to(tl.int64) + frame = t + if HAS_POSITIONS: + frame = tl.load(positions_ptr + row).to(tl.int64) + mask = mask & (frame >= 0) + dst = kv_ptr + kv_row * stride_kvn + heads[:, None] * stride_kvh + frame * stride_kvt + dims[None, :] + tl.store(dst, k, mask=mask) + + v = tl.load(src + 2 * inner, mask=mask, other=0.0) + tl.store(dst + HEAD_DIM, v, mask=mask) + + +def qkv_head_layer_norm( + qkv: torch.Tensor, + kv: torch.Tensor, + *, + num_heads: int, + head_dim: int, + q_norm: torch.nn.LayerNorm, + k_norm: torch.nn.LayerNorm, + rows: torch.Tensor | None = None, + positions: torch.Tensor | None = None, +) -> torch.Tensor: + """Split packed ``qkv``, LayerNorm q/k per head, write q and interleaved ``kv``.""" + batch, frames, width = (int(dim) for dim in qkv.shape) + inner = num_heads * head_dim + if width != 3 * inner or qkv.stride(2) != 1: + raise ValueError(f"qkv_head_layer_norm: qkv must be (N, T, 3*{inner}) with contiguous channels") + if ( + kv.dim() != 4 + or (rows is None and int(kv.shape[0]) != batch) + or int(kv.shape[1]) != num_heads + or (positions is None and int(kv.shape[2]) < frames) + or int(kv.shape[3]) != 2 * head_dim + or kv.stride(3) != 1 + ): + raise ValueError(f"qkv_head_layer_norm: kv must be (N, H, >= T, 2*D), got {tuple(kv.shape)}") + if rows is not None and tuple(rows.shape) != (batch,): + raise ValueError(f"qkv_head_layer_norm: rows must be ({batch},), got {tuple(rows.shape)}") + if positions is not None and (tuple(positions.shape) != (batch, frames) or not positions.is_contiguous()): + raise ValueError(f"qkv_head_layer_norm: positions must be contiguous ({batch}, {frames})") + q_out = torch.empty((batch, num_heads, frames, head_dim), device=qkv.device, dtype=qkv.dtype) + + if not _use_triton(qkv): + heads = qkv.view(batch, frames, 3, num_heads, head_dim) + q_out.copy_(q_norm(heads[:, :, 0].transpose(1, 2))) + keys = k_norm(heads[:, :, 1].transpose(1, 2)) + values = heads[:, :, 2].transpose(1, 2) + if rows is None and positions is None: + kv[:, :, :frames, :head_dim].copy_(keys) + kv[:, :, :frames, head_dim:].copy_(values) + return q_out + pool_rows = rows.tolist() if rows is not None else range(batch) + for n, pool_row in enumerate(pool_rows): + frame_index = positions[n] if positions is not None else torch.arange(frames, device=kv.device) + written = frame_index >= 0 + destination = frame_index[written].to(kv.device) + kv[pool_row, :, destination, :head_dim] = keys[n][:, written] + kv[pool_row, :, destination, head_dim:] = values[n][:, written] + return q_out + if batch * frames == 0: + return q_out + q_weight, q_bias, q_eps = q_norm.weight, q_norm.bias, float(q_norm.eps) + k_weight, k_bias, k_eps = k_norm.weight, k_norm.bias, float(k_norm.eps) + if (q_weight is None) != (q_bias is None) or (k_weight is None) != (k_bias is None): + raise ValueError("qkv_head_layer_norm: LayerNorm affine needs both weight and bias") + _qkv_head_layer_norm_kernel[(batch * frames,)]( + q_out, + kv, + qkv, + q_weight if q_weight is not None else qkv, + q_bias if q_bias is not None else qkv, + k_weight if k_weight is not None else qkv, + k_bias if k_bias is not None else qkv, + rows if rows is not None else qkv, + positions if positions is not None else qkv, + frames, + qkv.stride(0), + qkv.stride(1), + q_out.stride(0), + q_out.stride(1), + q_out.stride(2), + kv.stride(0), + kv.stride(1), + kv.stride(2), + q_eps, + k_eps, + HEADS=num_heads, + HEAD_DIM=head_dim, + BLOCK_H=triton.next_power_of_2(num_heads), + BLOCK_D=triton.next_power_of_2(head_dim), + HAS_Q_AFFINE=q_weight is not None, + HAS_K_AFFINE=k_weight is not None, + HAS_ROWS=rows is not None, + HAS_POSITIONS=positions is not None, + num_warps=4, + ) + return q_out diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py b/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py new file mode 100644 index 000000000..bc35f9152 --- /dev/null +++ b/vllm_omni/model_executor/models/minicpmo_4_5/ops/fused_residual_layer_norm.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +# ruff: noqa: N803 + +"""Fused ``act(layer_norm(residual + gate * y) * weight + bias)`` (optional conv taps).""" + +import torch +import torch.nn.functional as F +from vllm.triton_utils import HAS_TRITON, tl, triton + +_ACTIVATIONS = {None: 0, "mish": 1} + + +def _use_triton(tensor: torch.Tensor) -> bool: + """Triton on CUDA tensors; everything else runs the native path.""" + return HAS_TRITON and tensor.is_cuda + + +@triton.jit +def _residual_layer_norm_kernel( + out_ptr, + residual_out_ptr, + x_ptr, + y_ptr, + gate_ptr, + y_bias_ptr, + weight_ptr, + bias_ptr, + frames, + channels, + stride_on, + stride_out_t, + stride_rn, + stride_rt, + stride_xn, + stride_xt, + stride_yn, + stride_yt, + eps, + HAS_X: tl.constexpr, + HAS_Y: tl.constexpr, + TAPS: tl.constexpr, + HAS_GATE: tl.constexpr, + HAS_Y_BIAS: tl.constexpr, + HAS_WEIGHT: tl.constexpr, + HAS_BIAS: tl.constexpr, + STORE_RESIDUAL: tl.constexpr, + ACTIVATION: tl.constexpr, + BLOCK_C: tl.constexpr, +): + row = tl.program_id(0) + n = row // frames + t = row % frames + cols = tl.arange(0, BLOCK_C) + mask = cols < channels + + value = tl.zeros((BLOCK_C,), dtype=tl.float32) + if HAS_Y: + y = tl.zeros((BLOCK_C,), dtype=tl.float32) + for k in tl.static_range(TAPS): + y += tl.load( + y_ptr + n * stride_yn + (t + k) * stride_yt + k * channels + cols, + mask=mask, + other=0.0, + ).to(tl.float32) + if HAS_Y_BIAS: + y += tl.load(y_bias_ptr + cols, mask=mask, other=0.0).to(tl.float32) + if HAS_GATE: + y = tl.load(gate_ptr + cols, mask=mask, other=0.0).to(tl.float32) * y + value = y + if HAS_X: + value += tl.load(x_ptr + n * stride_xn + t * stride_xt + cols, mask=mask, other=0.0).to(tl.float32) + if STORE_RESIDUAL: + tl.store(residual_out_ptr + n * stride_rn + t * stride_rt + cols, value, mask=mask) + + mean = tl.sum(value, axis=0) / channels + centered = tl.where(mask, value - mean, 0.0) + variance = tl.sum(centered * centered, axis=0) / channels + normalized = centered * tl.rsqrt(variance + eps) + if HAS_WEIGHT: + normalized = normalized * tl.load(weight_ptr + cols, mask=mask, other=1.0).to(tl.float32) + if HAS_BIAS: + normalized = normalized + tl.load(bias_ptr + cols, mask=mask, other=0.0).to(tl.float32) + if ACTIVATION == 1: + e = tl.exp(tl.minimum(normalized, 20.0)) + numerator = e * (e + 2.0) + factor = tl.where(normalized > 20.0, 1.0, numerator / (numerator + 2.0)) + normalized = normalized * factor + tl.store(out_ptr + n * stride_on + t * stride_out_t + cols, normalized, mask=mask) + + +def _check_rows(name: str, tensor: torch.Tensor, rows: int, frames: int, width: int) -> None: + if tensor.dim() != 3 or int(tensor.shape[0]) != rows or int(tensor.shape[1]) < frames: + raise ValueError(f"residual_layer_norm: {name} must be (N, >= T, C), got {tuple(tensor.shape)}") + if int(tensor.shape[2]) != width or (width > 1 and tensor.stride(2) != 1): + raise ValueError(f"residual_layer_norm: {name} needs {width} contiguous channels, got {tuple(tensor.shape)}") + + +def _native(residual, y, gate, weight, bias, eps, taps, y_bias, activation, residual_out, out, frames, channels): + value = None + if y is not None: + value = sum(y[:, k : k + frames, k * channels : (k + 1) * channels] for k in range(taps)) + if y_bias is not None: + value = value + y_bias + if gate is not None: + value = gate * value + if residual is not None: + value = residual[:, :frames] if value is None else residual[:, :frames] + value + if residual_out is not None: + residual_out.copy_(value) + result = F.layer_norm(value, (channels,), None, None, eps) + if weight is not None: + result = result * weight + if bias is not None: + result = result + bias + if activation == "mish": + result = F.mish(result) + out.copy_(result) + return out + + +def residual_layer_norm( + residual: torch.Tensor | None, + y: torch.Tensor | None = None, + *, + gate: torch.Tensor | None = None, + weight: torch.Tensor | None = None, + bias: torch.Tensor | None = None, + eps: float = 1e-6, + taps: int = 1, + y_bias: torch.Tensor | None = None, + activation: str | None = None, + residual_out: torch.Tensor | None = None, + out: torch.Tensor | None = None, +) -> torch.Tensor: + """``act(layer_norm(residual + gate * y) * weight + bias)`` in one pass per row.""" + if activation not in _ACTIVATIONS: + raise ValueError(f"residual_layer_norm: unsupported activation {activation!r}") + if residual is None and y is None: + raise ValueError("residual_layer_norm needs a residual or a y") + reference = residual if residual is not None else y + rows = int(reference.shape[0]) + if residual is not None: + frames, channels = int(residual.shape[1]), int(residual.shape[2]) + _check_rows("residual", residual, rows, frames, channels) + else: + channels = int(y.shape[2]) // taps + frames = int(y.shape[1]) - (taps - 1) + if y is not None: + _check_rows("y", y, rows, frames + taps - 1, taps * channels) + if residual_out is not None: + _check_rows("residual_out", residual_out, rows, frames, channels) + if out is None: + out = torch.empty((rows, frames, channels), device=reference.device, dtype=reference.dtype) + _check_rows("out", out, rows, frames, channels) + + if not _use_triton(reference): + return _native( + residual, y, gate, weight, bias, eps, taps, y_bias, activation, residual_out, out, frames, channels + ) + if rows * frames == 0: + return out + + def strides(tensor: torch.Tensor | None) -> tuple[int, int]: + return (tensor.stride(0), tensor.stride(1)) if tensor is not None else (0, 0) + + _residual_layer_norm_kernel[(rows * frames,)]( + out, + *(t if t is not None else reference for t in (residual_out, residual, y, gate, y_bias, weight, bias)), + frames, + channels, + *strides(out), + *strides(residual_out), + *strides(residual), + *strides(y), + float(eps), + HAS_X=residual is not None, + HAS_Y=y is not None, + TAPS=int(taps), + HAS_GATE=gate is not None, + HAS_Y_BIAS=y_bias is not None, + HAS_WEIGHT=weight is not None, + HAS_BIAS=bias is not None, + STORE_RESIDUAL=residual_out is not None, + ACTIVATION=_ACTIVATIONS[activation], + BLOCK_C=triton.next_power_of_2(channels), + num_warps=4 if channels <= 1024 else 8, + ) + return out diff --git a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py index 8e3c56b32..9523d3412 100644 --- a/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py +++ b/vllm_omni/model_executor/models/step_audio2/step_audio2_token2wav.py @@ -95,6 +95,22 @@ class _ConvRNNF0Predictor(nn.Module): return self.classifier(self.condnet(x).transpose(1, 2)).squeeze(-1).abs() +def drop_upstream_chunk_att_buffers(flow: nn.Module) -> int: + """Drop upstream cosyvoice2's unused chunk-streaming ``att_cache_buffer`` tensors (~2.1 GiB); returns bytes freed. + + Only ``flow.setup_cache`` / ``flow.inference_chunk`` read them, and then + fail instead of computing on a stand-in. + """ + released = 0 + decoder = getattr(flow, "decoder", None) + for module in (decoder, getattr(decoder, "estimator", None)): + buffer = module._buffers.get("att_cache_buffer") if isinstance(module, nn.Module) else None + if isinstance(buffer, torch.Tensor): + released += buffer.numel() * buffer.element_size() + del module._buffers["att_cache_buffer"] + return released + + def _build_hift() -> HiFTGenerator: return HiFTGenerator( sampling_rate=24000, @@ -115,9 +131,13 @@ class StepAudio2Token2WavCore(nn.Module): float16: bool = False, device: str = "cuda", n_timesteps: int = DEFAULT_TOKEN2WAV_CONFIG.n_timesteps, + drop_upstream_chunk_att_buffers: bool = False, ): super().__init__() self.model_path = model_path + # A backend that never runs ``flow.setup_cache`` / ``flow.inference_chunk`` + # (MiniCPM-o's ``BatchedToken2Wav``) can drop their 2.1 GiB of buffers. + self._drop_chunk_att_buffers = bool(drop_upstream_chunk_att_buffers) self.float16 = float16 self.device = torch.device(device) self.n_timesteps = n_timesteps @@ -176,6 +196,10 @@ class StepAudio2Token2WavCore(nn.Module): with open(f"{self.model_path}/flow.yaml") as f: configs = load_hyperpyyaml(f) self._flow = configs["flow"] + if self._drop_chunk_att_buffers: + # Before ``.to(self.device)``: the buffers never reach the device. + released = drop_upstream_chunk_att_buffers(self._flow) + logger.info("Dropped the upstream flow chunk attention buffers (%.1f MiB)", released / 2**20) if self.float16: self._flow.half() self._flow.load_state_dict(