diff --git a/.buildkite/cuda/test-ready.yml b/.buildkite/cuda/test-ready.yml index 3b5213d9e..ee93db4f5 100644 --- a/.buildkite/cuda/test-ready.yml +++ b/.buildkite/cuda/test-ready.yml @@ -205,6 +205,16 @@ steps: - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_minicpmo_4_5_ready.json -m 'core_model and cuda' mirror_hardwares: h100_1 + - label: "TTS · Higgs v3 Full Response Test" + source_file_dependencies: + - tests/e2e/online_serving/test_higgs_audio_v3.py + - vllm_omni/model_executor/models/higgs_audio_v3/ + - vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py + - vllm_omni/deploy/higgs_multimodal_qwen3.yaml + commands: + - timeout 20m pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py::test_full_response_v1_preserves_audio -m 'core_model and cuda' --run-level 'core_model' + mirror_hardwares: h100_1 + - label: "TTS · Qwen3-TTS CustomVoice Test" source_file_dependencies: tts_qwen3_tts_function commands: diff --git a/tests/e2e/online_serving/test_higgs_audio_v3.py b/tests/e2e/online_serving/test_higgs_audio_v3.py index e9ba68909..3c5abcf08 100644 --- a/tests/e2e/online_serving/test_higgs_audio_v3.py +++ b/tests/e2e/online_serving/test_higgs_audio_v3.py @@ -315,3 +315,34 @@ class TestHiggsAudioV3OnlineVoiceClone: "timeout": DEFAULT_SPEECH_TIMEOUT_S, } ) + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.tts +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.parametrize( + "omni_server", + [ + OmniServerParams( + model=MODEL, + stage_config_path=STAGE_CONFIG, + server_args=[*SERVER_ARGS, "--no-async-chunk"], + env_dict=SERVER_ENV, + ) + ], + indirect=True, +) +def test_full_response_v1_preserves_audio(omni_server, online_client) -> None: + """Non-chunk V1 must retain codec rows through the terminal talker step.""" + online_client.send_audio_speech_request( + { + "model": omni_server.model, + "input": "The quick brown fox jumps over the lazy dog.", + "stream": False, + "response_format": "wav", + "min_audio_bytes": _MIN_AUDIO_BYTES, + "timeout": DEFAULT_SPEECH_TIMEOUT_S, + }, + request_num=3, + ) diff --git a/tests/model_executor/models/higgs_audio_v3/test_async_audio_mode.py b/tests/model_executor/models/higgs_audio_v3/test_async_audio_mode.py new file mode 100644 index 000000000..345688986 --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_async_audio_mode.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_model_metadata_reset(enabled): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration as C, + ) + + t = C.__new__(C) + torch.nn.Module.__init__(t) + t.config = SimpleNamespace(audio_async_prompt_mode=enabled) + t._set_last_step_query_start_loc = lambda x: None + t._sync_decode_state_with_batch = lambda x: None + t.update_decode_step_metadata(audio_prompt_mode_rows=2) + assert t._step_audio_mode_rows == (2 if enabled else 0) + t.update_decode_step_metadata() + assert t._step_audio_mode_rows == 0 + + +@pytest.mark.parametrize("bad_state", [False, True]) +def test_actual_state_guard_and_eos(bad_state): + import importlib.util + from pathlib import Path + + import vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker as mod + + root = Path(mod.__file__).parents[4] + spec = importlib.util.spec_from_file_location( + "helper_fixture", root / "tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py" + ) + assert spec is not None and spec.loader is not None + helpers = importlib.util.module_from_spec(spec) + spec.loader.exec_module(helpers) + t = helpers.TestSamplerMethods()._make_batched_sampler_talker(2) + t._restore_terminal_audio_rows = ( + mod.HiggsAudioV3TalkerForConditionalGeneration._restore_terminal_audio_rows.__get__(t) + ) + t.config = SimpleNamespace(audio_async_prompt_mode=True, audio_full_sample_graph=False) + t._resolve_token_ids = lambda: None + t._audio_continuation_id = 99999 + t._eos_token_id = 151671 + t._last_logits_hidden = torch.zeros(2, 16) + t._last_step_input_ids = torch.tensor([99999, 12345 if bad_state else 151671]) + t._last_step_query_start_loc = None + t._decode_has_codes = torch.tensor([True, False]) + t._decode_generation_done = torch.tensor([False, False]) + t._decode_delay_count = torch.zeros(2, dtype=torch.long) + t._decode_eoc_countdown = torch.full((2,), -1, dtype=torch.long) + t._fast_audio_direct_rows = 0 + t._step_audio_tail_rows = 0 + t._step_audio_mode_rows = 2 + t._fast_audio_sampler_gpu_fallback_reason = lambda **kw: None + t._audio_codebook_logits_from_rows = lambda hidden, rows, all_rows=False: torch.zeros(2, 8, 1026) + t._apply_delay_pattern_masking_batched = lambda *a, **kw: None + t._sample_audio_codes = lambda logits, *a, **kw: torch.zeros(logits.shape[0], dtype=torch.long) + seen_codes = {} + t._update_delay_state_batched = lambda *a, **kw: seen_codes.update(kw) + + def call(): + return mod.HiggsAudioV3TalkerForConditionalGeneration.sample( + t, torch.zeros(2, 200000), SimpleNamespace(no_penalties=True) + ) + + if bad_state: + with pytest.raises(RuntimeError, match="tail mismatch"): + call() + else: + assert call().sampled_token_ids.tolist() == [[99999], [151671]] + assert seen_codes["code_row_mask"].tolist() == [True, False] diff --git a/tests/model_executor/models/higgs_audio_v3/test_async_payload.py b/tests/model_executor/models/higgs_audio_v3/test_async_payload.py new file mode 100644 index 000000000..8371b4080 --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_async_payload.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration as C, +) +from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_two_inflight_steps_keep_owned_codes_and_masks(): + t = C.__new__(C) + torch.nn.Module.__init__(t) + t.num_codebooks = 2 + t.use_async_omni_output = True + first = torch.tensor([[10, 11, 1, 0], [20, 21, 0, 0], [30, 31, 1, 1], [40, 41, 1, 0]]) + t._finish_audio_staging(first, 4, torch.device("cpu"), 4) + first.zero_() # Next graph replay overwrites original GPU storage. + a = t.post_sample_multimodal_outputs( + req_ids=["b", "a", "c", "invalid"], invalid_req_indices=[3], multimodal_outputs=None + ) + second = torch.tensor([[50, 51, 1, 0]]) + t._finish_audio_staging(second, 1, torch.device("cpu"), 1) + b = t.post_sample_multimodal_outputs(req_ids=["new"], invalid_req_indices=[], multimodal_outputs=None) + second.zero_() + assert t.post_sample_multimodal_outputs(req_ids=[], invalid_req_indices=[], multimodal_outputs=None) is None + # Finish callbacks out of order. They cannot read current model state. + rb = t.finalize_multimodal_outputs_from_cpu_snapshot(b)["codes"]["audio"] + ra = t.finalize_multimodal_outputs_from_cpu_snapshot(a)["codes"]["audio"] + assert [x.tolist() for x in ra] == [[[10, 11]], [], [[30, 31]], []] + assert rb[0].tolist() == [[50, 51]] + a["_higgs_audio_snapshot"].zero_() + assert ra[0].tolist() == [[10, 11]] + assert t._fast_audio_direct_rows == 0 + assert t.finalize_multimodal_outputs_from_cpu_snapshot(None) is None + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_whole_output_explicit_opt_in(enabled): + r = object.__new__(GPUARModelRunner) + r.use_async_scheduling = True + r.speculative_config = None + r.model_config = SimpleNamespace(async_chunk=False, enable_return_routed_experts=False) + r.model = SimpleNamespace(use_async_omni_output=True, supports_async_whole_payload=enabled, has_postprocess=False) + assert r._should_use_async_omni_output() == enabled + r.use_async_scheduling = False + assert not r._should_use_async_omni_output() + + +@pytest.mark.parametrize("async_output", [False, True]) +def test_terminal_and_invalid_rows_preserve_full_response_codes(async_output): + from vllm_omni.distributed.omni_connectors.model_runner.omni_connector_payload_transport import ( + _OmniConnectorPayloadTransportMixin, + ) + from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import talker2code2wav_full_payload + + model = C.__new__(C) + torch.nn.Module.__init__(model) + model.num_codebooks = 8 + model.use_async_omni_output = async_output + model._last_audio_host_staging = None + model._last_audio_staging_event = None + transport = _OmniConnectorPayloadTransportMixin() + transport._pending_full_payload_send = {} + transport._full_payload_replace_keys_cached = frozenset() + expected = torch.arange(80, dtype=torch.int32).reshape(10, 8) + # Terminal and invalid rows must append zero rows, never replace earlier codes. + for codes, valid, done, invalid in [ + *[(row, 1, 0, []) for row in expected], + (torch.zeros(8), 0, 1, []), + (torch.ones(8), 1, 0, [0]), + ]: + staging = torch.cat((codes.to(torch.int32), torch.tensor([valid, done], dtype=torch.int32))).reshape(1, 10) + model._finish_audio_staging(staging, 1, torch.device("cpu"), 1) + payload = model.post_sample_multimodal_outputs( + req_ids=["request"], invalid_req_indices=invalid, multimodal_outputs=None + ) + payload = model.finalize_multimodal_outputs_from_cpu_snapshot(payload) + audio = payload["codes"]["audio"][0] + assert audio.shape == (int(bool(valid) and not invalid), 8) + transport.accumulate_full_payload_output("request", {"codes.audio": audio}, None) + full, _ = transport._materialize_full_payload_entry(transport._pending_full_payload_send["request"]) + torch.testing.assert_close(full["codes.audio"], expected) + assert talker2code2wav_full_payload(None, full, None)["codes"]["audio"].numel() == 16 diff --git a/tests/model_executor/models/higgs_audio_v3/test_codec_batching.py b/tests/model_executor/models/higgs_audio_v3/test_codec_batching.py new file mode 100644 index 000000000..7041df772 --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_codec_batching.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Codec batching must preserve per-request order and convolution boundaries.""" + +import pytest +import torch + +from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import HiggsAudioV3Code2Wav +from vllm_omni.transformers_utils.configs.higgs_audio_v3 import HiggsAudioV3Config + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def codec(monkeypatch): + model = HiggsAudioV3Code2Wav(config=HiggsAudioV3Config()) + calls = [] + + def decode(codes): + calls.append(tuple(codes.shape)) + if (codes < 0).any(): + raise ValueError("invalid codec row") + # A time-dependent decoder makes both trim boundaries observable. + pcm = codes[:, :1].float().repeat_interleave(model.hop_length, dim=-1) + return pcm + + monkeypatch.setattr(model, "decode_codes", decode) + return model, calls + + +def test_ragged_codec_batch_preserves_request_order_and_trims(codec): + model, calls = codec + codes = [ + torch.arange(6).expand(8, -1), + torch.arange(10, 14).expand(8, -1), + torch.arange(20, 26).expand(8, -1), + ] + output = model( + input_ids=torch.cat([row.reshape(-1) for row in codes]), + seq_token_counts=[row.numel() for row in codes], + runtime_additional_information=[ + {"meta": {"left_context_size": 1, "right_holdback_size": 1}}, + {"meta": {"left_context_size": 0, "right_holdback_size": 1}}, + {"meta": {"left_context_size": 1, "right_holdback_size": 1}}, + ], + ) + assert calls == [(2, 8, 6), (1, 8, 4)] + for actual, values in zip( + output.multimodal_outputs["model_outputs"], ([1, 2, 3, 4], [10, 11, 12], [21, 22, 23, 24]) + ): + expected = torch.tensor(values, dtype=torch.float32).repeat_interleave(model.hop_length) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("bad_length,bad_value", [(0, 0), (7, 1), (24, -1)]) +def test_invalid_codec_request_does_not_discard_valid_peer(codec, bad_length, bad_value): + model, _ = codec + valid = torch.arange(3).expand(8, -1).reshape(-1) + bad = torch.full((bad_length,), bad_value, dtype=torch.long) + output = model(input_ids=torch.cat([bad, valid]), seq_token_counts=[bad.numel(), valid.numel()]) + audio = output.multimodal_outputs["model_outputs"] + assert len(audio) == 2 + assert audio[0].numel() == 0 + expected = torch.arange(3, dtype=torch.float32).repeat_interleave(model.hop_length) + torch.testing.assert_close(audio[1], expected, rtol=0, atol=0) + + +def test_native_payload_overrides_control_slots_and_empty_terminal(codec): + model, calls = codec + codes = torch.arange(6).expand(8, -1).reshape(-1) + output = model( + input_ids=torch.zeros(2, dtype=torch.long), + seq_token_counts=[1, 1], + runtime_additional_information=[ + {"codes": {"audio": codes}, "meta": {"left_context_size": 1, "right_holdback_size": 1}}, + {"codes": {"audio": torch.empty(0, dtype=torch.long)}, "meta": {"finished": True}}, + ], + ) + assert calls == [(1, 8, 6)] + wavs = output.multimodal_outputs["model_outputs"] + torch.testing.assert_close(wavs[0], torch.arange(1, 5).float().repeat_interleave(model.hop_length)) + assert wavs[1].numel() == 0 diff --git a/tests/model_executor/models/higgs_audio_v3/test_codec_graph_shapes.py b/tests/model_executor/models/higgs_audio_v3/test_codec_graph_shapes.py new file mode 100644 index 000000000..65b6241c3 --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_codec_graph_shapes.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.higgs_audio_v3.full_sample_graph import run_dense_sample +from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_code2wav import ( + HiggsAudioV3Code2WavForConditionalGeneration as Codec, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_matched_streaming_graphs_cover_every_exact_batch(): + model = SimpleNamespace( + config=SimpleNamespace( + codec_graph_single_max_frames=33, + codec_graph_batch_sizes=list(range(1, 65)), + codec_graph_frame_sizes=[8, 33], + ) + ) + shapes = Codec._decode_graph_shapes(model) + assert len(shapes) == 159 + assert all((b, f) in shapes for b in range(1, 65) for f in (8, 33)) + assert all((1, f) in shapes for f in range(1, 34)) + assert (2, 32) not in shapes + + +def test_default_graph_shapes_are_preserved(): + shapes = Codec._decode_graph_shapes(SimpleNamespace(config=SimpleNamespace())) + assert len(shapes) == 162 + assert (1, 150) in shapes and (16, 54) in shapes + + +def test_invalid_graph_shape_rejected(): + with pytest.raises(ValueError): + Codec._decode_graph_shapes(SimpleNamespace(config=SimpleNamespace(codec_graph_batch_sizes=[0]))) + + +@pytest.mark.parametrize("use_v2", [False, True]) +@pytest.mark.parametrize("sampler", [None, False, True]) +@pytest.mark.parametrize("codec", [None, False, True]) +def test_graph_defaults_follow_runner_and_preserve_explicit_overrides(use_v2, sampler, codec): + from vllm_omni.transformers_utils.configs.higgs_audio_v3 import HiggsAudioV3Config + + config = HiggsAudioV3Config(audio_full_sample_graph=sampler, codec_cuda_graph=codec) + assert config.resolve_graph_defaults(use_v2_model_runner=use_v2) == ( + use_v2 if sampler is None else sampler, + use_v2 if codec is None else codec, + ) + # Resolving a stage must not bake its defaults into shared model config. + assert config.audio_full_sample_graph is sampler + assert config.codec_cuda_graph is codec + + +def test_sampler_cache_miss_does_not_capture_or_advance_rng(): + model = SimpleNamespace(_dense_sample_graphs={}) + torch.manual_seed(7) + before = torch.random.get_rng_state() + assert run_dense_sample(model, torch.zeros(2, 4), None, None) is None + assert torch.equal(before, torch.random.get_rng_state()) + assert not model._dense_sample_graphs diff --git a/tests/model_executor/models/higgs_audio_v3/test_codec_snake_kernel.py b/tests/model_executor/models/higgs_audio_v3/test_codec_snake_kernel.py new file mode 100644 index 000000000..208c841aa --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_codec_snake_kernel.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The fused DAC Snake / decoder forward is bit-identical to transformers' DacDecoder.""" + +import copy + +import pytest +import torch +import torch.nn as nn + +from tests.helpers.mark import hardware_test +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model] + +requires_cuda = pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="requires CUDA") + +_DAC_CONFIG = { + "codebook_dim": 8, + "codebook_size": 1024, + "decoder_hidden_size": 256, + "downsampling_ratios": [8, 5, 4, 2, 3], + "encoder_hidden_size": 16, + "hidden_size": 64, + "hop_length": 960, + "n_codebooks": 9, + "sampling_rate": 16000, + "upsampling_ratios": [8, 5, 4, 2, 3], +} + + +def _dac_decoder_pair(device, dtype): + from transformers import DacConfig, DacModel + + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import ( + _fuse_dac_decoder, + adjust_conv_transpose_output_padding, + ) + + torch.manual_seed(0) + reference = DacModel(DacConfig(**_DAC_CONFIG)).decoder + adjust_conv_transpose_output_padding(reference) + with torch.no_grad(): + for name, param in reference.named_parameters(): + if name.endswith("alpha"): + param.copy_(torch.rand_like(param) * 2 + 0.05) + elif name.endswith("bias"): + param.normal_(0.0, 0.1) + reference = reference.to(device, dtype).eval() + fused = copy.deepcopy(reference) + assert _fuse_dac_decoder(fused) == 1 + 7 * len(reference.block) + fused.load_state_dict(reference.state_dict()) + return reference, fused + + +@requires_cuda +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("shape", [(1, 64, 33), (3, 1024, 8), (2, 96, 5000)]) +@torch.inference_mode() +def test_fused_snake_matches_transformers_snake(dtype, shape): + from transformers.models.dac.modeling_dac import Snake1d + + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _FusedDacSnake1d + + torch.manual_seed(0) + reference = Snake1d(shape[1]).to("cuda", dtype) + reference.alpha.copy_(torch.rand_like(reference.alpha) * 4 + 0.05) + fused = _FusedDacSnake1d(shape[1]).to("cuda", dtype) + fused.load_state_dict(reference.state_dict()) + x = (torch.randn(shape, device="cuda") * 3).to(dtype) + residual = (torch.randn(shape, device="cuda") * 3).to(dtype) + bias = (torch.randn(shape[1], device="cuda") * 3).to(dtype) + assert torch.equal(fused(x), reference(x)) + + total, out = fused.fused(x, bias, residual, write_sum=True) + expected_total = residual + (x + bias.reshape(1, -1, 1)) + assert torch.equal(total, expected_total) + assert torch.equal(out, reference(expected_total)) + assert fused.fused(x, bias)[0] is None + assert torch.equal(fused.fused(x, bias)[1], reference(x + bias.reshape(1, -1, 1))) + + +@requires_cuda +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@pytest.mark.parametrize("batch,frames", [(1, 33), (4, 8)]) +def test_fused_dac_decoder_matches_transformers_decoder(dtype, batch, frames): + reference, fused = _dac_decoder_pair("cuda", dtype) + x = torch.randn(batch, _DAC_CONFIG["hidden_size"], frames, device="cuda").to(dtype) + with torch.inference_mode(): + assert torch.equal(fused(x), reference(x)) + + # Reloaded weights must not reuse the cached Snake reciprocals. + with torch.no_grad(): + for module in (reference, fused): + module.snake1.alpha.mul_(0.5) + with torch.inference_mode(): + assert torch.equal(fused(x), reference(x)) + + +@pytest.mark.cpu +def test_fuse_dac_decoder_keeps_state_dict_and_falls_back_on_cpu(): + from transformers.models.dac.modeling_dac import Snake1d + + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _FusedDacSnake1d + + reference, fused = _dac_decoder_pair("cpu", torch.float32) + assert not any(isinstance(module, Snake1d) for module in fused.modules()) + assert set(fused.state_dict()) == set(reference.state_dict()) + x = torch.randn(2, _DAC_CONFIG["hidden_size"], 5) + with torch.no_grad(): + assert torch.equal(fused(x), reference(x)) + + decoder = nn.Sequential(Snake1d(8), nn.Conv1d(8, 8, 3)) + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _fuse_dac_decoder + + assert _fuse_dac_decoder(decoder) == 1 + assert isinstance(decoder[0], _FusedDacSnake1d) + assert "forward" not in vars(decoder) + + +@requires_cuda +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float32]) +@torch.inference_mode() +def test_fused_decoder_graph_replays_changed_inputs(dtype): + reference, fused = _dac_decoder_pair("cuda", dtype) + x = torch.randn(2, _DAC_CONFIG["hidden_size"], 8, device="cuda").to(dtype) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + fused(x) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = fused(x) + for _ in range(2): + x.normal_() + graph.replay() + assert torch.equal(output, reference(x)) + + +@requires_cuda +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_fused_snake_uses_eager_on_non_nvidia_platform(monkeypatch): + from transformers.models.dac.modeling_dac import Snake1d + + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _FusedDacSnake1d + + reference = Snake1d(8).cuda() + fused = _FusedDacSnake1d(8).cuda() + fused.load_state_dict(reference.state_dict()) + x = torch.randn(2, 8, 33, device="cuda") + monkeypatch.setattr(current_omni_platform, "is_cuda", lambda: False) + assert not fused.can_fuse(x) + assert torch.equal(fused(x), reference(x)) + + +@requires_cuda +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@torch.inference_mode() +def test_fused_snake_preserves_float64_precision(): + from transformers.models.dac.modeling_dac import Snake1d + + from vllm_omni.model_executor.models.higgs_audio_v2.higgs_audio_decoder import _FusedDacSnake1d + + reference = Snake1d(8).to("cuda", torch.float64) + fused = _FusedDacSnake1d(8).to("cuda", torch.float64) + fused.load_state_dict(reference.state_dict()) + x = torch.randn(2, 8, 33, device="cuda", dtype=torch.float64) + assert not fused.can_fuse(x) + assert torch.equal(fused(x), reference(x)) diff --git a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py index 3d43078ce..206d9922a 100755 --- a/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/model_executor/models/higgs_audio_v3/test_higgs_audio_v3.py @@ -24,6 +24,27 @@ from vllm_omni.entrypoints.openai.tts_adapters.higgs_audio_v3 import HiggsAudioV pytestmark = [pytest.mark.core_model, pytest.mark.cpu] +def test_external_decode_enters_backbone_compile_boundary(): + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration, + ) + + class Backbone(torch.nn.Module): + def forward(self, input_ids, positions, inputs_embeds): + assert input_ids is None + return inputs_embeds + positions[:, None] + + @property + def layers(self): + pytest.fail("External decode bypassed the Qwen3Model compilation boundary") + + talker = SimpleNamespace(_use_external_decode_cudagraph=True, model=Backbone()) + hidden = torch.randn(2, 4) + positions = torch.tensor([11, 27]) + actual = HiggsAudioV3TalkerForConditionalGeneration._run_qwen3_layers_eager(talker, positions, hidden) + torch.testing.assert_close(actual, hidden + positions[:, None]) + + # ---- AC-1: Configuration ---- @@ -284,7 +305,7 @@ class TestSamplerMethods: """Create a minimal talker-like object with sampler/masking methods.""" from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod - class FakeTalker: + class FakeTalker(SimpleNamespace): num_codebooks = 8 codebook_size = 1026 @@ -297,9 +318,10 @@ class TestSamplerMethods: """Create a fake talker with GPU-resident sampler state helpers.""" from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod - class FakeTalker: + class FakeTalker(SimpleNamespace): num_codebooks = 8 codebook_size = 1026 + config = type("Config", (), {"audio_full_sample_graph": False})() t = FakeTalker() t._decode_last_codes = torch.arange(num_rows * t.num_codebooks, dtype=torch.long).view( @@ -337,6 +359,7 @@ class TestSamplerMethods: "_get_audio_host_staging_buffer", "_apply_delay_pattern_masking_batched", "_update_delay_state_batched", + "_finish_audio_staging", "_prefill_row_mask", "_audio_seed_mask_from_step_input", ): @@ -426,9 +449,12 @@ class TestSamplerMethods: def test_sampling_metadata_invalid_temperature_is_sanitized(self, monkeypatch): t = self._make_minimal_talker() seen = [] - monkeypatch.setattr( - torch, "multinomial", lambda probs, num_samples: seen.append(probs) or probs.argmax(-1, True) - ) + + def sample(probs, num_samples): + seen.append(probs) + return probs.argmax(-1, True) + + monkeypatch.setattr(torch, "multinomial", sample) logits = torch.tensor([[0.0, 2.0], [3.0, 0.0]]) result = t._sample_audio_codes( @@ -444,9 +470,12 @@ class TestSamplerMethods: def test_sampling_metadata_is_expanded_per_request_and_vq(self, monkeypatch): t = self._make_minimal_talker() seen = [] - monkeypatch.setattr( - torch, "multinomial", lambda probs, num_samples: seen.append(probs) or probs.argmax(-1, True) - ) + + def sample(probs, num_samples): + seen.append(probs) + return probs.argmax(-1, True) + + monkeypatch.setattr(torch, "multinomial", sample) cb_logits = torch.tensor([[[4.0, 3.0, 2.0, 1.0]] * 2] * 2) # [2 requests, 2 VQs, 4 vocab] logits = cb_logits.reshape(-1, cb_logits.shape[-1]) metadata = self._sampling_metadata( @@ -840,7 +869,7 @@ class TestFeedbackMethods: """postprocess() should return codes from _last_audio_codes.""" from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod - class FakeTalker: + class FakeTalker(SimpleNamespace): _last_audio_codes = torch.tensor([[100, 200, 300, 400, 500, 600, 700, 800]]) _last_audio_code_valid = [True] _last_audio_host_staging = None @@ -859,7 +888,7 @@ class TestFeedbackMethods: """postprocess() should skip rows with -1 (no audio).""" from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod - class FakeTalker: + class FakeTalker(SimpleNamespace): _last_audio_codes = torch.tensor([[-1, -1, -1, -1, -1, -1, -1, -1]]) _last_audio_code_valid = [False] _last_audio_host_staging = None @@ -876,7 +905,7 @@ class TestFeedbackMethods: """postprocess() should advance cursor by 1 per call.""" from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod - class FakeTalker: + class FakeTalker(SimpleNamespace): _last_audio_codes = torch.tensor( [[100, 200, 300, 400, 500, 600, 700, 800], [-1, -1, -1, -1, -1, -1, -1, -1]] ) @@ -901,7 +930,7 @@ class TestReferenceAudioSubstitution: def _make_talker(self, query_start_loc): from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod - class FakeTalker: + class FakeTalker(SimpleNamespace): _last_step_query_start_loc = query_start_loc _apply_ref_audio_substitution = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_ref_audio_substitution @@ -981,7 +1010,7 @@ class TestAudioFeedback: embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16) torch.nn.init.ones_(embed.weight) - class FakeTalker: + class FakeTalker(SimpleNamespace): num_codebooks = 8 codebook_size = 1026 _audio_continuation_id = 99999 # fake audio token diff --git a/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py b/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py new file mode 100644 index 000000000..e0210295a --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_mrv2_state.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace +from unittest.mock import Mock + +import numpy as np +import pytest +import torch +from vllm.sampling_params import SamplingParams + +from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration, +) +from vllm_omni.model_executor.models.higgs_audio_v3.model_state import HiggsModelState +from vllm_omni.worker_v2.model_states.intermediate_buffer import OmniIntermediateBuffer + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def make_state(): + state = object.__new__(HiggsModelState) + state.sampling_slots = None + state.request_slots = {} + state.device = torch.device("cpu") + state.rope_state = state.prompt_embeds_state = state._static_inputs_embeds = None + state.requests = {} + state.generators = {} + state.metadata_key = state.metadata = None + state.intermediate_buffer = OmniIntermediateBuffer(4) + state.model = SimpleNamespace( + on_requests_finished=Mock(), + _resolve_token_ids=Mock(), + _audio_continuation_id=99, + update_decode_step_metadata=Mock(), + ) + return state + + +def request(name, seed, **kwargs): + return SimpleNamespace( + req_id=name, prompt_token_ids=[1, 2, 99], mm_features=[], sampling_params=SamplingParams(seed=seed, **kwargs) + ) + + +@pytest.mark.parametrize("slots", [False, True]) +def test_sampling_follows_request_identity_and_reinitializes_reused_slots(slots): + s = make_state() + if slots: + s.sampling_slots = { + name: torch.zeros(4, dtype=torch.long if name == "top_k" else torch.float32) + for name in ("temperature", "top_p", "top_k", "zeros", "ones") + } + s.sampling_slots["ones"].fill_(1) + s.add_request(0, request("a", 42, temperature=0.5)) + s.add_request(2, request("b", 99, temperature=0.9)) + a = s.generators["a"] + b = s.generators["b"] + first = s.sampling_metadata(["a", "b"]) + assert first.generators == {0: a, 1: b} + torch.testing.assert_close(first.temperature, torch.tensor([0.5, 0.9])) + swapped = s.sampling_metadata(["b", "a"]) + assert swapped.generators == {0: b, 1: a} + torch.testing.assert_close(swapped.temperature, torch.tensor([0.9, 0.5])) + s.remove_request("a") + s.add_request(0, request("new", 7)) + assert "a" not in s.generators and "a" not in s.requests + s.model.on_requests_finished.assert_called_once_with({"a"}) + assert s.sampling_metadata(["new"]).generators[0].initial_seed() == 7 + + +@pytest.mark.parametrize( + "kwargs", + [ + {"repetition_penalty": 1.2}, + {"logprobs": 1}, + {"min_p": 0.1}, + {"frequency_penalty": 0.2}, + {"prompt_logprobs": 1}, + {"min_tokens": 1}, + ], +) +def test_unsupported_sampling_is_rejected_before_admission(kwargs): + s = make_state() + with pytest.raises(ValueError, match="temperature/top-k/top-p"): + s.add_request(0, request("a", 42, **kwargs)) + assert not s.requests + + +def test_partial_prefill_not_classified_as_audio_and_slot_order_is_authoritative(): + s = make_state() + s.add_request(0, request("a", 42)) + s.add_request(2, request("b", 99)) + batch = SimpleNamespace( + req_ids=["b", "a"], + idx_mapping_np=np.array([2, 0]), + num_scheduled_tokens=np.array([1, 1]), + positions=torch.tensor([2, 1]), + query_start_loc=torch.tensor([0, 1, 2]), + num_reqs=2, + num_tokens=2, + ) + states = SimpleNamespace(num_computed_prefill_tokens=np.array([1, 0, 2, 0])) + s.run_preprocess(batch, {"input_ids": torch.tensor([99, 2])}, states) + assert s.model.update_decode_step_metadata.call_args.kwargs["audio_prompt_mode_rows"] == 0 + states.num_computed_prefill_tokens[0] = 2 + s.run_preprocess(batch, {"input_ids": torch.tensor([99, 99])}, states) + assert s.model.update_decode_step_metadata.call_args.kwargs["audio_prompt_mode_rows"] == 2 + + +def test_finalizer_filters_unsampled_prefill_and_keeps_owned_audio(): + s = make_state() + s.model = object.__new__(HiggsAudioV3TalkerForConditionalGeneration) + torch.nn.Module.__init__(s.model) + s.model.num_codebooks = 2 + snapshot = torch.tensor([[11, 12, 1, 0], [21, 22, 1, 0], [31, 32, 0, 1]]) + payload = {"_higgs_audio_snapshot": snapshot, "_higgs_invalid_rows": ()} + result = s.finalize_audio_snapshot(payload, [1, 0, 1]) + snapshot.zero_() + assert [x.tolist() for x in result["codes"]["audio"]] == [[[11, 12]], [], []] + assert payload["_higgs_invalid_rows"] == () + + +def test_upstream_warmup_retains_generic_text_sampling(): + s = make_state() + s.add_request(0, request("_warmup_0_", 0, repetition_penalty=1.2, logprobs=1)) + batch = SimpleNamespace(req_ids=["_warmup_0_"]) + sampler, rejection = s.custom_sampler(Mock()) + expected = (object(), torch.ones(1), torch.zeros(1)) + standard = Mock(return_value=expected) + output = sampler.sample_step(None, batch, None, None, standard) + assert output.sampler_output is expected[0] and output.multimodal_outputs is None + assert rejection is None + standard.assert_called_once_with(None, batch, None) + s.remove_request("_warmup_0_") + assert not s.requests + + +def test_static_embeddings_exclude_padding_and_preserve_storage(): + s = make_state() + embed = torch.nn.Embedding(100, 3) + s.model.model = SimpleNamespace(embed_tokens=embed) + s.model._apply_audio_feedback = Mock(side_effect=lambda x, ids: x + 2) + s.model._apply_ref_audio_substitution = Mock(side_effect=lambda x, *args: x + 5) + s._static_inputs_embeds = torch.full((8, 3), 999.0) + ptr = s._static_inputs_embeds.data_ptr() + ids = torch.tensor([1, 2, 3, 0]) + batch = SimpleNamespace(num_tokens=3, num_reqs=3, num_tokens_after_padding=4, has_prefill=False) + s._prepare_static_embeddings(batch, {"input_ids": ids}) + assert s._static_inputs_embeds.data_ptr() == ptr + torch.testing.assert_close(s._static_inputs_embeds[:3], embed(ids[:3]) + 2) + assert s._static_inputs_embeds[3].count_nonzero() == 0 + assert s.model._apply_audio_feedback.call_args.args[1].numel() == 3 + s.model._apply_ref_audio_substitution.assert_not_called() + batch.has_prefill = True + batch.positions = torch.arange(4) + s._prepare_static_embeddings(batch, {"input_ids": ids, "model_intermediate_buffer": [{}]}) + torch.testing.assert_close(s._static_inputs_embeds[:3], embed(ids[:3]) + 7) + + +def test_mixed_feedback_updates_only_single_token_spans_with_codes(): + s = make_state() + s.model._ensure_decode_state_capacity = Mock() + s.model._decode_last_codes = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8]]) + s.model._decode_has_codes = torch.tensor([True, True, False, True]) + s.model.multimodal_embedding = lambda codes: codes.sum(1)[:, None].expand(-1, 3) + batch = SimpleNamespace(num_tokens=7, num_reqs=4, query_start_loc=torch.tensor([0, 1, 4, 5, 7])) + embeds = torch.arange(21).reshape(7, 3).float() + output = s._apply_audio_feedback(embeds, torch.ones(7).long(), batch) + expected = embeds.clone() + expected[0] = 3 + torch.testing.assert_close(output, expected) + assert embeds[0].tolist() == [0, 1, 2] + + +def test_direct_payload_owns_rows_and_filters_partial_prefill(): + s = make_state() + s.direct_payload = True + s.model.num_codebooks = 2 + snapshot = torch.tensor([[11, 12, 1, 0], [21, 22, 1, 0], [31, 32, 0, 1]]) + out = s.finalize_audio_snapshot({"_higgs_audio_snapshot": snapshot}, [1, 0, 1]) + snapshot.zero_() + assert out.inter_stage[0]["codes.audio"].tolist() == [[11, 12]] + assert out.inter_stage[1:] == [None, None] + + +def test_registered_audio_sampler_preserves_rows_counts_and_owned_output(mocker): + state = make_state() + metadata, payload = object(), {"_higgs_audio_snapshot": torch.ones(2, 4)} + state.sampling_metadata = Mock(return_value=metadata) + state.model.compute_logits = Mock(return_value=torch.ones(2, 8)) + state.model.sample = Mock( + return_value=SimpleNamespace(sampled_token_ids=torch.tensor([[9], [8]]), logprobs_tensors=None) + ) + state.model.post_sample_multimodal_outputs = Mock(return_value=payload) + count, rejected = torch.tensor([1, 0]), torch.tensor([0, 0]) + mocker.patch( + "vllm_omni.model_executor.models.higgs_audio_v3.model_state.get_num_sampled_and_rejected", + return_value=(count, rejected), + ) + batch = SimpleNamespace( + req_ids=["b", "a"], + num_reqs=2, + num_draft_tokens=0, + idx_mapping=torch.tensor([3, 1]), + logits_indices=torch.tensor([2, 0]), + seq_lens=None, + cu_num_logits=None, + ) + hidden = torch.arange(12).reshape(3, 4).float() + standard = Mock(side_effect=AssertionError("must not sample twice")) + sampler, rejection = state.custom_sampler(Mock()) + result = sampler.sample_step(hidden, batch, SimpleNamespace(prefill_len=SimpleNamespace(gpu=None)), None, standard) + state.sampling_metadata.assert_called_once() + torch.testing.assert_close(state.model.compute_logits.call_args.args[0], hidden[[2, 0]]) + assert state.model.compute_logits.call_args.args[1] is metadata + assert result.num_sampled is count and result.num_rejected is rejected + assert result.multimodal_outputs is payload and result.owns_multimodal_outputs + assert not result.include_hidden_states and result.finalize_multimodal == state.finalize_audio_snapshot + assert rejection is None + standard.assert_not_called() + + +@pytest.mark.parametrize("draft,grammar", [(1, None), (0, object())]) +def test_registered_audio_sampler_rejects_unsupported_modes(draft, grammar): + sampler, _ = make_state().custom_sampler(Mock()) + batch = SimpleNamespace(req_ids=["real-request"], num_draft_tokens=draft) + with pytest.raises(ValueError, match="grammar or speculative"): + sampler.sample_step(None, batch, None, grammar, Mock()) diff --git a/tests/model_executor/models/higgs_audio_v3/test_native_payload.py b/tests/model_executor/models/higgs_audio_v3/test_native_payload.py new file mode 100644 index 000000000..47911d99f --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_native_payload.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( + talker2code2wav, + talker2code2wav_full_payload, + talker2code2wav_token_only, +) +from vllm_omni.outputs.mm_outputs import MultimodalPayload + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("wrapped", [False, True]) +@pytest.mark.parametrize("rows", [0, 4, 8, 9, 30]) +@pytest.mark.parametrize("offset", [0, -16, 1016]) +def test_native_full_payload_preserves_legacy_dedelay_and_tail(rows, offset, wrapped): + audio = torch.arange(rows * 8).reshape(rows, 8) + offset + payload = {"codes": {"audio": audio}} + if wrapped: + payload = MultimodalPayload(metadata=payload) + out = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=payload)]) + legacy = talker2code2wav([out])[0]["prompt_token_ids"] + native = talker2code2wav_full_payload(None, {"codes.audio": audio}, None) + assert native["codes"]["audio"].tolist() == legacy + # Independent expected layout: book q reads delayed rows q + frame. + frames = max(rows - 7, 0) if rows >= 8 else 0 + kept = frames - 1 if frames >= 2 else frames + expected = [] + for q in range(8): + for frame in range(kept): + value = (q + frame) * 8 + q + offset + expected.append(value if 0 <= value < 1024 else 0) + assert legacy == expected + assert native["meta"]["finished"] + assert talker2code2wav_token_only([out])[0]["prompt_token_ids"] == legacy + + +def test_native_control_slot_without_audio_payload(): + out = SimpleNamespace(finished=True, outputs=[SimpleNamespace(multimodal_output=None)]) + assert talker2code2wav_token_only([out])[0]["prompt_token_ids"] == [0] diff --git a/tests/model_executor/models/higgs_audio_v3/test_sampling_kernels.py b/tests/model_executor/models/higgs_audio_v3/test_sampling_kernels.py new file mode 100644 index 000000000..aebabd1d2 --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_sampling_kernels.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused Higgs sampling kernels must be bit-identical to the eager code they replace.""" + +import pytest +import torch +import torch.nn.functional as F + +from tests.helpers.mark import hardware_test +from vllm_omni.platforms import current_omni_platform + +pytestmark = [pytest.mark.core_model, pytest.mark.skipif(not current_omni_platform.is_cuda(), reason="requires CUDA")] + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("numel", [8 * 1026, 1000, 300000]) +def test_batched_generator_noise_matches_per_request_exponential(numel): + from vllm_omni.utils.seeded_exponential import fill_exponential_rows + + device = torch.device("cuda") + seeds = [42, 42, None, 2**63 + 5] + + def generators(): + gens = [None if s is None else torch.Generator(device=device).manual_seed(s) for s in seeds] + resumed = gens[1] + assert resumed is not None + resumed.set_offset(resumed.get_offset() + 4) # a request that already sampled + return gens + + default = torch.cuda.default_generators[device.index or 0] + start = default.get_offset() + expected_gens = generators() + expected = torch.empty(len(seeds), numel, device=device) + for row, generator in enumerate(expected_gens): + expected[row].exponential_(generator=generator) + expected_default = default.get_offset() + + default.set_offset(start) + actual_gens = generators() + actual = fill_exponential_rows(torch.empty_like(expected), actual_gens) + assert torch.equal(actual, expected) + assert default.get_offset() == expected_default + for got, want in zip(actual_gens, expected_gens): + if want is not None: + assert got.get_offset() == want.get_offset() + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_fused_decode_embeddings_match_feedback_and_zero_padding(): + from vllm_omni.model_executor.models.higgs_audio_v3.step_kernels import decode_embeddings + + torch.manual_seed(0) + vocab, hidden, books, audio_vocab = 1000, 96, 8, 17 + text = torch.randn(vocab, hidden, device="cuda", dtype=torch.bfloat16) + audio = torch.randn(books * audio_vocab, hidden, device="cuda", dtype=torch.bfloat16) + ids = torch.randint(0, vocab, (5,), device="cuda") + ids[0] = -1 + codes = torch.randint(0, audio_vocab, (8, books), device="cuda") + has = torch.tensor([True, False, True, False, True, True, False, False], device="cuda") + offsets = torch.arange(books, device="cuda") * audio_vocab + expected = torch.where( + has[:5, None], F.embedding(codes[:5] + offsets, audio).sum(-2), F.embedding(ids.clamp_min(0), text) + ) + out = torch.full((8, hidden), 3.0, device="cuda", dtype=torch.bfloat16) + decode_embeddings(out, ids, text, audio, codes, has, 7, audio_vocab) + assert torch.equal(out[:5], expected) + assert torch.count_nonzero(out[5:7]) == 0 + assert torch.all(out[7] == 3) + + +@hardware_test(res={"cuda": "L4"}, num_cards=1) +@pytest.mark.parametrize("prompt_mode", [False, True]) +def test_step_tail_guard_matches_eager_restore_reset_and_validation(prompt_mode): + from vllm_omni.model_executor.models.higgs_audio_v3.step_kernels import step_tail_guard + + audio, eos = 7, 9 + # Rows: decode audio, decode EOS, 3-token prefill ending in audio, decode audio. + ids = torch.tensor([audio, eos, 1, 2, audio, audio], device="cuda") + qsl = torch.tensor([0, 1, 2, 5, 6], device="cuda", dtype=torch.int32) + has = torch.tensor([True, True, True, False], device="cuda") + done = torch.tensor([False, False, True, False], device="cuda") + delay = torch.tensor([3, 4, 5, 6], device="cuda") + eoc = torch.tensor([1, 2, 3, 4], device="cuda") + error = torch.zeros(1, device="cuda", dtype=torch.int32) + step_tail_guard(ids, qsl, has, done, delay, eoc, error, audio_id=audio, eos_id=eos, prompt_mode=prompt_mode) + prefill = torch.tensor([False, False, True, False], device="cuda") + tail = ids[(qsl[1:] - 1).long()] + want_done = torch.tensor([False, False, True, False], device="cuda") + if prompt_mode: + want_done |= tail == eos + want_done &= ~prefill + assert torch.equal(done, want_done) + assert torch.equal(has, torch.tensor([True, True, False, False], device="cuda")) + assert delay.tolist() == [3, 4, 0, 6] and eoc.tolist() == [1, 2, -1, 4] + # The EOS row is valid only as a terminal prompt-mode row. + assert int(error.item()) == (0 if prompt_mode else 1) diff --git a/tests/model_executor/models/higgs_audio_v3/test_text_head_skip.py b/tests/model_executor/models/higgs_audio_v3/test_text_head_skip.py new file mode 100644 index 000000000..56d84df5d --- /dev/null +++ b/tests/model_executor/models/higgs_audio_v3/test_text_head_skip.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + HiggsAudioV3TalkerForConditionalGeneration, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("fallback", [None, "prefill", "logprobs", "allowed", "bad_words", "penalty", "inactive"]) +def test_text_head_only_skipped_for_direct_audio_decode(fallback): + hidden = torch.randn(3, 4) + calls = [] + + def logits_processor(head, h): + calls.append(h) + return torch.ones(3, 7) + + model = SimpleNamespace( + _audio_continuation_id=5, + _last_step_input_ids=torch.ones(4 if fallback == "prefill" else 3, dtype=torch.long), + _fast_audio_direct_rows=0 if fallback == "inactive" else 3, + _text_vocab_size=7, + lm_head=None, + logits_processor=logits_processor, + ) + metadata = SimpleNamespace( + no_penalties=fallback != "penalty", + max_num_logprobs=1 if fallback == "logprobs" else None, + allowed_token_ids_mask=torch.ones(3, 7) if fallback == "allowed" else None, + bad_words_token_ids={0: [1]} if fallback == "bad_words" else {}, + ) + logits = HiggsAudioV3TalkerForConditionalGeneration.compute_logits(model, hidden, metadata) + assert logits.shape == (3, 7) + assert model._last_logits_hidden is hidden + assert len(calls) == (fallback is not None) diff --git a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py index d62abd60f..f22a2c693 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py +++ b/tests/model_executor/models/minicpmo_4_5/test_cuda_graph_wrapper.py @@ -215,6 +215,8 @@ class _MiniDiT(nn.Module): cnn_cache_buffer[b_idx] = x[:, -2:, :].transpose(1, 2).contiguous() dt = x.shape[1] att_cache_buffer[b_idx][:, :, :dt, :] = x.unsqueeze(1) + if att_b is not None: + att_cache_buffer[b_idx][:, :, dt:, :] = att_b x = self.final_layer(x) x = x.transpose(1, 2) return x @@ -248,6 +250,9 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes( with torch.inference_mode(): for _, chunk_size, old_att_len in ((2, 10, 0), (2, 10, 5)): inputs = _cfm_inputs(2, chunk_size, old_att_len) + # A cache output must overwrite every row, including the old tail. + inputs[4].fill_(float("nan")) + inputs[5].fill_(float("nan")) eager_inputs = tuple(v.clone() for v in inputs) with torch.no_grad(): @@ -268,6 +273,8 @@ def test_cfm_graph_replay_matches_eager_for_uncached_and_cached_shapes( torch.testing.assert_close(graph_result, eager_result, rtol=1e-4, atol=1e-5) torch.testing.assert_close(graph_cnn, eager_inputs[4], rtol=1e-4, atol=1e-5) torch.testing.assert_close(graph_att, eager_inputs[5], rtol=1e-4, atol=1e-5) + assert torch.isfinite(graph_att).all() + torch.testing.assert_close(graph_att[:, :, :, chunk_size:, :], inputs[3]) wrapper._flush() diff --git a/tests/model_executor/stage_input_processors/test_higgs_audio_v3_async_chunk_batch.py b/tests/model_executor/stage_input_processors/test_higgs_audio_v3_async_chunk_batch.py new file mode 100644 index 000000000..5808ff185 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_higgs_audio_v3_async_chunk_batch.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""The batched Higgs async-chunk builder must emit exactly what the per-request one does.""" + +from collections import defaultdict +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( + talker2code2wav_async_chunk, + talker2code2wav_async_chunk_batch, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _manager(): + extra = { + "codec_chunk_frames": 5, + "codec_left_context_frames": 3, + "codec_right_holdback_frames": 1, + "initial_codec_chunk_frames": 2, + } + return SimpleNamespace(connector=SimpleNamespace(config={"extra": extra}), code_prompt_token_ids=defaultdict(list)) + + +class _Request: + def __init__(self, rid, finish_step): + self.external_req_id = rid + self.finish_step = finish_step + self.step = 0 + + def is_finished(self): + return self.step >= self.finish_step + + +def _payload_items(payload): + if payload is None: + return None + meta = payload.meta + return ( + payload.codes.audio.tolist(), + None if meta is None else (meta.left_context_size, meta.right_holdback_size, bool(meta.finished)), + ) + + +def test_batch_async_chunk_matches_per_request_builder(): + torch.manual_seed(0) + single, batched = _manager(), _manager() + finish = [23, 9, 31] + requests_single = [_Request(f"r{i}", f) for i, f in enumerate(finish)] + requests_batch = [_Request(f"r{i}", f) for i, f in enumerate(finish)] + for step in range(32): + outputs: list[dict | None] = [] + flags: list[bool] = [] + for i, request in enumerate(requests_single): + if step > request.finish_step: + outputs.append(None) + elif step == 7 and i == 2: + outputs.append(None) # a step without an emitted row + else: + outputs.append({"codes": {"audio": torch.randint(0, 1026, (1, 8))}}) + flags.append(step == request.finish_step) + for request in requests_single + requests_batch: + request.step = step + active = [i for i, request in enumerate(requests_single) if step <= request.finish_step] + want = [ + _payload_items(talker2code2wav_async_chunk(single, outputs[i], requests_single[i], flags[i])) + for i in active + ] + got = talker2code2wav_async_chunk_batch( + batched, + [outputs[i] for i in active], + [requests_batch[i] for i in active], + [flags[i] for i in active], + ) + assert [_payload_items(p) for p in got] == want + assert single.higgs_v3_emitted_frames == batched.higgs_v3_emitted_frames == {} diff --git a/tests/worker/test_higgs_forward_metadata_scope.py b/tests/worker/test_higgs_forward_metadata_scope.py new file mode 100644 index 000000000..e97ce5ee6 --- /dev/null +++ b/tests/worker/test_higgs_forward_metadata_scope.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import pytest +import torch +from vllm.v1.worker.gpu_model_runner import GPUModelRunner + +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.mark.parametrize("fail", [False, True]) +def test_runner_closes_metadata_scope_on_return_and_error(monkeypatch, fail): + events = [] + + def begin(**kw): + events.append("begin") + + def end(): + events.append("end") + + def forward(*args, **kw): + assert events == ["begin"] + events.extend(["warmup", "capture"]) + if fail: + raise RuntimeError("capture failed") + return torch.ones(2) + + r = object.__new__(OmniGPUModelRunner) + r.model = SimpleNamespace( + supports_omni_decode_step_metadata=True, update_decode_step_metadata=begin, finish_decode_step_forward=end + ) + r.input_batch = SimpleNamespace(req_ids=["a", "b"]) + r._build_model_kwargs_extra = lambda: {} + monkeypatch.setattr(GPUModelRunner, "_model_forward", forward) + if fail: + with pytest.raises(RuntimeError, match="capture failed"): + r._model_forward() + else: + torch.testing.assert_close(r._model_forward(), torch.ones(2)) + assert events == ["begin", "warmup", "capture", "end"] + + +@pytest.mark.parametrize("requires_tails", [False, True]) +@pytest.mark.parametrize("async_scheduling", [False, True]) +def test_cpu_tail_metadata_requires_model_capability(monkeypatch, requires_tails, async_scheduling): + observed = {} + runner = object.__new__(OmniGPUModelRunner) + runner.model = SimpleNamespace( + supports_omni_decode_step_metadata=True, + requires_cpu_input_tail_ids=requires_tails, + update_decode_step_metadata=lambda **kwargs: observed.update(kwargs), + ) + ids = torch.tensor([10, 11, 12]) + runner.use_async_scheduling = async_scheduling + runner.input_ids = SimpleNamespace(gpu=ids, cpu=ids) + runner.query_start_loc = SimpleNamespace(cpu=torch.tensor([0, 2, 3])) + runner.input_batch = SimpleNamespace(req_ids=["a", "b"]) + runner._build_model_kwargs_extra = lambda: {} + monkeypatch.setattr(GPUModelRunner, "_model_forward", lambda *args, **kwargs: torch.ones(2)) + runner._model_forward(input_ids=ids) + expected = [11, 12] if requires_tails and not async_scheduling else None + assert observed["cpu_input_tail_ids"] == expected + + +@pytest.mark.parametrize("generation", [False, True]) +@pytest.mark.parametrize("enabled,fail", [(False, False), (True, False), (True, True)]) +def test_auxiliary_capture_finishes_before_worker_readiness(monkeypatch, generation, enabled, fail): + from vllm.v1.worker.gpu_worker import Worker + + from vllm_omni.worker.base import OmniGPUWorkerBase + from vllm_omni.worker.gpu_generation_worker import GPUGenerationWorker + + events: list[str] = [] + + def capture(): + assert events == ["warmup"] + events.append("capture") + if fail: + raise RuntimeError("auxiliary capture failed") + + monkeypatch.setattr(Worker, "compile_or_warm_up_model", lambda _: events.append("warmup")) + worker = object.__new__(GPUGenerationWorker if generation else OmniGPUWorkerBase) + worker.use_v2_model_runner = generation + model = SimpleNamespace(capture_auxiliary_graphs=capture) if enabled else SimpleNamespace() + worker.model_runner = SimpleNamespace(model=model, profile_run=lambda: events.append("warmup")) + if fail: + with pytest.raises(RuntimeError, match="auxiliary capture failed"): + worker.compile_or_warm_up_model() + else: + worker.compile_or_warm_up_model() + events.append("ready") + assert events == ["warmup"] + (["capture"] if enabled else []) + ([] if fail else ["ready"]) diff --git a/tests/worker_v2/test_capture_query_bound_gpu.py b/tests/worker_v2/test_capture_query_bound_gpu.py new file mode 100644 index 000000000..93c53f313 --- /dev/null +++ b/tests/worker_v2/test_capture_query_bound_gpu.py @@ -0,0 +1,102 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from vllm.config import 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] + + +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_full_graph_prefill_uses_real_flash_attention_backend(): + if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 9: + pytest.skip("requires FA3 on Hopper or newer") + _check_prefill() + + +def _check_prefill(): + torch.manual_seed(17) + device = torch.device("cuda:0") + c = VllmConfig() + c.model_config = SimpleNamespace( + get_num_attention_heads=lambda p: 32, + rswa_window=None, + is_diffusion=False, + is_mm_prefix_lm=False, + hf_config=SimpleNamespace(model_type="qwen3"), + ) + c.scheduler_config.max_num_seqs = 64 + c.compilation_config.cudagraph_mode = CUDAGraphMode.FULL + c.compilation_config.max_cudagraph_capture_size = 256 + c.attention_config.flash_attn_version = 3 + 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.max_model_len = 256 + 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(32, 128, 128**-0.5, 8, None, None, "auto") + assert impl.vllm_flash_attn_version == 3 + buffers = InputBuffers(64, 256, device) + batch = InputBatch.make_dummy(64, 256, buffers) + tables = torch.arange(1024, dtype=torch.int32, device=device).reshape(64, 16) + slots = torch.zeros((1, 256), dtype=torch.int64, device=device) + cache = torch.randn(1024, 8, 16, 256, device=device, dtype=torch.bfloat16) + q = torch.randn(256, 32, 128, device=device, dtype=torch.bfloat16) + k = v = torch.empty(256, 8, 128, device=device, dtype=torch.bfloat16) + layer = SimpleNamespace( + _q_scale=torch.ones((), device=device), + _k_scale=torch.ones((), device=device), + _v_scale=torch.ones((), device=device), + ) + batch.query_start_loc.copy_(torch.arange(65, device=device, dtype=torch.int32) * 4) + batch.seq_lens.fill_(4) + prepare = OmniModelState.prepare_attn + md = prepare(state, batch, CUDAGraphMode.NONE, (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) + # Same captured 64-request shape: one 212-token prompt, 44 single-token + # requests, then zero-length padded rows. No backend is mocked. + lengths = np.array([212] + [1] * 44 + [0] * 19, dtype=np.int32) + batch.query_start_loc.copy_(torch.tensor(np.r_[0, np.cumsum(lengths)], device=device, dtype=torch.int32)) + batch.seq_lens.copy_(torch.tensor(lengths, device=device)) + from dataclasses import replace + + runtime = replace( + batch, + max_query_len=212, + num_scheduled_tokens=lengths, + query_start_loc_np=np.r_[0, np.cumsum(lengths)], + seq_lens_cpu_upper_bound=torch.tensor(lengths), + ) + eager_md = state.prepare_attn(runtime, CUDAGraphMode.NONE, (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.NONE, (tables,), slots, [[group]], kvconfig) + out.fill_(float("nan")) + graph.replay() + torch.accelerator.synchronize() + assert md.max_query_len == 256 + torch.testing.assert_close(out[:212], expected[:212], atol=0.002, rtol=0.002) diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_mrv2_h200.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_mrv2_h200.yaml new file mode 100644 index 000000000..a1cbd760e --- /dev/null +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_mrv2_h200.yaml @@ -0,0 +1,169 @@ +# Higgs v3 Model Runner V2 streaming profile, validated on one H200. +# Select an available physical GPU with CUDA_VISIBLE_DEVICES before launch. +# First audio: 8 frames (320 ms); following chunks: 25 frames (1 s). +# Codec graph windows are exact: 8 first frames, 25 + 8 left context. +# AR and codec capacities are 64; tested client concurrencies are 32 and 64. +# Keep model weights and sampling parameters identical for comparisons. +# Server flags used in measurements: --disable-log-stats --disable-uvicorn-access-log. +# Auxiliary graphs are captured after KV sizing and are outside the utilization +# budget. Leave at least 10 GiB free for codec graph pools on this H200 profile; +# inspect both allocated and reserved capture logs before raising stage budgets. +# Sampler graphs cover tensor-valued temperature/top_k/top_p. Disabling top-k +# or top-p for an entire batch (None metadata) uses eager sampling. +async_chunk: true +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + codec_streaming: true + codec_chunk_frames: 25 + codec_left_context_frames: 8 + codec_right_holdback_frames: 0 + initial_codec_chunk_frames: 8 + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 5000 + connector_get_max_wait: 500 +stages: +- stage_id: 0 + max_num_seqs: 64 + gpu_memory_utilization: 0.6 + enforce_eager: false + attention_backend: FLASH_ATTN + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: true + compilation_config: + cudagraph_capture_sizes: + - 1 + - 2 + - 4 + - 8 + - 16 + - 24 + - 32 + - 40 + - 48 + - 56 + - 64 + - 96 + - 128 + - 192 + - 256 + - 384 + - 512 + cudagraph_mode: FULL + cudagraph_num_of_warmups: 1 + max_cudagraph_capture_size: 512 + cudagraph_copy_inputs: true + max_num_batched_tokens: 4096 + max_model_len: 8192 + devices: '0' + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 1.0 + top_p: 0.95 + top_k: 50 + max_tokens: 2048 + seed: 42 + repetition_penalty: 1.0 + hf_overrides: + audio_async_prompt_mode: true + audio_async_payload: true + audio_mrv2_static_inputs: true + audio_mrv2_slot_parameters: true + audio_mrv2_direct_payload: true + audio_mrv2_sample_graph_buckets: true + audio_mrv2_async_state_indices: true + audio_mrv2_fused_state: true +- stage_id: 1 + max_num_seqs: 64 + gpu_memory_utilization: 0.25 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + async_scheduling: false + max_num_batched_tokens: 65536 + max_model_len: 65536 + devices: '0' + input_connectors: + from_stage_0: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + repetition_penalty: 1.0 + hf_overrides: + codec_graph_single_max_frames: 33 + codec_graph_batch_sizes: + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + - 8 + - 9 + - 10 + - 11 + - 12 + - 13 + - 14 + - 15 + - 16 + - 17 + - 18 + - 19 + - 20 + - 21 + - 22 + - 23 + - 24 + - 25 + - 26 + - 27 + - 28 + - 29 + - 30 + - 31 + - 32 + - 33 + - 34 + - 35 + - 36 + - 37 + - 38 + - 39 + - 40 + - 41 + - 42 + - 43 + - 44 + - 45 + - 46 + - 47 + - 48 + - 49 + - 50 + - 51 + - 52 + - 53 + - 54 + - 55 + - 56 + - 57 + - 58 + - 59 + - 60 + - 61 + - 62 + - 63 + - 64 + codec_graph_frame_sizes: + - 8 + - 33 +cuda_mps: true +model_runner: v2 diff --git a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py index 76ca7eaa5..7d2b16d50 100644 --- a/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py +++ b/vllm_omni/distributed/omni_connectors/model_runner/omni_connector_payload_transport.py @@ -1280,11 +1280,12 @@ class _OmniConnectorPayloadTransportMixin(_OmniConnectorRuntimeMixin): payload_kwarg: pooling_output, "request": request, } - supports_is_finished = getattr( - self, - "_custom_process_supports_is_finished", - self._custom_process_supports_is_finished_kwarg(), - ) + # The signature check is cached at initialization; never re-inspect + # the hook for every request of every model step. + supports_is_finished = getattr(self, "_custom_process_supports_is_finished", None) + if supports_is_finished is None: + supports_is_finished = self._custom_process_supports_is_finished_kwarg() + self._custom_process_supports_is_finished = supports_is_finished is_finished_fn = getattr(request, "is_finished", None) if callable(is_finished_fn): try: diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py index 74cdbb2db..77cb17983 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_decoder.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """HiggsAudio codec decoder kernel for higgs-audio v2. This module hosts the parameter-side building blocks for the higgs-audio-v2 @@ -16,11 +16,14 @@ from __future__ import annotations import json import os +import types from typing import Any import torch import torch.nn as nn +from vllm_omni.platforms import current_omni_platform + __all__ = [ "HiggsAudioVQLayer", "HiggsAudioRVQ", @@ -99,6 +102,138 @@ class _Snake1d(nn.Module): return _snake(x, self.alpha) +class _FusedDacSnake1d(nn.Module): + """Drop-in for ``transformers`` ``DacModel`` Snake1d with a fused CUDA kernel. + + Same ``alpha`` parameter, so checkpoints load unchanged. Inference on CUDA + runs one bit-identical kernel instead of about seven elementwise launches; + everything else keeps the transformers expression. + """ + + def __init__(self, channels: int): + super().__init__() + self.alpha = nn.Parameter(torch.ones(1, channels, 1)) + self._inverse: tuple[tuple[int, int], torch.Tensor] | None = None + + def can_fuse(self, x: torch.Tensor) -> bool: + # The kernel uses NVIDIA PTX and float32 intermediates. ROCm tensors + # also report is_cuda; retain the eager path there and for float64. + return ( + current_omni_platform.is_cuda() + and x.is_cuda + and x.dim() == 3 + and not torch.is_grad_enabled() + and x.dtype in (torch.float16, torch.bfloat16, torch.float32) + and x.dtype == self.alpha.dtype + ) + + def fused( + self, + x: torch.Tensor, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + write_sum: bool = False, + ) -> tuple[torch.Tensor | None, torch.Tensor]: + """Snake of ``residual + (x + bias)`` in one kernel; see ``fused_snake1d``.""" + from .snake_kernel import fused_snake1d + + alpha = self.alpha.detach() + try: + key = (alpha.data_ptr(), alpha._version) + except RuntimeError: # inference tensors carry no version counter + key = None + if key is not None and self._inverse is not None and self._inverse[0] == key: + inverse = self._inverse[1] + else: + # The eager (alpha + 1e-9).reciprocal(), computed once per weight + # update in the same dtype instead of on every call. + inverse = (alpha + 1e-9).reciprocal().reshape(-1).contiguous() + self._inverse = None if key is None else (key, inverse) + return fused_snake1d(x, alpha.reshape(-1), inverse, bias=bias, residual=residual, write_sum=write_sum) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.can_fuse(x): + return self.fused(x)[1] + shape = x.shape + x = x.reshape(shape[0], shape[1], -1) + x = x + (self.alpha + 1e-9).reciprocal() * torch.sin(self.alpha * x).pow(2) + return x.reshape(shape) + + +def _conv_without_bias(conv: nn.Module, x: torch.Tensor) -> torch.Tensor: + # cuDNN convolutions add the bias as a separate kernel anyway, so leaving + # it to the next fused Snake changes no bits. + if isinstance(conv, nn.ConvTranspose1d): + return nn.functional.conv_transpose1d( + x, conv.weight, None, conv.stride, conv.padding, conv.output_padding, conv.groups, conv.dilation + ) + return conv._conv_forward(x, conv.weight, None) + + +def _fused_dac_decoder_forward(decoder: nn.Module, hidden_state: torch.Tensor) -> torch.Tensor: + """transformers ``DacDecoder.forward`` with conv biases and residual adds + folded into the Snake kernels (5 elementwise passes per residual unit -> 2). + + A value is carried as ``residual + (raw + bias)`` until the next Snake + consumes it, and only materialized where a residual connection needs it. + """ + if not decoder.snake1.can_fuse(hidden_state): + return type(decoder).forward(decoder, hidden_state) + raw, bias, residual = _conv_without_bias(decoder.conv1, hidden_state), decoder.conv1.bias, None + for block in decoder.block: + _, act = block.snake1.fused(raw, bias, residual) + raw, bias, residual = _conv_without_bias(block.conv_t1, act), block.conv_t1.bias, None + for unit in (block.res_unit1, block.res_unit2, block.res_unit3): + residual, act = unit.snake1.fused(raw, bias, residual, write_sum=True) + _, act = unit.snake2.fused(_conv_without_bias(unit.conv1, act), unit.conv1.bias) + raw, bias = _conv_without_bias(unit.conv2, act), unit.conv2.bias + _, act = decoder.snake1.fused(raw, bias, residual) + return decoder.tanh(decoder.conv2(act)) + + +def _fuse_dac_decoder(decoder: nn.Module) -> int: + """Swap transformers DAC Snake1d modules for fused ones (before weights + load) and, when the decoder has the stock layout, fuse its forward.""" + try: + from transformers.models.dac.modeling_dac import DacDecoder + from transformers.models.dac.modeling_dac import Snake1d as DacSnake1d + except ImportError: + return 0 + replaced = 0 + for module in list(decoder.modules()): + for name, child in list(module.named_children()): + if isinstance(child, DacSnake1d): + fused = _FusedDacSnake1d(int(child.alpha.shape[1])) + fused.to(device=child.alpha.device, dtype=child.alpha.dtype) + setattr(module, name, fused) + replaced += 1 + if type(decoder) is DacDecoder and _has_fusable_dac_layout(decoder): + decoder.forward = types.MethodType(_fused_dac_decoder_forward, decoder) + return replaced + + +def _has_fusable_dac_layout(decoder: nn.Module) -> bool: + convs = [decoder.conv1] + for block in decoder.block: + units = (block.res_unit1, block.res_unit2, block.res_unit3) + convs.append(block.conv_t1) + for unit in units: + conv1, conv2 = unit.conv1, unit.conv2 + # Length-preserving convs, so the residual add needs no crop. + if conv1.stride[0] != 1 or 2 * conv1.padding[0] != conv1.dilation[0] * (conv1.kernel_size[0] - 1): + return False + if conv2.stride[0] != 1 or conv2.kernel_size[0] != 1 or conv2.padding[0] != 0: + return False + convs += [conv1, conv2] + if not all(isinstance(m, _FusedDacSnake1d) for m in (block.snake1, *(u.snake1 for u in units))): + return False + if not all(isinstance(u.snake2, _FusedDacSnake1d) for u in units): + return False + return isinstance(decoder.snake1, _FusedDacSnake1d) and all( + conv.bias is not None and conv.padding_mode == "zeros" for conv in convs + ) + + def _wn_conv1d(*args: Any, **kwargs: Any) -> nn.Module: from torch.nn.utils import weight_norm @@ -248,6 +383,7 @@ def build_higgs_audio_acoustic_decoder( dac_cfg = DacConfig(**tokenizer_config["acoustic_model_config"]) dac_model = DacModel(dac_cfg) decoder = dac_model.decoder.to(device) + _fuse_dac_decoder(decoder) adjust_conv_transpose_output_padding(decoder) if hasattr(decoder, "tanh"): decoder.tanh = nn.Identity() diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/snake_kernel.py b/vllm_omni/model_executor/models/higgs_audio_v2/snake_kernel.py new file mode 100644 index 000000000..e965c18bb --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v2/snake_kernel.py @@ -0,0 +1,113 @@ +# ruff: noqa: N803 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused DAC Snake activation: ``x + 1 / (alpha + 1e-9) * sin(alpha * x) ** 2``. + +The eager expression launches about seven elementwise kernels per call, each +reading and writing the whole activation; in the codec's CUDA graphs they are +most of its GPU time. This kernel reads ``x`` once and rounds every +intermediate to the tensor dtype exactly where the eager ops do, so outputs are +bit-identical to the unfused expression. It can also absorb the preceding conv +bias add and residual add, which are otherwise two more full passes each. +""" + +import torch +from vllm.triton_utils import tl, tldevice, triton + + +@triton.jit +def _mul_rn(a, b): + # Explicit round-to-nearest ops are never contracted into FMA, so each + # product/sum rounds exactly where the separate eager kernels round. + return tl.inline_asm_elementwise("mul.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1) + + +@triton.jit +def _add_rn(a, b): + return tl.inline_asm_elementwise("add.rn.f32 $0, $1, $2;", "=r,r,r", [a, b], dtype=tl.float32, is_pure=True, pack=1) + + +@triton.jit +def _dac_snake_kernel( + x_ptr, + bias_ptr, + residual_ptr, + alpha_ptr, + inverse_ptr, + sum_ptr, + out_ptr, + stride_b, + stride_c, + t_len, + HAS_BIAS: tl.constexpr, + HAS_RESIDUAL: tl.constexpr, + WRITE_SUM: tl.constexpr, + BLOCK_T: tl.constexpr, +): + batch = tl.program_id(0).to(tl.int64) + channel = tl.program_id(1).to(tl.int64) + t = tl.program_id(2) * BLOCK_T + tl.arange(0, BLOCK_T) + mask = t < t_len + offset = batch * stride_b + channel * stride_c + t + x = tl.load(x_ptr + offset, mask=mask, other=0.0) + dtype = x.dtype + xf = x.to(tl.float32) + if HAS_BIAS: + # The conv bias add that cuDNN convolutions run as a separate kernel. + xf = _add_rn(xf, tl.load(bias_ptr + channel).to(tl.float32)).to(dtype).to(tl.float32) + if HAS_RESIDUAL: + residual = tl.load(residual_ptr + offset, mask=mask, other=0.0).to(tl.float32) + xf = _add_rn(residual, xf).to(dtype).to(tl.float32) + if WRITE_SUM: + tl.store(sum_ptr + offset, xf.to(dtype), mask=mask) + alpha = tl.load(alpha_ptr + channel).to(tl.float32) + inverse = tl.load(inverse_ptr + channel).to(tl.float32) + scaled = _mul_rn(alpha, xf).to(dtype).to(tl.float32) + # libdevice sinf, as the eager CUDA kernel (tl.sin is the approximation). + sine = tldevice.sin(scaled).to(dtype).to(tl.float32) + square = _mul_rn(sine, sine).to(dtype).to(tl.float32) + term = _mul_rn(inverse, square).to(dtype).to(tl.float32) + tl.store(out_ptr + offset, _add_rn(xf, term).to(dtype), mask=mask) + + +def fused_snake1d( + x: torch.Tensor, + alpha: torch.Tensor, + inverse: torch.Tensor, + *, + bias: torch.Tensor | None = None, + residual: torch.Tensor | None = None, + write_sum: bool = False, +) -> tuple[torch.Tensor | None, torch.Tensor]: + """Snake of ``residual + (x + bias)``, rounded as the separate eager ops. + + x, residual: (B, C, T); alpha, inverse, bias: (C,), all in x's dtype. + Returns ``(sum, snake(sum))``; ``sum`` is only materialized when + ``write_sum`` is set (for the next residual connection), else None. + """ + x = x.contiguous() + if residual is not None: + residual = residual.contiguous() + batch, channels, t_len = x.shape + out = torch.empty_like(x) + total = torch.empty_like(x) if write_sum else None + if x.numel() == 0: + return total, out + block = min(triton.next_power_of_2(t_len), 2048) + _dac_snake_kernel[(batch, channels, triton.cdiv(t_len, block))]( + x, + x if bias is None else bias, + x if residual is None else residual, + alpha, + inverse, + out if total is None else total, + out, + x.stride(0), + x.stride(1), + t_len, + HAS_BIAS=bias is not None, + HAS_RESIDUAL=residual is not None, + WRITE_SUM=write_sum, + BLOCK_T=block, + ) + return total, out diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/full_sample_graph.py b/vllm_omni/model_executor/models/higgs_audio_v3/full_sample_graph.py new file mode 100644 index 000000000..b318e736c --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/full_sample_graph.py @@ -0,0 +1,332 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project + +"""Dense decode sampler graphs captured and parity-checked during startup. + +Capture compares eager and captured state transitions with identical noise. +Serving only replays warmed shapes and preserves request torch RNG draws. CPU payload staging +remains outside capture. This module does not handle mixed prefill batches. +""" + +import copy +import time + +import torch +from vllm.logger import init_logger +from vllm.v1.outputs import SamplerOutput + +logger = init_logger(__name__) + +_STATE = ( + "_decode_last_codes", + "_decode_has_codes", + "_decode_delay_count", + "_decode_eoc_countdown", + "_decode_generation_done", +) + + +_BUCKETS = (1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 96, 128) +_FIELDS = ("temperature", "top_k", "top_p") + + +def _graph_key(model, batch, hidden, force_audio_inputs, metadata): + return ( + batch, + hidden.dtype, + hidden.device, + bool(force_audio_inputs), + bool(metadata.all_greedy), + tuple(getattr(metadata, f) is None for f in _FIELDS), + tuple(getattr(model, name).data_ptr() for name in _STATE), + ) + + +def _replay_captured(model, hidden, metadata, force_audio_inputs): + """Replay an already captured (bucket) graph writing inputs in place. + + The capture path below pads through temporaries on every call. Once a + graph exists, only the live rows are written into its static inputs: + padded rows keep finite values from earlier steps and own no request + state, forced-audio inputs are constants, sampling parameters are copied + only when the cached metadata changes, and request noise is drawn directly + into the graph's noise input. + """ + cache = getattr(model, "_dense_sample_graphs", None) + if not cache: + return None + batch = hidden.shape[0] + bucket = batch + if getattr(model.config, "audio_mrv2_sample_graph_buckets", False): + bucket = next((n for n in _BUCKETS if n >= batch), batch) + model._ensure_decode_state_capacity(bucket, hidden.device) + key = _graph_key(model, bucket, hidden, force_audio_inputs, metadata) + entry = cache.get(key) + if entry is None: + return None + graph, static_hidden, static_input_ids, static_noise, static_metadata, output, staging, keepalive = entry + books = model.num_codebooks + static_hidden[:batch].copy_(hidden) + if not force_audio_inputs: + static_input_ids[:batch].copy_(model._last_step_input_ids) + if batch < bucket: + static_input_ids[batch:].fill_(model._audio_continuation_id) + if metadata.generators: + from vllm_omni.utils.seeded_exponential import fill_exponential_rows + + fill_exponential_rows(static_noise[: batch * books], [metadata.generators.get(i) for i in range(batch)]) + else: + static_noise[: batch * books].exponential_() + sources = getattr(model, "_dense_sample_metadata_sources", None) + if sources is None: + sources = model._dense_sample_metadata_sources = {} + if sources.get(key) is not metadata: + for name in _FIELDS: + value = getattr(metadata, name) + if value is not None: + getattr(static_metadata, name)[:batch].copy_(value) + sources[key] = metadata + graph.replay() + model._finish_audio_staging(staging[:batch], batch, hidden.device, batch) + return SamplerOutput(sampled_token_ids=output.sampled_token_ids[:batch].clone(), logprobs_tensors=None) + + +def run_dense_sample(model, hidden, logits, metadata, force_audio_inputs=False, active_batch=None): + batch = hidden.shape[0] + if active_batch is None: + replayed = _replay_captured(model, hidden, metadata, force_audio_inputs) + if replayed is not None: + return replayed + # Serving never captures a new shape. Unsupported shapes/metadata use the + # regular sampler; startup owns capture and its parity checks. + if not getattr(model, "_warming_sample_graphs", False): + return None + if active_batch is None and getattr(model.config, "audio_mrv2_sample_graph_buckets", False): + bucket = next((n for n in (1, 2, 4, 8, 16, 24, 32, 40, 48, 56, 64, 96, 128) if n >= batch), batch) + if bucket != batch: + # Dummy rows own no request state. The batch-state gather replaces + # these rows before any later request can use them. + model._ensure_decode_state_capacity(bucket, hidden.device) + padded_hidden = hidden.new_zeros((bucket, hidden.shape[1])) + padded_hidden[:batch].copy_(hidden) + padded_metadata = copy.copy(metadata) + for name in ("temperature", "top_k", "top_p"): + value = getattr(metadata, name) + if value is not None: + padded = value.new_ones(bucket) + padded[:batch].copy_(value) + setattr(padded_metadata, name, padded) + original_ids = model._last_step_input_ids + padded_ids = original_ids.new_full((bucket,), model._audio_continuation_id) + if not force_audio_inputs: + padded_ids[:batch].copy_(original_ids) + model._last_step_input_ids = padded_ids + try: + output = run_dense_sample( + model, + padded_hidden, + logits.new_empty((bucket, logits.shape[1])), + padded_metadata, + force_audio_inputs=force_audio_inputs, + active_batch=batch, + ) + finally: + model._last_step_input_ids = original_ids + return SamplerOutput(sampled_token_ids=output.sampled_token_ids[:batch], logprobs_tensors=None) + actual = batch if active_batch is None else active_batch + effective_ids = ( + torch.full((batch,), model._audio_continuation_id, dtype=model._last_step_input_ids.dtype, device=hidden.device) + if force_audio_inputs + else model._last_step_input_ids + ) + books = model.num_codebooks + vocab = model.modality_head.vocab_size + noise = torch.empty((batch * books, vocab), dtype=torch.float32, device=hidden.device) + if actual < batch: + noise[actual * books :].fill_(1) + if metadata.generators: + from vllm_omni.utils.seeded_exponential import fill_exponential_rows + + fill_exponential_rows(noise[: actual * books], [metadata.generators.get(i) for i in range(actual)]) + else: + noise[: actual * books].exponential_() + fields = _FIELDS + # force_audio_inputs controls capture-time Python branches (including + # the mixed-tail assertion). Never reuse that graph for real EOS inputs. + key = _graph_key(model, batch, hidden, force_audio_inputs, metadata) + cache = getattr(model, "_dense_sample_graphs", None) + if cache is None: + cache = model._dense_sample_graphs = {} + entry = cache.get(key) + if entry is None: + static_hidden = hidden.clone() + static_input_ids = effective_ids.clone() + static_noise = noise.clone() + static_metadata = copy.copy(metadata) + for name in fields: + value = getattr(metadata, name) + setattr(static_metadata, name, None if value is None else value.clone()) + # Ensure the reusable staging exists before capture. Each graph keeps + # its own reference if a later larger batch grows the model buffer. + staging = model._get_audio_gpu_staging_buffer(batch, hidden.device) + state = {name: getattr(model, name).clone() for name in _STATE} + original_ids = model._last_step_input_ids + original_qsl = model._last_step_query_start_loc + original_mode_rows = getattr(model, "_step_audio_mode_rows", 0) + original_tail_rows = getattr(model, "_step_audio_tail_rows", 0) + original_direct_rows = getattr(model, "_fast_audio_direct_rows", 0) + # Dense decode has one token per row. Own this input instead of + # capturing a view of the runner's resizable staging buffer. + static_qsl = torch.arange(batch + 1, dtype=torch.int32, device=hidden.device) + + def restore(): + for name, value in state.items(): + getattr(model, name).copy_(value) + + try: + model._in_audio_sample_graph = True + # Outer eligibility and real-input validation already ran. The + # owned graph inputs represent exactly one audio-state row each. + # Normalize capture-time host dispatch; never capture an outer + # mixed-batch assertion or a stock-sampler branch into this graph. + model._step_audio_mode_rows = 0 + model._step_audio_tail_rows = 0 + model._fast_audio_direct_rows = batch + model._audio_graph_noise = static_noise + model._last_step_input_ids = static_input_ids + model._last_step_query_start_loc = static_qsl + model._last_logits_hidden = static_hidden + reference = model.sample(logits, static_metadata) + expected_ids = reference.sampled_token_ids.clone() + expected_staging = staging.clone() + expected_state = {name: getattr(model, name).clone() for name in _STATE} + restore() + model._last_logits_hidden = static_hidden + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + output = model.sample(logits, static_metadata) + # CUDA capture records work; it does not execute the transition. + # Validate only after replay from the preserved initial state. + restore() + graph.replay() + torch.testing.assert_close(output.sampled_token_ids, expected_ids, atol=0, rtol=0) + torch.testing.assert_close(staging, expected_staging, atol=0, rtol=0) + for name, expected in expected_state.items(): + torch.testing.assert_close(getattr(model, name), expected, atol=0, rtol=0) + restore() + logger.debug("Higgs sample graph parity checked for batch=%d", batch) + finally: + model._in_audio_sample_graph = False + model._step_audio_mode_rows = original_mode_rows + model._step_audio_tail_rows = original_tail_rows + model._fast_audio_direct_rows = original_direct_rows + model._audio_graph_noise = None + model._last_step_input_ids = original_ids + model._last_step_query_start_loc = original_qsl + model._last_logits_hidden = None + # Graphs reference storage, not Python owners. Model caches and state + # can grow after capture; keep their original allocations alive. + keepalive = [getattr(model, name) for name in _STATE] + for name in ("_row_index_cache", "_codebook_index_cache", "_boc_frame_cache"): + keepalive.extend(getattr(model, name).values()) + keepalive.append(static_qsl) + entry = (graph, static_hidden, static_input_ids, static_noise, static_metadata, output, staging, keepalive) + cache[key] = entry + graph, static_hidden, static_input_ids, static_noise, static_metadata, output, staging, keepalive = entry + static_hidden.copy_(hidden) + static_input_ids.copy_(effective_ids) + static_noise.copy_(noise) + for name in fields: + value = getattr(metadata, name) + if value is not None: + getattr(static_metadata, name).copy_(value) + graph.replay() + model._finish_audio_staging(staging[:actual], actual, hidden.device, actual) + return SamplerOutput(sampled_token_ids=output.sampled_token_ids.clone(), logprobs_tensors=None) + + +@torch.inference_mode() +def capture_sample_graphs(model, max_num_reqs): + """Warm the MRv2 sampler without consuming request RNG or decode state.""" + if not model._sample_graph_enabled or getattr(model, "_sample_graphs_ready", False): + return + from vllm.v1.sample.logits_processor import LogitsProcessors + from vllm.v1.sample.metadata import SamplingMetadata + + device = model.model.embed_tokens.weight.device + if device.type != "cuda": + return + start = time.perf_counter() + memory_before = torch.accelerator.memory_allocated(device) + reserved_before = torch.accelerator.memory_reserved(device) + model._resolve_token_ids() + buckets = ( + sorted({next((b for b in _BUCKETS if b >= n), n) for n in range(1, max_num_reqs + 1)}) + if getattr(model.config, "audio_mrv2_sample_graph_buckets", False) + else range(1, max_num_reqs + 1) + ) + model._ensure_decode_state_capacity(max(buckets), device) + model._get_audio_gpu_staging_buffer(max(buckets), device) + saved_state = {name: getattr(model, name).clone() for name in _STATE} + fields = ( + "_last_step_input_ids", + "_last_step_query_start_loc", + "_last_logits_hidden", + "_step_audio_mode_rows", + "_step_audio_tail_rows", + "_fast_audio_direct_rows", + ) + saved = {name: getattr(model, name, None) for name in fields} + try: + model._warming_sample_graphs = True + with torch.random.fork_rng(devices=[device.index]): + for batch in buckets: + hidden = model.model.embed_tokens.weight.new_zeros((batch, model.model.embed_tokens.weight.shape[1])) + logits = hidden.new_empty((batch, model.config.text_config.vocab_size)) + model._last_step_input_ids = torch.full( + (batch,), model._audio_continuation_id, device=device, dtype=torch.long + ) + model._last_step_query_start_loc = torch.arange(batch + 1, device=device, dtype=torch.int32) + ones = torch.ones(batch, device=device) + zeros = torch.zeros_like(ones) + for greedy in (False, True): + metadata = SamplingMetadata( + temperature=zeros if greedy else ones, + all_greedy=greedy, + all_random=not greedy, + top_p=ones, + top_k=torch.full((batch,), -1, device=device, dtype=torch.long), + generators={}, + max_num_logprobs=None, + no_penalties=True, + prompt_token_ids=None, + frequency_penalties=zeros, + presence_penalties=zeros, + repetition_penalties=ones, + output_token_ids=[[] for _ in range(batch)], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ) + for forced in (False, True): + run_dense_sample(model, hidden, logits, metadata, force_audio_inputs=forced) + finally: + for name, value in saved_state.items(): + getattr(model, name).copy_(value) + for name, value in saved.items(): + setattr(model, name, value) + model._warming_sample_graphs = False + model._clear_last_audio_outputs() + model._async_audio_gpu_staging = None + model._dense_sample_metadata_sources = {} + torch.accelerator.synchronize() + model._sample_graphs_ready = True + logger.info( + "Higgs sampler captured %d CUDA graphs in %.2f s " + "(allocated delta %.1f MiB, reserved delta %.1f MiB, reserved total %.1f MiB)", + len(model._dense_sample_graphs), + time.perf_counter() - start, + (torch.accelerator.memory_allocated(device) - memory_before) / 2**20, + (torch.accelerator.memory_reserved(device) - reserved_before) / 2**20, + torch.accelerator.memory_reserved(device) / 2**20, + ) diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py index 45b5d68be..dda33eb18 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_code2wav.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage 1 (codec decoder) for higgs-audio v3. Reuses higgs-audio-v2's RVQ + DAC codec decoder but loads weights from the @@ -10,6 +10,7 @@ prefix) rather than from a standalone tokenizer repo. from __future__ import annotations import os +import time from typing import Any import torch @@ -57,6 +58,9 @@ class HiggsAudioV3Code2Wav(nn.Module): input_modalities = "audio" + supports_native_payload_input = True + batched_gpu_staging_keys = {("codes", "audio")} + def __init__( self, config: HiggsAudioV3Config | None = None, @@ -81,6 +85,11 @@ class HiggsAudioV3Code2Wav(nn.Module): self._model_path = None self.vllm_config = None + _, self._codec_graph_enabled = self.config.resolve_graph_defaults( + use_v2_model_runner=bool( + getattr(getattr(self.vllm_config, "model_config", None), "use_v2_model_runner", False) + ) + ) self.sample_rate: int = int(self.config.sample_rate) self.num_codebooks: int = int(self.config.num_codebooks) self.num_real_codes: int = int(self.config.num_real_codes) @@ -97,6 +106,8 @@ class HiggsAudioV3Code2Wav(nn.Module): self.quantizer: HiggsAudioRVQ | None = None self.fc2: nn.Linear | None = None self.acoustic_decoder: nn.Module | None = None + self._decode_graphs: dict[tuple[int, int], tuple[Any, torch.Tensor, torch.Tensor]] = {} + self._decode_graphs_initialized = False self._loaded: bool = False # Do NOT eagerly load from standalone tokenizer here. # load_weights() will try bundled V3 codec first, then fall back. @@ -452,6 +463,30 @@ class HiggsAudioV3Code2Wav(nn.Module): self._ensure_codec_loaded() codes = self._validate_codes(audio_codes) + if codes.is_cuda and self._codec_graph_enabled: + batch, _, frames = codes.shape + entry = self._decode_graphs.get((batch, frames)) + if entry is not None: + graph, static_codes, static_audio = entry + static_codes[:batch].copy_(codes) + graph.replay() + # The next group may replay the same graph. Own the output + # before reusing its storage or shared graph memory pool. + return static_audio[:batch].clone() + return self._decode_codes_impl(codes) + + @torch.inference_mode() + def capture_auxiliary_graphs(self) -> None: + """Capture exact codec shapes before the generation worker is ready.""" + if not self._codec_graph_enabled or self._decode_graphs_initialized: + return + if not self._loaded: + self._ensure_codec_loaded() + device = self.fc2.weight.device + if device.type == "cuda": + self._capture_decode_graphs(device) + + def _decode_codes_impl(self, codes: torch.Tensor) -> torch.Tensor: rvq_codes = codes.transpose(0, 1).long() quantized = self.quantizer.decode(rvq_codes) quantized = quantized.to(dtype=self.fc2.weight.dtype) @@ -464,6 +499,55 @@ class HiggsAudioV3Code2Wav(nn.Module): audio = audio.unsqueeze(1) return audio.to(codes.device) + def _decode_graph_shapes(self) -> list[tuple[int, int]]: + single_max = int(getattr(self.config, "codec_graph_single_max_frames", 150)) + batches = getattr(self.config, "codec_graph_batch_sizes", (2, 4, 8, 16)) + frames = getattr(self.config, "codec_graph_frame_sizes", (5, 29, 54)) + if single_max < 1 or any(int(n) < 1 for n in (*batches, *frames)): + raise ValueError("Higgs codec graph dimensions must be positive") + return sorted({(1, n) for n in range(1, single_max + 1)} | {(int(b), int(f)) for b in batches for f in frames}) + + @torch.inference_mode() + def _capture_decode_graphs(self, device: torch.device) -> None: + # Never pad time: convolution boundary conditions must stay exact. + # Batch size also stays exact: padding can change convolution + # algorithms and rounding. Uncovered shapes retain the eager path. + start = time.perf_counter() + memory_before = torch.accelerator.memory_allocated(device) + reserved_before = torch.accelerator.memory_reserved(device) + shapes = self._decode_graph_shapes() + previous_device = current_omni_platform.current_device() + current_omni_platform.set_device(device) + try: + stream = torch.cuda.Stream(device=device) + stream.wait_stream(torch.cuda.current_stream(device)) + pool = torch.cuda.graph_pool_handle() + for batch, frames in shapes: + codes = torch.zeros((batch, self.num_codebooks, frames), device=device, dtype=torch.long) + stream.wait_stream(torch.cuda.current_stream(device)) + with torch.cuda.stream(stream): + for _ in range(3): + self._decode_codes_impl(codes) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, pool=pool, stream=stream): + audio = self._decode_codes_impl(codes) + self._decode_graphs[batch, frames] = (graph, codes, audio) + torch.cuda.current_stream(device).wait_stream(stream) + finally: + current_omni_platform.set_device(torch.device(device.type, previous_device)) + self._decode_graphs_initialized = True + torch.accelerator.synchronize() + logger.info( + "Higgs codec captured %d exact-frame CUDA graphs in %.2f s " + "(allocated delta %.1f MiB, reserved delta %.1f MiB, reserved total %.1f MiB)", + len(self._decode_graphs), + time.perf_counter() - start, + (torch.accelerator.memory_allocated(device) - memory_before) / 2**20, + (torch.accelerator.memory_reserved(device) - reserved_before) / 2**20, + torch.accelerator.memory_reserved(device) / 2**20, + ) + @torch.inference_mode() def forward_chunk( self, @@ -518,6 +602,18 @@ class HiggsAudioV3Code2Wav(nn.Module): ids = input_ids.reshape(-1).to(dtype=torch.long) request_ids_list = self._split_request_ids(ids, kwargs.get("seq_token_counts")) + # Native MRV2 uses token placeholders for scheduling. The connector + # payload owns the actual codebook-major codec sequence, including an + # empty terminal payload which must not decode a placeholder token. + if runtime_additional_information is not None: + for i, info in enumerate(runtime_additional_information): + if i >= len(request_ids_list): + break + codes = info.get("codes", {}) if isinstance(info, dict) else {} + audio = codes.get("audio") if isinstance(codes, dict) else None + if isinstance(audio, torch.Tensor): + request_ids_list[i] = audio.reshape(-1).to(device=ids.device, dtype=torch.long) + left_context_size = [0] * len(request_ids_list) right_holdback_size = [0] * len(request_ids_list) if runtime_additional_information is not None: @@ -530,11 +626,14 @@ class HiggsAudioV3Code2Wav(nn.Module): if "right_holdback_size" in meta: right_holdback_size[i] = int(meta["right_holdback_size"]) - wavs: list[torch.Tensor] = [] + wavs: list[torch.Tensor] = [empty] * len(request_ids_list) + # Requests with identical codec windows are independent along B. Keep + # time lengths and overlap trims exact rather than padding the temporal + # axis, which could change the convolution boundary conditions. + groups: dict[tuple[int, int, int], list[tuple[int, torch.Tensor]]] = {} for i, req_ids in enumerate(request_ids_list): n = int(req_ids.numel()) if n == 0: - wavs.append(empty) continue if n % self.num_codebooks != 0: logger.warning( @@ -542,23 +641,43 @@ class HiggsAudioV3Code2Wav(nn.Module): n, self.num_codebooks, ) - wavs.append(empty) continue frames = n // self.num_codebooks - codes_qf = req_ids.reshape(self.num_codebooks, frames) - codes_bqf = codes_qf.unsqueeze(0) + key = (frames, left_context_size[i], right_holdback_size[i]) + groups.setdefault(key, []).append((i, req_ids.reshape(self.num_codebooks, frames))) + + for (_, left, right), group in groups.items(): + codes_bqf = group[0][1].unsqueeze(0) if len(group) == 1 else torch.stack([codes for _, codes in group]) try: pcm = self.forward_chunk( codes_bqf, - left_context_size=left_context_size[i], - right_holdback_size=right_holdback_size[i], + left_context_size=left, + right_holdback_size=right, hop_length=self.hop_length, ) except ValueError as exc: - logger.warning("HiggsAudioV3Code2Wav: decode skipped (%s)", exc) - wavs.append(empty) + if len(group) == 1: + logger.warning("HiggsAudioV3Code2Wav: decode skipped (%s)", exc) + continue + # An invalid row must not discard valid peers. This slow path + # retains the previous per-request validation/error behavior. + for index, codes in group: + try: + row_pcm = self.forward_chunk( + codes.unsqueeze(0), + left_context_size=left, + right_holdback_size=right, + hop_length=self.hop_length, + ) + wavs[index] = row_pcm[0, 0].to(torch.float32) + except ValueError as row_exc: + logger.warning("HiggsAudioV3Code2Wav: decode skipped (%s)", row_exc) continue - wavs.append(pcm.squeeze(0).squeeze(0).to(torch.float32).cpu()) + pcm = pcm.to(torch.float32) + for row, (index, _) in enumerate(group): + # The generation runner owns the batched pinned D2H copy and + # its completion event; avoid blocking once per request here. + wavs[index] = pcm[row, 0] return OmniOutput( text_hidden_states=None, diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py index ba9c9eb2d..c6e0fc511 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/higgs_audio_v3_talker.py @@ -108,12 +108,26 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): prefer_model_sampler: bool = True # Tell the runner to call postprocess() to emit per-step audio codes. have_multimodal_outputs: bool = True - has_postprocess: bool = True + has_postprocess: bool = False skips_model_sampler_output_token_history: bool = True postprocess_uses_hidden_states: bool = False postprocess_uses_multimodal_outputs: bool = False postprocess_uses_req_infos: bool = False supports_omni_query_start_loc: bool = True + requires_cpu_input_tail_ids: bool = True + # The codec stage consumes audio codes, never backbone hidden states. + omni_pooler_payload_include_hidden: bool = False + + def capture_auxiliary_graphs(self) -> None: + from .full_sample_graph import capture_sample_graphs + + scheduler = self.vllm_config.scheduler_config + capture_sample_graphs(self, scheduler.max_num_seqs) + + def create_omni_model_state(self, vllm_config, encoder_cache, device): + from .model_state import HiggsModelState + + return HiggsModelState(vllm_config, self, encoder_cache, device) def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() @@ -122,6 +136,17 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self.config = hf_config else: self.config = HiggsAudioV3Config(**hf_config.to_dict()) + self._sample_graph_enabled, _ = self.config.resolve_graph_defaults( + use_v2_model_runner=bool(getattr(vllm_config.model_config, "use_v2_model_runner", False)) + ) + self.use_async_omni_output = getattr(self.config, "audio_async_payload", False) is True + self.supports_async_whole_payload = self.use_async_omni_output + if self.use_async_omni_output: + if not vllm_config.scheduler_config.async_scheduling: + raise ValueError("audio_async_payload requires async scheduling") + if getattr(self.config, "audio_async_prompt_mode", False) is not True: + raise ValueError("audio_async_payload requires audio prompt mode classification") + self._async_audio_gpu_staging = None model_config = getattr(vllm_config, "model_config", None) compilation_config = getattr(vllm_config, "compilation_config", None) cudagraph_mode = getattr(compilation_config, "cudagraph_mode", None) @@ -345,10 +370,26 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): inputs_embeds: torch.Tensor | None = None, omni_query_start_loc: torch.Tensor | None = None, req_ids: list[str] | None = None, + cpu_input_tail_ids: list[int] | None = None, + audio_prompt_mode_rows: int = 0, **_: Any, ) -> None: """Update per-step metadata before runner forward or CUDA graph replay.""" _ = positions, inputs_embeds + self._step_audio_mode_rows = ( + audio_prompt_mode_rows if getattr(self.config, "audio_async_prompt_mode", False) is True else 0 + ) + self._step_audio_tail_rows = 0 + if cpu_input_tail_ids: + self._resolve_token_ids() + if ( + cpu_input_tail_ids + and self._audio_continuation_id is not None + and req_ids is not None + and len(cpu_input_tail_ids) == len(req_ids) + and all(token == self._audio_continuation_id for token in cpu_input_tail_ids) + ): + self._step_audio_tail_rows = len(cpu_input_tail_ids) if input_ids is not None: self._last_step_input_ids = input_ids if self._use_external_decode_cudagraph and input_ids.is_cuda and input_ids.ndim >= 1: @@ -359,6 +400,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): if req_ids is not None: self._sync_decode_state_with_batch(req_ids) + def finish_decode_step_forward(self) -> None: + self._decode_step_metadata_from_runner = False + def _ensure_state_pool_capacity(self, num_rows: int, device: torch.device) -> None: if not hasattr(self, "_state_pool_indices"): self._state_pool_indices: dict[str, int] = {} @@ -387,17 +431,53 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self._state_pool_eoc_countdown = resize(self._state_pool_eoc_countdown, -1) self._state_pool_generation_done = resize(self._state_pool_generation_done, False) + def _state_indices(self, rows: list[int], device: torch.device) -> torch.Tensor: + # A fresh pinned allocation is owned by PyTorch's host allocator until + # its asynchronous copy completes; never overwrite a shared CPU staging + # buffer while the previous decode step may still be reading it. + if device.type == "cuda" and getattr(self.config, "audio_mrv2_async_state_indices", False): + return torch.tensor(rows, dtype=torch.long, pin_memory=True).to(device, non_blocking=True) + return torch.tensor(rows, dtype=torch.long, device=device) + def _sync_decode_state_with_batch(self, req_ids: list[str]) -> None: device = self._decode_has_codes.device self._ensure_state_pool_capacity(len(req_ids), device) previous_req_ids = getattr(self, "_last_batch_req_ids", []) + if ( + device.type == "cuda" + and getattr(self.config, "audio_mrv2_fused_state", False) + and (previous_req_ids != req_ids or any(r not in self._state_pool_indices for r in req_ids)) + ): + from .mrv2_state_kernels import sync_state + + previous_rows = [row for row, req_id in enumerate(previous_req_ids) if req_id in self._state_pool_indices] + pool_rows = [self._state_pool_indices[previous_req_ids[row]] for row in previous_rows] + current_rows = [] + for req_id in req_ids: + fresh = req_id not in self._state_pool_indices + if fresh: + slot = ( + self._free_state_pool_indices.pop() + if self._free_state_pool_indices + else len(self._state_pool_indices) + ) + self._ensure_state_pool_capacity(slot + 1, device) + self._state_pool_indices[req_id] = slot + slot = self._state_pool_indices[req_id] + current_rows.append(-slot - 1 if fresh else slot) + self._ensure_decode_state_capacity(len(req_ids), device) + if req_ids or previous_rows: + sync_state(self, previous_rows, pool_rows, current_rows) + self._last_batch_req_ids = list(req_ids) + return + if previous_req_ids != req_ids or any(req_id not in self._state_pool_indices for req_id in req_ids): previous_rows = [row for row, req_id in enumerate(previous_req_ids) if req_id in self._state_pool_indices] if previous_rows: pool_rows = [self._state_pool_indices[previous_req_ids[row]] for row in previous_rows] - src = torch.tensor(previous_rows, dtype=torch.long, device=device) - dst = torch.tensor(pool_rows, dtype=torch.long, device=device) + src = self._state_indices(previous_rows, device) + dst = self._state_indices(pool_rows, device) self._state_pool_has_codes[dst] = self._decode_has_codes[src] self._state_pool_last_codes[dst] = self._decode_last_codes[src] self._state_pool_delay_count[dst] = self._decode_delay_count[src] @@ -420,10 +500,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self._state_pool_generation_done[pool_row] = False self._ensure_decode_state_capacity(len(req_ids), device) - pool_rows = torch.tensor( + pool_rows = self._state_indices( [self._state_pool_indices[req_id] for req_id in req_ids], - dtype=torch.long, - device=device, + device, ) self._decode_has_codes[: len(req_ids)] = self._state_pool_has_codes[pool_rows] self._decode_last_codes[: len(req_ids)] = self._state_pool_last_codes[pool_rows] @@ -450,8 +529,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): inputs_embeds: torch.Tensor | None = None, **kwargs: Any, ) -> torch.Tensor: + # A runner invocation may call forward repeatedly while warming and + # capturing a graph. Keep its metadata authoritative for that scope. metadata_from_runner = self._decode_step_metadata_from_runner - self._decode_step_metadata_from_runner = False info_dicts = kwargs.get("model_intermediate_buffer") if info_dicts is None: @@ -501,7 +581,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): and inputs_embeds is None and not self._is_single_token_decode_step(int(hidden_states.shape[0])) ) - if is_prefill and info_dicts: + if is_prefill and info_dicts and inputs_embeds is None: # Voice clone: replace marked prompt positions with ref audio embeddings. hidden_states = self._apply_ref_audio_substitution(hidden_states, input_ids, positions, info_dicts) @@ -515,6 +595,11 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): return self._run_qwen3_layers_eager(positions, hidden_states) def _run_qwen3_layers_eager(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + if self._use_external_decode_cudagraph: + # Enter Qwen3Model's compilation boundary before the runner + # captures the decode graph. Audio feedback stays outside it. + return self.model(input_ids=None, positions=positions, inputs_embeds=hidden_states) + residual: torch.Tensor | None = None for layer in self.model.layers: hidden_states, residual = layer(positions, hidden_states, residual) @@ -619,6 +704,25 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: self._last_logits_hidden = hidden_states + ids = getattr(self, "_last_step_input_ids", None) + rows = int(hidden_states.shape[0]) + if ( + self._audio_continuation_id is not None + and sampling_metadata is not None + and getattr(sampling_metadata, "no_penalties", False) + and ( + (isinstance(ids, torch.Tensor) and ids.numel() == rows and self._fast_audio_direct_rows == rows) + or getattr(self, "_step_audio_tail_rows", 0) == rows + or getattr(self, "_step_audio_mode_rows", 0) == rows + ) + and getattr(sampling_metadata, "max_num_logprobs", None) is None + and getattr(sampling_metadata, "allowed_token_ids_mask", None) is None + and not getattr(sampling_metadata, "bad_words_token_ids", None) + ): + # sample() already selects continuation/EOS directly in this + # state, using the audio head rather than the text head. Keep + # the normal vocabulary shape for runner-side processors. + return hidden_states.new_empty((rows, self._text_vocab_size), dtype=torch.float32) return self.logits_processor(self.lm_head, hidden_states) def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: @@ -656,7 +760,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): ``audio_placeholder_positions`` ([T] absolute prompt positions). Legacy callers that still submit -100 token IDs remain supported. """ - if not info_dicts: + if not info_dicts or not any( + isinstance(info, dict) and info.get("audio_input_ids") is not None for info in info_dicts + ): return hidden_states flat_ids = input_ids.reshape(-1) @@ -788,10 +894,6 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): if self._decode_active_audio_count == 0 and not external_decode_graph: return hidden_states - req_rows, token_positions = self._decode_request_token_positions(num_tokens, hidden_states.device) - if int(req_rows.numel()) == 0: - return hidden_states - num_reqs = self._step_request_count(num_tokens) self._ensure_decode_state_capacity(num_reqs, hidden_states.device) @@ -800,6 +902,9 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): has_codes = self._decode_has_codes[:num_tokens].unsqueeze(-1) current_hidden = hidden_states else: + req_rows, token_positions = self._decode_request_token_positions(num_tokens, hidden_states.device) + if int(req_rows.numel()) == 0: + return hidden_states codes_slice = self._decode_last_codes.index_select(0, req_rows) has_codes = self._decode_has_codes.index_select(0, req_rows).unsqueeze(-1) current_hidden = hidden_states.index_select(0, token_positions) @@ -932,6 +1037,44 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): return None return tail_ids == int(audio_id) + def _step_tail_guard(self, num_rows: int, *, prompt_mode: bool, decode_only: bool) -> None: + from .step_kernels import step_tail_guard + + error = getattr(self, "_decode_tail_guard_error", None) + device = self._decode_has_codes.device + if error is None or error.device != device: + error = self._decode_tail_guard_error = torch.zeros(1, dtype=torch.int32, device=device) + step_tail_guard( + self._last_step_input_ids.reshape(-1), + None if decode_only else self._last_step_query_start_loc, + self._decode_has_codes[:num_rows], + self._decode_generation_done[:num_rows], + self._decode_delay_count[:num_rows], + self._decode_eoc_countdown[:num_rows], + error, + audio_id=int(self._audio_continuation_id), + eos_id=int(self._eos_token_id), + prompt_mode=prompt_mode, + ) + torch._assert_async(error[0] == 0, "CPU/GPU Higgs graph audio tail mismatch") + + def _restore_terminal_audio_rows(self, num_rows): + # Async scheduling can execute one final in-flight EOS row after + # request retirement has released its state slot. EOS is terminal: + # never resurrect that row as a new audio stream with a forced input. + eos_id = self._eos_token_id + if eos_id is None: + return + ids = self._last_step_input_ids.reshape(-1) + qsl = self._last_step_query_start_loc + if isinstance(qsl, torch.Tensor) and qsl.numel() == num_rows + 1: + tail = ids.index_select(0, (qsl[1:].long() - 1).clamp_min(0)) + elif ids.numel() == num_rows: + tail = ids + else: + raise RuntimeError("Cannot validate terminal audio rows") + self._decode_generation_done[:num_rows].logical_or_(tail == eos_id) + def _fast_audio_sampler_gpu_fallback_reason( self, *, @@ -997,6 +1140,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): and accumulate per-request state. """ self._resolve_token_ids() + self._async_audio_gpu_staging = None audio_id = self._audio_continuation_id @@ -1020,6 +1164,56 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self._ensure_decode_state_capacity(num_rows, hidden.device) ids = getattr(self, "_last_step_input_ids", None) decode_only = isinstance(ids, torch.Tensor) and int(ids.numel()) == num_rows + prompt_mode_rows = getattr(self, "_step_audio_mode_rows", 0) == num_rows + graph_forced_audio = getattr(self, "_step_audio_tail_rows", 0) == num_rows or prompt_mode_rows + use_sample_graph = bool( + getattr(self, "_sample_graph_enabled", False) + and not getattr(self, "_in_audio_sample_graph", False) + and hidden.is_cuda + and (graph_forced_audio or (decode_only and self._fast_audio_direct_rows == num_rows)) + and getattr(sampling_metadata, "no_penalties", False) + and self._fast_audio_sampler_gpu_fallback_reason( + logits=logits, sampling_metadata=sampling_metadata, num_rows=num_rows + ) + is None + ) + qsl = self._last_step_query_start_loc + fused_guard = ( + use_sample_graph + and graph_forced_audio + and self._eos_token_id is not None + and (decode_only or (isinstance(qsl, torch.Tensor) and int(qsl.numel()) == num_rows + 1)) + ) + if fused_guard: + # One launch restores terminal EOS rows, resets new prefill rows + # and validates the audio tail (~20 small eager kernels before). + self._step_tail_guard(num_rows, prompt_mode=prompt_mode_rows, decode_only=decode_only) + elif prompt_mode_rows: + self._restore_terminal_audio_rows(num_rows) + if use_sample_graph: + from .full_sample_graph import run_dense_sample + + if not fused_guard and graph_forced_audio: + actual_audio = self._audio_seed_mask_from_step_input(num_rows, hidden.device) + if getattr(self, "_step_audio_mode_rows", 0) == num_rows: + # Reset new prefill rows before accepting old active state. + if not decode_only: + self._reset_decode_state_rows( + self._prefill_row_mask(num_rows, hidden.device), num_rows, hidden.device + ) + actual_audio = ( + actual_audio | self._decode_has_codes[:num_rows] | self._decode_generation_done[:num_rows] + ) + torch._assert_async(actual_audio.all(), "CPU/GPU Higgs graph audio tail mismatch") + if not decode_only: + if not fused_guard: + pfmask = self._prefill_row_mask(num_rows, hidden.device) + self._reset_decode_state_rows(pfmask, num_rows, hidden.device) + self._fast_audio_direct_rows = num_rows + sampled = run_dense_sample(self, hidden, logits, sampling_metadata, force_audio_inputs=graph_forced_audio) + if sampled is not None: + return sampled + self._last_logits_hidden = None if not decode_only: self._fast_audio_direct_rows = 0 pfmask = self._prefill_row_mask(num_rows, hidden.device) @@ -1042,6 +1236,17 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): if fast_fallback_reason is None: direct_audio_batch = decode_only and self._fast_audio_direct_rows == num_rows + mixed_direct = ( + getattr(self, "_step_audio_tail_rows", 0) == num_rows + or getattr(self, "_step_audio_mode_rows", 0) == num_rows + ) and getattr(sampling_metadata, "no_penalties", False) + if mixed_direct: + # Guard the CPU classification against runner/device drift. + torch._assert_async( + (audio_mask if getattr(self, "_step_audio_mode_rows", 0) == num_rows else prev_audio_mask).all(), + "CPU/GPU Higgs audio tail mismatch", + ) + direct_audio_batch = True if direct_audio_batch: self._fast_audio_direct_rows = num_rows @@ -1194,6 +1399,62 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self._boc_frame_cache[key] = buf return buf[:num_rows] + def post_sample_multimodal_outputs(self, *, req_ids, invalid_req_indices, multimodal_outputs): + """Export current-step codes by current batch row, with owned CPU storage. + + Forward-time payloads refer to previous request-buffer rows, which may + be reordered or overwritten when the sampler reuses its host staging. + Export only after sampling, before the runner takes its output snapshot. + """ + if getattr(self, "use_async_omni_output", False): + staging = self._async_audio_gpu_staging + self._async_audio_gpu_staging = None + self._fast_audio_direct_rows = 0 + if staging is None: + return None + if staging.shape[0] != len(req_ids): + raise RuntimeError("Higgs async snapshot row mismatch") + return {"_higgs_audio_snapshot": staging, "_higgs_invalid_rows": tuple(invalid_req_indices)} + staging = self._last_audio_host_staging + if staging is None: + self._fast_audio_direct_rows = 0 + return None + if staging.shape[0] != len(req_ids): + raise RuntimeError("Higgs audio staging rows do not match the sampled request batch") + event = self._last_audio_staging_event + if event is not None: + event.synchronize() + self._last_audio_staging_event = None + codes = staging[:, : self.num_codebooks].clone() + valid = staging[:, self.num_codebooks].tolist() + done = staging[:, self.num_codebooks + 1].tolist() + invalid = set(invalid_req_indices) + self._fast_audio_direct_rows = len(req_ids) if not invalid and all(v or d for v, d in zip(valid, done)) else 0 + # Keep the codec-row rank on terminal/invalid steps: a 1-D empty + # value would replace previously accumulated full-response rows. + empty = codes.new_empty((0, self.num_codebooks)) + rows = [codes[i : i + 1] if valid[i] and i not in invalid else empty for i in range(len(req_ids))] + return {"codes": {"audio": rows}} + + def finalize_multimodal_outputs_from_cpu_snapshot(self, payload): + """Pure per-step conversion; never reads mutable decode/request state.""" + if not payload or "_higgs_audio_snapshot" not in payload: + return payload + staging = payload["_higgs_audio_snapshot"] + if staging.device.type != "cpu": + raise RuntimeError("Higgs snapshot finalized before D2H") + codes = staging[:, : self.num_codebooks].clone() + valid = staging[:, self.num_codebooks].tolist() + invalid = set(payload["_higgs_invalid_rows"]) + # Keep the codec-row rank on terminal/invalid steps: a 1-D empty + # value would replace previously accumulated full-response rows. + empty = codes.new_empty((0, self.num_codebooks)) + return { + "codes": { + "audio": [codes[i : i + 1] if valid[i] and i not in invalid else empty for i in range(staging.shape[0])] + } + } + # ------------------------------------------------------------------ postprocess def postprocess( self, @@ -1450,6 +1711,20 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): staging[:, num_codebooks].copy_(valid_full.to(torch.int32)) staging[:, num_codebooks + 1].copy_(done_full.to(torch.int32)) + self._finish_audio_staging(staging, num_rows, device, num_rows if all_rows else int(rows.numel())) + + def _finish_audio_staging(self, staging, num_rows, device, active_rows): + if getattr(self, "_in_audio_sample_graph", False): + return + if getattr(self, "use_async_omni_output", False): + # Own this step's GPU values before the next graph replay. The + # runner copies the snapshot on its output stream and waits there. + self._async_audio_gpu_staging = staging[:num_rows].clone() + self._last_audio_host_staging = None + self._last_audio_staging_event = None + self._last_audio_codes = None + self._decode_active_audio_count = active_rows + return host = self._get_audio_host_staging_buffer(num_rows) host.copy_(staging, non_blocking=True) if device.type == "cuda": @@ -1459,7 +1734,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): else: self._last_audio_staging_event = None - self._decode_active_audio_count = num_rows if all_rows else int(rows.numel()) + self._decode_active_audio_count = active_rows # Host staging is the authoritative postprocess path. Keeping another # full GPU codes buffer alive adds fill/copy work on the hot path. self._last_audio_codes = None @@ -1514,7 +1789,24 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): fallback = x.argmax(dim=-1) safe_x = torch.where(all_masked.unsqueeze(-1), torch.zeros_like(x), x) probs = safe_x.softmax(dim=-1) - if sampling_metadata.generators: + if getattr(self, "_audio_graph_noise", None) is not None: + torch._assert_async(torch.isfinite(probs).all(), "Audio sampling probabilities must be finite") + sampled = probs.div(self._audio_graph_noise).argmax(dim=-1) + elif probs.is_cuda: + # Match torch.multinomial's single-sample exponential race, but + # validate the whole batch asynchronously. The original per-request + # calls each copy validation scalars to the host before drawing. + torch._assert_async(torch.isfinite(probs).all(), "Audio sampling probabilities must be finite") + q = torch.empty_like(probs) + if sampling_metadata.generators: + from vllm_omni.utils.seeded_exponential import fill_exponential_rows + + rows = int(probs.shape[0]) // num_codebooks + fill_exponential_rows(q, [sampling_metadata.generators.get(i) for i in range(rows)]) + else: + q.exponential_() + sampled = probs.div(q).argmax(dim=-1) + elif sampling_metadata.generators: sampled = torch.cat( [ torch.multinomial( diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/model_state.py b/vllm_omni/model_executor/models/higgs_audio_v3/model_state.py new file mode 100644 index 000000000..d275ea7f8 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/model_state.py @@ -0,0 +1,320 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Higgs delay-pattern sampling and owned audio output on Model Runner V2.""" + +from typing import Any + +import numpy as np +import torch +from vllm.logger import init_logger +from vllm.v1.sample.logits_processor import LogitsProcessors +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.worker.gpu.input_batch import get_num_sampled_and_rejected +from vllm.v1.worker.gpu.sample.output import SamplerOutput + +from vllm_omni.worker_v2.model_states.omni_model_state import OmniModelState +from vllm_omni.worker_v2.omni_sampler import OmniSampler, OmniSamplingOutput + +logger = init_logger(__name__) + + +class HiggsModelState(OmniModelState): + def __init__(self, vllm_config, model, encoder_cache, device): + super().__init__(vllm_config, model, encoder_cache, device) + model._use_external_decode_cudagraph = True + # Fix decode-state addresses before any sampling graph can capture them. + model._ensure_decode_state_capacity(max(self.max_num_tokens, 128), device) + logger.info("Higgs MRV2 custom sampler and owned audio snapshots enabled") + if getattr(model.config, "audio_mrv2_static_inputs", False): + self._static_inputs_embeds = torch.zeros( + (self.max_num_tokens, model.model.embed_tokens.weight.shape[1]), + dtype=self.dtype, + device=device, + ) + self.direct_payload = bool(getattr(model.config, "audio_mrv2_direct_payload", False)) + self.sampling_slots = None + self.request_slots = {} + if getattr(model.config, "audio_mrv2_slot_parameters", False): + capacity = self.scheduler_config.max_num_seqs + self.sampling_slots = { + "temperature": torch.zeros(capacity, device=device), + "top_p": torch.ones(capacity, device=device), + "top_k": torch.full((capacity,), -1, dtype=torch.long, device=device), + "zeros": torch.zeros(capacity, device=device), + "ones": torch.ones(capacity, device=device), + } + self.requests = {} + self.generators = {} + self._ensure_slot_capacity(self.scheduler_config.max_num_seqs) + self.metadata_key = None + self.metadata = None + logger.info( + "Higgs MRV2 options: static=%s slot_parameters=%s direct_payload=%s", + self._static_inputs_embeds is not None, + self.sampling_slots is not None, + self.direct_payload, + ) + if not model.use_async_omni_output: + raise ValueError("Higgs MRV2 requires audio_async_payload; use higgs_multimodal_qwen3_mrv2_h200.yaml") + + def _ensure_slot_capacity(self, slots): + """Per-slot prompt facts for vectorized step classification.""" + current = getattr(self, "prompt_len_np", None) + if current is not None and len(current) >= slots: + return + size = max(slots, 2 * (0 if current is None else len(current))) + old_len, old_tail = current, getattr(self, "prompt_audio_tail_np", None) + self.prompt_len_np = np.zeros(size, dtype=np.int64) + self.prompt_audio_tail_np = np.zeros(size, dtype=bool) + if current is not None: + self.prompt_len_np[: len(old_len)] = old_len + self.prompt_audio_tail_np[: len(old_tail)] = old_tail + + def add_request(self, req_index, new_req_data): + params = new_req_data.sampling_params + is_warmup = str(new_req_data.req_id).startswith("_warmup_") + if not is_warmup and ( + params.frequency_penalty != 0 + or params.presence_penalty != 0 + or params.repetition_penalty != 1 + or params.logprobs is not None + or params.prompt_logprobs is not None + or params.min_tokens != 0 + or params.allowed_token_ids + or params.bad_words + or params.logit_bias + or params.min_p != 0 + ): + raise ValueError("Higgs MRV2 currently supports temperature/top-k/top-p sampling without text penalties") + super().add_request(req_index, new_req_data) + self.requests[new_req_data.req_id] = new_req_data + self.request_slots[new_req_data.req_id] = req_index + if self.sampling_slots is not None: + for name in ("temperature", "top_p", "top_k"): + self.sampling_slots[name][req_index : req_index + 1].fill_(getattr(params, name)) + self.metadata_key = None + prompt = new_req_data.prompt_token_ids or [] + self._ensure_slot_capacity(req_index + 1) + self.prompt_len_np[req_index] = len(prompt) + self.model._resolve_token_ids() + self.prompt_audio_tail_np[req_index] = bool(prompt) and prompt[-1] == self.model._audio_continuation_id + if params.seed is not None: + self.generators[new_req_data.req_id] = torch.Generator(device=self.device).manual_seed(params.seed) + + def remove_request(self, req_index_or_id): + index = self._resolve_req_index(req_index_or_id) + if index is not None: + req_id = self.intermediate_buffer.buffers[index].get("req_id") + if req_id is not None: + self.model.on_requests_finished({req_id}) + self.requests.pop(req_id, None) + self.request_slots.pop(req_id, None) + self.generators.pop(req_id, None) + self.metadata_key = None + super().remove_request(req_index_or_id) + + def run_preprocess(self, input_batch, model_inputs, req_states=None, mtp_batch_descriptor_dispatcher=None): + model = self.model + model._resolve_token_ids() + # Every row must have an audio-continuation prompt whose prefill ends + # within this step; one vectorized check replaces a per-row loop. + n = input_batch.num_reqs + slots = input_batch.idx_mapping_np[:n] + end = req_states.num_computed_prefill_tokens[slots] + input_batch.num_scheduled_tokens[:n] + audio_mode = bool(np.all(self.prompt_audio_tail_np[slots] & (end >= self.prompt_len_np[slots]))) + model.update_decode_step_metadata( + input_ids=model_inputs["input_ids"][: input_batch.num_tokens], + positions=input_batch.positions, + omni_query_start_loc=input_batch.query_start_loc[: input_batch.num_reqs + 1], + req_ids=input_batch.req_ids, + audio_prompt_mode_rows=input_batch.num_reqs if audio_mode else 0, + ) + + if self._static_inputs_embeds is not None: + self._prepare_static_embeddings(input_batch, model_inputs) + + def _prepare_static_embeddings(self, input_batch, model_inputs): + # Feedback and reference substitution are request-dependent and belong + # outside the captured backbone. Use only actual tokens: graph padding + # must not change the one-token-per-request decode classification. + n = input_batch.num_tokens + ids = model_inputs["input_ids"][:n] + if self._fused_decode_embeddings(input_batch, ids): + return + safe_ids = torch.where(ids < 0, torch.zeros_like(ids), ids) + embeds = self.model.model.embed_tokens(safe_ids) + info = model_inputs.get("model_intermediate_buffer") + if input_batch.has_prefill and info: + embeds = self.model._apply_ref_audio_substitution( + embeds, + ids, + input_batch.positions[:n], + info, + ) + embeds = self._apply_audio_feedback(embeds, ids, input_batch) + self._static_inputs_embeds[:n].copy_(embeds) + padded = input_batch.num_tokens_after_padding + self._static_inputs_embeds[n:padded].zero_() + + def _fused_decode_embeddings(self, input_batch, ids) -> bool: + """Pure decode steps: token/feedback embeddings and padding in one launch.""" + model = self.model + if ( + input_batch.has_prefill + or input_batch.num_tokens != input_batch.num_reqs + or not ids.is_cuda + or not getattr(model, "_use_external_decode_cudagraph", False) + ): + return False + text = getattr(model.model.embed_tokens, "weight", None) + audio = getattr(model.multimodal_embedding, "weight", None) + if text is None or audio is None or text.dtype != self._static_inputs_embeds.dtype: + return False + from .step_kernels import decode_embeddings + + n = int(ids.numel()) + model._ensure_decode_state_capacity(n, ids.device) + decode_embeddings( + self._static_inputs_embeds, + ids, + text, + audio, + model._decode_last_codes, + model._decode_has_codes, + int(input_batch.num_tokens_after_padding), + int(model.multimodal_embedding.vocab_size), + ) + return True + + def _apply_audio_feedback(self, embeds, ids, input_batch): + if input_batch.num_tokens == input_batch.num_reqs: + return self.model._apply_audio_feedback(embeds, ids) + # Every scheduled request has a nonempty span. Gather its first token + # for all rows, then mask on-device. Boolean indexing would first ask + # the GPU how many rows matched, synchronizing every mixed prefill. + n = input_batch.num_reqs + qsl = input_batch.query_start_loc[: n + 1] + starts = qsl[:-1].long() + model = self.model + model._ensure_decode_state_capacity(n, embeds.device) + audio = model.multimodal_embedding(model._decode_last_codes[:n]).to(embeds.dtype) + use_audio = model._decode_has_codes[:n] & ((qsl[1:] - qsl[:-1]) == 1) + current = embeds.index_select(0, starts) + result = embeds.clone() + result.index_copy_(0, starts, torch.where(use_audio[:, None], audio, current)) + return result + + def run_postprocess(self, hidden_states, input_batch): + self.model.finish_decode_step_forward() + + def postprocess_model_output(self, model_output, input_batch, req_states): + return model_output, None + + def sampling_metadata(self, req_ids, slots=None): + key = tuple(req_ids) + if self.metadata_key != key: + params = [self.requests[r].sampling_params for r in req_ids] + n = len(params) + if self.sampling_slots is not None: + if slots is None: + slots = torch.tensor([self.request_slots[r] for r in req_ids], device=self.device) + temperature, top_p, top_k = ( + self.sampling_slots[name].index_select(0, slots) for name in ("temperature", "top_p", "top_k") + ) + zeros = self.sampling_slots["zeros"][:n] + ones = self.sampling_slots["ones"][:n] + else: + temperature = torch.tensor([p.temperature for p in params], device=self.device) + top_p = torch.tensor([p.top_p for p in params], device=self.device) + top_k = torch.tensor([p.top_k for p in params], device=self.device) + zeros = torch.zeros(n, device=self.device) + ones = torch.ones(n, device=self.device) + self.metadata = SamplingMetadata( + temperature=temperature, + all_greedy=all(p.temperature == 0 for p in params), + all_random=all(p.temperature != 0 for p in params), + top_p=top_p, + top_k=top_k, + generators={i: self.generators[r] for i, r in enumerate(req_ids) if r in self.generators}, + max_num_logprobs=None, + no_penalties=True, + prompt_token_ids=None, + frequency_penalties=zeros, + presence_penalties=zeros, + repetition_penalties=ones, + output_token_ids=[[] for _ in params], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ) + self.metadata_key = key + return self.metadata + + def custom_sampler(self, sampler): + return HiggsSampler(sampler, self), None + + def finalize_audio_snapshot(self, payload: dict[str, Any], num_sampled: list[int]): + if payload and "_higgs_audio_snapshot" in payload and getattr(self, "direct_payload", False): + from vllm_omni.worker_v2.output_snapshot import RequestOutputSnapshot + + snapshot = payload["_higgs_audio_snapshot"] + if snapshot.device.type != "cpu": + raise RuntimeError("Higgs snapshot finalized before D2H") + codes = snapshot[:, : self.model.num_codebooks].clone() + valid = snapshot[:, self.model.num_codebooks].tolist() + invalid = set(payload.get("_higgs_invalid_rows", ())) + # One split creates every row view; per-row slicing cost one + # dispatcher call per request on the output-materialization path. + rows = torch.split(codes, 1) + return RequestOutputSnapshot( + [ + {"codes.audio": rows[i]} if valid[i] and n and i not in invalid else None + for i, n in enumerate(num_sampled) + ] + ) + if payload and "_higgs_audio_snapshot" in payload: + payload = dict(payload) + payload["_higgs_invalid_rows"] = tuple(i for i, n in enumerate(num_sampled) if n == 0) + return self.model.finalize_multimodal_outputs_from_cpu_snapshot(payload) + + +class HiggsSampler(OmniSampler): + """Higgs audio sampling registered through the upstream sampler factory.""" + + def __init__(self, base_sampler, state: HiggsModelState): + super().__init__(base_sampler) + self.state = state + + def sample_step(self, hidden_states, input_batch, req_states, grammar_output, standard_sample): + # Upstream warmup exercises its text sampler, including grammar and + # penalties, with synthetic non-audio prompts. Keep that path intact. + if all(str(r).startswith("_warmup_") for r in input_batch.req_ids): + return super().sample_step(hidden_states, input_batch, req_states, grammar_output, standard_sample) + if grammar_output is not None or input_batch.num_draft_tokens: + raise ValueError("Higgs MRV2 does not support grammar or speculative sampling") + metadata = self.state.sampling_metadata(input_batch.req_ids, input_batch.idx_mapping[: input_batch.num_reqs]) + hidden = hidden_states[input_batch.logits_indices] + logits = self.state.model.compute_logits(hidden, metadata) + sampled = self.state.model.sample(logits, metadata) + count, rejected = get_num_sampled_and_rejected( + torch.ones(input_batch.num_reqs, dtype=torch.int32, device=hidden.device), + input_batch.seq_lens, + input_batch.cu_num_logits, + input_batch.idx_mapping, + req_states.prefill_len.gpu, + ) + output = SamplerOutput(sampled.sampled_token_ids, sampled.logprobs_tensors, None, count, rejected) + payload = self.state.model.post_sample_multimodal_outputs( + req_ids=input_batch.req_ids, + invalid_req_indices=[], + multimodal_outputs=None, + ) + return OmniSamplingOutput( + output, + count, + rejected, + payload, + include_hidden_states=False, + owns_multimodal_outputs=True, + finalize_multimodal=self.state.finalize_audio_snapshot, + ) diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/mrv2_state_kernels.py b/vllm_omni/model_executor/models/higgs_audio_v3/mrv2_state_kernels.py new file mode 100644 index 000000000..e880cc014 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/mrv2_state_kernels.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused movement of model-owned delay state when a request batch changes.""" + +import torch +from vllm.triton_utils import tl +from vllm.triton_utils import triton as tr + + +@tr.jit(do_not_specialize=["count"]) +def _move( + indices, + source_codes, + source_has, + source_delay, + source_eoc, + source_done, + dest_codes, + dest_has, + dest_delay, + dest_eoc, + dest_done, + count, + books: tl.constexpr, + gather: tl.constexpr, + block: tl.constexpr, +): + row = tl.program_id(0) + col = tl.arange(0, block) + if gather: + encoded = tl.load(indices + row) + fresh = encoded < 0 + src = tl.where(fresh, -encoded - 1, encoded) + dst = row + else: + src = tl.load(indices + row) + dst = tl.load(indices + count + row) + fresh = False + codes = tl.load(source_codes + src * books + col, col < books, 0) + has = tl.load(source_has + src) + delay = tl.load(source_delay + src) + eoc = tl.load(source_eoc + src) + done = tl.load(source_done + src) + tl.store(dest_codes + dst * books + col, tl.where(fresh, 0, codes), col < books) + tl.store(dest_has + dst, tl.where(fresh, 0, has)) + tl.store(dest_delay + dst, tl.where(fresh, 0, delay)) + tl.store(dest_eoc + dst, tl.where(fresh, -1, eoc)) + tl.store(dest_done + dst, tl.where(fresh, 0, done)) + + +def sync_state(model, previous_rows, pool_rows, current_rows): + # One owned pinned upload, ordered on the decode stream. Separate scatter + # and gather kernels provide a global ordering point for moved pool rows. + device = model._decode_has_codes.device + packed = torch.tensor(previous_rows + pool_rows + current_rows, dtype=torch.int64, pin_memory=True).to( + device, non_blocking=True + ) + names = ("last_codes", "has_codes", "delay_count", "eoc_countdown", "generation_done") + decode = [getattr(model, "_decode_" + name) for name in names] + pool = [getattr(model, "_state_pool_" + name) for name in names] + n = len(previous_rows) + block = tr.next_power_of_2(model.num_codebooks) + if n: + _move[(n,)](packed, *decode, *pool, n, model.num_codebooks, False, block) + if current_rows: + _move[(len(current_rows),)]( + packed[2 * n :], *pool, *decode, len(current_rows), model.num_codebooks, True, block + ) diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py b/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py index 9bfc411df..c044c907d 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/pipeline.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """higgs-audio v3 pipeline: Talker (text -> 8-codebook codec) -> Code2Wav (codec -> 24 kHz PCM). Two delivery modes are wired here: @@ -19,6 +19,9 @@ from vllm_omni.config.stage_config import ( StageExecutionType, StagePipelineConfig, ) +from vllm_omni.outputs.output_modality import TensorAccumulationStrategy, register_key_accumulation_strategy + +register_key_accumulation_strategy("codes.audio", TensorAccumulationStrategy.CONCAT_DIM0) _PROC = "vllm_omni.model_executor.stage_input_processors.higgs_audio_v3" @@ -34,6 +37,8 @@ HIGGS_AUDIO_V3_PIPELINE = PipelineConfig( execution_type=StageExecutionType.LLM_AR, input_sources=(), owns_tokenizer=True, + supports_native_mrv2_data_plane=True, + custom_process_next_stage_input_func=f"{_PROC}.talker2code2wav_full_payload", engine_output_type="latent", # stop_token_ids: the model-owned sampler forces eos at ramp-down # completion. Safety stops from the actual V3 checkpoint: @@ -54,7 +59,9 @@ HIGGS_AUDIO_V3_PIPELINE = PipelineConfig( final_output_type="audio", engine_output_type="audio", model_arch="HiggsAudioV3Code2WavForConditionalGeneration", - sync_process_input_func=f"{_PROC}.talker2code2wav", + sync_process_input_func=f"{_PROC}.talker2code2wav_token_only", + supports_native_mrv2_data_plane=True, + requires_full_payload_input=True, sampling_constraints={"detokenize": True}, ), ), diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/step_kernels.py b/vllm_omni/model_executor/models/higgs_audio_v3/step_kernels.py new file mode 100644 index 000000000..10f9e1214 --- /dev/null +++ b/vllm_omni/model_executor/models/higgs_audio_v3/step_kernels.py @@ -0,0 +1,150 @@ +# ruff: noqa: N803 +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project +"""Fused per-step Higgs kernels for the MRV2 decode loop. + +Each replaces a chain of small eager launches that ran on every model step +(terminal-row restore, prefill reset and audio-tail validation before dense +sampling; token and feedback embeddings on pure decode steps). +""" + +import torch +from vllm.triton_utils import tl, triton + + +@triton.jit +def _step_tail_guard_kernel( + ids_ptr, + qsl_ptr, + has_codes_ptr, + done_ptr, + delay_ptr, + eoc_ptr, + error_ptr, + num_rows, + audio_id, + eos_id, + HAS_QSL: tl.constexpr, + PROMPT_MODE: tl.constexpr, + BLOCK: tl.constexpr, +): + rows = tl.arange(0, BLOCK) + mask = rows < num_rows + if HAS_QSL: + start = tl.load(qsl_ptr + rows, mask=mask, other=0) + end = tl.load(qsl_ptr + rows + 1, mask=mask, other=1) + valid = end > 0 + tail = tl.load(ids_ptr + tl.maximum(end - 1, 0), mask=mask, other=audio_id) + prefill = mask & ((end - start) > 1) + else: + valid = mask + tail = tl.load(ids_ptr + rows, mask=mask, other=audio_id) + prefill = rows < 0 + done = tl.load(done_ptr + rows, mask=mask, other=0) != 0 + has = tl.load(has_codes_ptr + rows, mask=mask, other=0) != 0 + if PROMPT_MODE: + # An in-flight EOS row after retirement is terminal, never a new stream. + done = done | (tail == eos_id) + # New prefill rows start a fresh delay-pattern stream. + done = done & ~prefill + has = has & ~prefill + tl.store(done_ptr + rows, done, mask=mask) + tl.store(has_codes_ptr + rows, has, mask=mask) + tl.store(delay_ptr + rows, 0, mask=prefill) + tl.store(eoc_ptr + rows, -1, mask=prefill) + ok = (tail == audio_id) & valid + if PROMPT_MODE: + ok = ok | has | done + tl.store(error_ptr, tl.max((mask & ~ok).to(tl.int32), axis=0)) + + +def step_tail_guard( + ids: torch.Tensor, + qsl: torch.Tensor | None, + has_codes: torch.Tensor, + done: torch.Tensor, + delay: torch.Tensor, + eoc: torch.Tensor, + error: torch.Tensor, + *, + audio_id: int, + eos_id: int, + prompt_mode: bool, +) -> None: + """Fused terminal-row restore, prefill reset and audio-tail validation. + + One launch replaces the index/compare/reset chains that ran before every + dense sample, for pure decode and mixed prefill steps alike; ``error`` + becomes nonzero if any row is not an audio row. ``qsl`` (query start + locations) selects each row's last token; ``None`` means one token per row. + """ + rows = int(has_codes.numel()) + _step_tail_guard_kernel[(1,)]( + ids, + ids if qsl is None else qsl, + has_codes, + done, + delay, + eoc, + error, + rows, + audio_id, + eos_id, + HAS_QSL=qsl is not None, + PROMPT_MODE=prompt_mode, + BLOCK=max(16, triton.next_power_of_2(rows)), + ) + + +@triton.jit +def _decode_embeddings_kernel( + out_ptr, + ids_ptr, + text_weight_ptr, + audio_weight_ptr, + codes_ptr, + has_codes_ptr, + num_rows, + hidden, + audio_vocab, + BOOKS: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + cols = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = cols < hidden + out = out_ptr + row.to(tl.int64) * hidden + cols + if row >= num_rows: + # Graph padding rows own no request and must stay zero. + tl.store(out, tl.zeros((BLOCK,), dtype=out_ptr.dtype.element_ty), mask=mask) + return + if tl.load(has_codes_ptr + row) != 0: + # Audio feedback: sum of the previous frame's codebook embeddings, + # accumulated in float32 and rounded once like ``embedding(...).sum()``. + acc = tl.zeros((BLOCK,), dtype=tl.float32) + for book in range(BOOKS): + code = tl.load(codes_ptr + row * BOOKS + book).to(tl.int64) + book * audio_vocab + acc += tl.load(audio_weight_ptr + code * hidden + cols, mask=mask, other=0.0).to(tl.float32) + tl.store(out, acc.to(out_ptr.dtype.element_ty), mask=mask) + else: + token = tl.maximum(tl.load(ids_ptr + row), 0).to(tl.int64) + tl.store(out, tl.load(text_weight_ptr + token * hidden + cols, mask=mask), mask=mask) + + +def decode_embeddings( + out: torch.Tensor, + ids: torch.Tensor, + text_weight: torch.Tensor, + audio_weight: torch.Tensor, + codes: torch.Tensor, + has_codes: torch.Tensor, + padded_rows: int, + audio_vocab: int, +) -> None: + """One-token-per-row input embeddings with audio feedback, into ``out[:padded_rows]``.""" + rows, hidden = int(ids.numel()), int(out.shape[1]) + books = int(codes.shape[1]) + block = 1024 + _decode_embeddings_kernel[(padded_rows, triton.cdiv(hidden, block))]( + out, ids, text_weight, audio_weight, codes, has_codes, rows, hidden, audio_vocab, BOOKS=books, BLOCK=block + ) diff --git a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py index 353be03e8..8a7c66121 100644 --- a/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py +++ b/vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Stage-input processor for higgs-audio v3: Talker -> Code2Wav. Two adapters: @@ -21,6 +21,7 @@ to avoid corrupting valid tail content. from __future__ import annotations +from collections.abc import Mapping from typing import Any import torch @@ -95,6 +96,18 @@ def _filter_real_code_frames(audio_codes_qt: torch.Tensor) -> torch.Tensor: return frames[valid].t().contiguous() +def _whole_utterance_codec_codes(audio_rows: torch.Tensor) -> torch.Tensor: + """Convert complete delayed [T, Q] rows to flat codebook-major CPU codes.""" + codes = _revert_delay_pattern(audio_rows.to(device="cpu", dtype=torch.long).t().contiguous()) + # Substitute special/padding codes after de-delay. Clamping would map EOC + # to a valid codec entry and introduce audible artifacts. + codes = torch.where((codes < 0) | (codes >= _NUM_REAL_CODES), 0, codes) + # Preserve the existing removal of the residual ramp-down frame. + if codes.shape[-1] >= 2: + codes = codes[:, :-1] + return codes.reshape(-1).contiguous() + + def talker2code2wav( source_outputs: list[Any], prompt: Any = None, @@ -133,43 +146,16 @@ def talker2code2wav( f"got {audio_codes.shape[1]}. Audio codes shape: {tuple(audio_codes.shape)}" ) - # Transpose to [Q, T] for delay pattern reversal - codes_qt = audio_codes.transpose(0, 1).contiguous().cpu() - - # Step 1: Revert delay pattern try: - codes_qt = _revert_delay_pattern(codes_qt) + codec_codes = _whole_utterance_codec_codes(audio_codes) except ValueError as exc: logger.warning("Skipping invalid Higgs Audio v3 code sequence for Stage 1: %s", exc) code2wav_inputs.append(_empty_code2wav_prompt()) continue - # Step 2: Replace out-of-range codes (BOC=1024, EOC=1025, -1) with 0. - # Must use torch.where, NOT clamp: clamp(max=1023) turns 1025→1023 - # which is a valid codec code and decodes to audio artifacts. - # Matches sglang's: torch.where(codes >= codec_vocab, 0, codes) - codes_qt = torch.where( - (codes_qt >= _NUM_REAL_CODES) | (codes_qt < 0), - torch.zeros_like(codes_qt), - codes_qt, - ) - - # Step 3: Trim the last frame. After de-delay, the final frame - # contains residual ramp-down codes (EOC→0 substituted) that - # decode to a brief noise artifact at the end of the audio. - if codes_qt.shape[-1] >= 2: - codes_qt = codes_qt[:, :-1] - - if codes_qt.numel() == 0: - code2wav_inputs.append(_empty_code2wav_prompt()) - continue - - # Code2Wav expects codebook-major flat: [Q * num_frames] - codec_codes = codes_qt.reshape(-1).tolist() - code2wav_inputs.append( OmniTokensPrompt( - prompt_token_ids=codec_codes, + prompt_token_ids=codec_codes.tolist(), multi_modal_data=None, mm_processor_kwargs=None, additional_information=None, @@ -261,6 +247,13 @@ def talker2code2wav_async_chunk( elif not finished: return None + return _flush_async_chunk( + transfer_manager, request_id, finished, _async_chunk_config(transfer_manager), emitted_frames + ) + + +def _async_chunk_config(transfer_manager: Any) -> tuple[int, int, int, int]: + """Validated (chunk, left context, right holdback, initial chunk) frames.""" connector = getattr(transfer_manager, "connector", None) raw_cfg = getattr(connector, "config", {}) or {} cfg = raw_cfg.get("extra", raw_cfg) if isinstance(raw_cfg, dict) else {} @@ -281,7 +274,18 @@ def talker2code2wav_async_chunk( f"codec_right_holdback_frames={right_holdback_size_config}, " f"initial_codec_chunk_frames={configured_initial_chunk_size}" ) + return chunk_size, left_context_size_config, right_holdback_size_config, configured_initial_chunk_size + +def _flush_async_chunk( + transfer_manager: Any, + request_id: str, + finished: bool, + config: tuple[int, int, int, int], + emitted_frames: dict[str, int], +) -> OmniPayloadStruct | None: + """Emit the next sliding-window codec chunk once enough rows accumulated.""" + chunk_size, left_context_size_config, right_holdback_size_config, configured_initial_chunk_size = config n_rows = len(transfer_manager.code_prompt_token_ids[request_id]) de_delayed_total = max(0, n_rows - (_NUM_CODEBOOKS - 1)) emitted = int(emitted_frames.get(request_id, 0)) @@ -385,3 +389,83 @@ def talker2code2wav_async_chunk( codes=CodesStruct(audio=codec_codes), meta=meta, ) + + +def talker2code2wav_async_chunk_batch( + transfer_manager: Any, + pooling_outputs: list[OmniPayload | None], + requests: list[Any], + is_finished: list[bool], +) -> list[OmniPayloadStruct | None]: + """One talker step of async-chunk payloads with one host conversion per row. + + Same accumulation and flush semantics as ``talker2code2wav_async_chunk``; + the chunk config is resolved once per step and each emitted ``[1, Q]`` + CPU row is converted with a single ``tolist`` instead of several small + tensor operations per request. + """ + if not (len(pooling_outputs) == len(requests) == len(is_finished)): + raise ValueError("batch codec inputs must have identical lengths") + emitted_frames = getattr(transfer_manager, "higgs_v3_emitted_frames", None) + if emitted_frames is None: + emitted_frames = {} + transfer_manager.higgs_v3_emitted_frames = emitted_frames + config = None + rows_by_request = transfer_manager.code_prompt_token_ids + payloads: list[OmniPayloadStruct | None] = [] + for pooling_output, request, finished_flag in zip(pooling_outputs, requests, is_finished): + request_id = request.external_req_id + finished = bool(finished_flag or request.is_finished()) + if isinstance(pooling_output, dict): + codes = pooling_output.get("codes") + audio = codes.get("audio") if isinstance(codes, dict) else None + if isinstance(audio, torch.Tensor) and audio.numel(): + if audio.ndim == 2 and audio.device.type == "cpu" and not audio.is_floating_point(): + row = audio.tolist()[-1] + if len(row) != _NUM_CODEBOOKS: + raise ValueError( + f"talker emit row has {len(row)} codebooks; expected {_NUM_CODEBOOKS} " + "for higgs_audio_v3 async_chunk." + ) + else: + row = _extract_last_step_row(pooling_output) + row = None if row is None else row.cpu().tolist() + if row is not None: + rows_by_request[request_id].append(row) + elif not finished: + payloads.append(None) + continue + if config is None: + config = _async_chunk_config(transfer_manager) + payloads.append(_flush_async_chunk(transfer_manager, request_id, finished, config, emitted_frames)) + return payloads + + +def talker2code2wav_token_only(source_outputs, prompt=None, _requires_multimodal_data=False): + """Control slots for the native full-payload consumer.""" + result = [] + for output in source_outputs: + if not output.finished: + continue + mm = getattr(output.outputs[0], "multimodal_output", None) + if isinstance(mm, Mapping) and isinstance(mm.get("codes", {}).get("audio"), torch.Tensor): + result.extend(talker2code2wav([output], prompt, _requires_multimodal_data)) + else: + result.append(OmniTokensPrompt(prompt_token_ids=[0])) + return result + + +def talker2code2wav_full_payload(transfer_manager, pooling_output, request): + """De-delay the accumulated raw rows once before native codec delivery.""" + del transfer_manager, request + audio = pooling_output.get("codes.audio") if isinstance(pooling_output, dict) else None + if audio is None and isinstance(pooling_output, dict): + audio = pooling_output.get("codes", {}).get("audio") + codes = torch.empty(0, dtype=torch.long) + if isinstance(audio, torch.Tensor) and audio.numel(): + if audio.numel() % _NUM_CODEBOOKS: + raise ValueError("Higgs payload must contain complete codebook rows") + rows = audio.reshape(-1, _NUM_CODEBOOKS).to(device="cpu", dtype=torch.long) + if rows.shape[0] >= _NUM_CODEBOOKS: + codes = _whole_utterance_codec_codes(rows) + return {"codes": {"audio": codes}, "meta": {"finished": torch.tensor(True)}} diff --git a/vllm_omni/transformers_utils/configs/higgs_audio_v3.py b/vllm_omni/transformers_utils/configs/higgs_audio_v3.py index f6db731c7..cc6b4638a 100644 --- a/vllm_omni/transformers_utils/configs/higgs_audio_v3.py +++ b/vllm_omni/transformers_utils/configs/higgs_audio_v3.py @@ -1,5 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Configuration class for higgs-audio v3 (HiggsMultimodalQwen3) in vllm-omni. ``HiggsAudioV3Config.from_pretrained(model_path)`` returns a config with @@ -95,6 +95,8 @@ class HiggsAudioV3Config(PretrainedConfig): audio_continuation_id: int | None = None, enable_flashinfer_api_unwrap: bool = True, enable_mlp_cudagraph: bool = True, + audio_full_sample_graph: bool | None = None, + codec_cuda_graph: bool | None = None, **kwargs: Any, ) -> None: # Legacy perf knob removed: Higgs v3 scheduler tokens now come from @@ -131,9 +133,18 @@ class HiggsAudioV3Config(PretrainedConfig): self.audio_continuation_id = audio_continuation_id self.enable_flashinfer_api_unwrap = bool(enable_flashinfer_api_unwrap) self.enable_mlp_cudagraph = bool(enable_mlp_cudagraph) + self.audio_full_sample_graph = audio_full_sample_graph + self.codec_cuda_graph = codec_cuda_graph super().__init__(**kwargs) + def resolve_graph_defaults(self, *, use_v2_model_runner: bool) -> tuple[bool, bool]: + """Resolve sampler/codec graphs without mutating shared checkpoint config.""" + return ( + use_v2_model_runner if self.audio_full_sample_graph is None else self.audio_full_sample_graph, + use_v2_model_runner if self.codec_cuda_graph is None else self.codec_cuda_graph, + ) + @classmethod def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs: Any) -> HiggsAudioV3Config: """Load config and resolve special token IDs from the checkpoint tokenizer. diff --git a/vllm_omni/worker/base.py b/vllm_omni/worker/base.py index 6ee27c719..6969cf79f 100644 --- a/vllm_omni/worker/base.py +++ b/vllm_omni/worker/base.py @@ -13,6 +13,7 @@ from contextlib import AbstractContextManager, nullcontext import torch from vllm.logger import init_logger from vllm.utils.mem_utils import format_gib, memory_profiling +from vllm.v1.worker.gpu_worker import CompilationTimes from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm_omni.diffusion.data import ( @@ -39,6 +40,17 @@ class OmniGPUWorkerBase(GPUWorker): for custom trace naming, background gzip, and trace path collection. """ + def _capture_auxiliary_graphs(self) -> None: + """Let opt-in models warm valid inputs before the worker becomes ready.""" + capture = getattr(self.model_runner.model, "capture_auxiliary_graphs", None) + if callable(capture): + capture() + + def compile_or_warm_up_model(self) -> CompilationTimes: + result = super().compile_or_warm_up_model() + self._capture_auxiliary_graphs() + return result + def load_model(self, *args, **kwargs): with self._maybe_get_memory_pool_context("weights"): res = super().load_model(*args, **kwargs) @@ -225,7 +237,7 @@ class OmniGPUWorkerBase(GPUWorker): logger.info(f"[LLM Worker {self.rank}] Wake-up complete.") return True - def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK: + def handle_sleep_task(self, task: OmniSleepTask) -> OmniACK | None: "Handle deterministic Sleep command from the main process" try: if isinstance(task, dict): @@ -279,7 +291,7 @@ class OmniGPUWorkerBase(GPUWorker): pass return OmniACK(task_id=task.task_id, status="ERROR", error_msg=str(e)) - def handle_wake_task(self, task: OmniWakeTask) -> OmniACK: + def handle_wake_task(self, task: OmniWakeTask) -> OmniACK | None: "Handle deterministic Wakeup command from the main process" try: if isinstance(task, dict): diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index ad3e50d52..330355c64 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -1655,7 +1655,9 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex model_config = getattr(self, "model_config", None) if model_config is None: model_config = getattr(getattr(self, "vllm_config", None), "model_config", None) - if not bool(getattr(model_config, "async_chunk", False)): + if not bool(getattr(model_config, "async_chunk", False)) and not self._model_omni_flag( + getattr(self, "model", None), "supports_async_whole_payload" + ): return False if bool(getattr(model_config, "enable_return_routed_experts", False)): return False @@ -2110,6 +2112,14 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex scheduler_output.total_num_scheduled_tokens, ) + post_sample_mm = getattr(getattr(self, "model", None), "post_sample_multimodal_outputs", None) + if callable(post_sample_mm): + multimodal_outputs = post_sample_mm( + req_ids=req_ids_output_copy, + invalid_req_indices=invalid_req_indices, + multimodal_outputs=multimodal_outputs, + ) + multimodal_outputs = self._run_post_sample_talker_mtp( req_ids=req_ids_output_copy, valid_sampled_token_ids=valid_sampled_token_ids, @@ -2184,12 +2194,18 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin, Duplex if output_tensor_snapshot.async_payload is not None: with record_function_or_nullcontext("omni_async_output:wait_cpu_payload"): output_tensor_snapshot.async_payload.wait() + mm_snapshot = output_tensor_snapshot.multimodal_outputs + finalize_snapshot = getattr( + getattr(self, "model", None), "finalize_multimodal_outputs_from_cpu_snapshot", None + ) + if callable(finalize_snapshot): + mm_snapshot = finalize_snapshot(mm_snapshot) with record_function_or_nullcontext("omni_output_builder:total"): output = self._build_omni_model_runner_output_from_snapshot( scheduler_output=scheduler_output_snapshot, hidden_states=output_tensor_snapshot.hidden_states, staged_hidden_states_cpu=output_tensor_snapshot.staged_hidden_states_cpu, - multimodal_outputs=output_tensor_snapshot.multimodal_outputs, + multimodal_outputs=mm_snapshot, req_ids_output_copy=req_ids_output_snapshot, req_id_to_index_output_copy=req_id_to_index_output_snapshot, valid_sampled_token_ids=valid_sampled_token_ids_snapshot, diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 0e939957e..9e1a0b818 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -126,4 +126,5 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): start = time.perf_counter() self.model_runner.profile_run() + self._capture_auxiliary_graphs() return CompilationTimes(language_model=time.perf_counter() - start, encoder=0.0) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 9cd81fb1a..d8af6b201 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -2148,7 +2148,19 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): model_kwargs_extra = self._build_model_kwargs_extra() update_decode_metadata = getattr(self.model, "update_decode_step_metadata", None) if getattr(self.model, "supports_omni_decode_step_metadata", False) and callable(update_decode_metadata): + cpu_input_tail_ids = None + if ( + getattr(self.model, "requires_cpu_input_tail_ids", False) + and not self.use_async_scheduling + and input_ids is not None + and input_ids.data_ptr() == self.input_ids.gpu.data_ptr() + ): + count = len(self.input_batch.req_ids) + tails = self.query_start_loc.cpu[1 : count + 1].to(torch.long) - 1 + if count and bool((tails >= 0).all()) and bool((tails < input_ids.numel()).all()): + cpu_input_tail_ids = self.input_ids.cpu.index_select(0, tails).tolist() update_decode_metadata( + cpu_input_tail_ids=cpu_input_tail_ids, input_ids=input_ids, positions=positions, inputs_embeds=inputs_embeds, @@ -2156,14 +2168,19 @@ class OmniGPUModelRunner(PrefixCacheRunnerMixin, GPUModelRunner): req_ids=self.input_batch.req_ids, ) - model_output = super()._model_forward( - input_ids=input_ids, - positions=positions, - intermediate_tensors=intermediate_tensors, - inputs_embeds=inputs_embeds, - **model_kwargs, - **model_kwargs_extra, - ) + try: + model_output = super()._model_forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + **model_kwargs_extra, + ) + finally: + finish_decode_step = getattr(self.model, "finish_decode_step_forward", None) + if callable(finish_decode_step): + finish_decode_step() # CUDAGraphWrapper's weak_ref_tensors preserves the fields but turns # NamedTuple outputs into plain tuples. Restore the Omni envelope # before a model adapter or extraction discards its multimodal fields. diff --git a/vllm_omni/worker_v2/model_states/__init__.py b/vllm_omni/worker_v2/model_states/__init__.py index df327a129..be7e19dbb 100644 --- a/vllm_omni/worker_v2/model_states/__init__.py +++ b/vllm_omni/worker_v2/model_states/__init__.py @@ -31,6 +31,9 @@ def init_omni_model_state( capabilities; otherwise delegates to the upstream v2 factory. """ + factory = getattr(model, "create_omni_model_state", None) + if factory is not None: + return factory(vllm_config, encoder_cache, device) uses_omni_lifecycle = any( getattr(model, flag, False) is True for flag in ("has_preprocess", "has_postprocess", "have_multimodal_outputs") ) 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 08b7acdd0..6700e26b5 100644 --- a/vllm_omni/worker_v2/model_states/omni_model_state.py +++ b/vllm_omni/worker_v2/model_states/omni_model_state.py @@ -17,6 +17,7 @@ import inspect import threading import types from collections.abc import Callable +from dataclasses import replace from typing import Any, cast import numpy as np @@ -26,10 +27,12 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.logger import init_logger from vllm.v1.core.sched.output import NewRequestData +from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache from vllm.v1.worker.gpu.model_states.default import DefaultModelState from vllm.v1.worker.gpu.states import RequestState +from vllm.v1.worker.utils import AttentionGroup from vllm_omni.model_executor.models.output_templates import OmniOutput, OwnedBatchTensor from vllm_omni.platforms import current_omni_platform @@ -413,6 +416,36 @@ class OmniModelState(DefaultModelState): # Input preparation # ------------------------------------------------------------------ + def prepare_attn( + self, + input_batch: InputBatch, + cudagraph_mode: CUDAGraphMode, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + for_capture: bool = False, + ubatch_idx: int = 0, + ) -> dict[str, Any]: + if for_capture and input_batch.max_query_len is None: + # 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. + # 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( + input_batch, + cudagraph_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=for_capture, + ubatch_idx=ubatch_idx, + ) + def prepare_inputs_embeds( self, scheduled_encoder_inputs: dict[str, list[int]],