[Model] Optimize CosyVoice3 packed streaming with bounded HiFT graphs

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

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: 75,
+        EnvironmentVariableCategory.MODEL_SPECIFIC: 77,
         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: 46,
+        ModelEnvironmentVariableDisposition.PROMOTE: 48,
         ModelEnvironmentVariableDisposition.REQUEST_SCOPE: 5,
         ModelEnvironmentVariableDisposition.EXTERNAL: 0,
         ModelEnvironmentVariableDisposition.INTERNALIZE: 17,

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

vllm_omni/config/environment_variable_inventory.py

@@ -148,11 +148,13 @@ _MODEL_PROMOTE = (
     "COSMOS3_SOUND_TOKENIZER_CONFIG_PATH",
     "COSMOS3_SOUND_TOKENIZER_PATH",
     "COSYVOICE3_BATCH_FLOW",
+    "COSYVOICE3_CACHED_ISTFT",
     "COSYVOICE3_ESTIMATOR_ONNX",
     "COSYVOICE3_FLOW_CUDAGRAPH",
     "COSYVOICE3_FLOW_GRAPH",
     "COSYVOICE3_FLOW_GRAPH_MAX_TOTAL",
     "COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS",
+    "COSYVOICE3_HIFT_GRAPH",
     "COSYVOICE3_PACKED_STREAMING",
     "COSYVOICE3_REFERENCE_PREFETCH",
     "COSYVOICE3_REFERENCE_WORKERS",

vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml

@@ -0,0 +1,22 @@
+# Opt-in single-Hopper streaming throughput profile: AR32 / codec32, Flow10.
+# Use --hf-overrides '{"cosyvoice3_sampling_mode": "standard"}' and stream=true.
+# F0 and HiFT stay FP32; only deterministic HiFT decode is graphed.
+base_config: cosyvoice3_packed_streaming.yaml
+stages:
+  - stage_id: 0
+    max_num_seqs: 32
+    default_sampling_params:
+      max_tokens: 2048
+      repetition_penalty: 1.21
+      temperature: 0.7
+      top_k: 20
+      top_p: 0.8
+  - stage_id: 1
+    max_num_seqs: 32
+    env:
+      COSYVOICE3_TRT: "0"
+      COSYVOICE3_BATCH_FLOW: "1"
+      COSYVOICE3_FULL_RESPONSE_OPTIMIZATIONS: "0"
+      COSYVOICE3_PACKED_STREAMING: "1"
+      COSYVOICE3_CACHED_ISTFT: "1"
+      COSYVOICE3_HIFT_GRAPH: "1"

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py

@@ -966,7 +966,20 @@ class CausalHiFTGenerator(HiFTGenerator):
         self.conv_pre_look_right = conv_pre_look_right
         self.f0_predictor = f0_predictor
 
+    def enable_decode_graphs(self) -> None:
+        """Opt into exact-shape CUDA replay after frozen weights are loaded."""
+        from .hift_graph import HiFTDecodeGraphs
+
+        self.enable_cached_istft()
+        self._decode_graphs = HiFTDecodeGraphs(self._decode_eager)
+
     def decode(self, x: torch.Tensor, s: torch.Tensor = torch.zeros(1, 1, 0), finalize: bool = True) -> torch.Tensor:
+        runner = getattr(self, "_decode_graphs", None)
+        if runner is not None:
+            return runner.run(x, s, finalize)
+        return self._decode_eager(x, s, finalize=finalize)
+
+    def _decode_eager(self, x: torch.Tensor, s: torch.Tensor, finalize: bool = True) -> torch.Tensor:
         s_stft_real, s_stft_imag = self._stft(s.squeeze(1))
         if finalize is True:
             x = self.conv_pre(x)

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py

@@ -0,0 +1,85 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
+"""Bounded CUDA graph replay for HiFT's deterministic waveform decoder.
+
+F0, sine/noise generation, phase carry, and per-request streaming state remain
+outside the graph. Each replay consumes current mel and source tensors and
+returns an owned waveform, because multiple requests can reuse one shape.
+"""
+
+from collections import OrderedDict
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class _DecodeGraph:
+    graph: torch.cuda.CUDAGraph
+    mel: torch.Tensor
+    source: torch.Tensor
+    output: torch.Tensor
+
+
+class HiFTDecodeGraphs:
+    """Capture recurrent exact shapes, with at most 32 resident graphs.
+
+    ``decode`` must have warmed its ISTFT envelope before capture and remain
+    deterministic for its mel, source and finalize arguments. Uncommon shapes,
+    CPU execution and nested CUDA capture use that same eager implementation.
+    """
+
+    def __init__(self, decode: Callable[..., torch.Tensor], max_graphs: int = 32) -> None:
+        self.decode = decode
+        self.max_graphs = max_graphs
+        self.graphs: dict[tuple, _DecodeGraph] = {}
+        self.seen: OrderedDict[tuple, int] = OrderedDict()
+        self.pool = None
+
+    @torch.inference_mode()
+    def run(self, mel: torch.Tensor, source: torch.Tensor, finalize: bool) -> torch.Tensor:
+        """Replay a recurring layout without aliasing another request's output."""
+        if not mel.is_cuda or not source.is_cuda or torch.cuda.is_current_stream_capturing():
+            return self.decode(mel, source, finalize=finalize)
+        autocast = torch.is_autocast_enabled("cuda")
+        key = (
+            mel.shape,
+            source.shape,
+            mel.device,
+            mel.dtype,
+            source.dtype,
+            finalize,
+            torch.backends.cudnn.allow_tf32,
+            torch.backends.cuda.matmul.allow_tf32,
+            autocast,
+            torch.get_autocast_dtype("cuda") if autocast else None,
+        )
+        state = self.graphs.get(key)
+        if state is None:
+            if len(self.graphs) >= self.max_graphs:
+                return self.decode(mel, source, finalize=finalize)
+            count = self.seen.pop(key, 0) + 1
+            self.seen[key] = count
+            if len(self.seen) > 128:
+                self.seen.popitem(last=False)
+            if count < 3:
+                return self.decode(mel, source, finalize=finalize)
+            static_mel, static_source = mel.clone(), source.clone()
+            stream = torch.cuda.Stream(device=mel.device)
+            stream.wait_stream(torch.cuda.current_stream(mel.device))
+            with torch.cuda.stream(stream):
+                self.decode(static_mel, static_source, finalize=finalize)
+            torch.cuda.current_stream(mel.device).wait_stream(stream)
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, pool=self.pool):
+                output = self.decode(static_mel, static_source, finalize=finalize)
+            if self.pool is None:
+                self.pool = graph.pool()
+            state = _DecodeGraph(graph, static_mel, static_source, output)
+            self.graphs[key] = state
+        state.mel.copy_(mel)
+        state.source.copy_(source)
+        state.graph.replay()
+        # A later request may replay this shape before the caller consumes it.
+        return state.output.clone()

vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py

@@ -170,7 +170,7 @@ class PackedDiT:
             t,
             rows,
             attention,
-            max_seqlen_q=attention.page_table.shape[1],
+            max_seqlen_q=attention.max_seqlen_q,
             block_modulation=block_modulation,
             final_modulation=final_modulation,
             rope=rope,
@@ -771,7 +771,9 @@ class FlowGraphRunner:
         attention.heads, attention.head_dim = self.heads, self.head_dim
         attention.cache_seqlens, attention.cu_seqlens_q = state.cache_seqlens, state.cu_seqlens_q
         attention.page_table = state.page_table
-        attention.max_seqlen_q = state.page_table.shape[1]
+        # Streaming queries cover one training chunk, even when their KV
+        # prefix and the packed graph allocation contain thousands of frames.
+        attention.max_seqlen_q = estimator.dit.static_chunk_size if key[0] else state.page_table.shape[1]
         mark_packed_compile_metadata(rows, attention)
         mu_cfg = torch.cat((state.mu, torch.zeros_like(state.mu)), dim=1)
         cond_cfg = torch.cat((state.cond, torch.zeros_like(state.cond)), dim=1)

vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py

@@ -769,4 +769,11 @@ class CosyVoice3Code2Wav(nn.Module):
         if f0_folded == 0:
             logger.warning("HiFT F0 predictor had no weight-norm layers to fold; check config drift")
         self.hift.eval()
+        if device.type == "cuda" and cosyvoice3_packed_streaming_enabled():
+            if os.getenv("COSYVOICE3_CACHED_ISTFT", "0") == "1":
+                self.hift.enable_cached_istft()
+                logger.info("CosyVoice3 streaming HiFT: cached ISTFT enabled")
+            if os.getenv("COSYVOICE3_HIFT_GRAPH", "0") == "1":
+                self.hift.enable_decode_graphs()
+                logger.info("CosyVoice3 streaming HiFT: decode CUDA graphs enabled")
         logger.info(f"Loaded hift weights from {hift_path}")