From 97f2de55fa6dd112d57f57b42882148621a00fbb Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Fri, 2 Oct 2026 17:00:11 +0800 Subject: [PATCH] [Bugfix] Restore Fish Speech and CosyVoice3 audio quality at high concurrency Track Fish Speech's absolute emitted-frame frontier across adaptive chunk size changes and flush the complete final tail without replaying frames. Preserve the uniform query bound of separate MRV2 decode graphs so FA3 capture matches the runtime GQA schedule, while retaining the mixed FULL prefill bound. Exclude graph padding rows from multimodal query boundaries. Add frame-integrity and real FA3 graph regressions and route them to the H100 ready pipeline. Signed-off-by: Sy03 <1370724210@qq.com> --- .buildkite/cuda/test-ready.yml | 10 ++ .../test_fish_speech_async_chunk.py | 79 ++++++++++++++- .../worker_v2/test_capture_query_bound_gpu.py | 1 + .../worker_v2/test_decode_query_bound_gpu.py | 99 +++++++++++++++++++ tests/worker_v2/test_omni_model_state.py | 33 +++++++ .../stage_input_processors/fish_speech.py | 85 ++++++++-------- .../model_states/omni_model_state.py | 10 +- 7 files changed, 270 insertions(+), 47 deletions(-) create mode 100644 tests/worker_v2/test_decode_query_bound_gpu.py diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index ee93db4f56b..8f3ca6bdcf1 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -25,6 +25,16 @@ steps: - "pytest -sv tests/model_executor/models/cosyvoice3/test_cosyvoice3_full_batch.py -k packed -m 'core_model and cuda' --run-level=core_model" mirror_hardwares: h100_1 + - label: "Simple · MRV2 FA3 Graph Bounds Test" + timeout_in_minutes: 15 + source_file_dependencies: + - tests/worker_v2/test_capture_query_bound_gpu.py + - tests/worker_v2/test_decode_query_bound_gpu.py + - vllm_omni/worker_v2/model_states/omni_model_state.py + commands: + - "pytest -sv tests/worker_v2/test_capture_query_bound_gpu.py tests/worker_v2/test_decode_query_bound_gpu.py -m 'core_model and cuda and H100' --run-level=core_model" + mirror_hardwares: h100_1 + - label: "Simple · Diffusion Test" timeout_in_minutes: 30 commands: diff --git a/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py b/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py index 132615b1428..c6a68c898cc 100644 --- a/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py +++ b/tests/model_executor/stage_input_processors/test_fish_speech_async_chunk.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + from collections import defaultdict from types import SimpleNamespace @@ -18,7 +21,7 @@ def _make_transfer_manager( left_context: int = 1, initial_chunk_frames: int = 0, single_initial_chunk: bool = False, - backlog_chunk_frames: int = 0, + backlog_chunk_frames: int | str = 0, backlog_load_threshold: float = 0.75, max_num_seqs: int = 1, ): @@ -227,3 +230,77 @@ def test_invalid_backlog_config_fails_loudly(): {"audio_codes": torch.tensor([[1, 2, 3]])}, request, ) + + +@pytest.mark.parametrize("single_initial_chunk", [False, True]) +@pytest.mark.parametrize("initial_chunk_frames", [0, 4, 10]) +@pytest.mark.parametrize("tensor_payload", [False, True]) +@pytest.mark.parametrize("pattern", ["steady", "drop_at_finish", "rise_midstream", "oscillate"]) +def test_backlog_changes_preserve_every_frame_once(single_initial_chunk, initial_chunk_frames, tensor_payload, pattern): + tm = _make_transfer_manager( + tensor_payload=tensor_payload, + chunk_frames=25, + left_context=25, + initial_chunk_frames=initial_chunk_frames, + single_initial_chunk=single_initial_chunk, + backlog_chunk_frames=50, + max_num_seqs=64, + ) + request = _make_request() + emitted = [] + for frame in range(1, 141): + high = ( + pattern == "steady" + or (pattern == "drop_at_finish" and frame < 140) + or (pattern == "rise_midstream" and frame >= 37) + or (pattern == "oscillate" and (frame // 13) % 2 == 0) + ) + others = 63 if high else 1 + tm.code_prompt_token_ids = defaultdict(list, {"req": tm.code_prompt_token_ids["req"]}) + for idx in range(others): + tm.code_prompt_token_ids[f"other-{idx}"] = [None] * 50 + out = slow_ar_to_dac_decoder_async_chunk( + tm, + {"audio_codes": torch.tensor([[frame, frame + 1000]]), "audio_code_valid": torch.tensor([True])}, + request, + ) + if out is not None: + codes = out.codes.audio.reshape(2, -1) + emitted.extend(codes[0, out.meta.left_context_size :].tolist()) + tm.put_req_chunk["req"] += 1 + # EOS may arrive in a separate callback after the last valid frame. + out = slow_ar_to_dac_decoder_async_chunk(tm, {}, request, is_finished=True) + if out is not None and out.codes.audio.numel(): + emitted.extend(out.codes.audio.reshape(2, -1)[0, out.meta.left_context_size :].tolist()) + assert emitted == list(range(1, 141)) + assert slow_ar_to_dac_decoder_async_chunk(tm, {}, request, is_finished=True) is None + + +def test_final_backlog_drop_flushes_entire_pending_tail(): + tm = _make_transfer_manager( + tensor_payload=True, + chunk_frames=25, + left_context=25, + initial_chunk_frames=4, + single_initial_chunk=True, + backlog_chunk_frames=50, + max_num_seqs=64, + ) + request = _make_request() + emitted = [] + for frame in range(1, 41): + if frame < 40: + for idx in range(63): + tm.code_prompt_token_ids[f"other-{idx}"] = [None] * 30 + else: + for idx in range(32, 63): + tm.code_prompt_token_ids.pop(f"other-{idx}") + out = slow_ar_to_dac_decoder_async_chunk( + tm, + {"audio_codes": torch.tensor([[frame]])}, + request, + is_finished=frame == 40, + ) + if out is not None: + emitted.extend(out.codes.audio[0, out.meta.left_context_size :].tolist()) + assert emitted == list(range(1, 41)) diff --git a/tests/worker_v2/test_capture_query_bound_gpu.py b/tests/worker_v2/test_capture_query_bound_gpu.py index 93c53f31344..c341a575b93 100644 --- a/tests/worker_v2/test_capture_query_bound_gpu.py +++ b/tests/worker_v2/test_capture_query_bound_gpu.py @@ -44,6 +44,7 @@ def _check_prefill(): spec = FullAttentionSpec(block_size=16, num_kv_heads=8, head_size=128, dtype=torch.bfloat16) kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)]) state = OmniModelState.__new__(OmniModelState) + state.vllm_config = c state.max_model_len = 256 state.supports_mm_inputs = False with set_current_vllm_config(c): diff --git a/tests/worker_v2/test_decode_query_bound_gpu.py b/tests/worker_v2/test_decode_query_bound_gpu.py new file mode 100644 index 00000000000..114e9f343e0 --- /dev/null +++ b/tests/worker_v2/test_decode_query_bound_gpu.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Real FA3 graph replay must preserve the decode GQA layout.""" + +from dataclasses import dataclass, replace + +import numpy as np +import pytest +import torch +from transformers import Qwen2Config +from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend, FlashAttentionImpl +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.utils import AttentionGroup + +from tests.helpers.mark import hardware_test +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState + +pytestmark = [pytest.mark.core_model] + + +@dataclass +class _AttentionScales: + _q_scale: torch.Tensor + _k_scale: torch.Tensor + _v_scale: torch.Tensor + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize("graph_mode", [CUDAGraphMode.FULL_AND_PIECEWISE, CUDAGraphMode.FULL_DECODE_ONLY]) +def test_full_decode_graph_matches_fa3_aot_schedule(graph_mode, tmp_path): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + pytest.skip("requires FA3 on Hopper or newer") + torch.manual_seed(17) + device = torch.device("cuda:0") + c = VllmConfig() + Qwen2Config( + architectures=["Qwen2ForCausalLM"], + hidden_size=1024, + intermediate_size=2048, + num_hidden_layers=1, + num_attention_heads=16, + num_key_value_heads=2, + max_position_embeddings=2048, + ).save_pretrained(tmp_path) + c.model_config = ModelConfig(model=str(tmp_path), skip_tokenizer_init=True, dtype="bfloat16") + c.scheduler_config.max_num_seqs = 64 + c.compilation_config.cudagraph_mode = graph_mode + c.compilation_config.max_cudagraph_capture_size = 128 + c.attention_config.flash_attn_version = 3 + spec = FullAttentionSpec(block_size=16, num_kv_heads=2, head_size=64, dtype=torch.bfloat16) + kvconfig = KVCacheConfig(num_blocks=1024, kv_cache_tensors=[], kv_cache_groups=[KVCacheGroupSpec(["attn"], spec)]) + state = OmniModelState.__new__(OmniModelState) + state.vllm_config = c + state.max_model_len = 2048 + state.supports_mm_inputs = False + with set_current_vllm_config(c): + group = AttentionGroup(FlashAttentionBackend, ["attn"], spec, 0) + group.create_metadata_builders(c, device) + impl = FlashAttentionImpl(16, 64, 64**-0.5, 2, None, None, "auto") + assert impl.vllm_flash_attn_version == 3 + buffers = InputBuffers(64, 128, device) + batch = InputBatch.make_dummy(8, 8, buffers, max_query_len=None) + tables = torch.arange(1024, dtype=torch.int32, device=device).reshape(8, 128) + slots = torch.zeros((1, 8), dtype=torch.int64, device=device) + cache = torch.randn(1024, 2, 16, 128, device=device, dtype=torch.bfloat16) + q = torch.randn(8, 16, 64, device=device, dtype=torch.bfloat16) + k = v = torch.empty(8, 2, 64, device=device, dtype=torch.bfloat16) + layer = _AttentionScales( + _q_scale=torch.ones((), device=device), + _k_scale=torch.ones((), device=device), + _v_scale=torch.ones((), device=device), + ) + md = state.prepare_attn(batch, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig, for_capture=True)[ + "attn" + ] + out = torch.full_like(q, float("nan")) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + impl.forward(layer, q, k, v, cache, md, out) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + impl.forward(layer, q, k, v, cache, md, out) + lengths = np.arange(37, 45, dtype=np.int32) + batch.seq_lens.copy_(torch.tensor(lengths, device=device)) + runtime = replace(batch, max_query_len=1, seq_lens_cpu_upper_bound=torch.tensor(lengths)) + eager_md = state.prepare_attn(runtime, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig)["attn"] + expected = torch.empty_like(q) + impl.forward(layer, q, k, v, cache, eager_md, expected) + state.prepare_attn(runtime, CUDAGraphMode.FULL, (tables,), slots, [[group]], kvconfig) + graph.replay() + torch.accelerator.synchronize() + assert md.max_query_len == 1 + torch.testing.assert_close(out, expected, atol=0.002, rtol=0.002) diff --git a/tests/worker_v2/test_omni_model_state.py b/tests/worker_v2/test_omni_model_state.py index 4edfed603c1..e85790c2010 100644 --- a/tests/worker_v2/test_omni_model_state.py +++ b/tests/worker_v2/test_omni_model_state.py @@ -4,13 +4,17 @@ graph/eager paths with per-request seed independence, async snapshot ownership.""" from contextlib import nullcontext +from dataclasses import replace from types import SimpleNamespace from unittest.mock import MagicMock, patch import numpy as np import pytest import torch +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers +from vllm.v1.worker.gpu.mm.encoder_runner import EncoderRunner from vllm.v1.worker.gpu.model_states.default import DefaultModelState +from vllm.v1.worker.gpu.states import RequestState from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState, _make_safe_get_rope @@ -655,3 +659,32 @@ def preprocess(input_ids, input_embeds, **info): state.run_preprocess(batch, {"input_ids": torch.tensor([1]), "inputs_embeds": embeds}) assert seen == (["r1"] if prefilling else []) assert torch.equal(embeds, torch.ones(1, 4)) + + +def test_mm_embeddings_exclude_zero_length_graph_padding_rows(): + state = object.__new__(OmniModelState) + state.supports_mm_inputs = True + state.mm_pruner = None + state.prompt_embeds_state = None + state.encoder_runner = MagicMock(spec=EncoderRunner, inputs_embeds=torch.zeros(8, 2)) + state.execute_mm_encoder = lambda _: None + state.gather_mm_embeddings = lambda _: ([torch.ones(1, 2)], torch.ones(5, dtype=torch.bool)) + + def embed(input_ids, *, query_start_loc, multimodal_embeddings, is_multimodal): + assert query_start_loc == [0, 2, 5] + return input_ids[:, None].expand(-1, 2).float() + + state.model = MagicMock(supports_embed_input_ids_query_start_loc=True, embed_input_ids=embed) + buffers = InputBuffers(4, 8, torch.device("cpu")) + batch = replace( + InputBatch.make_dummy(2, 5, buffers), + num_reqs=2, + num_reqs_after_padding=4, + num_tokens=5, + num_tokens_after_padding=8, + input_ids=torch.arange(8), + query_start_loc_np=np.array([0, 2, 5, 5, 5], dtype=np.int32), + ) + result = state.prepare_inputs_embeds({}, batch, MagicMock(spec=RequestState)) + torch.testing.assert_close(result[:5], torch.arange(5).float()[:, None].expand(-1, 2)) + assert result.shape == (8, 2) diff --git a/vllm_omni/model_executor/stage_input_processors/fish_speech.py b/vllm_omni/model_executor/stage_input_processors/fish_speech.py index f428c34a454..cfcb277eabb 100644 --- a/vllm_omni/model_executor/stage_input_processors/fish_speech.py +++ b/vllm_omni/model_executor/stage_input_processors/fish_speech.py @@ -1,3 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + """Stage input processor for Fish Speech S2 Pro: Slow AR → DAC Decoder.""" from collections.abc import Mapping @@ -122,9 +125,9 @@ def slow_ar_to_dac_decoder_async_chunk( ) -> OmniPayloadStruct | None: """Async streaming processor: emit code chunks as they are produced. - Accumulates per-step codes and emits fixed-size chunks with left context - overlap for smooth audio transitions, analogous to - ``talker2code2wav_async_chunk`` in Qwen3 TTS. + Accumulates per-step codes and emits contiguous chunks with left context + overlap. The per-request output frontier remains absolute when the backlog + policy changes the chunk size; a final callback flushes every pending frame. """ request_id = request.external_req_id finished = bool(is_finished or request.is_finished()) @@ -174,49 +177,41 @@ def slow_ar_to_dac_decoder_async_chunk( ) return None - single_initial_chunk = _cfg_bool( - cfg, - "fish_speech_single_initial_chunk", - False, + # Keep the absolute output frontier in the request payload so that changing + # the backlog chunk size cannot reinterpret frames already emitted. Both + # connector runtimes reclaim request_payload on completion/cancellation. + request_payload = getattr(transfer_manager, "request_payload", None) + if request_payload is None: + request_payload = {} + transfer_manager.request_payload = request_payload + state = request_payload.setdefault(request_id, {}).setdefault( + "_fish_speech_async_state", {"emitted_frames": 0, "terminal_sent": False} ) - use_first_chunk = initial_chunk_size > 0 and initial_chunk_size < steady_chunk_size - - if single_initial_chunk and use_first_chunk: - if length <= initial_chunk_size: - if not finished and length < initial_chunk_size: - return None - context_length = length if finished and length < initial_chunk_size else initial_chunk_size - else: - adjusted = length - initial_chunk_size - if adjusted <= 0: - return None - if not finished and adjusted % steady_chunk_size != 0: - return None - chunk_length = adjusted % steady_chunk_size - context_length = chunk_length if chunk_length != 0 else steady_chunk_size - end_index = min(length, left_context_size_config + context_length) - left_context_size = max(0, int(end_index - context_length)) - window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] - elif initial_chunk_size > 0 and length <= chunk_size: - already_sent = transfer_manager.put_req_chunk[request_id] * initial_chunk_size - pending = length - already_sent - if pending <= 0: - return None - if pending < initial_chunk_size and not finished: - return None - context_length = min(pending, initial_chunk_size) - left_context_size = max(0, length - context_length) - window_frames = transfer_manager.code_prompt_token_ids[request_id][:length] - else: - initial_coverage = (chunk_size // initial_chunk_size) * initial_chunk_size if initial_chunk_size > 0 else 0 - adjusted = length - initial_coverage - chunk_length = adjusted % steady_chunk_size - if chunk_length != 0 and not finished: + if state["terminal_sent"]: + return None + emitted_frames = state["emitted_frames"] + pending = length - emitted_frames + if pending <= 0: + if not finished: return None - context_length = chunk_length if chunk_length != 0 else steady_chunk_size - end_index = min(length, left_context_size_config + context_length) - left_context_size = max(0, int(end_index - context_length)) - window_frames = transfer_manager.code_prompt_token_ids[request_id][-end_index:] + state["terminal_sent"] = True + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) + + single_initial_chunk = _cfg_bool(cfg, "fish_speech_single_initial_chunk", False) + initial_coverage = ( + (initial_chunk_size if single_initial_chunk else (chunk_size // initial_chunk_size) * initial_chunk_size) + if initial_chunk_size > 0 + else 0 + ) + target_chunk_size = initial_chunk_size if emitted_frames < initial_coverage else steady_chunk_size + if not finished and pending < target_chunk_size: + return None + + left_context_size = min(emitted_frames, left_context_size_config) + window_frames = transfer_manager.code_prompt_token_ids[request_id][emitted_frames - left_context_size : length] # Pack into codebook-major codes. The tensor path avoids expanding codec # indices into Python ints across the connector boundary; Stage1 schedules @@ -228,6 +223,8 @@ def slow_ar_to_dac_decoder_async_chunk( else: code_predictor_codes = codes_qf.reshape(-1) + state["emitted_frames"] = length + state["terminal_sent"] = finished return OmniPayloadStruct( codes=CodesStruct(audio=code_predictor_codes), meta=MetaStruct( diff --git a/vllm_omni/worker_v2/model_states/omni_model_state.py b/vllm_omni/worker_v2/model_states/omni_model_state.py index f20a0f2db89..d5a6ce313ef 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -427,12 +427,18 @@ def prepare_attn( for_capture: bool = False, ubatch_idx: int = 0, ) -> dict[str, Any]: - if for_capture and input_batch.max_query_len is None: + if ( + for_capture + and input_batch.max_query_len is None + and not self.vllm_config.compilation_config.cudagraph_mode.separate_routine() + ): # vLLM 0.30 distributes dummy tokens evenly across requests. For # an unconstrained FULL graph that split is not a query-length # bound: replay may put the entire token bucket in one request. # Attention launch parameters are fixed at capture, so use the # bucket's worst-case query length, not the dummy per-row length. + # Separate decode graphs must retain their uniform query bound: + # FA3's captured GQA layout must match the runtime AOT schedule. # Keep explicit bounds (e.g. varlen decode) and runtime metadata. input_batch = replace(input_batch, max_query_len=input_batch.num_tokens) return super().prepare_attn( @@ -464,7 +470,7 @@ def prepare_inputs_embeds( mm_embeds, is_mm_embed = self.gather_mm_embeddings(input_batch) kwargs: dict[str, Any] = {"multimodal_embeddings": mm_embeds, "is_multimodal": is_mm_embed} if mm_embeds: - kwargs["query_start_loc"] = input_batch.query_start_loc_np.tolist() + kwargs["query_start_loc"] = input_batch.query_start_loc_np[: input_batch.num_reqs + 1].tolist() embeds = self.model.embed_input_ids(input_batch.input_ids[: input_batch.num_tokens], **kwargs) inputs_embeds = self.encoder_runner.inputs_embeds inputs_embeds[: embeds.shape[0]] = embeds