Compare commits
4
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
118f646c87 | ||
|
|
7d5a6bc5cb | ||
|
|
8361765b33 | ||
|
|
b3c25a608d |
No files matched your search
@@ -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,
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
# 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)
|
||||
saved: list[tuple[torch.Tensor, torch.Tensor]] = []
|
||||
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)
|
||||
for previous, snapshot in saved:
|
||||
assert torch.equal(previous, snapshot)
|
||||
if hift._decode_graphs.graphs:
|
||||
saved.append((actual, actual.clone()))
|
||||
assert len(hift._decode_graphs.graphs) == 1
|
||||
hift._decode_graphs.max_graphs = 1
|
||||
mel = torch.randn(1, 80, 31, device="cuda", generator=generator)
|
||||
source = torch.randn(1, 1, 31 * 480, device="cuda", generator=generator)
|
||||
actual = hift.decode(mel, source, finalize=finalize)
|
||||
expected = hift._decode_eager(mel, source, finalize=finalize)
|
||||
assert torch.equal(actual, expected)
|
||||
assert len(hift._decode_graphs.graphs) == 1
|
||||
@@ -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",
|
||||
|
||||
@@ -0,0 +1,24 @@
|
||||
# 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.
|
||||
# HiFT graphs capture on the third use of each exact shape, adding latency then.
|
||||
# The codec graph pool is not capped by gpu_memory_utilization; leave memory headroom.
|
||||
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"
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
@@ -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)
|
||||
|
||||
@@ -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}")
|
||||
Reference in new issue
Block a user