[Fix][HiggsAudioV3] Fix ramp-down off-by-one crash and buffer state bugs in Stage0 talker (#4219)

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

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

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]

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.