[Model] Optimize CosyVoice3 packed streaming with bounded HiFT graphs

Signed-off-by: Sy03 <1370724210@qq.com>

tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py

@@ -108,6 +108,26 @@ def test_packed_stream_noise_growth_preserves_prefix_and_global_rng():
     assert model._stream_position_noise(0, like).shape == (1, 80, 0)
 
 
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_packed_stream_launch_uses_query_bound(monkeypatch):
+    """A long KV prefix must not inflate the chunk-causal query launch grid."""
+    from vllm_omni.model_executor.models.cosyvoice3.code2wav_core import packed_dit
+
+    bound = []
+
+    def forward(estimator, x, *args, max_seqlen_q, **kwargs):
+        bound.append(max_seqlen_q)
+        return x
+
+    monkeypatch.setattr(packed_dit, "forward_packed_tensor_geometry", forward)
+    estimator = object.__new__(packed_dit.PackedDiT)
+    rows = packed_dit.pack_rows([50], torch.device("cpu"))
+    attention = packed_dit.RaggedRowAttention(rows, heads=16, head_dim=64, chunk_size=50, width=4096)
+    hidden = torch.empty(1, 50, 80)
+    assert estimator.forward_full(hidden, None, None, None, None, None, attention) is hidden
+    assert bound == [50]
+
+
 @hardware_test(res={"cuda": "H100"}, num_cards=1)
 @torch.inference_mode()
 def test_packed_stream_attention_matches_chunk_causal_reference():

tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py

@@ -0,0 +1,50 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""HiFT graphs consume live inputs and return independently owned waveforms."""
+
+import pytest
+import torch
+
+from tests.helpers.mark import hardware_test
+from tests.model_executor.models.cosyvoice3.test_cosyvoice3_incremental_hift import _make_hift
+
+pytestmark = [pytest.mark.core_model]
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.parametrize("finalize", [False, True])
+@torch.inference_mode()
+def test_hift_graph_preserves_current_and_previous_request(finalize):
+    hift = _make_hift().cuda().eval()
+    hift.remove_weight_norm()
+    hift.f0_predictor.remove_weight_norm()
+    hift.enable_decode_graphs()
+    generator = torch.Generator(device="cuda").manual_seed(17)
+    first = None
+    first_saved = None
+    for _ in range(5):
+        mel = torch.randn(1, 80, 30, device="cuda", generator=generator)
+        source = torch.randn(1, 1, 30 * 480, device="cuda", generator=generator)
+        expected = hift._decode_eager(mel, source, finalize=finalize)
+        actual = hift.decode(mel, source, finalize=finalize)
+        torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6)
+        if first is None and hift._decode_graphs.graphs:
+            first, first_saved = actual, actual.clone()
+        elif first is not None:
+            assert torch.equal(first, first_saved)
+    assert len(hift._decode_graphs.graphs) == 1
+
+
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@torch.inference_mode()
+def test_hift_graph_capacity_keeps_eager_fallback():
+    hift = _make_hift().cuda().eval()
+    hift.remove_weight_norm()
+    hift.enable_decode_graphs()
+    hift._decode_graphs.max_graphs = 0
+    mel = torch.randn(1, 80, 30, device="cuda")
+    source = torch.randn(1, 1, 30 * 480, device="cuda")
+    actual = hift.decode(mel, source, finalize=True)
+    expected = hift._decode_eager(mel, source, finalize=True)
+    assert torch.equal(actual, expected)
+    assert not hift._decode_graphs.graphs