diff --git a/tests/config/test_environment_variables.py b/tests/config/test_environment_variables.py index e1f6aa651..eb3db91ec 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: 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, diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py index 3722bbcc7..3753a75c6 100644 --- a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py +++ b/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(): diff --git a/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py b/tests/model_executor/models/cosyvoice3/test_cosyvoice3_hift_graph.py new file mode 100644 index 000000000..1e1fcce0c --- /dev/null +++ b/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 diff --git a/vllm_omni/config/environment_variable_inventory.py b/vllm_omni/config/environment_variable_inventory.py index ee7f599f2..06dcb3ef4 100644 --- a/vllm_omni/config/environment_variable_inventory.py +++ b/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", diff --git a/vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml b/vllm_omni/deploy/cosyvoice3_packed_streaming_optimized_standard.yaml new file mode 100644 index 000000000..48bcc562b --- /dev/null +++ b/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" diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py index a353675bf..a1e31c070 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hifigan.py +++ b/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) diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/hift_graph.py new file mode 100644 index 000000000..588497310 --- /dev/null +++ b/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() diff --git a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py b/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py index ff6535b87..579f7d32e 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/code2wav_core/packed_dit.py +++ b/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) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py index 6fd13c173..96ec400c6 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3_code2wav.py +++ b/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}")