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

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

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)