[Bugfix] Restore Fish Speech and CosyVoice3 audio quality at high concurrency (#8422)

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

.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:

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))

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):

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)

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
@@ -661,3 +665,32 @@ def test_identity_preprocess_skips_only_decode_rows(prefilling):
     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)

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 state["terminal_sent"]:
+        return None
+    emitted_frames = state["emitted_frames"]
+    pending = length - emitted_frames
+    if pending <= 0:
+        if not finished:
+            return None
+        state["terminal_sent"] = True
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)),
+            meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)),
+        )
 
-    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:
-            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:]
+    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(

vllm_omni/worker_v2/model_states/omni_model_state.py

@@ -448,12 +448,18 @@ class OmniModelState(DefaultModelState):
         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(
@@ -485,7 +491,7 @@ class OmniModelState(DefaultModelState):
         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