diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index a324d6bbf..dd6d5ec31 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -568,6 +568,55 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test" + key: nightly-tts-performance-higgs-audio-v3 + depends_on: upload-nightly-pipeline + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" + timeout_in_minutes: 180 + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - export VLLM_USE_DEEP_GEMM="0" + - export VLLM_MOE_USE_DEEP_GEMM="0" + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below. - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test" key: nightly-diffusion-x2iat-group @@ -1463,6 +1512,7 @@ steps: - nightly-omni-performance-vllm-text - nightly-omni-performance-multi-replicas - nightly-tts-performance + - nightly-tts-performance-higgs-audio-v3 - nightly-diffusion-x2iat-performance-qwen-image - nightly-diffusion-x2iat-performance-qwen-image-edit - nightly-diffusion-x2iat-performance-qwen-image-edit-2511 @@ -1476,6 +1526,7 @@ steps: - export DEFAULT_INPUT_DIR=tests/dfx/perf/results - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3 - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text diff --git a/benchmarks/tts/model_configs.yaml b/benchmarks/tts/model_configs.yaml index 0a26bc0cb..9acf4e47e 100644 --- a/benchmarks/tts/model_configs.yaml +++ b/benchmarks/tts/model_configs.yaml @@ -33,6 +33,15 @@ models: voice_clone: task_type: Base + bosonai/higgs-audio-v3-tts-4b: + supported_tasks: [default_voice, voice_clone] + backend: openai-audio-speech + endpoint: /v1/audio/speech + trust_remote_code: true + task_extra_body: + default_voice: {} + voice_clone: {} + openbmb/VoxCPM2: supported_tasks: [voice_clone, default_voice] trust_remote_code: true diff --git a/recipes/BosonAI/Higgs-Audio-V3-TTS.md b/recipes/BosonAI/Higgs-Audio-V3-TTS.md index 07f252bef..b924ac472 100644 --- a/recipes/BosonAI/Higgs-Audio-V3-TTS.md +++ b/recipes/BosonAI/Higgs-Audio-V3-TTS.md @@ -52,6 +52,7 @@ vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`). +It matches the high-throughput profile. **Offline batch inference:** @@ -103,8 +104,28 @@ python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client - Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity. - Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`). - `max_num_seqs=16` for both stages. - - `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1). + - Stage 0 and Stage 1 default to the same device (`0`) for single-GPU serving. + - Stage 0 intentionally keeps `enforce_eager=true`. This preserves the Higgs-specific local MLP CUDA graph path, which is the current high-throughput default. + - Stage 1 remains `enforce_eager=true` for the codec decoder. +- Deploy profiles: + - High throughput: `vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml`. + Use this for medium/high concurrency. The auto-discovered + `higgs_multimodal_qwen3.yaml` is kept as a compatibility/default alias for + this profile. + - Low latency: `vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml`. + Use this for low-concurrency serving (for example c1-c4). It sets Stage 0 + `enforce_eager=false` and enables vLLM `FULL_DECODE_ONLY` CUDA graph through + YAML `compilation_config`; no environment variable is required. + - Profile details: `vllm_omni/deploy/README_higgs_audio_v3.md`. +- Performance note: + - Do not switch the auto-discovered default Stage 0 profile to vLLM + `FULL_DECODE_ONLY` CUDA graph without an end-to-end throughput and + audio-quality revalidation. On the H20 SeedTTS c16/full-dataset benchmark, + the eager Stage 0 path with Higgs local MLP CUDA graph reproduced + ~35 audio_s/s (`1088/1088` OK per run, three runs around 134-136s). A + separate H20 low-concurrency smoke after the FULL_DECODE audio-feedback + capture fix showed lower c1/c4 latency, so the FULL_DECODE profile is kept + as an explicit low-latency option rather than the throughput default. - Known limitations: - - CUDA graph for stage 0 helps batch=1 (-13% RTF) but not batch>1 due to model-owned sampler running outside graph. - Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture). - - Async streaming (chunk-based) not yet implemented; sync-only pipeline. + - Stage 0 full-decode CUDA graph is experimental; sampler, delay-state updates, staging, and request postprocess remain outside the graph. diff --git a/tests/dfx/perf/tests/test_higgs_audio_v3.json b/tests/dfx/perf/tests/test_higgs_audio_v3.json new file mode 100644 index 000000000..6fdaec580 --- /dev/null +++ b/tests/dfx/perf/tests/test_higgs_audio_v3.json @@ -0,0 +1,74 @@ +[ + { + "test_name": "test_higgs_audio_v3", + "server_params": { + "model": "bosonai/higgs-audio-v3-tts-4b", + "stage_overrides": { + "0": { + "gpu_memory_utilization": 0.3 + }, + "1": { + "gpu_memory_utilization": 0.2 + } + }, + "extra_cli_args": [ + "--deploy-config", + "vllm_omni/deploy/higgs_multimodal_qwen3.yaml", + "--trust-remote-code" + ] + }, + "benchmark_params": [ + { + "task": "default_voice", + "eval_phase": "latency", + "dataset_name": "seed-tts-text", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", + "num_prompts": [ + 20 + ], + "max_concurrency": [ + 1 + ], + "seed_tts_locale": "en", + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "median_audio_ttfp_ms": [ + 3000 + ], + "median_audio_rtf": [ + 0.5 + ] + } + }, + { + "task": "default_voice", + "eval_phase": "throughput", + "dataset_name": "seed-tts-text", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "dataset_path": "benchmarks/build_dataset/seed_tts_smoke", + "num_prompts": [ + 500 + ], + "max_concurrency": [ + 16 + ], + "seed_tts_locale": "en", + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "median_audio_ttfp_ms": [ + 5000 + ], + "median_audio_rtf": [ + 1.2 + ], + "audio_throughput": [ + 30 + ] + } + } + ] + } +] diff --git a/tests/distributed/omni_connectors/test_adapter_and_flow.py b/tests/distributed/omni_connectors/test_adapter_and_flow.py index c2be5003c..548d16eb6 100644 --- a/tests/distributed/omni_connectors/test_adapter_and_flow.py +++ b/tests/distributed/omni_connectors/test_adapter_and_flow.py @@ -145,7 +145,7 @@ def test_shm_connector_flow(mocker: MockerFixture): Using real SharedMemoryConnector (inline mode for simplicity). """ # 1. Setup Connector - config = {"shm_threshold_bytes": 1024} # Large threshold to use inline + config = {"shm_threshold_bytes": 1024, "inline_small_payloads": True} connector = SharedMemoryConnector(config) connectors_map = {("0", "1"): connector} diff --git a/tests/distributed/omni_connectors/test_basic_connectors.py b/tests/distributed/omni_connectors/test_basic_connectors.py index 4d0661be8..1518e98b1 100644 --- a/tests/distributed/omni_connectors/test_basic_connectors.py +++ b/tests/distributed/omni_connectors/test_basic_connectors.py @@ -75,18 +75,21 @@ def test_create_unknown_connector(): @pytest.fixture def shm_connector(): - config = {"shm_threshold_bytes": 100} # Small threshold for testing + config = {"shm_threshold_bytes": 100, "inline_small_payloads": True} return SharedMemoryConnector(config) def test_put_get_inline(shm_connector): - """Test transfer for small data (currently always uses SHM).""" + """Test inline transfer for small data.""" data = {"small": "data"} success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data) assert success is True - assert "shm" in metadata + assert "inline_bytes" in metadata + assert "shm" not in metadata assert "size" in metadata + assert shm_connector._metrics["inline_writes"] == 1 + assert shm_connector._metrics["shm_writes"] == 0 # Retrieve retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata) diff --git a/tests/distributed/omni_connectors/test_shm_connector.py b/tests/distributed/omni_connectors/test_shm_connector.py index e702318e3..73928aaf6 100644 --- a/tests/distributed/omni_connectors/test_shm_connector.py +++ b/tests/distributed/omni_connectors/test_shm_connector.py @@ -13,7 +13,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @pytest.fixture() def connector(): - c = SharedMemoryConnector({"shm_threshold_bytes": 64}) + c = SharedMemoryConnector({"shm_threshold_bytes": 0}) yield c c.close() diff --git a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py index 627e22ffb..066d0c6a0 100644 --- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py @@ -6,6 +6,10 @@ AC-5 (delay pattern), AC-7 (stage processor), and AC-10 (registry) without requiring the actual checkpoint or GPU. """ +import asyncio +import time +from collections import OrderedDict + import pytest import torch @@ -232,7 +236,7 @@ class TestDelayPatternBehavior: class TestSamplerMethods: - """Test the actual _sample_audio_codes and _apply_delay_pattern_masking methods.""" + """Test the actual sampler and batched delay masking methods.""" def _make_minimal_talker(self): """Create a minimal talker-like object with sampler/masking methods.""" @@ -241,18 +245,65 @@ class TestSamplerMethods: class FakeTalker: num_codebooks = 8 codebook_size = 1026 - _audio_state = {} t = FakeTalker() t._sample_audio_codes = mod.HiggsAudioV3TalkerForConditionalGeneration._sample_audio_codes.__get__(t) - t._apply_delay_pattern_masking = ( - mod.HiggsAudioV3TalkerForConditionalGeneration._apply_delay_pattern_masking.__get__(t) + return t + + def _make_batched_sampler_talker(self, num_rows=4): + """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: + num_codebooks = 8 + codebook_size = 1026 + + t = FakeTalker() + t._decode_last_codes = torch.arange(num_rows * t.num_codebooks, dtype=torch.long).view( + num_rows, t.num_codebooks ) + t._decode_has_codes = torch.tensor([True, False, True, True])[:num_rows].clone() + t._decode_delay_count = torch.tensor([0, 2, 7, 4], dtype=torch.int32)[:num_rows].clone() + t._decode_eoc_countdown = torch.tensor([-1, -1, 3, 1], dtype=torch.int32)[:num_rows].clone() + t._decode_generation_done = torch.tensor([False, False, False, True])[:num_rows].clone() + t._decode_active_audio_count = 0 + t._codebook_index_cache = {} + t._row_index_cache = {} + t._boc_frame_cache = {} + t._last_audio_codes_buffer = None + t._last_audio_host_staging = None + t._last_audio_gpu_staging = None + t._last_audio_staging_event = None + t._audio_staging_events = [] + t._audio_staging_event_cursor = 0 + t._last_audio_codes = None + t._last_audio_code_valid = [] + t._postprocess_cursor = 0 + t._postprocess_audio_rows = 0 + t._postprocess_audio_active_rows = 0 + + cls = mod.HiggsAudioV3TalkerForConditionalGeneration + for name in ( + "_sample_audio_codes", + "_ensure_decode_state_capacity", + "_get_row_indices", + "_get_codebook_indices", + "_get_boc_frames", + "_get_audio_codes_buffer", + "_get_audio_gpu_staging_buffer", + "_get_audio_host_staging_buffer", + "_apply_delay_pattern_masking_batched", + "_update_delay_state_batched", + "_prefill_row_mask", + "_audio_seed_mask_from_step_input", + ): + setattr(t, name, getattr(cls, name).__get__(t)) + t._device_cache_key = cls._device_cache_key return t def test_sample_respects_mask(self): """Tokens masked to -inf must never be sampled.""" - t = self._make_minimal_talker() + t = self._make_batched_sampler_talker(num_rows=1) logits = torch.full((10, 1026), float("-inf")) # Only token 500 is allowed for each row logits[:, 500] = 0.0 @@ -262,7 +313,7 @@ class TestSamplerMethods: def test_sample_all_masked_falls_back_to_argmax(self): """All-masked row should fall back to argmax (least-negative logit).""" - t = self._make_minimal_talker() + t = self._make_batched_sampler_talker(num_rows=1) logits = torch.full((2, 1026), float("-inf")) # Row 0: all masked # Row 1: only token 42 allowed @@ -275,11 +326,13 @@ class TestSamplerMethods: """During delay phase, codebooks beyond delay_count must have only BOC allowed.""" from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID - t = self._make_minimal_talker() + t = self._make_batched_sampler_talker(num_rows=1) # Simulate delay_count=2 for batch row 0 - t._audio_state[0] = {"num_delay": 2, "num_remaining_delays": None} + t._ensure_decode_state_capacity(1, torch.device("cpu")) + t._decode_delay_count[0] = 2 + t._decode_eoc_countdown[0] = -1 cb_logits = torch.zeros(1, 8, 1026) # [1 audio row, 8 codebooks, 1026 vocab] - t._apply_delay_pattern_masking(cb_logits, [0]) + t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long)) # CBs 3-7 should have only BOC allowed (everything else -inf) for q in range(3, 8): row = cb_logits[0, q] @@ -292,10 +345,12 @@ class TestSamplerMethods: """Active codebooks during delay should have BOC disallowed.""" from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID - t = self._make_minimal_talker() - t._audio_state[0] = {"num_delay": 3, "num_remaining_delays": None} + t = self._make_batched_sampler_talker(num_rows=1) + t._ensure_decode_state_capacity(1, torch.device("cpu")) + t._decode_delay_count[0] = 3 + t._decode_eoc_countdown[0] = -1 cb_logits = torch.zeros(1, 8, 1026) - t._apply_delay_pattern_masking(cb_logits, [0]) + t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long)) # CBs 0-3 are active: BOC should be -inf for q in range(4): assert cb_logits[0, q, BOC_ID].item() == float("-inf") @@ -304,10 +359,12 @@ class TestSamplerMethods: """Only codebook 0 should allow EOC during normal generation.""" from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID - t = self._make_minimal_talker() - t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": None} + t = self._make_batched_sampler_talker(num_rows=1) + t._ensure_decode_state_capacity(1, torch.device("cpu")) + t._decode_delay_count[0] = 8 + t._decode_eoc_countdown[0] = -1 cb_logits = torch.zeros(1, 8, 1026) - t._apply_delay_pattern_masking(cb_logits, [0]) + t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long)) # CB0 should keep EOC assert cb_logits[0, 0, EOC_ID].item() != float("-inf") # CB1-7 should have EOC masked @@ -321,11 +378,13 @@ class TestSamplerMethods: EOC_ID, ) - t = self._make_minimal_talker() + t = self._make_batched_sampler_talker(num_rows=1) # Ramp-down with 4 remaining delays: lock CBs 0-3 to EOC - t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": 4} + t._ensure_decode_state_capacity(1, torch.device("cpu")) + t._decode_delay_count[0] = 8 + t._decode_eoc_countdown[0] = 4 cb_logits = torch.zeros(1, 8, 1026) - t._apply_delay_pattern_masking(cb_logits, [0]) + t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long)) # CBs 0-3 locked: only EOC allowed for q in range(4): row = cb_logits[0, q] @@ -337,6 +396,195 @@ class TestSamplerMethods: assert cb_logits[0, q, BOC_ID].item() == float("-inf") assert cb_logits[0, q, EOC_ID].item() == float("-inf") + def test_batched_delay_masking_all_rows_matches_sparse_path(self): + """The all-row fast path must preserve sparse path masking semantics.""" + torch.manual_seed(0) + sparse = self._make_batched_sampler_talker() + all_rows = self._make_batched_sampler_talker() + rows = torch.arange(4, dtype=torch.long) + logits_sparse = torch.randn(4, 8, 1026) + logits_all_rows = logits_sparse.clone() + + sparse._apply_delay_pattern_masking_batched(logits_sparse, rows, all_rows=False) + all_rows._apply_delay_pattern_masking_batched(logits_all_rows, rows, all_rows=True) + + assert torch.equal(logits_all_rows, logits_sparse) + + def test_delay_state_update_all_rows_matches_sparse_path(self): + """The all-row fast path must update state and staging like index_copy_.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( + EOC_ID, + ) + + sparse = self._make_batched_sampler_talker() + all_rows = self._make_batched_sampler_talker() + rows = torch.arange(4, dtype=torch.long) + codes = torch.tensor( + [ + [10, 11, 12, 13, 14, 15, 16, 17], + [20, 21, 22, 23, 24, 25, 26, 27], + [30, 31, EOC_ID, 33, 34, 35, 36, 37], + [40, 41, 42, 43, 44, 45, 46, 47], + ], + dtype=torch.long, + ) + code_row_mask = torch.tensor([True, False, True, True]) + + sparse._update_delay_state_batched( + codes.clone(), rows, 4, torch.device("cpu"), code_row_mask=code_row_mask, all_rows=False + ) + all_rows._update_delay_state_batched( + codes.clone(), rows, 4, torch.device("cpu"), code_row_mask=code_row_mask, all_rows=True + ) + + assert torch.equal(all_rows._decode_delay_count, sparse._decode_delay_count) + assert torch.equal(all_rows._decode_eoc_countdown, sparse._decode_eoc_countdown) + assert torch.equal(all_rows._decode_generation_done, sparse._decode_generation_done) + assert torch.equal(all_rows._decode_last_codes, sparse._decode_last_codes) + assert torch.equal(all_rows._decode_has_codes, sparse._decode_has_codes) + assert all_rows._last_audio_codes is None + assert sparse._last_audio_codes is None + assert torch.equal(all_rows._last_audio_host_staging, sparse._last_audio_host_staging) + + def test_direct_audio_sampler_does_not_override_done_tokens(self): + """Done rows are represented directly in GPU sampled tokens.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + t = self._make_batched_sampler_talker(num_rows=2) + 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, 99999]) + t._last_step_query_start_loc = None + t._decode_has_codes = torch.tensor([True, True]) + t._decode_generation_done = torch.tensor([False, True]) + 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 = 2 + t._last_audio_done_flags = None + t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: 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 cb_logits, audio_rows, all_rows=False: None + t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long) + t._update_delay_state_batched = lambda *args, **kwargs: None + t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t) + + sampler_output = t.sample(torch.zeros(2, 200000), sampling_metadata=object()) + + assert not getattr( + mod.HiggsAudioV3TalkerForConditionalGeneration, "supports_sampled_token_ids_cpu_override", False + ) + assert sampler_output.sampled_token_ids.tolist() == [[99999], [151671]] + + def test_direct_audio_sampler_uses_gpu_tokens_for_active_rows(self): + """Direct audio sampling leaves scheduler-visible tokens to the GPU tensor.""" + from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod + + t = self._make_batched_sampler_talker(num_rows=2) + 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, 99999]) + t._last_step_query_start_loc = None + t._decode_has_codes = torch.tensor([True, True]) + 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 = 2 + t._last_audio_done_flags = [0, 0] + t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: 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 cb_logits, audio_rows, all_rows=False: None + t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long) + t._update_delay_state_batched = lambda *args, **kwargs: None + t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t) + + sampler_output = t.sample(torch.zeros(2, 200000), sampling_metadata=object()) + + assert sampler_output.sampled_token_ids.tolist() == [[99999], [99999]] + + def test_terminal_rampdown_frame_is_emitted_before_done(self): + """The final ramp-down frame must be emitted, not replaced by a done marker.""" + from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID + + t = self._make_batched_sampler_talker(num_rows=1) + t._decode_delay_count[0] = 8 + t._decode_eoc_countdown[0] = 1 + t._decode_generation_done[0] = False + t._decode_has_codes[0] = True + codes = torch.tensor([[10, 11, 12, 13, 14, 15, 16, 17]], dtype=torch.long) + + t._update_delay_state_batched( + codes, + torch.tensor([0], dtype=torch.long), + 1, + torch.device("cpu"), + code_row_mask=torch.tensor([True]), + all_rows=True, + ) + + staged = t._last_audio_host_staging[0] + assert staged[t.num_codebooks].item() == 1 # valid + assert staged[t.num_codebooks + 1].item() == 1 # done + assert not (staged[: t.num_codebooks] == -1).any() + assert staged[: t.num_codebooks - 1].eq(EOC_ID).all() + assert t._decode_generation_done[0].item() is True + assert t._decode_has_codes[0].item() is False + + def test_decode_state_growth_preserves_inflight_rows(self): + """Growing decode buffers must copy old rows and initialize new rows inactive.""" + t = self._make_batched_sampler_talker(num_rows=2) + old_codes = t._decode_last_codes.clone() + old_has = t._decode_has_codes.clone() + old_delay = t._decode_delay_count.clone() + old_rem = t._decode_eoc_countdown.clone() + old_done = t._decode_generation_done.clone() + + t._ensure_decode_state_capacity(5, torch.device("cpu")) + + assert t._decode_last_codes.shape[0] >= 5 + assert torch.equal(t._decode_last_codes[:2], old_codes) + assert torch.equal(t._decode_has_codes[:2], old_has) + assert torch.equal(t._decode_delay_count[:2], old_delay) + assert torch.equal(t._decode_eoc_countdown[:2], old_rem) + assert torch.equal(t._decode_generation_done[:2], old_done) + assert not t._decode_has_codes[2:5].any() + assert not t._decode_generation_done[2:5].any() + assert t._decode_eoc_countdown[2:5].eq(-1).all() + + def test_mixed_batch_prefill_mask_targets_request_rows(self): + """Mixed prefill/decode must reset only prefill request rows.""" + t = self._make_batched_sampler_talker(num_rows=3) + t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5]) + + mask = t._prefill_row_mask(3, torch.device("cpu")) + + assert torch.equal(mask, torch.tensor([False, True, False])) + + def test_mixed_batch_seed_mask_maps_token_tails_to_request_rows(self): + """Seed detection should return one flag per request, not per token.""" + t = self._make_batched_sampler_talker(num_rows=3) + t._audio_continuation_id = 99999 + t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5]) + t._last_step_input_ids = torch.tensor([99999, 11, 12, 13, 99999]) + + mask = t._audio_seed_mask_from_step_input(3, torch.device("cpu")) + + assert torch.equal(mask, torch.tensor([True, False, True])) + + def test_mixed_batch_seed_mask_allows_prefill_tail_audio_token(self): + """A prefill span ending in <|audio|> starts audio-code generation.""" + t = self._make_batched_sampler_talker(num_rows=3) + t._audio_continuation_id = 99999 + t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5]) + t._last_step_input_ids = torch.tensor([99999, 11, 12, 99999, 99999]) + + mask = t._audio_seed_mask_from_step_input(3, torch.device("cpu")) + + assert torch.equal(mask, torch.tensor([True, True, True])) + # ---- AC-6: Feedback Method Tests ---- @@ -348,10 +596,14 @@ class TestFeedbackMethods: class FakeTalker: _last_audio_codes = torch.tensor([[100, 200, 300, 400, 500, 600, 700, 800]]) + _last_audio_code_valid = [True] + _last_audio_host_staging = None + _last_audio_staging_event = None _postprocess_cursor = 0 t = FakeTalker() t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t) + t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t) result = t.postprocess(torch.zeros(1, 64)) assert "codes" in result assert "audio" in result["codes"] @@ -363,10 +615,14 @@ class TestFeedbackMethods: class FakeTalker: _last_audio_codes = torch.tensor([[-1, -1, -1, -1, -1, -1, -1, -1]]) + _last_audio_code_valid = [False] + _last_audio_host_staging = None + _last_audio_staging_event = None _postprocess_cursor = 0 t = FakeTalker() t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t) + t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t) result = t.postprocess(torch.zeros(1, 64)) assert result == {} @@ -378,10 +634,14 @@ class TestFeedbackMethods: _last_audio_codes = torch.tensor( [[100, 200, 300, 400, 500, 600, 700, 800], [-1, -1, -1, -1, -1, -1, -1, -1]] ) + _last_audio_code_valid = [True, False] + _last_audio_host_staging = None + _last_audio_staging_event = None _postprocess_cursor = 0 t = FakeTalker() t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t) + t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t) r1 = t.postprocess(torch.zeros(1, 64)) assert "codes" in r1 r2 = t.postprocess(torch.zeros(1, 64)) @@ -405,14 +665,23 @@ class TestAudioFeedback: codebook_size = 1026 _audio_continuation_id = 99999 # fake audio token _last_step_query_start_loc = None - _audio_state = {} + _decode_active_audio_count = 0 + _decode_last_codes = torch.zeros(4, 8, dtype=torch.long) + _decode_has_codes = torch.zeros(4, dtype=torch.bool) + _use_external_decode_cudagraph = False multimodal_embedding = embed model = type("M", (), {"embed_tokens": lambda self, ids: torch.zeros(ids.shape[0], 16)})() t = FakeTalker() - t._decode_last_codes = torch.zeros(256, 8, dtype=torch.long) - t._decode_has_codes = torch.zeros(256, dtype=torch.bool) t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t) + t._is_single_token_decode_step = ( + mod.HiggsAudioV3TalkerForConditionalGeneration._is_single_token_decode_step.__get__(t) + ) + t._step_request_count = mod.HiggsAudioV3TalkerForConditionalGeneration._step_request_count.__get__(t) + t._decode_request_token_positions = ( + mod.HiggsAudioV3TalkerForConditionalGeneration._decode_request_token_positions.__get__(t) + ) + t._ensure_decode_state_capacity = lambda min_bs, device: None return t def test_text_positions_unchanged(self): @@ -428,9 +697,11 @@ class TestAudioFeedback: """Audio position with state should have its embedding replaced.""" t = self._make_feedback_talker() audio_id = t._audio_continuation_id - t._audio_state[1] = {"last_codes": torch.zeros(8, dtype=torch.long)} t._decode_last_codes[1] = torch.zeros(8, dtype=torch.long) t._decode_has_codes[1] = True + t._decode_active_audio_count = 1 + t._ensure_decode_state_capacity = lambda min_bs, device: None + t._last_step_query_start_loc = torch.tensor([0, 1, 2, 3]) input_ids = torch.tensor([1, audio_id, 3]) hidden = torch.zeros(3, 16) result = t._apply_audio_feedback(hidden, input_ids) @@ -444,12 +715,80 @@ class TestAudioFeedback: """Audio position without state should not be replaced.""" t = self._make_feedback_talker() audio_id = t._audio_continuation_id - # No state for row 0 — _decode_has_codes[0] defaults to False + # No state for row 0 input_ids = torch.tensor([audio_id]) hidden = torch.zeros(1, 16) result = t._apply_audio_feedback(hidden, input_ids) assert torch.equal(result, hidden) + def test_prefill_span_does_not_receive_audio_feedback(self): + """A prefill span can have numel == hidden rows but must not be treated as decode.""" + t = self._make_feedback_talker() + t._last_step_input_ids = torch.tensor([1, 2, 3, 4]) + t._last_step_query_start_loc = torch.tensor([0, 4]) + t._decode_last_codes[0] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[0] = True + t._decode_active_audio_count = 1 + hidden = torch.zeros(4, 16) + + result = t._apply_audio_feedback(hidden, t._last_step_input_ids) + + assert torch.equal(result, hidden) + + def test_multi_request_decode_receives_audio_feedback(self): + """Concurrent decode has numel > 1 but each request span is exactly one token.""" + t = self._make_feedback_talker() + t._last_step_input_ids = torch.tensor([10, 11, 12, 13]) + t._last_step_query_start_loc = torch.tensor([0, 1, 2, 3, 4]) + t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[2] = True + t._decode_active_audio_count = 1 + hidden = torch.zeros(4, 16) + + result = t._apply_audio_feedback(hidden, t._last_step_input_ids) + + assert result[2].abs().sum() > 0 + assert result[0].abs().sum() == 0 + assert result[1].abs().sum() == 0 + assert result[3].abs().sum() == 0 + + def test_external_decode_graph_capture_keeps_feedback_ops_without_query_start_loc(self): + """FULL_DECODE capture has no request qsl but still needs dense decode feedback ops.""" + t = self._make_feedback_talker() + t._use_external_decode_cudagraph = True + t._last_step_input_ids = torch.tensor([10, 11, 12, 13]) + t._last_step_query_start_loc = None + t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[2] = True + hidden = torch.zeros(4, 16) + + result = t._apply_audio_feedback(hidden, t._last_step_input_ids) + + assert result[2].abs().sum() > 0 + assert result[0].abs().sum() == 0 + assert result[1].abs().sum() == 0 + assert result[3].abs().sum() == 0 + + def test_mixed_prefill_decode_receives_feedback_only_on_decode_rows(self): + """Mixed batches should skip prefill spans but keep audio feedback for decode spans.""" + t = self._make_feedback_talker() + t._last_step_input_ids = torch.tensor([99999, 21, 22, 23, 99999]) + t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5]) + t._decode_last_codes[0] = torch.zeros(8, dtype=torch.long) + t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long) + t._decode_has_codes[0] = True + t._decode_has_codes[2] = True + t._decode_active_audio_count = 2 + hidden = torch.zeros(5, 16) + + result = t._apply_audio_feedback(hidden, t._last_step_input_ids) + + assert result[0].abs().sum() > 0 + assert result[4].abs().sum() > 0 + assert result[1].abs().sum() == 0 + assert result[2].abs().sum() == 0 + assert result[3].abs().sum() == 0 + # ---- AC-8: Codec Strictness ---- @@ -775,6 +1114,23 @@ class TestStageInputProcessor: with pytest.raises(ValueError, match="Not enough frames"): _revert_delay_pattern(codes) + def test_talker2code2wav_skips_too_few_frames_without_crashing(self): + from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( + talker2code2wav, + ) + + class Output: + multimodal_output = {"codes": {"audio": torch.zeros(5, 8, dtype=torch.long)}} + + class TalkerOutput: + finished = True + outputs = [Output()] + + result = talker2code2wav([TalkerOutput()]) + + assert len(result) == 1 + assert result[0]["prompt_token_ids"] == [] + def test_filter_real_code_frames(self): from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( _filter_real_code_frames, @@ -826,8 +1182,14 @@ class TestRegistry: # __file__ = tests/unit/higgs_audio_v3/test_*.py → 4 dirnames to repo root repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) - yaml_path = os.path.join(repo_root, "vllm_omni", "deploy", "higgs_multimodal_qwen3.yaml") - assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}" + deploy_dir = os.path.join(repo_root, "vllm_omni", "deploy") + for name in ( + "higgs_multimodal_qwen3.yaml", + "higgs_multimodal_qwen3_high_throughput.yaml", + "higgs_multimodal_qwen3_low_latency.yaml", + ): + yaml_path = os.path.join(deploy_dir, name) + assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}" # ---- AC-3: Prompt Builder ---- @@ -914,298 +1276,89 @@ class TestPromptBuilder: assert 151704 not in ids # <|ref_text|> -# ---- Regression: ramp-down done-before-emit off-by-one ---- +class TestVoiceCloneReferenceCache: + def test_audio_tokenizer_dir_env_accepts_parent_or_subdir(self, tmp_path, monkeypatch): + from vllm_omni.model_executor.models.higgs_audio_v2 import higgs_audio_v2_tokenizer as tok + parent = tmp_path / "OmniVoice" + subdir = parent / "audio_tokenizer" + subdir.mkdir(parents=True) + (subdir / "config.json").write_text("{}", encoding="utf-8") -def _simulate_delay_state_machine( - sampled_codes_per_step: list[torch.Tensor], -) -> list[torch.Tensor]: - """Reproduce the delay-pattern state machine from sample(). + monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent)) + assert tok._resolve_audio_tokenizer_dir() == str(subdir) - Feeds ``sampled_codes_per_step`` (each [Q]) through the exact same - logic as HiggsAudioV3TalkerForConditionalGeneration.sample() and - returns the list of emitted frames (non -1 codes). - """ - from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import ( - BOC_ID, - EOC_ID, - ) + monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir)) + assert tok._resolve_audio_tokenizer_dir() == str(subdir) - num_codebooks = 8 - bos = BOC_ID - eos_stream = EOC_ID + def test_higgs_v3_ref_code_cache_returns_clone(self): + from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech - state: dict = { - "num_delay": 0, - "num_remaining_delays": None, - "audio_out_ids": None, - "last_codes": torch.full((num_codebooks,), bos, dtype=torch.long), - "should_terminate": False, - } + serving = object.__new__(OmniOpenAIServingSpeech) + serving._higgs_audio_v3_ref_code_cache = OrderedDict() + serving._higgs_audio_v3_ref_code_cache_bytes = 0 + serving._higgs_audio_v3_ref_code_inflight = {} - emitted: list[torch.Tensor] = [] - for step_codes in sampled_codes_per_step: - if state["should_terminate"]: - break + codes = torch.arange(16, dtype=torch.long).reshape(2, 8) + serving._put_higgs_audio_v3_ref_codes("ref-a", codes) - num_delay = state["num_delay"] - num_remaining_delays = state["num_remaining_delays"] - this_codes = step_codes.clone() + cached = serving._get_higgs_audio_v3_ref_codes("ref-a") + assert cached is not None + cached.fill_(0) - if num_delay + 1 < num_codebooks: - this_codes[num_delay + 1 :] = bos - num_delay += 1 + cached_again = serving._get_higgs_audio_v3_ref_codes("ref-a") + assert cached_again is not None + assert torch.equal(cached_again, codes) - if num_remaining_delays is not None: - this_codes[: num_codebooks - num_remaining_delays] = eos_stream - num_remaining_delays -= 1 - else: - eos_positions = (this_codes == eos_stream).nonzero(as_tuple=False).reshape(-1) - if eos_positions.numel() > 0: - last_eos_idx = int(eos_positions[-1].item()) - this_codes[: last_eos_idx + 1] = eos_stream - num_remaining_delays = num_codebooks - last_eos_idx - 1 + def test_higgs_v3_ref_code_inflight_deduplicates_concurrent_encode(self): + from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech - if num_remaining_delays is not None and num_remaining_delays <= 0: - state["num_delay"] = 0 - state["num_remaining_delays"] = None - state["should_terminate"] = True - continue + serving = object.__new__(OmniOpenAIServingSpeech) + serving._higgs_audio_v3_ref_code_cache = OrderedDict() + serving._higgs_audio_v3_ref_code_cache_bytes = 0 + serving._higgs_audio_v3_ref_code_inflight = {} + calls = 0 - state["num_delay"] = num_delay - state["num_remaining_delays"] = num_remaining_delays - state["last_codes"] = this_codes.clone() - emitted.append(this_codes) + def encode_reference_audio(_wav, _sr): + nonlocal calls + calls += 1 + time.sleep(0.05) + return torch.arange(16, dtype=torch.long).reshape(2, 8) - return emitted + def apply_delay_pattern(codes): + return codes + 1 + async def run(): + return await asyncio.gather( + *[ + serving._resolve_higgs_audio_v3_ref_codes( + "ref-a", + object(), + 24000, + encode_reference_audio, + apply_delay_pattern, + ) + for _ in range(3) + ] + ) -class TestRampDownOffByOne: - """Regression tests for the done-before-emit off-by-one bug. + results = asyncio.run(run()) + assert calls == 1 + assert sum(int(inflight_wait) for _, _, inflight_wait in results) == 2 + for codes, cache_hit, _ in results: + assert cache_hit is False + assert torch.equal(codes, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1) - When EOC fires early, the ramp-down must emit enough frames so that - Stage1's _revert_delay_pattern receives T >= Q = 8. - """ - - def test_eoc_at_step0_terminates(self): - """EOC on cb0 at step 0 terminates; T < Q is handled by Stage1 defensive try/except.""" - from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID - - Q = 8 - steps = [] - for s in range(20): - codes = torch.randint(0, 1024, (Q,)) - if s == 0: - codes[0] = EOC_ID - steps.append(codes) - - emitted = _simulate_delay_state_machine(steps) - assert len(emitted) == Q - 1, f"EOC at step 0: expected {Q - 1} emitted frames, got {len(emitted)}" - - def test_eoc_at_step1_emits_enough_frames(self): - """EOC on cb0 at step 1 must produce >= Q emitted frames.""" - from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID - - Q = 8 - steps = [] - for s in range(20): - codes = torch.randint(0, 1024, (Q,)) - if s == 1: - codes[0] = EOC_ID - steps.append(codes) - - emitted = _simulate_delay_state_machine(steps) - assert len(emitted) >= Q, f"EOC at step 1: expected >= {Q} emitted frames, got {len(emitted)}" - - def test_eoc_after_full_rampup_emits_enough(self): - """EOC after all codebooks are active produces T >> Q.""" - from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID - - Q = 8 - eoc_step = 10 - steps = [] - for s in range(30): - codes = torch.randint(0, 1024, (Q,)) - if s == eoc_step: - codes[0] = EOC_ID - steps.append(codes) - - emitted = _simulate_delay_state_machine(steps) - assert len(emitted) >= Q, f"EOC at step {eoc_step}: expected >= {Q} emitted frames, got {len(emitted)}" - - def test_last_rampdown_frame_not_all_negative(self): - """The last emitted ramp-down frame must contain real codes, not -1.""" - from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID - - Q = 8 - steps = [] - for s in range(20): - codes = torch.randint(0, 1024, (Q,)) - if s == 1: - codes[0] = EOC_ID - steps.append(codes) - - emitted = _simulate_delay_state_machine(steps) - last = emitted[-1] - assert (last >= 0).all(), f"Last ramp-down frame contains negative values: {last.tolist()}" - - def test_emitted_frames_revertible_by_stage1(self): - """Emitted frames from EOC-at-step-1 must be revertible by _revert_delay_pattern.""" - from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _revert_delay_pattern, + cached, cache_hit, inflight_wait = asyncio.run( + serving._resolve_higgs_audio_v3_ref_codes( + "ref-a", + object(), + 24000, + encode_reference_audio, + apply_delay_pattern, + ) ) - - Q = 8 - steps = [] - for s in range(20): - codes = torch.randint(0, 1024, (Q,)) - if s == 1: - codes[0] = EOC_ID - steps.append(codes) - - emitted = _simulate_delay_state_machine(steps) - # Stack to [T, Q], transpose to [Q, T] - codes_qt = torch.stack(emitted, dim=0).t().contiguous() - assert codes_qt.shape[0] == Q - assert codes_qt.shape[1] >= Q - - # Must not raise - reverted = _revert_delay_pattern(codes_qt) - assert reverted.shape[0] == Q - assert reverted.shape[1] == codes_qt.shape[1] - Q + 1 - - -class TestStage1DefensiveHandling: - """Stage1 must not crash the server when receiving T < Q frames.""" - - def test_revert_delay_pattern_raises_on_insufficient_frames(self): - """_revert_delay_pattern correctly raises ValueError for T < Q.""" - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - _revert_delay_pattern, - ) - - codes = torch.zeros(8, 7) # T=7 < Q=8 - with pytest.raises(ValueError, match="Not enough frames"): - _revert_delay_pattern(codes) - - def test_talker2code2wav_survives_insufficient_frames(self): - """talker2code2wav must not crash when audio_codes has T < Q frames.""" - from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import ( - talker2code2wav, - ) - - # Build a fake talker output with only 7 frames (T=7 < Q=8) - class FakeOutput: - def __init__(self): - self.multimodal_output = { - "codes": {"audio": torch.zeros(7, 8, dtype=torch.long)}, - } - - class FakeTalkerOutput: - finished = True - outputs = [FakeOutput()] - - # Must not raise — should return gracefully - result = talker2code2wav([FakeTalkerOutput()]) - assert len(result) == 1 - # Result is an OmniTokensPrompt with empty prompt_token_ids - out = result[0] - token_ids = out.prompt_token_ids if hasattr(out, "prompt_token_ids") else out.get("prompt_token_ids", []) - assert token_ids == [] - - -# ---- Regression: buffer growth + prefill guard ---- - - -class TestBufferGrowthPreservesState: - """Bug #1: buffer growth in _apply_audio_feedback must not wipe existing state.""" - - def _make_feedback_talker_with_buffers(self): - from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod - - embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16) - torch.nn.init.ones_(embed.weight) - - class FakeTalker: - num_codebooks = 8 - codebook_size = 1026 - multimodal_embedding = embed - - t = FakeTalker() - t._decode_last_codes = torch.zeros(4, 8, dtype=torch.long) - t._decode_has_codes = torch.zeros(4, dtype=torch.bool) - t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t) - return t - - def test_buffer_growth_preserves_existing_codes(self): - """Growing the buffer must copy old _decode_last_codes data.""" - t = self._make_feedback_talker_with_buffers() - # Set state for slot 2 - t._decode_last_codes[2] = torch.tensor([100, 200, 300, 400, 500, 600, 700, 800]) - t._decode_has_codes[2] = True - - # Trigger growth by calling with bs=8 > current capacity 4 - hidden = torch.randn(8, 16) - t._apply_audio_feedback(hidden, torch.zeros(8, dtype=torch.long)) - - # After growth, slot 2 data must still be present - assert t._decode_has_codes[2].item() is True - assert t._decode_last_codes[2, 0].item() == 100 - - def test_buffer_growth_new_slots_are_inactive(self): - """New slots from buffer growth must have has_codes=False.""" - t = self._make_feedback_talker_with_buffers() - t._decode_has_codes[1] = True - - hidden = torch.randn(8, 16) - t._apply_audio_feedback(hidden, torch.zeros(8, dtype=torch.long)) - - # Slots 4-7 (new) must be inactive - assert not t._decode_has_codes[4].item() - assert not t._decode_has_codes[7].item() - # Slot 1 (old) must still be active - assert t._decode_has_codes[1].item() - - -class TestPrefillGuard: - """Bug #2: _apply_audio_feedback must NOT run during prefill. - - The prefill detection uses attn_metadata.max_query_len (== 1 for decode, - > 1 for prefill). This correctly handles concurrent decode (N requests, - each contributing 1 token → numel=N but max_query_len=1). - """ - - @staticmethod - def _is_prefill_via_max_query_len(max_query_len: int | None, input_ids: torch.Tensor) -> bool: - """Mirror the actual detection logic in forward().""" - if max_query_len is not None: - return int(max_query_len) > 1 - return input_ids is not None and int(input_ids.numel()) > 1 - - def test_prefill_detected_via_max_query_len(self): - """Prefill: max_query_len > 1 → is_prefill=True.""" - input_ids = torch.tensor([1, 2, 3, 4, 5]) - assert self._is_prefill_via_max_query_len(5, input_ids) is True - - def test_single_decode_detected_via_max_query_len(self): - """Single-request decode: max_query_len=1 → is_prefill=False.""" - input_ids = torch.tensor([42]) - assert self._is_prefill_via_max_query_len(1, input_ids) is False - - def test_concurrent_decode_not_misclassified(self): - """Multi-request decode: numel=4 but max_query_len=1 → is_prefill=False. - - This is the regression caught by reviewer: with c=4 concurrent decode - requests, input_ids.numel()==4 but each request contributes exactly 1 - token, so max_query_len==1 and audio feedback must still apply. - """ - input_ids = torch.tensor([42, 43, 44, 45]) # 4 concurrent decode - assert self._is_prefill_via_max_query_len(1, input_ids) is False - - def test_fallback_when_max_query_len_unavailable(self): - """When attn_metadata is unavailable, fall back to numel heuristic.""" - prefill_ids = torch.tensor([1, 2, 3]) - assert self._is_prefill_via_max_query_len(None, prefill_ids) is True - decode_ids = torch.tensor([42]) - assert self._is_prefill_via_max_query_len(None, decode_ids) is False + assert calls == 1 + assert cache_hit is True + assert inflight_wait is False + assert torch.equal(cached, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1) diff --git a/vllm_omni/deploy/README_higgs_audio_v3.md b/vllm_omni/deploy/README_higgs_audio_v3.md new file mode 100644 index 000000000..ea1363069 --- /dev/null +++ b/vllm_omni/deploy/README_higgs_audio_v3.md @@ -0,0 +1,53 @@ +# Higgs-Audio V3 Deploy Profiles + +Higgs-Audio V3 has two Stage 0 graph profiles. They are intentionally separate +because the graph paths are mutually exclusive. + +## High Throughput + +Use `higgs_multimodal_qwen3_high_throughput.yaml` for medium/high concurrency +serving. This profile keeps Stage 0 `enforce_eager: true`, which preserves the +Higgs-specific local MLP CUDA graph path. It is the default production profile +for throughput-oriented serving. + +`higgs_multimodal_qwen3.yaml` is kept as the auto-discovered default deploy +config for `model_type=higgs_multimodal_qwen3`, and matches the high-throughput +profile. + +```bash +vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ + --omni --trust-remote-code \ + --deploy-config vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml +``` + +## Low Latency + +Use `higgs_multimodal_qwen3_low_latency.yaml` for low-concurrency serving +(for example c1-c4) where Stage 0 decode launch overhead dominates. This profile +sets Stage 0 `enforce_eager: false` and explicitly enables vLLM +`FULL_DECODE_ONLY` CUDA graph: + +```yaml +compilation_config: + cudagraph_capture_sizes: [1, 2, 4, 8, 16] + cudagraph_mode: FULL_DECODE_ONLY + cudagraph_num_of_warmups: 1 +``` + +FULL_DECODE is controlled by deploy configuration, not by an environment +variable. When this external decode graph is active, the Higgs talker disables +the local MLP CUDA graph automatically. + +```bash +vllm-omni serve bosonai/higgs-audio-v3-tts-4b \ + --omni --trust-remote-code \ + --deploy-config vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml +``` + +## Notes + +- Stage 1 remains `enforce_eager: true` in both profiles. +- Keep `VLLM_USE_DEEP_GEMM=0` and `VLLM_MOE_USE_DEEP_GEMM=0` for this model + unless DeepGEMM support is revalidated. +- Revalidate end-to-end throughput and audio quality before changing the default + auto-discovered config. diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml index c7016537b..5fc44a295 100644 --- a/vllm_omni/deploy/higgs_multimodal_qwen3.yaml +++ b/vllm_omni/deploy/higgs_multimodal_qwen3.yaml @@ -1,5 +1,11 @@ # Higgs-audio v3 deploy: Stage 0 (talker) -> Stage 1 (code2wav) via shared memory. # +# Auto-discovered default for model_type=higgs_multimodal_qwen3. This file is +# the high-throughput profile and intentionally keeps Stage 0 enforce_eager=true +# so the Higgs-specific local MLP CUDA graph path remains active. For low +# concurrency latency-sensitive serving, use +# higgs_multimodal_qwen3_low_latency.yaml explicitly. +# # Plain TTS only (no voice cloning). async_chunk streams codec chunks from # Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's # worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size @@ -32,6 +38,7 @@ stages: max_num_seqs: 16 gpu_memory_utilization: 0.6 enforce_eager: true + attention_backend: FLASHINFER trust_remote_code: true # Stage-0 prefix caching is on by default. The talker's HS opt-out and # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml new file mode 100644 index 000000000..86427be13 --- /dev/null +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml @@ -0,0 +1,80 @@ +# Higgs-audio v3 deploy, high-throughput profile. +# +# This is the default production profile for medium/high concurrency serving. +# Stage 0 keeps enforce_eager=true so Higgs' local MLP CUDA graph remains active. +# Do not enable vLLM FULL_DECODE graph here: it replaces the local MLP graph +# path instead of stacking on top of it. +# +# Plain TTS only (no voice cloning). async_chunk streams codec chunks from +# Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's +# worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size +# so the per-chunk codec windows stitch into a coherent PCM stream. +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + # Stream PCM bytes back to the client as Stage 1 emits each chunk. + codec_streaming: true + # Streaming chunk shape for talker2code2wav_async_chunk: + # codec_chunk_frames steady-state new frames per emit + # codec_left_context_frames sliding past context to avoid leading-edge artifacts + # codec_right_holdback_frames trailing frames held back so the codec + # has future context for emit-region tail samples + # initial_codec_chunk_frames smaller first chunk to minimize TTFA + codec_chunk_frames: 25 + codec_left_context_frames: 25 + codec_right_holdback_frames: 4 + initial_codec_chunk_frames: 1 + 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: true + attention_backend: FLASHINFER + trust_remote_code: true + # Stage-0 prefix caching is on by default. The talker's HS opt-out and + # deferred ``codes.audio`` mm-output write (see higgs_audio_v3_talker.py) + # keep per-step bookkeeping negligible; any shared text prefix across + # concurrent requests is a net throughput win. Stage 1 (codec decoder) + # has no KV cache and stays disabled. + enable_prefix_caching: true + async_scheduling: false + 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 + + - stage_id: 1 + max_num_seqs: 16 + 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 diff --git a/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml new file mode 100644 index 000000000..8210db4a8 --- /dev/null +++ b/vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml @@ -0,0 +1,85 @@ +# Higgs-audio v3 deploy, low-latency profile. +# +# Use this profile for low-concurrency serving (for example c1-c4) where Stage 0 +# decode launch overhead dominates. Stage 0 enables vLLM FULL_DECODE_ONLY CUDA +# graph through YAML configuration: +# +# enforce_eager: false +# compilation_config.cudagraph_mode: FULL_DECODE_ONLY +# +# This is not controlled by an environment variable. The Higgs talker disables +# its local MLP CUDA graph automatically when the external decode graph is active, +# because the two graph paths are mutually exclusive. +# +# Plain TTS only (no voice cloning). async_chunk streams codec chunks from +# Stage 0 to Stage 1 as they are produced (TTFA approaches one AR loop's +# worth of work). Stage 1 honors meta.left_context_size + meta.right_holdback_size +# so the per-chunk codec windows stitch into a coherent PCM stream. +async_chunk: true + +connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + # Stream PCM bytes back to the client as Stage 1 emits each chunk. + codec_streaming: true + # Streaming chunk shape for talker2code2wav_async_chunk: + # codec_chunk_frames steady-state new frames per emit + # codec_left_context_frames sliding past context to avoid leading-edge artifacts + # codec_right_holdback_frames trailing frames held back so the codec + # has future context for emit-region tail samples + # initial_codec_chunk_frames smaller first chunk to minimize TTFA + codec_chunk_frames: 25 + codec_left_context_frames: 25 + codec_right_holdback_frames: 4 + initial_codec_chunk_frames: 1 + 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: 16 + gpu_memory_utilization: 0.6 + enforce_eager: false + attention_backend: FLASHINFER + trust_remote_code: true + enable_prefix_caching: true + async_scheduling: false + compilation_config: + cudagraph_capture_sizes: [1, 2, 4, 8, 16] + cudagraph_mode: FULL_DECODE_ONLY + cudagraph_num_of_warmups: 1 + 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 + + - stage_id: 1 + max_num_seqs: 16 + 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 diff --git a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py index 6cf5c2f15..63466d116 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/shm_connector.py @@ -29,6 +29,7 @@ class SharedMemoryConnector(OmniConnectorBase): self.stage_id = config.get("stage_id", -1) self.device = config.get("device", "cuda:0") self.threshold = int(config.get("shm_threshold_bytes", 65536)) + self.inline_small_payloads = bool(config.get("inline_small_payloads", False)) self._pending_keys: set[str] = set() self._metrics = { "puts": 0, @@ -53,9 +54,11 @@ class SharedMemoryConnector(OmniConnectorBase): payload = self.serialize_obj(data) size = len(payload) - # Currently, we always use SHM. - if True: - # Use Shared Memory + # The legacy async-chunk adapter transfers only the connector key + # across stages; it cannot deliver inline metadata to the receiver. + # Keep key-addressed SHM as the default and only inline payloads + # when the caller explicitly enables the metadata-aware path. + if size >= self.threshold or not self.inline_small_payloads: lock_file = f"/dev/shm/shm_{put_key}_lockfile.lock" with open(lock_file, "wb+") as lockf: fcntl.flock(lockf, fcntl.LOCK_EX) @@ -67,9 +70,8 @@ class SharedMemoryConnector(OmniConnectorBase): self._pending_keys.add(put_key) self._metrics["shm_writes"] += 1 else: - # Inline - pass bytes directly to avoid double serialization of the object - # We already serialized it to check size, so we pass the bytes. - # The Queue will pickle these bytes (fast), avoiding re-serializing the complex object. + # Inline small payloads to avoid the mmap + file-lock overhead + # that dominates codec chunk transfers. metadata = {"inline_bytes": payload, "size": size} self._metrics["inline_writes"] += 1 diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 9127ac7db..c6709ba20 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -78,7 +78,7 @@ _MING_TTS_MODEL_STAGES = {"ming_tts"} _MOSS_TTS_MODEL_STAGES = {"moss_tts_nano"} _MOSS_TTS_FULL_MODEL_STAGES = {"moss_tts", "moss_tts_codec"} _HIGGS_AUDIO_V2_TTS_MODEL_STAGES = {"higgs_audio_v2"} -_HIGGS_AUDIO_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"} +_HIGGS_V3_TTS_MODEL_STAGES = {"higgs_audio_v3"} _GLM_TTS_MODEL_STAGES = {"glm_tts"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES @@ -87,7 +87,7 @@ _TTS_MODEL_STAGES: set[str] = ( | _COSYVOICE3_TTS_MODEL_STAGES | _OMNIVOICE_TTS_MODEL_STAGES | _HIGGS_AUDIO_V2_TTS_MODEL_STAGES - | _HIGGS_AUDIO_V3_TTS_MODEL_STAGES + | _HIGGS_V3_TTS_MODEL_STAGES | _COVO_AUDIO_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES | _MING_TTS_MODEL_STAGES @@ -123,6 +123,8 @@ _REF_AUDIO_MIN_DURATION = 1.0 # seconds _REF_AUDIO_MAX_DURATION = 30.0 # seconds _REF_AUDIO_RESOLVE_CACHE_MAX_ENTRIES = 256 _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES = 256 * 1024 * 1024 +_HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES = 256 +_HIGGS_V3_REF_CODE_CACHE_MAX_BYTES = 64 * 1024 * 1024 _QWEN3_TTS_REF_AUDIO_CACHE_KEY = "_qwen3_tts_ref_audio_cache_key" _TTS_MAX_INSTRUCTIONS_LENGTH = 500 _TTS_MAX_NEW_TOKENS_MIN = 1 @@ -288,6 +290,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): self._ref_audio_resolve_cache_max_bytes = _REF_AUDIO_RESOLVE_CACHE_MAX_BYTES self._ref_audio_model_artifact_ready: set[str] = set() self._request_ref_audio_artifact_keys: dict[str, str] = {} + self._higgs_audio_v3_ref_code_cache: OrderedDict[str, tuple[torch.Tensor, int]] = OrderedDict() + self._higgs_audio_v3_ref_code_cache_bytes = 0 + self._higgs_audio_v3_ref_code_inflight: dict[str, asyncio.Task[torch.Tensor]] = {} self._speaker_cache = get_speaker_cache() self._last_upload_ts = 0 self._upload_lock = asyncio.Lock() @@ -646,7 +651,7 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "moss_tts" if model_stage in _HIGGS_AUDIO_V2_TTS_MODEL_STAGES: return "higgs_audio_v2" - if model_stage in _HIGGS_AUDIO_V3_TTS_MODEL_STAGES: + if model_stage in _HIGGS_V3_TTS_MODEL_STAGES: return "higgs_audio_v3" if model_stage in _GLM_TTS_MODEL_STAGES: return "glm_tts" @@ -2018,17 +2023,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): apply_delay_pattern, encode_reference_audio, ) - from vllm_omni.model_executor.models.higgs_audio_v3.ref_audio_cache import ( - cached_encode_reference_audio, - ) wav_list, sr = await self._resolve_ref_audio(request.ref_audio) + artifact_key = self._get_resolved_ref_audio_artifact_key(request.ref_audio) wav = np.asarray(wav_list, dtype=np.float32) - # Cache the codec encode by content hash so identical ref_audio - # payloads skip the encode and reuse the same [num_frames, - # num_codebooks] codes tensor across requests. - ref_codes_raw = await asyncio.to_thread(cached_encode_reference_audio, wav, int(sr), encode_reference_audio) - ref_codes_delayed = apply_delay_pattern(ref_codes_raw) + ref_codes_delayed, cache_hit, inflight_wait = await self._resolve_higgs_audio_v3_ref_codes( + artifact_key, + wav, + int(sr), + encode_reference_audio, + apply_delay_pattern, + ) + del cache_hit, inflight_wait prompt_ids = adapter.build_prompt( request.input, @@ -2044,6 +2050,68 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): } return prompt + async def _resolve_higgs_audio_v3_ref_codes( + self, + artifact_key: str | None, + wav: np.ndarray, + sr: int, + encode_reference_audio, + apply_delay_pattern, + ) -> tuple[torch.Tensor, bool, bool]: + ref_codes_delayed = self._get_higgs_audio_v3_ref_codes(artifact_key) + if ref_codes_delayed is not None: + return ref_codes_delayed, True, False + if not artifact_key: + ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) + return apply_delay_pattern(ref_codes_raw), False, False + + task = self._higgs_audio_v3_ref_code_inflight.get(artifact_key) + if task is not None: + return (await task).clone(), False, True + + async def _encode_and_cache() -> torch.Tensor: + ref_codes_raw = await asyncio.to_thread(encode_reference_audio, wav, sr) + delayed = apply_delay_pattern(ref_codes_raw) + self._put_higgs_audio_v3_ref_codes(artifact_key, delayed) + cached = self._get_higgs_audio_v3_ref_codes(artifact_key) + return cached if cached is not None else delayed.detach().to("cpu", dtype=torch.long).contiguous() + + task = asyncio.create_task(_encode_and_cache()) + self._higgs_audio_v3_ref_code_inflight[artifact_key] = task + try: + return (await task).clone(), False, False + finally: + if self._higgs_audio_v3_ref_code_inflight.get(artifact_key) is task: + self._higgs_audio_v3_ref_code_inflight.pop(artifact_key, None) + + def _get_higgs_audio_v3_ref_codes(self, artifact_key: str | None) -> torch.Tensor | None: + if not artifact_key: + return None + cached = self._higgs_audio_v3_ref_code_cache.get(artifact_key) + if cached is None: + return None + self._higgs_audio_v3_ref_code_cache.move_to_end(artifact_key) + return cached[0].clone() + + def _put_higgs_audio_v3_ref_codes(self, artifact_key: str, codes: torch.Tensor) -> None: + if _HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES <= 0 or _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES <= 0 or not artifact_key: + return + cached_codes = codes.detach().to("cpu", dtype=torch.long).contiguous() + size = int(cached_codes.numel() * cached_codes.element_size()) + if size > _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES: + return + previous = self._higgs_audio_v3_ref_code_cache.pop(artifact_key, None) + if previous is not None: + self._higgs_audio_v3_ref_code_cache_bytes -= previous[1] + self._higgs_audio_v3_ref_code_cache[artifact_key] = (cached_codes, size) + self._higgs_audio_v3_ref_code_cache_bytes += size + while len(self._higgs_audio_v3_ref_code_cache) > _HIGGS_V3_REF_CODE_CACHE_MAX_ENTRIES: + _, (_, old_size) = self._higgs_audio_v3_ref_code_cache.popitem(last=False) + self._higgs_audio_v3_ref_code_cache_bytes -= old_size + while self._higgs_audio_v3_ref_code_cache_bytes > _HIGGS_V3_REF_CODE_CACHE_MAX_BYTES: + _, (_, old_size) = self._higgs_audio_v3_ref_code_cache.popitem(last=False) + self._higgs_audio_v3_ref_code_cache_bytes -= old_size + async def _resolve_higgs_audio_v3_adapter(self): """Lazy-load the tokenizer adapter for higgs_audio_v3.""" cached = getattr(self, "_higgs_audio_v3_adapter", None) diff --git a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py index a51d692b9..cacf82ee5 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py +++ b/vllm_omni/model_executor/models/higgs_audio_v2/higgs_audio_v2_tokenizer.py @@ -162,6 +162,54 @@ def _build_delay_pattern(codes: torch.Tensor) -> torch.Tensor: _ENCODER_CACHE: Any | None = None _K2_OMNIVOICE_REPO = "k2-fsa/OmniVoice" _K2_OMNIVOICE_SUBDIR = "audio_tokenizer" +_AUDIO_TOKENIZER_PATH_ENVS = ( + "HIGGS_AUDIO_TOKENIZER_PATH", + "HIGGS_AUDIO_V2_TOKENIZER_PATH", +) + + +def _normalize_audio_tokenizer_dir(path: str) -> str | None: + if not path: + return None + expanded = os.path.abspath(os.path.expanduser(path)) + if os.path.isfile(os.path.join(expanded, "config.json")): + return expanded + nested = os.path.join(expanded, _K2_OMNIVOICE_SUBDIR) + if os.path.isfile(os.path.join(nested, "config.json")): + return nested + return None + + +def _resolve_audio_tokenizer_dir() -> str | None: + """Resolve the Higgs codec encoder dir without requiring online HF access. + + Voice cloning runs in the API server process. Production H20 jobs are often + offline, so prefer explicit local directories and already-populated HF cache + before falling back to ``snapshot_download``. + """ + for env_name in _AUDIO_TOKENIZER_PATH_ENVS: + candidate = _normalize_audio_tokenizer_dir(os.getenv(env_name, "")) + if candidate is not None: + return candidate + + from huggingface_hub import try_to_load_from_cache + + cached_config = try_to_load_from_cache( + repo_id=_K2_OMNIVOICE_REPO, + filename=f"{_K2_OMNIVOICE_SUBDIR}/config.json", + ) + if isinstance(cached_config, str) and os.path.isfile(cached_config): + return os.path.dirname(cached_config) + + from huggingface_hub.constants import HF_HUB_CACHE + + safe = _K2_OMNIVOICE_REPO.replace("/", "--") + snapshots_dir = os.path.join(HF_HUB_CACHE, f"models--{safe}", "snapshots") + if os.path.isdir(snapshots_dir): + for rev in os.listdir(snapshots_dir): + candidate = _normalize_audio_tokenizer_dir(os.path.join(snapshots_dir, rev)) + if candidate is not None: + return candidate def _load_audio_tokenizer(): @@ -180,14 +228,17 @@ def _load_audio_tokenizer(): Only the ``audio_tokenizer/`` subdirectory (~806 MB) is downloaded. """ - from huggingface_hub import snapshot_download from transformers import HiggsAudioV2TokenizerModel - repo_path = snapshot_download( - _K2_OMNIVOICE_REPO, - allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], - ) - audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR) + audio_tokenizer_dir = _resolve_audio_tokenizer_dir() + if audio_tokenizer_dir is None: + from huggingface_hub import snapshot_download + + repo_path = snapshot_download( + _K2_OMNIVOICE_REPO, + allow_patterns=[f"{_K2_OMNIVOICE_SUBDIR}/*"], + ) + audio_tokenizer_dir = os.path.join(repo_path, _K2_OMNIVOICE_SUBDIR) device = "cuda" if torch.cuda.is_available() else "cpu" model = HiggsAudioV2TokenizerModel.from_pretrained(audio_tokenizer_dir).to(device) return model.eval() diff --git a/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py b/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py index bed0467c6..f6db731c7 100644 --- a/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py +++ b/vllm_omni/model_executor/models/higgs_audio_v3/configuration_higgs_audio_v3.py @@ -93,8 +93,13 @@ class HiggsAudioV3Config(PretrainedConfig): tts_token_id: int | None = None, text_token_id: int | None = None, audio_continuation_id: int | None = None, + enable_flashinfer_api_unwrap: bool = True, + enable_mlp_cudagraph: bool = True, **kwargs: Any, ) -> None: + # Legacy perf knob removed: Higgs v3 scheduler tokens now come from + # SamplerOutput.sampled_token_ids only. + kwargs.pop("enable_cpu_token_override", None) self.audio_token_id = audio_token_id self.mel_per_sample = mel_per_sample @@ -124,6 +129,8 @@ class HiggsAudioV3Config(PretrainedConfig): self.tts_token_id = tts_token_id self.text_token_id = text_token_id 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) super().__init__(**kwargs) 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 2baf70b4d..40936f85d 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 @@ -35,6 +35,7 @@ __all__ = [ logger = init_logger(__name__) + # Prefix under which codec weights are stored in the v3 checkpoint. _CODEC_PREFIX = "tied.embedding.modality_embeddings.0.model." 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 5e20245b0..9d4e05356 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 @@ -33,6 +33,8 @@ from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.qwen3 import Qwen3Model +from vllm.platforms import current_platform +from vllm.v1.outputs import SamplerOutput from vllm_omni.model_executor.models.higgs_audio_v3.configuration_higgs_audio_v3 import ( HiggsAudioV3Config, @@ -43,6 +45,7 @@ __all__ = ["HiggsAudioV3TalkerForConditionalGeneration"] logger = init_logger(__name__) + # Delay pattern constants BOC_ID = 1024 # beginning of codebook EOC_ID = 1025 # end of codebook @@ -67,12 +70,14 @@ class HiggsFusedMultiTextEmbedding(nn.Module): self.weight = nn.Parameter(torch.empty(num_codebooks * vocab_size, hidden_size)) self.num_codebooks = num_codebooks self.vocab_size = vocab_size + self.register_buffer( + "_offsets", + torch.arange(num_codebooks, dtype=torch.long) * vocab_size, + persistent=False, + ) def forward(self, codes: torch.Tensor) -> torch.Tensor: - N = self.num_codebooks - V = self.vocab_size - offsets = torch.arange(N, device=codes.device, dtype=codes.dtype) * V - fused_ids = codes + offsets + fused_ids = codes + self._offsets return F.embedding(fused_ids, self.weight).sum(dim=-2) @@ -102,15 +107,33 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): # Tell the runner to call postprocess() to emit per-step audio codes. have_multimodal_outputs: bool = True has_postprocess: bool = True + 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 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() - hf_config = vllm_config.model_config.hf_config if isinstance(hf_config, HiggsAudioV3Config): self.config = hf_config else: self.config = HiggsAudioV3Config(**hf_config.to_dict()) + model_config = getattr(vllm_config, "model_config", None) + compilation_config = getattr(vllm_config, "compilation_config", None) + cudagraph_mode = getattr(compilation_config, "cudagraph_mode", None) + mode_name = getattr(cudagraph_mode, "name", str(cudagraph_mode)).upper() + self._use_external_decode_cudagraph = ( + not bool(getattr(model_config, "enforce_eager", True)) + and cudagraph_mode is not None + and "NONE" not in mode_name + ) + if self._use_external_decode_cudagraph: + self.config.enable_mlp_cudagraph = False + assert not (self._use_external_decode_cudagraph and self.config.enable_mlp_cudagraph), ( + "Higgs Audio v3 local MLP graph and external decode CUDA graph are mutually exclusive" + ) self.vllm_config = vllm_config self.num_codebooks = int(self.config.num_codebooks) @@ -135,6 +158,8 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): vllm_config=backbone_vllm_config, prefix=f"{prefix}.model" if prefix else "model", ) + if self.config.enable_flashinfer_api_unwrap: + self._maybe_unwrap_flashinfer_api_wrappers() if self._backbone_config.tie_word_embeddings: self.lm_head = self.model.embed_tokens @@ -146,6 +171,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): ) self.logits_processor = LogitsProcessor(self._backbone_config.vocab_size) + self._text_vocab_size = int(self._backbone_config.vocab_size) # Audio continuation token ID — resolved lazily from tokenizer. # This is the <|audio|> token that serves as the LM-level continuation @@ -154,22 +180,49 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self._eos_token_id: int | None = None self._resolved_tokens = False - # Per-request audio state keyed by batch row index. - # Reset per slot via _slot_output_len tracking (same pattern as v2). - self._audio_state: dict[int, dict[str, Any]] = {} - self._slot_output_len: dict[int, int] = {} self._last_logits_hidden: torch.Tensor | None = None self._last_step_input_ids: torch.Tensor | None = None self._last_step_query_start_loc: torch.Tensor | None = None - self._last_first_audio_after_start: torch.Tensor | None = None + self._last_step_query_start_loc_buffer: torch.Tensor | None = None + self.supports_omni_decode_step_metadata = True + self._decode_step_metadata_from_runner = False self._last_audio_codes: torch.Tensor | None = None + self._last_audio_code_valid: list[bool] = [] self._postprocess_cursor: int = 0 + self._last_audio_codes_buffer: torch.Tensor | None = None + self._last_audio_host_staging: torch.Tensor | None = None + self._last_audio_gpu_staging: torch.Tensor | None = None + self._last_audio_staging_event: torch.cuda.Event | None = None + self._last_audio_valid_flags: list[int] | None = None + self._last_audio_done_flags: list[int] | None = None + self._audio_staging_events: list[torch.cuda.Event] = [] + self._audio_staging_event_cursor: int = 0 + self._row_index_cache: dict[tuple[str, int], torch.Tensor] = {} + self._codebook_index_cache: dict[tuple[str, int], torch.Tensor] = {} + self._boc_frame_cache: dict[tuple[str, int], torch.Tensor] = {} + self._fast_audio_direct_rows: int = 0 + scheduler_config = getattr(vllm_config, "scheduler_config", None) + self._mlp_cudagraph_max_batch = max(1, int(getattr(scheduler_config, "max_num_seqs", 16))) + self._postprocess_audio_rows: int = 0 + self._postprocess_audio_active_rows: int = 0 + self._mlp_graphs: dict[tuple[int, int], dict[str, Any]] = {} + self._mlp_graph_disabled: set[tuple[int, int]] = set() # Pre-allocated decode-step audio feedback buffers (CUDA-graph safe). # Populated by sample(), read by forward() via torch.where (no dict). - max_bs = 64 # safe upper bound; will grow if needed - self._decode_last_codes = torch.zeros(max_bs, self.num_codebooks, dtype=torch.long) - self._decode_has_codes = torch.zeros(max_bs, dtype=torch.bool) + max_bs = max(64, self._mlp_cudagraph_max_batch) + self.register_buffer( + "_decode_last_codes", + torch.zeros(max_bs, self.num_codebooks, dtype=torch.long), + persistent=False, + ) + self.register_buffer("_decode_has_codes", torch.zeros(max_bs, dtype=torch.bool), persistent=False) + # Keep delay/ramp counters as int64 to match arange/codebook-index + # arithmetic in the decode hot path without per-step dtype casts. + self.register_buffer("_decode_delay_count", torch.zeros(max_bs, dtype=torch.long), persistent=False) + self.register_buffer("_decode_eoc_countdown", torch.full((max_bs,), -1, dtype=torch.long), persistent=False) + self.register_buffer("_decode_generation_done", torch.zeros(max_bs, dtype=torch.bool), persistent=False) + self._decode_active_audio_count: int = 0 # PrefixCache opt-outs (mirror qwen3_tts pattern): # 1. The talker only consumes the last token's hidden state, so the @@ -183,6 +236,58 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self.requires_full_prefix_cached_hidden_states = False self.deferred_prefix_cache_mm_keys = {"codes.audio"} + def _maybe_unwrap_flashinfer_api_wrappers(self) -> None: + """Remove FlashInfer trace wrappers only from wrappers owned by this model.""" + patched: list[str] = [] + saw_flashinfer = False + seen: set[int] = set() + stack: list[Any] = [] + for layer in getattr(self.model, "layers", []): + impl = getattr(getattr(getattr(layer, "self_attn", None), "attn", None), "impl", None) + if impl is not None: + stack.append(impl) + + while stack and len(seen) < 512: + obj = stack.pop() + obj_id = id(obj) + if obj_id in seen: + continue + seen.add(obj_id) + + cls = obj.__class__ + if "flashinfer" in cls.__module__.lower(): + saw_flashinfer = True + for method_name in ("run", "paged_run", "plan"): + method = getattr(obj, method_name, None) + original = getattr(method, "__wrapped__", None) + if original is None: + continue + try: + bound = original.__get__(obj, cls) if hasattr(original, "__get__") else original + setattr(obj, method_name, bound) + except Exception: + continue + patched.append(f"{cls.__name__}.{method_name}") + + obj_dict = getattr(obj, "__dict__", None) + if not isinstance(obj_dict, dict): + continue + for value in obj_dict.values(): + if isinstance(value, dict): + stack.extend(value.values()) + elif isinstance(value, (list, tuple, set)): + stack.extend(value) + elif not isinstance(value, (bool, int, float, str, bytes, type(None), torch.Tensor)): + stack.append(value) + + if patched: + logger.info("HiggsAudioV3Talker: unwrapped FlashInfer API wrappers: %s", sorted(set(patched))) + elif saw_flashinfer: + logger.warning( + "HiggsAudioV3Talker: FlashInfer backend was detected but no API wrappers were unwrapped; " + "FlashInfer internals may have changed and wrapper-unwrapping performance benefits may be unavailable." + ) + def _resolve_token_ids(self) -> None: """Resolve <|audio|> and eos token IDs. @@ -230,6 +335,24 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): except Exception as exc: logger.warning("Failed to resolve token IDs from tokenizer: %s", exc) + def update_decode_step_metadata( + self, + *, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + omni_query_start_loc: torch.Tensor | None = None, + **_: Any, + ) -> None: + """Update per-step metadata before runner forward or CUDA graph replay.""" + _ = positions, inputs_embeds + 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: + self._ensure_decode_state_capacity(int(input_ids.shape[0]), input_ids.device) + self._set_last_step_query_start_loc(omni_query_start_loc) + self._decode_step_metadata_from_runner = True + # ------------------------------------------------------------------ forward def forward( self, @@ -239,6 +362,13 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): inputs_embeds: torch.Tensor | None = None, **kwargs: Any, ) -> torch.Tensor: + 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: + info_dicts = kwargs.get("runtime_additional_information") + if inputs_embeds is None: # Mask -100 placeholders to 0 before embedding. Use torch.where # (no Python data-dependent branch) so this is CUDA-graph safe. @@ -247,69 +377,180 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): else: hidden_states = inputs_embeds - if input_ids is not None: + if input_ids is not None and not metadata_from_runner: self._last_step_input_ids = input_ids - # Stash query_start_loc and max_query_len for prefill detection - _max_query_len = None - try: - from vllm.forward_context import get_forward_context + # Stash query_start_loc for audio-state row mapping and max_query_len + # for prefill detection. Some attention backends do not expose + # query_start_loc, so prefer backend metadata when available and + # otherwise use the runner-supplied buffer. + max_query_len = None + if not metadata_from_runner: + try: + fallback_qsl = kwargs.get("omni_query_start_loc") + from vllm.forward_context import get_forward_context - attn_metadata = get_forward_context().attn_metadata - if isinstance(attn_metadata, dict) and attn_metadata: - attn = next(iter(attn_metadata.values())) - else: - attn = attn_metadata - qsl = getattr(attn, "query_start_loc", None) - if isinstance(qsl, torch.Tensor): - self._last_step_query_start_loc = qsl.detach().clone() - else: + attn_metadata = get_forward_context().attn_metadata + if isinstance(attn_metadata, dict) and attn_metadata: + attn = next(iter(attn_metadata.values())) + else: + attn = attn_metadata + qsl = getattr(attn, "query_start_loc", None) + self._set_last_step_query_start_loc(qsl if isinstance(qsl, torch.Tensor) else fallback_qsl) + max_query_len = getattr(attn, "max_query_len", None) + except Exception: self._last_step_query_start_loc = None - _max_query_len = getattr(attn, "max_query_len", None) - except Exception: - self._last_step_query_start_loc = None - # Detect prefill vs decode using attn_metadata.max_query_len. - # max_query_len == 1 means pure decode (even with N concurrent - # requests, each contributes exactly 1 token). numel > 1 is NOT - # reliable because decode with N>1 concurrent requests gives - # input_ids.numel() == N. - if _max_query_len is not None: - is_prefill = int(_max_query_len) > 1 + # Prefill-only operations: ref audio substitution and audio feedback + # require Python dict/list ops that break CUDA graph capture. + # Prefer max_query_len where available: concurrent decode has + # input_ids.numel() > 1 but max_query_len == 1. + if max_query_len is not None: + is_prefill = int(max_query_len) > 1 else: - is_prefill = input_ids is not None and inputs_embeds is None and int(input_ids.numel()) > 1 - if is_prefill: + is_prefill = ( + input_ids is not None + and inputs_embeds is None + and not self._is_single_token_decode_step(int(hidden_states.shape[0])) + ) + if is_prefill and info_dicts: # Voice clone: replace -100 placeholder positions with ref audio embeddings - info_dicts = kwargs.get("model_intermediate_buffer") - if info_dicts is None: - info_dicts = kwargs.get("runtime_additional_information") hidden_states = self._apply_ref_audio_substitution(hidden_states, input_ids, info_dicts) - # Audio feedback: replace continuation token embeddings with audio - # embeddings from the last decoded frame (CUDA-graph safe). + # Audio feedback at decode: replace continuation token embeddings if input_ids is not None and inputs_embeds is None: hidden_states = self._apply_audio_feedback(hidden_states, input_ids) + if self.config.enable_mlp_cudagraph: + return self._run_qwen3_layers_mlp_graph(positions, hidden_states) + + return self._run_qwen3_layers_eager(positions, hidden_states) + + def _run_qwen3_layers_eager(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: residual: torch.Tensor | None = None for layer in self.model.layers: hidden_states, residual = layer(positions, hidden_states, residual) - norm_out = self.model.norm(hidden_states, residual) if isinstance(norm_out, tuple): norm_out = norm_out[0] return norm_out + def _run_qwen3_layers_mlp_graph(self, positions: torch.Tensor, hidden_states: torch.Tensor) -> torch.Tensor: + residual: torch.Tensor | None = None + batch = int(hidden_states.shape[0]) if hidden_states.ndim >= 1 else 0 + use_graph = ( + torch.cuda.is_available() + and hidden_states.is_cuda + and hidden_states.ndim == 2 + and positions.ndim == 1 + and batch > 0 + and batch <= self._mlp_cudagraph_max_batch + and int(positions.shape[0]) == batch + and self._is_decode_only_graph_batch(batch) + ) + for layer_idx, layer in enumerate(self.model.layers): + if residual is None: + residual = hidden_states + hidden_states = layer.input_layernorm(hidden_states) + else: + hidden_states, residual = layer.input_layernorm(hidden_states, residual) + hidden_states = layer.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + if use_graph: + graph_out = self._run_qwen3_mlp_graph(layer_idx, layer, hidden_states, residual) + if graph_out is not None: + hidden_states, residual = graph_out + continue + hidden_states, residual = layer.post_attention_layernorm(hidden_states, residual) + hidden_states = layer.mlp(hidden_states) + norm_out = self.model.norm(hidden_states, residual) + if isinstance(norm_out, tuple): + norm_out = norm_out[0] + return norm_out + + def _run_qwen3_mlp_graph( + self, + layer_idx: int, + layer: nn.Module, + hidden_states: torch.Tensor, + residual: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor] | None: + key = (int(layer_idx), int(hidden_states.shape[0])) + if key in self._mlp_graph_disabled: + return None + if hidden_states.ndim != 2 or residual.ndim != 2 or hidden_states.shape != residual.shape: + self._mlp_graph_disabled.add(key) + return None + if torch.cuda.is_current_stream_capturing(): + return None + + entry = self._mlp_graphs.get(key) + if entry is not None: + entry["static_hidden"].copy_(hidden_states) + entry["static_residual"].copy_(residual) + entry["graph"].replay() + return entry["static_output"], entry["static_residual_out"] + + try: + static_hidden = torch.empty_like(hidden_states) + static_residual = torch.empty_like(residual) + static_hidden.copy_(hidden_states) + static_residual.copy_(residual) + graph = torch.cuda.CUDAGraph() + with torch.inference_mode(), torch.cuda.graph(graph, pool=current_platform.get_global_graph_pool()): + normed, static_residual_out = layer.post_attention_layernorm(static_hidden, static_residual) + static_output = layer.mlp(normed) + self._mlp_graphs[key] = { + "graph": graph, + "static_hidden": static_hidden, + "static_residual": static_residual, + "static_output": static_output, + "static_residual_out": static_residual_out, + } + logger.info( + "HiggsAudioV3Talker: captured MLP CUDA graph for layer=%d decode batch=%d", + layer_idx, + int(hidden_states.shape[0]), + ) + return static_output, static_residual_out + except Exception as exc: + self._mlp_graph_disabled.add(key) + logger.warning( + "HiggsAudioV3Talker: MLP CUDA graph disabled for layer=%d batch=%d: %s", + layer_idx, + int(hidden_states.shape[0]), + exc, + ) + return None + + def _is_decode_only_graph_batch(self, batch: int) -> bool: + q_start = self._last_step_query_start_loc + return isinstance(q_start, torch.Tensor) and int(q_start.numel()) == batch + 1 + def compute_logits(self, hidden_states: torch.Tensor, sampling_metadata: Any = None) -> torch.Tensor: self._last_logits_hidden = hidden_states return self.logits_processor(self.lm_head, hidden_states, sampling_metadata) def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: - safe_ids = input_ids - if input_ids is not None and (input_ids < 0).any(): - safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids) + safe_ids = torch.where(input_ids < 0, torch.zeros_like(input_ids), input_ids) text_embed = self.model.embed_tokens(safe_ids) return self._apply_audio_feedback(text_embed, input_ids) + def _set_last_step_query_start_loc(self, query_start_loc: Any) -> None: + if not isinstance(query_start_loc, torch.Tensor): + self._last_step_query_start_loc = None + return + source = query_start_loc.detach() + numel = int(source.numel()) + buf = self._last_step_query_start_loc_buffer + if buf is None or buf.device != source.device or buf.dtype != source.dtype or int(buf.numel()) < numel: + buf = torch.empty(max(numel, 17), dtype=source.dtype, device=source.device) + self._last_step_query_start_loc_buffer = buf + buf[:numel].copy_(source.reshape(-1)) + self._last_step_query_start_loc = buf[:numel] + # ------------------------------------------------------------------ ref audio substitution def _apply_ref_audio_substitution( self, @@ -323,7 +564,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): Called at prefill to inject voice clone reference. ``info_dicts`` is a list of per-request dicts from ``model_intermediate_buffer``, each containing ``audio_input_ids`` ([T, N] delayed codes) and - ``audio_input_ids_mask`` ([T] bool mask). + ``audio_input_ids_mask`` ([T] bool mask). """ if not info_dicts: return hidden_states @@ -412,30 +653,210 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): For prefill: audio feedback is not needed (ref audio substitution handles the prefill path separately). """ - bs = hidden_states.shape[0] + num_tokens = int(hidden_states.shape[0]) + external_decode_graph = bool(getattr(self, "_use_external_decode_cudagraph", False)) + if self._decode_active_audio_count == 0 and not external_decode_graph: + return hidden_states - # Ensure buffers are on the right device and large enough - if self._decode_last_codes.device != hidden_states.device: - self._decode_last_codes = self._decode_last_codes.to(hidden_states.device) - self._decode_has_codes = self._decode_has_codes.to(hidden_states.device) - if bs > self._decode_last_codes.shape[0]: - old_size = self._decode_last_codes.shape[0] - new_size = max(bs, old_size * 2) - new_codes = torch.zeros(new_size, self.num_codebooks, dtype=torch.long, device=hidden_states.device) - new_has = torch.zeros(new_size, dtype=torch.bool, device=hidden_states.device) - new_codes[:old_size] = self._decode_last_codes - new_has[:old_size] = self._decode_has_codes - self._decode_last_codes = new_codes - self._decode_has_codes = new_has + req_rows, token_positions = self._decode_request_token_positions(num_tokens, hidden_states.device) + if int(req_rows.numel()) == 0: + return hidden_states - # Compute audio embeddings from last_codes for ALL rows (graph-safe) - codes_slice = self._decode_last_codes[:bs] # [bs, N] - has_codes = self._decode_has_codes[:bs].unsqueeze(-1) # [bs, 1] + num_reqs = self._step_request_count(num_tokens) + self._ensure_decode_state_capacity(num_reqs, hidden_states.device) + + if self._is_single_token_decode_step(num_tokens): + codes_slice = self._decode_last_codes[:num_tokens] + has_codes = self._decode_has_codes[:num_tokens].unsqueeze(-1) + current_hidden = hidden_states + else: + 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) audio_embeds = self.multimodal_embedding(codes_slice) # [bs, D] audio_embeds = audio_embeds.to(dtype=hidden_states.dtype) + replaced = torch.where(has_codes, audio_embeds, current_hidden) - # Select: where has_codes, use audio embed; else keep text embed - return torch.where(has_codes, audio_embeds, hidden_states) + if self._is_single_token_decode_step(num_tokens): + return replaced + new_hidden = hidden_states.clone() + new_hidden.index_copy_(0, token_positions, replaced) + return new_hidden + + def _is_single_token_decode_step(self, num_rows: int) -> bool: + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor): + return int(q_start.numel()) == int(num_rows) + 1 + ids = getattr(self, "_last_step_input_ids", None) + if not isinstance(ids, torch.Tensor) or int(ids.numel()) != int(num_rows): + return False + if bool(getattr(self, "_use_external_decode_cudagraph", False)): + return True + return int(num_rows) == 1 + + def _step_request_count(self, num_tokens: int) -> int: + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) >= 2: + return int(q_start.numel()) - 1 + return int(num_tokens) + + def _decode_request_token_positions( + self, + num_tokens: int, + device: torch.device, + ) -> tuple[torch.Tensor, torch.Tensor]: + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) >= 2: + q = q_start.to(device=device, dtype=torch.long) + starts = q[:-1] + spans = q[1:] - q[:-1] + req_rows = torch.arange(int(q.numel()) - 1, dtype=torch.long, device=device) + decode_mask = (spans == 1) & (starts >= 0) & (starts < int(num_tokens)) + return req_rows[decode_mask], starts[decode_mask] + + ids = getattr(self, "_last_step_input_ids", None) + if isinstance(ids, torch.Tensor) and int(ids.numel()) == int(num_tokens): + if bool(getattr(self, "_use_external_decode_cudagraph", False)): + rows = torch.arange(int(num_tokens), dtype=torch.long, device=device) + return rows, rows + if int(num_tokens) == 1: + row = torch.zeros(1, dtype=torch.long, device=device) + return row, row + empty = torch.empty(0, dtype=torch.long, device=device) + return empty, empty + + def _prefill_row_mask(self, num_rows: int, device: torch.device) -> torch.Tensor: + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) >= 2: + q = q_start.to(device=device, dtype=torch.long) + return (q[1:] - q[:-1]) > 1 + return torch.ones(num_rows, dtype=torch.bool, device=device) + + def _reset_decode_state_rows(self, mask: torch.Tensor, num_rows: int, device: torch.device) -> None: + self._ensure_decode_state_capacity(num_rows, device) + row_mask = mask.to(device=device, dtype=torch.bool) + self._decode_has_codes[:num_rows].masked_fill_(row_mask, False) + self._decode_generation_done[:num_rows].masked_fill_(row_mask, False) + self._decode_delay_count[:num_rows].masked_fill_(row_mask, 0) + self._decode_eoc_countdown[:num_rows].masked_fill_(row_mask, -1) + + def _ensure_decode_state_capacity(self, num_rows: int, device: torch.device | None = None) -> None: + """Keep GPU-resident decode state tensors aligned and large enough.""" + if device is not None and self._decode_last_codes.device != device: + self._decode_last_codes = self._decode_last_codes.to(device) + self._decode_has_codes = self._decode_has_codes.to(device) + self._decode_delay_count = self._decode_delay_count.to(device) + self._decode_eoc_countdown = self._decode_eoc_countdown.to(device) + self._decode_generation_done = self._decode_generation_done.to(device) + + cur_rows = int(self._decode_last_codes.shape[0]) + if num_rows <= cur_rows: + return + + new_size = max(num_rows, cur_rows * 2) + state_device = self._decode_last_codes.device + + new_last_codes = torch.zeros(new_size, self.num_codebooks, dtype=torch.long, device=state_device) + new_last_codes[:cur_rows].copy_(self._decode_last_codes) + self._decode_last_codes = new_last_codes + + new_has_codes = torch.zeros(new_size, dtype=torch.bool, device=state_device) + new_has_codes[:cur_rows].copy_(self._decode_has_codes) + self._decode_has_codes = new_has_codes + + new_delay_count = torch.zeros(new_size, dtype=torch.long, device=state_device) + new_delay_count[:cur_rows].copy_(self._decode_delay_count) + self._decode_delay_count = new_delay_count + + new_eoc_countdown = torch.full((new_size,), -1, dtype=torch.long, device=state_device) + new_eoc_countdown[:cur_rows].copy_(self._decode_eoc_countdown) + self._decode_eoc_countdown = new_eoc_countdown + + new_generation_done = torch.zeros(new_size, dtype=torch.bool, device=state_device) + new_generation_done[:cur_rows].copy_(self._decode_generation_done) + self._decode_generation_done = new_generation_done + + def _audio_seed_mask_from_step_input(self, num_rows: int, device: torch.device) -> torch.Tensor | None: + """Return rows whose previous token is <|audio|> using current step input_ids. + + This is the hot-path replacement for scanning sampling_metadata + output_token_ids. In decode, vLLM feeds the previous sampled token as + the current single-token input for each row. + """ + audio_id = self._audio_continuation_id + ids = self._last_step_input_ids + if audio_id is None or not isinstance(ids, torch.Tensor) or int(ids.numel()) <= 0: + return None + + flat_ids = ids.reshape(-1).to(device=device) + q_start = self._last_step_query_start_loc + if isinstance(q_start, torch.Tensor) and int(q_start.numel()) == num_rows + 1: + q = q_start.to(device=device, dtype=torch.long) + tail_idx = q[1:] - 1 + valid_tail = tail_idx >= 0 + tail_ids = flat_ids.index_select(0, tail_idx.clamp_min(0)) + return (tail_ids == int(audio_id)) & valid_tail + if int(flat_ids.numel()) == num_rows: + tail_ids = flat_ids + else: + return None + return tail_ids == int(audio_id) + + def _fast_audio_sampler_gpu_fallback_reason( + self, + *, + logits: torch.Tensor, + sampling_metadata: Any, + num_rows: int, + ) -> str | None: + if logits is None or logits.ndim != 2 or int(logits.shape[0]) != num_rows: + return "invalid_logits" + if num_rows <= 0: + return "empty_batch" + if getattr(sampling_metadata, "max_num_logprobs", None) is not None: + return "logprobs" + if getattr(sampling_metadata, "allowed_token_ids_mask", None) is not None: + return "allowed_token_ids" + if bool(getattr(sampling_metadata, "bad_words_token_ids", None)): + return "bad_words" + return None + + def _clear_last_audio_outputs(self) -> None: + self._last_audio_codes = None + self._last_audio_code_valid = [] + self._last_audio_host_staging = None + self._last_audio_staging_event = None + self._decode_active_audio_count = 0 + + def _apply_audio_mode_bias_batched( + self, + logits: torch.Tensor, + audio_mask: torch.Tensor, + done_mask: torch.Tensor, + ) -> None: + """Force audio/EOS rows with GPU masks, avoiding metadata scans.""" + audio_id = self._audio_continuation_id + if audio_id is None or logits is None or logits.ndim != 2: + return + + vocab = int(logits.shape[-1]) + if audio_id < 0 or audio_id >= vocab: + return + + target = torch.full( + (int(logits.shape[0]),), + int(audio_id), + dtype=torch.long, + device=logits.device, + ) + if self._eos_token_id is not None and 0 <= int(self._eos_token_id) < vocab: + eos_target = torch.full_like(target, int(self._eos_token_id)) + target = torch.where(done_mask.to(device=logits.device), eos_target, target) + + target = target.unsqueeze(1) + target_values = logits.gather(1, target) + logits.masked_fill_(audio_mask.to(device=logits.device).unsqueeze(1), float("-inf")) + logits.scatter_(1, target, target_values) # ------------------------------------------------------------------ sampling def sample(self, logits: torch.Tensor, sampling_metadata: Any) -> Any: @@ -447,167 +868,198 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): """ self._resolve_token_ids() - sampler = getattr(self, "_stock_sampler", None) - if sampler is None: - from vllm.v1.sample.sampler import Sampler - - sampler = Sampler() - self._stock_sampler = sampler - audio_id = self._audio_continuation_id - num_codebooks = self.num_codebooks - # Slot reuse detection: check output_token_ids length for each row. - # If a slot's output length decreased (fresh request reusing a finished - # slot), drop stale _audio_state for that row. - output_ids = getattr(sampling_metadata, "output_token_ids", None) - if output_ids is not None: - for i in range(len(output_ids)): - cur_len = len(output_ids[i]) if output_ids[i] else 0 - prev_len = self._slot_output_len.get(i, 0) - if cur_len < prev_len: - # Fresh request reusing this slot — clear stale state - self._audio_state.pop(i, None) - self._slot_output_len[i] = 0 - self._slot_output_len[i] = cur_len + def run_stock_sampler() -> Any: + sampler = getattr(self, "_stock_sampler", None) + if sampler is None: + from vllm.v1.sample.sampler import Sampler - # Bias LM logits for audio continuation - self._apply_audio_mode_bias(logits, sampling_metadata) - sampler_output = sampler(logits=logits, sampling_metadata=sampling_metadata) + sampler = Sampler() + self._stock_sampler = sampler + return sampler(logits=logits, sampling_metadata=sampling_metadata) hidden = self._last_logits_hidden self._last_logits_hidden = None if hidden is None or audio_id is None: - self._last_audio_codes = None + sampler_output = run_stock_sampler() + self._clear_last_audio_outputs() return sampler_output - sampled = getattr(sampler_output, "sampled_token_ids", None) - if sampled is None: - self._last_audio_codes = None + num_rows = int(hidden.shape[0]) + 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 + if not decode_only: + self._fast_audio_direct_rows = 0 + pfmask = self._prefill_row_mask(num_rows, hidden.device) + self._reset_decode_state_rows(pfmask, num_rows, hidden.device) + prev_audio_mask = self._audio_seed_mask_from_step_input(num_rows, hidden.device) + if prev_audio_mask is None: + prev_audio_mask = torch.zeros(num_rows, dtype=torch.bool, device=hidden.device) + active_mask = self._decode_has_codes[:num_rows].to(device=hidden.device, dtype=torch.bool) + done_mask = self._decode_generation_done[:num_rows].to(device=hidden.device, dtype=torch.bool) + fast_fallback_reason = self._fast_audio_sampler_gpu_fallback_reason( + logits=logits, + sampling_metadata=sampling_metadata, + num_rows=num_rows, + ) + seed_mask = prev_audio_mask & ~active_mask & ~done_mask + audio_mask = prev_audio_mask | active_mask | done_mask + code_row_mask = (seed_mask | active_mask) & ~done_mask + gpu_stock_sampler_reasons = {"logprobs", "allowed_token_ids", "bad_words"} + use_gpu_audio_mode = fast_fallback_reason is None or fast_fallback_reason in gpu_stock_sampler_reasons + + if fast_fallback_reason is None: + direct_audio_batch = decode_only and self._fast_audio_direct_rows == num_rows + + if direct_audio_batch: + self._fast_audio_direct_rows = num_rows + audio_tokens = torch.full((num_rows,), int(audio_id), dtype=torch.int32, device=hidden.device) + if self._eos_token_id is not None: + eos_tokens = torch.full_like(audio_tokens, int(self._eos_token_id)) + audio_tokens = torch.where(done_mask, eos_tokens, audio_tokens) + sampled = audio_tokens.unsqueeze(-1) + sampler_output = SamplerOutput(sampled_token_ids=sampled, logprobs_tensors=None) + else: + self._apply_audio_mode_bias_batched(logits, audio_mask, done_mask) + sampler_output = run_stock_sampler() + sampled = getattr(sampler_output, "sampled_token_ids", None) + if sampled is None: + self._clear_last_audio_outputs() + return sampler_output + elif use_gpu_audio_mode: + self._fast_audio_direct_rows = 0 + self._apply_audio_mode_bias_batched(logits, audio_mask, done_mask) + sampler_output = run_stock_sampler() + sampled = getattr(sampler_output, "sampled_token_ids", None) + if sampled is None: + self._clear_last_audio_outputs() + return sampler_output + else: + self._fast_audio_direct_rows = 0 + sampler_output = run_stock_sampler() + self._clear_last_audio_outputs() return sampler_output + sampled_flat = sampled.reshape(-1) - if int(sampled_flat.numel()) != int(hidden.shape[0]): - self._last_audio_codes = None + if int(sampled_flat.numel()) != num_rows: + self._clear_last_audio_outputs() return sampler_output - is_audio = sampled_flat == audio_id + if use_gpu_audio_mode: + audio_row_tensor = self._get_row_indices(num_rows, hidden.device) + seed_mask_1d = seed_mask.to(device=hidden.device, dtype=torch.bool) + done_mask_1d = done_mask.to(device=hidden.device, dtype=torch.bool) + self._decode_generation_done[:num_rows].masked_fill_(done_mask_1d, False) + self._decode_has_codes[:num_rows].masked_fill_(done_mask_1d, False) + self._decode_delay_count[:num_rows].masked_fill_(done_mask_1d, 0) + self._decode_eoc_countdown[:num_rows].masked_fill_(done_mask_1d, -1) - # Handle first-after-<|audio|> transition: initialize state but do NOT - # skip sampling. Unlike v2 (which seeds an all-BOC frame and skips the - # first audio step), v3/sglang samples codebook 0 on the very first - # step with delay_count=0 (codebooks 1-7 masked to BOC). - first_after_start = self._last_first_audio_after_start - self._last_first_audio_after_start = None - - if isinstance(first_after_start, torch.Tensor) and first_after_start.numel() == is_audio.shape[0]: - first_after_start = first_after_start.to(is_audio.device) - init_rows = first_after_start & is_audio - if bool(init_rows.any()): - boc_frame = torch.full((num_codebooks,), BOC_ID, dtype=torch.long, device=hidden.device) - for bi in torch.nonzero(init_rows, as_tuple=False).reshape(-1).tolist(): - bi = int(bi) - self._slot_output_len[bi] = 0 - self._audio_state[bi] = { - "num_delay": 0, - "num_remaining_delays": None, - "audio_out_ids": None, - "last_codes": boc_frame.clone(), - "should_terminate": False, - } - # Update CUDA-graph-safe feedback buffers - self._decode_last_codes[bi] = boc_frame.to(self._decode_last_codes.device) - self._decode_has_codes[bi] = True - # Do NOT remove init_rows from is_audio — let them be sampled - - if not bool(is_audio.any()): - self._last_audio_codes = None - return sampler_output - - audio_row_indices = torch.nonzero(is_audio, as_tuple=False).reshape(-1).tolist() + boc_frames = self._get_boc_frames(num_rows, hidden.device) + seed_mask_2d = seed_mask_1d.unsqueeze(1) + self._decode_last_codes[:num_rows] = torch.where( + seed_mask_2d, + boc_frames, + self._decode_last_codes[:num_rows], + ) + self._decode_has_codes[:num_rows].masked_fill_(seed_mask_1d, True) + self._decode_delay_count[:num_rows].masked_fill_(seed_mask_1d, 0) + self._decode_eoc_countdown[:num_rows].masked_fill_(seed_mask_1d, -1) + self._decode_generation_done[:num_rows].masked_fill_(seed_mask_1d, False) + self._decode_active_audio_count = num_rows + else: + audio_row_tensor = self._get_row_indices(num_rows, hidden.device) # Per-codebook logits at audio positions - cb_logits = self._audio_codebook_logits(hidden, is_audio) + all_audio_rows = use_gpu_audio_mode + if use_gpu_audio_mode: + cb_logits = self._audio_codebook_logits_from_rows(hidden, audio_row_tensor, all_rows=True) + else: + cb_logits = self._audio_codebook_logits_from_rows(hidden, audio_row_tensor, all_rows=False) # Apply delay pattern masking BEFORE sampling - self._apply_delay_pattern_masking(cb_logits, audio_row_indices) + self._apply_delay_pattern_masking_batched(cb_logits, audio_row_tensor, all_rows=all_audio_rows) # Sample per-codebook cb_logits_2d = cb_logits.reshape(-1, cb_logits.shape[-1]) codes_2d = self._sample_audio_codes(cb_logits_2d) codes_flat = codes_2d.view(cb_logits.shape[0], cb_logits.shape[1]).to(torch.long) - # Update delay pattern state - eos_stream = EOC_ID - bos = BOC_ID - new_codes_flat: list[torch.Tensor] = [] - for local_i, batch_i in enumerate(audio_row_indices): - state = self._audio_state.setdefault( - int(batch_i), - { - "num_delay": 0, - "num_remaining_delays": None, - "audio_out_ids": None, - "last_codes": torch.full((num_codebooks,), bos, dtype=torch.long, device=hidden.device), - "should_terminate": False, - }, - ) - num_delay: int = state["num_delay"] - num_remaining_delays: int | None = state["num_remaining_delays"] - this_codes = codes_flat[local_i].clone() - - # Leading delay-pattern BOS pad - if num_delay + 1 < num_codebooks: - this_codes[num_delay + 1 :] = bos - num_delay += 1 - - # Trailing eos ramp-down - if num_remaining_delays is not None: - this_codes[: num_codebooks - num_remaining_delays] = eos_stream - num_remaining_delays -= 1 - else: - eos_positions = (this_codes == eos_stream).nonzero(as_tuple=False).reshape(-1) - if eos_positions.numel() > 0: - last_eos_idx = int(eos_positions[-1].item()) - this_codes[: last_eos_idx + 1] = eos_stream - num_remaining_delays = num_codebooks - last_eos_idx - 1 - - if num_remaining_delays is not None and num_remaining_delays <= 0: - # Ramp-down complete — terminate without emitting this frame. - state["num_delay"] = 0 - state["num_remaining_delays"] = None - state["should_terminate"] = True - new_codes_flat.append(torch.full_like(this_codes, -1)) - continue - - state["num_delay"] = num_delay - state["num_remaining_delays"] = num_remaining_delays - state["last_codes"] = this_codes.clone() - # Update CUDA-graph-safe feedback buffers - self._decode_last_codes[int(batch_i)] = this_codes.to(self._decode_last_codes.device) - self._decode_has_codes[int(batch_i)] = True - if state["audio_out_ids"] is None: - state["audio_out_ids"] = this_codes.unsqueeze(-1).clone() - else: - state["audio_out_ids"] = torch.cat([state["audio_out_ids"], this_codes.unsqueeze(-1)], dim=-1) - - new_codes_flat.append(this_codes) - - # Build full codes tensor [batch_size, num_codebooks] - # -1 marks "no audio code at this position" - codes_full = torch.full( - (int(sampled_flat.numel()), num_codebooks), - -1, - dtype=torch.long, - device=hidden.device, + self._update_delay_state_batched( + codes_flat, + audio_row_tensor, + num_rows, + hidden.device, + code_row_mask=code_row_mask, + all_rows=all_audio_rows, ) - if new_codes_flat: - stacked = torch.stack(new_codes_flat, dim=0).to(hidden.device) - codes_full[is_audio] = stacked - - self._last_audio_codes = codes_full - self._postprocess_cursor = 0 return sampler_output + def _get_audio_codes_buffer(self, num_rows: int, device: torch.device) -> torch.Tensor: + buf = self._last_audio_codes_buffer + if buf is None or buf.device != device or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + buf = torch.empty((rows, self.num_codebooks), dtype=torch.long, device=device) + self._last_audio_codes_buffer = buf + return buf[:num_rows] + + def _get_audio_gpu_staging_buffer(self, num_rows: int, device: torch.device) -> torch.Tensor: + width = self.num_codebooks + 2 + buf = self._last_audio_gpu_staging + if buf is None or buf.device != device or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + buf = torch.empty((rows, width), dtype=torch.int32, device=device) + self._last_audio_gpu_staging = buf + return buf[:num_rows] + + def _get_audio_host_staging_buffer(self, num_rows: int) -> torch.Tensor: + width = self.num_codebooks + 2 + buf = self._last_audio_host_staging + if buf is None or buf.device.type != "cpu" or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + pin_memory = torch.cuda.is_available() + buf = torch.empty((rows, width), dtype=torch.int32, device="cpu", pin_memory=pin_memory) + self._last_audio_host_staging = buf + return buf[:num_rows] + + def _next_audio_staging_event(self) -> torch.cuda.Event: + if not self._audio_staging_events: + self._audio_staging_events = [torch.cuda.Event(), torch.cuda.Event()] + event = self._audio_staging_events[self._audio_staging_event_cursor] + self._audio_staging_event_cursor = (self._audio_staging_event_cursor + 1) % len(self._audio_staging_events) + return event + + @staticmethod + def _device_cache_key(device: torch.device) -> tuple[str, int]: + return (device.type, -1 if device.index is None else int(device.index)) + + def _get_row_indices(self, num_rows: int, device: torch.device) -> torch.Tensor: + key = self._device_cache_key(device) + buf = self._row_index_cache.get(key) + if buf is None or buf.device != device or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + buf = torch.arange(rows, dtype=torch.long, device=device) + self._row_index_cache[key] = buf + return buf[:num_rows] + + def _get_codebook_indices(self, device: torch.device) -> torch.Tensor: + key = self._device_cache_key(device) + buf = self._codebook_index_cache.get(key) + if buf is None or buf.device != device or int(buf.shape[1]) != self.num_codebooks: + buf = torch.arange(self.num_codebooks, dtype=torch.long, device=device).view(1, self.num_codebooks) + self._codebook_index_cache[key] = buf + return buf + + def _get_boc_frames(self, num_rows: int, device: torch.device) -> torch.Tensor: + key = self._device_cache_key(device) + buf = self._boc_frame_cache.get(key) + if buf is None or buf.device != device or int(buf.shape[0]) < num_rows: + rows = max(num_rows, 16) + buf = torch.full((rows, self.num_codebooks), BOC_ID, dtype=torch.long, device=device) + self._boc_frame_cache[key] = buf + return buf[:num_rows] + # ------------------------------------------------------------------ postprocess def postprocess( self, @@ -621,73 +1073,269 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): by a running cursor (one row per request per step). """ _ = multimodal_outputs + return self._postprocess_impl() + + def _postprocess_impl(self) -> dict[str, Any]: + host_staging = self._last_audio_host_staging codes_full = self._last_audio_codes - if codes_full is None: + if host_staging is None and codes_full is None: return {} + event = self._last_audio_staging_event + if event is not None: + event.synchronize() + self._last_audio_staging_event = None cursor = int(self._postprocess_cursor) - if cursor >= int(codes_full.shape[0]): + num_rows = int(host_staging.shape[0]) if host_staging is not None else int(codes_full.shape[0]) + if cursor >= num_rows: self._postprocess_cursor = 0 return {} - slice_codes = codes_full[cursor : cursor + 1] self._postprocess_cursor = cursor + 1 - # Drop placeholder rows (-1) - audio_rows = slice_codes[:, 0] >= 0 - if not bool(audio_rows.any()): + if host_staging is not None: + if self._last_audio_valid_flags is None or self._last_audio_done_flags is None: + self._last_audio_valid_flags = host_staging[:num_rows, self.num_codebooks].tolist() + self._last_audio_done_flags = host_staging[:num_rows, self.num_codebooks + 1].tolist() + valid_flag = int(self._last_audio_valid_flags[cursor]) + done_flag = int(self._last_audio_done_flags[cursor]) + if valid_flag or done_flag: + self._postprocess_audio_active_rows += 1 + if cursor + 1 >= num_rows: + if self._postprocess_audio_rows == num_rows and self._postprocess_audio_active_rows == num_rows: + self._fast_audio_direct_rows = num_rows + self._postprocess_audio_rows = 0 + self._postprocess_audio_active_rows = 0 + if valid_flag == 0: + return {} + new_codes = host_staging[cursor : cursor + 1, : self.num_codebooks] + return {"codes": {"audio": new_codes}} + + if cursor >= len(self._last_audio_code_valid) or not self._last_audio_code_valid[cursor]: return {} - new_codes = slice_codes[audio_rows].to(torch.int32) + slice_codes = codes_full[cursor : cursor + 1] + new_codes = slice_codes.to(torch.int32) return {"codes": {"audio": new_codes}} # ------------------------------------------------------------------ helpers - def _audio_codebook_logits(self, hidden_states: torch.Tensor, audio_mask: torch.Tensor) -> torch.Tensor: - mask = audio_mask.reshape(-1).to(hidden_states.device) + def _audio_codebook_logits_from_rows( + self, hidden_states: torch.Tensor, audio_rows: torch.Tensor, *, all_rows: bool = False + ) -> torch.Tensor: hidden_flat = hidden_states.reshape(-1, hidden_states.shape[-1]) - if not mask.any(): + if audio_rows.numel() == 0: return torch.empty( (0, self.num_codebooks, self.codebook_size), dtype=hidden_states.dtype, device=hidden_states.device, ) - return self.modality_head.generate(hidden_flat[mask]) + if all_rows: + return self.modality_head.generate(hidden_flat[: int(audio_rows.numel())]) + return self.modality_head.generate(hidden_flat.index_select(0, audio_rows.to(hidden_flat.device))) - def _apply_delay_pattern_masking(self, cb_logits: torch.Tensor, audio_row_indices: list[int]) -> None: - """Mask per-codebook logits according to delay pattern state, in-place. + def _apply_delay_pattern_masking_batched( + self, cb_logits: torch.Tensor, audio_rows: torch.Tensor, *, all_rows: bool = False + ) -> None: + """Vectorized delay-pattern masking using GPU-resident sampler state.""" + if cb_logits.numel() == 0: + return - During delay phase: codebooks beyond delay_count only allow BOC. - During ramp-down: locked codebooks only allow EOC. - Normal generation: BOC disallowed; only cb0 allows EOC. - """ - bos_pre = BOC_ID - eos_pre = EOC_ID + rows = audio_rows.to(device=cb_logits.device, dtype=torch.long) + num_audio_rows = int(cb_logits.shape[0]) num_codebooks = self.num_codebooks - for local_i, batch_i in enumerate(audio_row_indices): - state = self._audio_state.get(int(batch_i)) - num_delay = int(state["num_delay"]) if state else 0 - num_rem = state.get("num_remaining_delays") if state else None + q = self._get_codebook_indices(cb_logits.device) + if all_rows: + delay = self._decode_delay_count[:num_audio_rows] + rem = self._decode_eoc_countdown[:num_audio_rows] + else: + delay = self._decode_delay_count.index_select(0, rows) + rem = self._decode_eoc_countdown.index_select(0, rows) + if delay.dtype != torch.long: + delay = delay.to(torch.long) + if rem.dtype != torch.long: + rem = rem.to(torch.long) + ramp = rem >= 0 - if num_rem is not None: - lock_until = num_codebooks - int(num_rem) - for q in range(num_codebooks): - row = cb_logits[local_i, q] - if q < lock_until: - mask = torch.full_like(row, float("-inf")) - mask[eos_pre] = row[eos_pre] - cb_logits[local_i, q] = mask - else: - cb_logits[local_i, q, bos_pre] = float("-inf") - cb_logits[local_i, q, eos_pre] = float("-inf") - else: - for q in range(num_codebooks): - row = cb_logits[local_i, q] - if q > num_delay: - mask = torch.full_like(row, float("-inf")) - mask[bos_pre] = row[bos_pre] - cb_logits[local_i, q] = mask - else: - cb_logits[local_i, q, bos_pre] = float("-inf") - if q != 0: - cb_logits[local_i, q, eos_pre] = float("-inf") + bos = BOC_ID + eos = EOC_ID + + # Ramp-down: codebooks already behind the EOC wave are locked to EOC. + lock_until = num_codebooks - rem + locked = ramp.unsqueeze(1) & (q < lock_until.unsqueeze(1)) + eoc_logits = cb_logits[:, :, eos].clone() + cb_logits.masked_fill_(locked.unsqueeze(-1), float("-inf")) + cb_logits[:, :, eos] = torch.where(locked, eoc_logits, cb_logits[:, :, eos]) + + # Ramp-down tail: active codebooks cannot emit stream specials. + ramp_tail = ramp.unsqueeze(1) & (q >= lock_until.unsqueeze(1)) + cb_logits[:, :, bos] = torch.where( + ramp_tail, + torch.full_like(cb_logits[:, :, bos], float("-inf")), + cb_logits[:, :, bos], + ) + cb_logits[:, :, eos] = torch.where( + ramp_tail, + torch.full_like(cb_logits[:, :, eos], float("-inf")), + cb_logits[:, :, eos], + ) + + # Normal generation: delayed codebooks are forced to BOC. + normal = ~ramp + allowed_until = torch.clamp(delay + 1, max=num_codebooks) + delayed = normal.unsqueeze(1) & (q >= allowed_until.unsqueeze(1)) + boc_logits = cb_logits[:, :, bos].clone() + cb_logits.masked_fill_(delayed.unsqueeze(-1), float("-inf")) + cb_logits[:, :, bos] = torch.where(delayed, boc_logits, cb_logits[:, :, bos]) + + # Normal active codebooks: BOC is disallowed; EOC is only allowed on cb0. + allowed = normal.unsqueeze(1) & (q < allowed_until.unsqueeze(1)) + cb_logits[:, :, bos] = torch.where( + allowed, + torch.full_like(cb_logits[:, :, bos], float("-inf")), + cb_logits[:, :, bos], + ) + nonzero_allowed = allowed & (q > 0) + cb_logits[:, :, eos] = torch.where( + nonzero_allowed, + torch.full_like(cb_logits[:, :, eos], float("-inf")), + cb_logits[:, :, eos], + ) + + def _update_delay_state_batched( + self, + codes_flat: torch.Tensor, + audio_rows: torch.Tensor, + num_rows: int, + device: torch.device, + *, + code_row_mask: torch.Tensor | None = None, + all_rows: bool = False, + ) -> None: + """Update delay/ramp-down state in batch and stage per-request outputs.""" + self._ensure_decode_state_capacity(num_rows, device) + if codes_flat.numel() == 0: + self._clear_last_audio_outputs() + return + + rows = audio_rows.to(device=device, dtype=torch.long) + num_audio_rows = int(codes_flat.shape[0]) + num_codebooks = self.num_codebooks + q = self._get_codebook_indices(device) + + if all_rows: + prev_delay = self._decode_delay_count[:num_audio_rows] + prev_rem = self._decode_eoc_countdown[:num_audio_rows] + prev_done = self._decode_generation_done[:num_audio_rows].to(torch.bool) + prev_has_codes = self._decode_has_codes[:num_audio_rows].to(torch.bool) + prev_codes = self._decode_last_codes[:num_audio_rows] + else: + prev_delay = self._decode_delay_count.index_select(0, rows) + prev_rem = self._decode_eoc_countdown.index_select(0, rows) + prev_done = self._decode_generation_done.index_select(0, rows).to(torch.bool) + prev_has_codes = self._decode_has_codes.index_select(0, rows).to(torch.bool) + prev_codes = self._decode_last_codes.index_select(0, rows) + if prev_delay.dtype != torch.long: + prev_delay = prev_delay.to(torch.long) + if prev_rem.dtype != torch.long: + prev_rem = prev_rem.to(torch.long) + if code_row_mask is None: + update_mask = torch.ones_like(prev_has_codes) + elif all_rows: + update_mask = code_row_mask.to(device=device, dtype=torch.bool)[:num_audio_rows] + else: + update_mask = code_row_mask.to(device=device, dtype=torch.bool).index_select(0, rows) + ramp = prev_rem >= 0 + + codes = codes_flat.to(device=device, dtype=torch.long).clone() + + # Leading BOC delay pad. + delay_active = (~ramp) & ((prev_delay + 1) < num_codebooks) + delay_mask = delay_active.unsqueeze(1) & (q > prev_delay.unsqueeze(1)) + codes = torch.where(delay_mask, torch.full_like(codes, BOC_ID), codes) + next_delay = torch.where(delay_active, prev_delay + 1, prev_delay) + + # Trailing EOC ramp-down. + lock_until = num_codebooks - prev_rem + ramp_mask = ramp.unsqueeze(1) & (q < lock_until.unsqueeze(1)) + codes = torch.where(ramp_mask, torch.full_like(codes, EOC_ID), codes) + ramp_next_rem = prev_rem - 1 + + # New EOC detection in normal mode. This avoids tensor->Python control + # flow; ties are impossible for the index because q is monotonic. + eos_mask = (~ramp).unsqueeze(1) & (codes == EOC_ID) + eos_idx_values = torch.where(eos_mask, q.expand_as(codes), torch.full_like(codes, -1)) + last_eos_idx = eos_idx_values.max(dim=1).values + has_eos = last_eos_idx >= 0 + eos_prefix = has_eos.unsqueeze(1) & (q <= last_eos_idx.unsqueeze(1)) + codes = torch.where(eos_prefix, torch.full_like(codes, EOC_ID), codes) + normal_next_rem = torch.where(has_eos, num_codebooks - last_eos_idx - 1, torch.full_like(last_eos_idx, -1)) + + next_rem = torch.where(ramp, ramp_next_rem, normal_next_rem) + done = ((next_rem >= 0) & (next_rem <= 0)) & update_mask + valid = update_mask + + next_delay = torch.where(done, torch.zeros_like(next_delay), next_delay) + next_rem = torch.where(done, torch.full_like(next_rem, -1), next_rem) + output_codes = codes + + write_delay = torch.where(update_mask, next_delay, prev_delay) + write_rem = torch.where(update_mask, next_rem, prev_rem) + write_done = torch.where(update_mask, done, prev_done) + write_codes = torch.where(update_mask.unsqueeze(1), codes, prev_codes) + write_has_codes = torch.where(update_mask, ~done, prev_has_codes) + + if all_rows: + self._decode_delay_count[:num_audio_rows].copy_(write_delay.to(dtype=self._decode_delay_count.dtype)) + self._decode_eoc_countdown[:num_audio_rows].copy_(write_rem.to(dtype=self._decode_eoc_countdown.dtype)) + self._decode_generation_done[:num_audio_rows].copy_(write_done) + self._decode_last_codes[:num_audio_rows].copy_(write_codes) + self._decode_has_codes[:num_audio_rows].copy_(write_has_codes) + else: + self._decode_delay_count.index_copy_(0, rows, write_delay.to(dtype=self._decode_delay_count.dtype)) + self._decode_eoc_countdown.index_copy_(0, rows, write_rem.to(dtype=self._decode_eoc_countdown.dtype)) + self._decode_generation_done.index_copy_(0, rows, write_done) + self._decode_last_codes.index_copy_(0, rows, write_codes) + self._decode_has_codes.index_copy_(0, rows, write_has_codes) + + staged_codes = torch.where(update_mask.unsqueeze(1), output_codes, torch.full_like(output_codes, -1)) + if all_rows: + valid_full = valid + done_full = done + staging = self._get_audio_gpu_staging_buffer(num_rows, device) + staging[:, :num_codebooks].copy_(staged_codes.to(torch.int32)) + staging[:, num_codebooks].copy_(valid_full.to(torch.int32)) + staging[:, num_codebooks + 1].copy_(done_full.to(torch.int32)) + else: + codes_full = self._get_audio_codes_buffer(num_rows, device) + codes_full.fill_(-1) + codes_full.index_copy_(0, rows, staged_codes) + valid_full = torch.zeros(num_rows, dtype=torch.bool, device=device) + done_full = torch.zeros(num_rows, dtype=torch.bool, device=device) + valid_full.index_copy_(0, rows, valid) + done_full.index_copy_(0, rows, done) + staging = self._get_audio_gpu_staging_buffer(num_rows, device) + staging[:, :num_codebooks].copy_(codes_full.to(torch.int32)) + staging[:, num_codebooks].copy_(valid_full.to(torch.int32)) + staging[:, num_codebooks + 1].copy_(done_full.to(torch.int32)) + + host = self._get_audio_host_staging_buffer(num_rows) + host.copy_(staging, non_blocking=True) + if device.type == "cuda": + event = self._next_audio_staging_event() + event.record(torch.cuda.current_stream(device)) + self._last_audio_staging_event = event + else: + self._last_audio_staging_event = None + + self._decode_active_audio_count = num_rows if all_rows else int(rows.numel()) + # 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 + self._last_audio_host_staging = host[:num_rows] + self._last_audio_valid_flags = None + self._last_audio_done_flags = None + self._last_audio_code_valid = [] + self._postprocess_cursor = 0 + self._postprocess_audio_rows = num_rows + self._postprocess_audio_active_rows = 0 def _sample_audio_codes(self, logits_2d: torch.Tensor) -> torch.Tensor: """Replicate upstream sampling: temperature → top-k → top-p → multinomial.""" @@ -709,106 +1357,12 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): # Detect all-masked rows BEFORE softmax (softmax of all-inf yields NaN). has_finite = torch.isfinite(x).any(dim=-1) all_masked = ~has_finite - if all_masked.any(): - # For all-masked rows, fall back to argmax (picks least-negative). - fallback = x.argmax(dim=-1) - if has_finite.any(): - probs = x[has_finite].softmax(dim=-1) - sampled = torch.multinomial(probs, num_samples=1).squeeze(-1) - result = fallback.clone() - result[has_finite] = sampled - return result - return fallback - probs = x.softmax(dim=-1) - return torch.multinomial(probs, num_samples=1).squeeze(-1) - - def _apply_audio_mode_bias(self, logits: torch.Tensor, sampling_metadata: Any) -> None: - """Detect <|audio|> transition, force continuation, force eos at ramp-down. - - Mirrors v2's _apply_audio_mode_bias: walks per-request to find the - previous token, and if it was <|audio|> or the continuation token, - forces the next emit to <|audio|>. On ramp-down completion, forces eos. - """ - if logits is None or logits.ndim != 2: - return - - audio_id = self._audio_continuation_id - eos_id = self._eos_token_id - if audio_id is None: - return - - num_rows = int(logits.shape[0]) - prompt_ids = getattr(sampling_metadata, "prompt_token_ids", None) - output_ids = getattr(sampling_metadata, "output_token_ids", None) - - # Fallback prev-token source from stashed input_ids - stash_ids = self._last_step_input_ids - stash_tail: list[int] | None = None - if isinstance(stash_ids, torch.Tensor) and stash_ids.numel() > 0: - q_start = self._last_step_query_start_loc - if isinstance(q_start, torch.Tensor) and int(q_start.numel()) == num_rows + 1: - q_start_cpu = q_start.detach().to("cpu").tolist() - tail_idx = [max(0, int(q_start_cpu[i + 1]) - 1) for i in range(num_rows)] - flat_ids = stash_ids.detach().to("cpu").tolist() - stash_tail = [int(flat_ids[idx]) if idx < len(flat_ids) else -1 for idx in tail_idx] - elif int(stash_ids.numel()) >= num_rows: - stash_tail = stash_ids[-num_rows:].detach().to("cpu").tolist() - - first_after_start = torch.zeros(num_rows, dtype=torch.bool, device=logits.device) - - for i in range(num_rows): - prev: int | None = None - if output_ids is not None and i < len(output_ids): - hist = output_ids[i] - if hist: - prev = int(hist[-1]) - if prev is None and prompt_ids is not None: - try: - p_i = prompt_ids[i] - if hasattr(p_i, "tolist"): - p_i = p_i.tolist() - if p_i: - prev = int(p_i[-1]) - except (IndexError, TypeError): - prev = None - if prev is None and stash_tail is not None and i < len(stash_tail): - prev = int(stash_tail[i]) - if prev is None: - continue - - # Only bias if previous token was <|audio|> (the continuation token) - if prev != audio_id: - continue - - # Check if this is the FIRST step after <|audio|> appears - # (i.e., transitioning from prompt to audio generation) - audio_state = self._audio_state.get(i) - if audio_state is None or audio_state.get("should_terminate"): - # No state yet, or stale state from a finished prior request - # reusing this slot — treat as first audio step - self._audio_state.pop(i, None) - # Clear CUDA-graph feedback buffers for this slot - self._decode_has_codes[i] = False - first_after_start[i] = True - - # Check for ramp-down termination - should_terminate = bool(isinstance(audio_state, dict) and audio_state.get("should_terminate")) - if should_terminate and eos_id is not None and 0 <= eos_id < int(logits.shape[-1]): - row = logits[i] - mask = torch.full_like(row, float("-inf")) - mask[eos_id] = row[eos_id] - logits[i].copy_(mask) - audio_state["should_terminate"] = False - continue - - # Force audio continuation token - row = logits[i] - mask = torch.full_like(row, float("-inf")) - if 0 <= audio_id < row.shape[-1]: - mask[audio_id] = row[audio_id] - logits[i].copy_(mask) - - self._last_first_audio_after_start = first_after_start + # Do not branch on tensor bools here; that synchronizes every audio step. + fallback = x.argmax(dim=-1) + safe_x = torch.where(all_masked.unsqueeze(-1), torch.zeros_like(x), x) + probs = safe_x.softmax(dim=-1) + sampled = torch.multinomial(probs, num_samples=1).squeeze(-1) + return torch.where(all_masked, fallback, sampled) # ------------------------------------------------------------------ omni output def make_omni_output(self, model_outputs: torch.Tensor | OmniOutput, **kwargs: Any) -> OmniOutput: 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 054b0d66a..6d71649e3 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 @@ -32,6 +32,7 @@ from vllm_omni.data_entry_keys import ( OmniPayload, OmniPayloadStruct, ) +from vllm_omni.inputs.data import OmniTokensPrompt __all__ = ["talker2code2wav", "talker2code2wav_async_chunk"] @@ -48,6 +49,15 @@ _DEFAULT_CODEC_LEFT_CONTEXT_FRAMES = 25 _DEFAULT_CODEC_RIGHT_HOLDBACK_FRAMES = 4 +def _empty_code2wav_prompt() -> Any: + return OmniTokensPrompt( + prompt_token_ids=[], + multi_modal_data=None, + mm_processor_kwargs=None, + additional_information=None, + ) + + def _revert_delay_pattern(audio_codes_qt: torch.Tensor) -> torch.Tensor: """Reverse the MusicGen-style delay pattern. @@ -91,9 +101,8 @@ def talker2code2wav( _requires_multimodal_data: bool = False, ) -> list[Any]: """Sync: collect all talker codes, revert delay pattern, filter, pass to code2wav.""" - from vllm_omni.inputs.data import OmniTokensPrompt - code2wav_inputs: list[OmniTokensPrompt] = [] + code2wav_inputs: list[Any] = [] for talker_output in source_outputs: if not talker_output.finished: continue @@ -103,14 +112,7 @@ def talker2code2wav( audio_codes = mm_codes.get("audio") if audio_codes is None or not isinstance(audio_codes, torch.Tensor) or audio_codes.numel() == 0: - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[], - multi_modal_data=None, - mm_processor_kwargs=None, - additional_information=None, - ) - ) + code2wav_inputs.append(_empty_code2wav_prompt()) continue audio_codes = audio_codes.to(torch.long) @@ -137,22 +139,9 @@ def talker2code2wav( # Step 1: Revert delay pattern try: codes_qt = _revert_delay_pattern(codes_qt) - except ValueError: - logger.warning( - "Skipping request: insufficient frames for delay pattern " - "reversal (T=%d < Q=%d). This is a request-level error, " - "not a server-fatal condition.", - codes_qt.shape[1], - codes_qt.shape[0], - ) - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[], - multi_modal_data=None, - mm_processor_kwargs=None, - additional_information=None, - ) - ) + 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. @@ -172,14 +161,7 @@ def talker2code2wav( codes_qt = codes_qt[:, :-1] if codes_qt.numel() == 0: - code2wav_inputs.append( - OmniTokensPrompt( - prompt_token_ids=[], - multi_modal_data=None, - mm_processor_kwargs=None, - additional_information=None, - ) - ) + code2wav_inputs.append(_empty_code2wav_prompt()) continue # Code2Wav expects codebook-major flat: [Q * num_frames] diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index c59b3b4c7..b68cf3800 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -163,6 +163,8 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): return output_token_ids def _sampling_metadata_for_model_sampler(self, sampling_metadata): + if getattr(self.model, "skips_model_sampler_output_token_history", False): + return sampling_metadata output_token_ids = self._build_model_sampler_output_token_ids() if output_token_ids == sampling_metadata.output_token_ids: return sampling_metadata @@ -568,7 +570,6 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): use_cascade_attn=cascade_attn_prefix_lens is not None, num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), ) - num_tokens_padded = batch_desc.num_tokens num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index abef62f31..120c7a565 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -57,6 +57,16 @@ class OmniGPUModelRunner(GPUModelRunner): # The Omni tensor prefix cache will be allocated # when we initialize the metadata builders if enabled self.omni_prefix_cache = None + self._sampled_token_ids_cpu_override = None + self._omni_query_start_loc_model_kwarg = False + + def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]: + override_fn = self._sampled_token_ids_cpu_override + if callable(override_fn): + sampled = override_fn(sampled_token_ids) + if sampled is not None: + return sampled + return super()._to_list(sampled_token_ids) def _omni_routed_experts_d2h(self, scheduler_output) -> None: """Issue routed-experts D2H copy matching upstream GPUModelRunner pattern. @@ -105,14 +115,13 @@ class OmniGPUModelRunner(GPUModelRunner): ) def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes): - """Override to fix scheduler_metadata buffer size for FA3 + CUDA graph. + """Initialize metadata builders and keep FA3 graph metadata buffers sized. - The upstream FlashAttentionMetadataBuilder pre-allocates - scheduler_metadata with (max_num_seqs + 1) entries, but FA3's - get_scheduler_metadata() can return up to - (max_num_seqs * max_num_splits + 1) entries, causing a RuntimeError - during CUDA graph capture. After calling the parent implementation - we resize any too-small buffers. + FlashAttentionMetadataBuilder can pre-allocate scheduler_metadata for + only max_num_seqs + 1 entries while FA3 with split scheduling may need + max_num_seqs * max_num_splits + 1 entries during CUDA graph capture. + This runner is shared across Omni models, so preserve the existing + workaround for non-Higgs models that still use FA3. """ super().initialize_metadata_builders(kv_cache_config, kernel_block_sizes) @@ -143,6 +152,14 @@ class OmniGPUModelRunner(GPUModelRunner): @instrument(span_name="Loading (GPU)") def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) + model = getattr(self, "model", None) + override_fn = None + if bool(getattr(model, "supports_sampled_token_ids_cpu_override", False)): + candidate = getattr(model, "consume_sampled_token_ids_cpu_override", None) + if callable(candidate): + override_fn = candidate + self._sampled_token_ids_cpu_override = override_fn + self._omni_query_start_loc_model_kwarg = bool(getattr(model, "supports_omni_query_start_loc", False)) self._maybe_enable_output_token_ids_for_model_sampler() self._init_talker_mtp() self._prewarm_attention_capture_workspaces() @@ -1295,6 +1312,13 @@ class OmniGPUModelRunner(GPUModelRunner): traceback.print_exc() + if self._omni_query_start_loc_model_kwarg: + try: + num_reqs = len(self.input_batch.req_ids) + model_kwargs_extra["omni_query_start_loc"] = self.query_start_loc.gpu[: num_reqs + 1] + except Exception as e: + logger.debug("[OMNI] Failed to attach query_start_loc: %s", e) + if getattr(self.model_config, "has_sampling_extra_args", False): extra_args_list: list[dict] = [] for req_id in self.input_batch.req_ids: @@ -1320,21 +1344,29 @@ class OmniGPUModelRunner(GPUModelRunner): # execute the custom postprocess function # TODO(Peiqi): do we have a more elegant way to do this? if hasattr(self.model, "has_postprocess") and self.model.has_postprocess: + postprocess_uses_hidden_states = getattr(self.model, "postprocess_uses_hidden_states", True) + postprocess_uses_multimodal_outputs = getattr(self.model, "postprocess_uses_multimodal_outputs", True) + postprocess_uses_req_infos = getattr(self.model, "postprocess_uses_req_infos", True) for req_index, req_id in enumerate(self.input_batch.req_ids): if req_ids_filter is not None and req_id not in req_ids_filter: continue - req_infos = self.model_intermediate_buffer.get(req_id, {}) - if combined_hidden_states: - # Combined hidden states contains all hidden states for every request - hidden_states_slice = combined_hidden_states[req_id] + req_infos = self.model_intermediate_buffer.get(req_id, {}) if postprocess_uses_req_infos else {} + if postprocess_uses_hidden_states: + if combined_hidden_states: + # Combined hidden states contains all hidden states for every request + hidden_states_slice = combined_hidden_states[req_id] + else: + start_offset = int(self.query_start_loc.cpu[req_index]) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + s, e = start_offset, start_offset + sched_tokens + # only consider to store data into update dict. + hidden_states_slice = hidden_states[s:e] else: - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - s, e = start_offset, start_offset + sched_tokens - # only consider to store data into update dict. - hidden_states_slice = hidden_states[s:e] + hidden_states_slice = hidden_states - if combined_multimodal_outputs: + if not postprocess_uses_multimodal_outputs: + mm_out = None + elif combined_multimodal_outputs: # NOTE this is a bit ugly, but the mm data is structured as a list of # keys mapping to request IDs, and if enabled, we will always have all # request IDs in every subdict, including for cache misses. @@ -1807,6 +1839,14 @@ class OmniGPUModelRunner(GPUModelRunner): ): """Inject omni-specific kwargs into forward and cache model output""" 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): + update_decode_metadata( + input_ids=input_ids, + positions=positions, + inputs_embeds=inputs_embeds, + omni_query_start_loc=model_kwargs_extra.get("omni_query_start_loc"), + ) model_output = super()._model_forward( input_ids=input_ids,