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