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 d73f0c6a5..627e22ffb 100644 --- a/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py +++ b/tests/unit/higgs_audio_v3/test_higgs_audio_v3.py @@ -410,6 +410,8 @@ class TestAudioFeedback: 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) return t @@ -427,6 +429,8 @@ class TestAudioFeedback: 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 input_ids = torch.tensor([1, audio_id, 3]) hidden = torch.zeros(3, 16) result = t._apply_audio_feedback(hidden, input_ids) @@ -440,7 +444,7 @@ 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 + # No state for row 0 — _decode_has_codes[0] defaults to False input_ids = torch.tensor([audio_id]) hidden = torch.zeros(1, 16) result = t._apply_audio_feedback(hidden, input_ids) @@ -908,3 +912,300 @@ class TestPromptBuilder: # Should not contain ref_audio or ref_text token IDs assert 151703 not in ids # <|ref_audio|> assert 151704 not in ids # <|ref_text|> + + +# ---- Regression: ramp-down done-before-emit off-by-one ---- + + +def _simulate_delay_state_machine( + sampled_codes_per_step: list[torch.Tensor], +) -> list[torch.Tensor]: + """Reproduce the delay-pattern state machine from sample(). + + 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, + ) + + num_codebooks = 8 + bos = BOC_ID + eos_stream = EOC_ID + + 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, + } + + emitted: list[torch.Tensor] = [] + for step_codes in sampled_codes_per_step: + if state["should_terminate"]: + break + + num_delay = state["num_delay"] + num_remaining_delays = state["num_remaining_delays"] + this_codes = step_codes.clone() + + if num_delay + 1 < num_codebooks: + this_codes[num_delay + 1 :] = bos + num_delay += 1 + + 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: + state["num_delay"] = 0 + state["num_remaining_delays"] = None + state["should_terminate"] = True + continue + + state["num_delay"] = num_delay + state["num_remaining_delays"] = num_remaining_delays + state["last_codes"] = this_codes.clone() + emitted.append(this_codes) + + return emitted + + +class TestRampDownOffByOne: + """Regression tests for the done-before-emit off-by-one bug. + + 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, + ) + + 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 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 4fc89341c..5e20245b0 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 @@ -250,7 +250,8 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): if input_ids is not None: self._last_step_input_ids = input_ids - # Stash query_start_loc for _apply_audio_mode_bias + # Stash query_start_loc and max_query_len for prefill detection + _max_query_len = None try: from vllm.forward_context import get_forward_context @@ -264,13 +265,19 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): self._last_step_query_start_loc = qsl.detach().clone() else: 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 - # Prefill-only operations: ref audio substitution and audio feedback - # require Python dict/list ops that break CUDA graph capture. - # Detect prefill (sequence length > batch size heuristic) vs decode. - is_prefill = input_ids is not None and inputs_embeds is None and int(input_ids.numel()) > 1 + # 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 + else: + is_prefill = input_ids is not None and inputs_embeds is None and int(input_ids.numel()) > 1 if is_prefill: # Voice clone: replace -100 placeholder positions with ref audio embeddings info_dicts = kwargs.get("model_intermediate_buffer") @@ -278,7 +285,8 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): info_dicts = kwargs.get("runtime_additional_information") hidden_states = self._apply_ref_audio_substitution(hidden_states, input_ids, info_dicts) - # Audio feedback at decode: replace continuation token embeddings + # Audio feedback: replace continuation token embeddings with audio + # embeddings from the last decoded frame (CUDA-graph safe). if input_ids is not None and inputs_embeds is None: hidden_states = self._apply_audio_feedback(hidden_states, input_ids) @@ -411,11 +419,14 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): 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]: - new_size = max(bs, self._decode_last_codes.shape[0] * 2) - self._decode_last_codes = torch.zeros( - new_size, self.num_codebooks, dtype=torch.long, device=hidden_states.device - ) - self._decode_has_codes = torch.zeros(new_size, dtype=torch.bool, device=hidden_states.device) + 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 # Compute audio embeddings from last_codes for ALL rows (graph-safe) codes_slice = self._decode_last_codes[:bs] # [bs, N] @@ -561,7 +572,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): 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 + # Ramp-down complete — terminate without emitting this frame. state["num_delay"] = 0 state["num_remaining_delays"] = None state["should_terminate"] = True @@ -578,6 +589,7 @@ class HiggsAudioV3TalkerForConditionalGeneration(nn.Module): 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] 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 95fbe74a5..054b0d66a 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 @@ -135,7 +135,25 @@ def talker2code2wav( codes_qt = audio_codes.transpose(0, 1).contiguous().cpu() # Step 1: Revert delay pattern - codes_qt = _revert_delay_pattern(codes_qt) + 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, + ) + ) + continue # Step 2: Replace out-of-range codes (BOC=1024, EOC=1025, -1) with 0. # Must use torch.where, NOT clamp: clamp(max=1023) turns 1025→1023 @@ -336,7 +354,21 @@ def talker2code2wav_async_chunk( return None codes_qt = torch.tensor(window_rows, dtype=torch.long).t().contiguous() - de_delayed = _revert_delay_pattern(codes_qt) + try: + de_delayed = _revert_delay_pattern(codes_qt) + except ValueError: + logger.warning( + "async_chunk: insufficient frames for delay pattern reversal " + "(T=%d < Q=%d, request=%s). Returning empty chunk.", + codes_qt.shape[1], + codes_qt.shape[0], + request_id, + ) + emitted_frames.pop(request_id, None) + return OmniPayloadStruct( + codes=CodesStruct(audio=torch.empty(0, dtype=torch.long)), + meta=MetaStruct(finished=torch.tensor(True, dtype=torch.bool)), + ) # Replace BOC=1024/EOC=1025 (and any negative pads) with 0; matches the # sync-path substitution. Clamp would turn 1025 into 1023 which is a # VALID codec code and decodes to audible artifacts.