[Performance][MiniCPM-o] Whole-Euler CUDA graphs for Code2Wav with a shared attention arena (#8007)

tests/model_executor/models/minicpmo_4_5/test_cfm_graph_bucketing.py

@@ -16,9 +16,15 @@ import torch
 
 from vllm_omni.model_executor.models.minicpmo_4_5.batched_token2wav import (
     _cfm_pad_frames,
-    _zero_padded_cnn_cache,
     _zero_padded_frames,
 )
+from vllm_omni.model_executor.models.minicpmo_4_5.cuda_graph_wrapper import (
+    _att_keep_ranges,
+    _build_capture_mask,
+    _capture_query_width,
+    _whole_euler_att_segments,
+    _zero_padded_cnn_cache,
+)
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -139,6 +145,75 @@ def test_varied_chunk_lengths_collapse_onto_few_widths():
     assert len(padded) < len(varied)
 
 
+def test_capture_query_width_collapses_decode_grid_onto_one_bucket():
+    """Whole-Euler capture bucket is the #7416 pad helper with a 64-wide decode grid."""
+    assert _capture_query_width(16, 0) == 16
+    assert _capture_query_width(16, 1) == 16
+    assert {_capture_query_width(w, 64) for w in (16, 32, 48, 50, 64)} == {64}
+    assert _capture_query_width(304, 64) == 320
+    assert _capture_query_width(320, 64) == 320
+
+
+def test_att_keep_ranges_match_the_streaming_trim():
+    """Past ``prompt_len + 100`` frames the cache keeps its first ``prompt_len`` and last 100."""
+    assert _att_keep_ranges(368, (300, 100)) == [(0, 368)]
+    assert _att_keep_ranges(400, (300, 100)) == [(0, 400)]
+    assert _att_keep_ranges(464, (300, 100)) == [(0, 300), (364, 100)]
+    assert _att_keep_ranges(464, None) == [(0, 464)]
+
+
+def test_whole_euler_att_segments_skip_capture_padding():
+    # Steady chunk: capture width == chunk width, so only the trim applies.
+    assert _whole_euler_att_segments(mel_width=64, query_cap=64, offset=400, keep=(300, 100)) == [
+        (0, 300),
+        (364, 100),
+    ]
+    # A 32-wide chunk on the 64 capture: the graph wrote [current 64 | cache 400];
+    # frames 32:64 are capture padding and the cache starts at 64.
+    assert _whole_euler_att_segments(mel_width=32, query_cap=64, offset=400, keep=None) == [(0, 32), (64, 400)]
+    assert _whole_euler_att_segments(mel_width=32, query_cap=64, offset=400, keep=(300, 100)) == [
+        (0, 32),
+        (64, 268),
+        (364, 100),
+    ]
+    # First chunk: no cache, only the tail of the capture width is dropped.
+    assert _whole_euler_att_segments(mel_width=304, query_cap=320, offset=0, keep=None) == [(0, 304)]
+
+
+def test_capture_mask_moves_the_cache_block_past_capture_padding():
+    width, cap, offset, valid = 32, 64, 8, 30
+    caller = torch.ones(2, width, width + offset, dtype=torch.bool)
+    caller[:, :, valid:width] = False
+    mask = _build_capture_mask(
+        attn_mask=caller,
+        batch_size=1,
+        query_cap=cap,
+        offset=offset,
+        mel_width=width,
+        mel_frames=valid,
+        device=torch.device("cpu"),
+    )
+    assert mask.shape == (2, cap, cap + offset)
+    assert mask[:, :, :valid].all()
+    assert not mask[:, :, valid:cap].any()
+    assert mask[:, :, cap:].all()
+    # Capture-padding query rows are never empty, so SDPA cannot produce NaN there.
+    assert mask.any(dim=-1).all()
+
+    implicit = _build_capture_mask(
+        attn_mask=None,
+        batch_size=1,
+        query_cap=cap,
+        offset=offset,
+        mel_width=width,
+        mel_frames=width,
+        device=torch.device("cpu"),
+    )
+    assert implicit[:, :, :width].all()
+    assert not implicit[:, :, width:cap].any()
+    assert implicit[:, :, cap:].all()
+
+
 def test_zero_padded_frames_survives_an_integration_step():
     """The padded region must stay zero after ``x = x + dt * velocity``.
 

tests/model_executor/models/minicpmo_4_5/test_code2wav_batching.py

@@ -172,6 +172,19 @@ def _config(minimum: int = 1, initial: int = 0, *, runtime_prompt_cache_size: in
     )
 
 
+@pytest.mark.parametrize(
+    ("extra", "max_num_seqs", "micro"), [({}, 6, 6), ({}, 64, 16), ({"micro_batch_size": 2}, 6, 2)]
+)
+def test_whole_euler_micro_batch_defaults_to_stage_concurrency(extra, max_num_seqs, micro):
+    config = _config()
+    config.model_config.stage_connector_config["extra"].update(extra)
+    config.scheduler_config = SimpleNamespace(max_num_seqs=max_num_seqs)
+
+    model = MiniCPMO45Code2Wav(vllm_config=config)
+
+    assert model._cfm_graph_config["micro_batch_size"] == micro
+
+
 def _model(initial: int = 0, minimum: int = 1, *, runtime_prompt_cache_size: int = 4, setup_cache_size: int = 1):
     token2wav = _FakeToken2Wav()
     backend = BatchedToken2Wav(token2wav, setup_cache_size=setup_cache_size)
@@ -213,17 +226,43 @@ def test_setup_cache_reuses_read_only_state_for_the_same_exact_batch():
         torch.testing.assert_close(cache[name], expected)
 
 
-def test_setup_cache_keys_batch_size_and_evicts_least_recent_entry():
+def test_setup_solves_the_prompt_once_for_every_batch_size():
     token2wav = _FakeToken2Wav()
     adapter = BatchedToken2Wav(token2wav, setup_cache_size=1)
     prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
 
-    adapter.setup_batch(prompt, 1)
-    adapter.setup_batch(prompt, 2)
-    adapter.setup_batch(prompt, 1)
+    single = adapter.setup_batch(prompt, 1)
+    pair = adapter.setup_batch(prompt, 2)
 
-    assert token2wav.flow.encoder.calls == [1, 2, 1]
-    assert list(adapter._setup_cache) == [(prompt.cache_key, 1, 0)]
+    assert token2wav.flow.encoder.calls == [1]
+    assert pair == [single[0], single[0]]
+    assert list(adapter._setup_cache) == [(prompt.cache_key, 0)]
+
+
+def test_setup_cache_evicts_least_recent_prompt():
+    token2wav = _FakeToken2Wav()
+    adapter = BatchedToken2Wav(token2wav, setup_cache_size=1)
+    first = adapter.prepare_prompt("first", "/fake/first.wav")
+    second = adapter.prepare_prompt("second", "/fake/second.wav")
+
+    adapter.setup_batch(first, 1)
+    adapter.setup_batch(second, 2)
+    adapter.setup_batch(first, 1)
+
+    assert token2wav.flow.encoder.calls == [1, 1, 1]
+    assert list(adapter._setup_cache) == [(first.cache_key, 0)]
+
+
+def test_setup_without_cache_still_solves_one_shared_row():
+    token2wav = _FakeToken2Wav()
+    adapter = BatchedToken2Wav(token2wav, setup_cache_size=0)
+    prompt = adapter.prepare_prompt("shared", "/fake/prompt.wav")
+
+    states = adapter.setup_batch(prompt, 3)
+
+    assert token2wav.flow.encoder.calls == [1]
+    assert states[0] is states[1] is states[2]
+    assert not adapter._setup_cache
 
 
 def test_evict_prompt_clears_features_and_setup_state():
@@ -481,9 +520,12 @@ def test_adapter_runs_true_batch_cfg_and_splits_request_caches():
     )
 
     assert token2wav.prompt_calls == 1
-    assert token2wav.flow.encoder.calls == [2, 2]
-    assert token2wav.flow.decoder.estimator.cfg_batches == [4, 4, 4, 4]
-    assert all(order == [1.0, 1.0, 0.0, 0.0] for order in token2wav.flow.decoder.estimator.speaker_order)
+    # Setup solves the shared prompt once; the chunk runs both requests together.
+    assert token2wav.flow.encoder.calls == [1, 2]
+    assert token2wav.flow.decoder.estimator.cfg_batches == [2, 2, 4, 4]
+    speaker_order = token2wav.flow.decoder.estimator.speaker_order
+    assert all(order == [1.0, 0.0] for order in speaker_order[:2])
+    assert all(order == [1.0, 1.0, 0.0, 0.0] for order in speaker_order[2:])
     assert token2wav.hift.calls == [2]
     assert len(audios) == 2
     cache0 = states[0].flow_cache["estimator_cnn_cache"]
@@ -706,6 +748,40 @@ def test_decode_cfm_enters_platform_sdpa_context(monkeypatch):
     assert entered == ["enter", "exit"]
 
 
+def test_eager_cfm_reads_the_step_sizes_to_the_host_once(monkeypatch):
+    """The eager Euler loop must not sync the stream for ``dt`` on every decode."""
+    adapter = BatchedToken2Wav(_FakeToken2Wav())
+    reads: list[str] = []
+
+    def recording(name):
+        original = getattr(torch.Tensor, name)
+
+        def read(self, *args, **kwargs):
+            # The fake estimator records its inputs with ``tolist``; skip those.
+            if sys._getframe(1).f_code.co_filename != __file__:
+                reads.append(name)
+            return original(self, *args, **kwargs)
+
+        return read
+
+    for name in ("item", "tolist"):
+        monkeypatch.setattr(torch.Tensor, name, recording(name))
+
+    def decode():
+        adapter._decode_cfm(
+            torch.ones((1, 1, 2)),
+            torch.ones((1, 1)),
+            torch.zeros((1, 1, 2)),
+            cnn_cache=None,
+            att_cache=None,
+        )
+
+    decode()
+    assert reads == ["tolist"]
+    decode()
+    assert reads == ["tolist"]
+
+
 def test_ragged_decode_bypasses_exact_shape_accelerators():
     adapter = BatchedToken2Wav(_FakeToken2Wav())
     _enable_fake_ragged_kernel(adapter)
@@ -963,6 +1039,44 @@ def test_estimator_cache_stack_split_round_trip_preserves_cfg_rows():
         )
 
 
+@pytest.mark.parametrize("attention_dtype", [torch.float32, torch.bfloat16])
+def test_split_flow_cache_trims_each_request_like_the_batch_trim(attention_dtype):
+    """Trimming per request while splitting equals trimming the stacked batch first."""
+    adapter = BatchedToken2Wav(_FakeToken2Wav(), bfloat16_attention_cache=attention_dtype == torch.bfloat16)
+    batch_size, prompt_len, frames = 3, 5, 12
+    torch.manual_seed(0)
+    stacked = {
+        "conformer_cnn_cache": torch.randn(batch_size, 2, 3),
+        "conformer_att_cache": torch.randn(2, batch_size, 2, 4, 6),
+        "estimator_cnn_cache": torch.randn(4, 2, 2 * batch_size, 3, 2),
+        "estimator_att_cache": torch.randn(4, 2, 2 * batch_size, 2, frames, 6),
+    }
+
+    restored = adapter._split_flow_cache(stacked, batch_size, estimator_att_keep=(prompt_len, 4))
+
+    att = stacked["estimator_att_cache"]
+    trimmed = torch.cat((att[..., :prompt_len, :], att[..., -4:, :]), dim=4)
+    for row, request in enumerate(restored):
+        expected = torch.cat(
+            (trimmed[:, :, row : row + 1], trimmed[:, :, batch_size + row : batch_size + row + 1]),
+            dim=2,
+        ).to(attention_dtype)
+        assert request["estimator_att_cache"].dtype == attention_dtype
+        torch.testing.assert_close(request["estimator_att_cache"], expected)
+        torch.testing.assert_close(
+            request["estimator_cnn_cache"],
+            torch.cat(
+                (
+                    stacked["estimator_cnn_cache"][:, :, row : row + 1],
+                    stacked["estimator_cnn_cache"][:, :, batch_size + row : batch_size + row + 1],
+                ),
+                dim=2,
+            ),
+        )
+    untrimmed = adapter._split_flow_cache(stacked, batch_size, estimator_att_keep=(frames, 4))
+    assert untrimmed[0]["estimator_att_cache"].shape[4] == frames
+
+
 def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep():
     token2wav = _FakeToken2Wav()
     adapter = BatchedToken2Wav(token2wav, bfloat16_attention_cache=True)
@@ -970,9 +1084,8 @@ def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep()
     states = adapter.setup_batch(prompt, 2)
 
     assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in states)
-    assert (
-        states[0].flow_cache["estimator_att_cache"].data_ptr() != states[1].flow_cache["estimator_att_cache"].data_ptr()
-    )
+    # Both requests start from the one shared prompt state.
+    assert states[0] is states[1]
 
     stacked = adapter._stack_flow_cache(states)
     assert stacked["estimator_att_cache"].dtype == torch.bfloat16
@@ -991,6 +1104,10 @@ def test_bfloat16_estimator_attention_cache_materializes_only_current_timestep()
         last_chunk=False,
     )
     assert all(state.flow_cache["estimator_att_cache"].dtype == torch.bfloat16 for state in next_states)
+    assert (
+        next_states[0].flow_cache["estimator_att_cache"].data_ptr()
+        != next_states[1].flow_cache["estimator_att_cache"].data_ptr()
+    )
     assert token2wav.flow.decoder.estimator.attention_cache_dtypes[-2:] == [torch.float32, torch.float32]
 
 
@@ -1117,9 +1234,17 @@ def test_initial_empty_segment_marker_initializes_stream_without_audio():
     assert "duplex" in model._states
 
 
-@pytest.mark.parametrize("setup_cache_size,expected_calls", [(0, [2, 2, 1]), (1, [2, 1])])
+@pytest.mark.parametrize("setup_cache_size,expected_calls", [(0, [1, 1, 1]), (1, [1])])
 def test_initial_empty_segment_markers_respect_initial_batch_limit(setup_cache_size, expected_calls):
     model, token2wav = _model(initial=2, setup_cache_size=setup_cache_size)
+    setup_batches: list[int] = []
+    setup_batch = model.backend.setup_batch
+
+    def record_setup(features, batch_size):
+        setup_batches.append(batch_size)
+        return setup_batch(features, batch_size)
+
+    model.backend.setup_batch = record_setup
     names = ["a", "b", "c", "d", "e"]
     boundaries = []
     for name in names:
@@ -1135,6 +1260,8 @@ def test_initial_empty_segment_markers_respect_initial_batch_limit(setup_cache_s
 
     output = _forward(model, boundaries)
 
+    assert setup_batches == [2, 2, 1]
+    # One shared prompt solve per setup, or one in all with the setup cache.
     assert token2wav.flow.encoder.calls == expected_calls
     assert [audio.numel() for audio in output.multimodal_outputs["model_outputs"]] == [0] * len(names)
     assert set(model._states) == set(names)
@@ -1631,8 +1758,8 @@ def test_setup_cache_misses_when_cfm_graph_padding_is_disabled(monkeypatch, disa
     assert unpadded[0] is not padded[0]
     assert reused[0] is unpadded[0]
     assert adapter.flow.encoder.calls == [1, 1]
-    assert padded_keys == [(prompt.cache_key, 1, 16)]
-    assert list(adapter._setup_cache) == [(prompt.cache_key, 1, 0)]
+    assert padded_keys == [(prompt.cache_key, 16)]
+    assert list(adapter._setup_cache) == [(prompt.cache_key, 0)]
 
 
 def test_padded_chunk_keeps_the_cross_chunk_caches_on_the_valid_boundary():

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -170,7 +170,7 @@ class TestDeployTopology:
         assert stages[1].yaml_engine_args["custom_process_next_stage_input_func"].endswith(expected_processor)
         assert "hf_overrides" not in stages[1].yaml_engine_args
         if filename == "minicpmo_4_5.yaml":
-            assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [4, 4, 4]
+            assert [stage.yaml_engine_args["max_num_seqs"] for stage in stages] == [16, 16, 16]
             memory_utilizations = [stage.yaml_engine_args["gpu_memory_utilization"] for stage in stages]
             assert memory_utilizations == [
                 0.55,