[Model] Add MiniMax Music 3 text-to-music (#6186)

Signed-off-by: Yueqian Lin <linyueqian@outlook.com>

tests/model_executor/models/test_minimax_music3_batch_payload.py

@@ -0,0 +1,132 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniMax Music 3 batched payload partitioning.
+
+The stage-0 payload is split per request by indexing list values: a list is
+indexed by batch row, anything else is handed to every request unchanged. So
+per-request metadata has to be carried as lists. Merging it into one dict made
+every request in a batch receive the last request's seed, window offset and
+terminal flag, which crosses seeds between songs, misplaces acoustic windows,
+and lets one request's terminal flag end another request's stream early.
+
+These drive ``make_omni_output`` directly with a hand-built batch, so no GPU
+and no engine are needed.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minimax_music3.talker import (
+    MiniMaxMusic3TalkerForConditionalGeneration as Talker,
+)
+from vllm_omni.utils.mm_outputs import to_payload_element
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+_WIDTH = 4
+
+
+def _info(*, seed: int, start: int, frames: int, last: bool) -> dict:
+    """One batch row's intermediate buffer, as ``postprocess`` leaves it."""
+    return {
+        "latent": torch.full((frames, _WIDTH), float(seed)),
+        "codes": {"audio": torch.full((frames, 8), seed)},
+        "meta": {
+            "audio_seed": seed,
+            "num_processed_tokens": start,
+            "codec_chunk_frames": frames,
+            "last_chunk": last,
+        },
+    }
+
+
+def _payload_for(output, row: int) -> dict:
+    """Partition the batch payload down to one request, as the runner does."""
+    flat = {}
+    for key, value in output.multimodal_outputs.items():
+        if isinstance(value, dict):
+            for qual, val in value.items():
+                flat[f"{key}.{qual}"] = val
+        else:
+            flat[key] = value
+    return {key: to_payload_element(value, row, 1, 1) for key, value in flat.items()}
+
+
+class TestBatchedMetadataPartitioning:
+    """Two requests in one batch must keep their own metadata."""
+
+    @staticmethod
+    def _two_request_batch():
+        infos = [
+            _info(seed=111, start=100, frames=3, last=False),
+            _info(seed=222, start=200, frames=5, last=True),
+        ]
+        talker = Talker.__new__(Talker)
+        return talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
+
+    def test_seeds_are_not_crossed(self) -> None:
+        """A song decoded with another request's seed is a different song."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.audio_seed"] == 111
+        assert _payload_for(output, 1)["meta.audio_seed"] == 222
+
+    def test_window_offsets_stay_with_their_request(self) -> None:
+        """The start frame decides where a window lands in the song."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.num_processed_tokens"] == 100
+        assert _payload_for(output, 1)["meta.num_processed_tokens"] == 200
+
+    def test_terminal_flag_does_not_leak_across_requests(self) -> None:
+        """One request finishing must not end another request's stream."""
+        output = self._two_request_batch()
+        assert _payload_for(output, 0)["meta.last_chunk"] is False
+        assert _payload_for(output, 1)["meta.last_chunk"] is True
+
+    def test_span_length_matches_its_own_latent(self) -> None:
+        output = self._two_request_batch()
+        for row, frames in ((0, 3), (1, 5)):
+            entry = _payload_for(output, row)
+            assert entry["meta.codec_chunk_frames"] == frames
+            assert entry["latent"].shape[0] == frames
+
+    def test_latent_and_codes_follow_the_same_row(self) -> None:
+        """The tensors were already per-request; pin that they stay aligned."""
+        output = self._two_request_batch()
+        for row, seed in ((0, 111), (1, 222)):
+            entry = _payload_for(output, row)
+            assert entry["latent"][0, 0].item() == pytest.approx(float(seed))
+            assert int(entry["codes.audio"][0, 0]) == seed
+
+
+class TestPartialBatches:
+    """Only some rows carry a span on any given step."""
+
+    def test_row_without_a_span_does_not_take_another_row_metadata(self) -> None:
+        infos = [{}, _info(seed=777, start=300, frames=2, last=True)]
+        talker = Talker.__new__(Talker)
+        output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=infos)
+
+        idle = _payload_for(output, 0)
+        assert idle["latent"].numel() == 0
+        assert idle["meta.audio_seed"] is None
+        assert idle["meta.last_chunk"] is None
+
+        active = _payload_for(output, 1)
+        assert active["meta.audio_seed"] == 777
+        assert active["meta.last_chunk"] is True
+
+    def test_batch_with_no_spans_emits_no_payload(self) -> None:
+        talker = Talker.__new__(Talker)
+        output = talker.make_omni_output(torch.zeros(2, _WIDTH), model_intermediate_buffer=[{}, {}])
+        assert output.multimodal_outputs is None
+
+    def test_consumed_entries_are_cleared_from_the_buffer(self) -> None:
+        """The runner never clears this buffer, so a span left behind would be
+        re-emitted on every later step."""
+        infos = [_info(seed=1, start=0, frames=2, last=False)]
+        talker = Talker.__new__(Talker)
+        talker.make_omni_output(torch.zeros(1, _WIDTH), model_intermediate_buffer=infos)
+        assert "latent" not in infos[0]
+        assert "meta" not in infos[0]

tests/model_executor/models/test_minimax_music3_frame_state.py

@@ -0,0 +1,214 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniMax Music 3 per-request AR frame bookkeeping.
+
+The talker emits overlapping windows while dropping frames it can no longer
+need, and flushes whatever is left as one span when the song ends. Getting the
+eviction boundary wrong by one frame silently corrupts the conditioning handed
+to the acoustic stage from the second window onward, which is audible but hard
+to attribute, so the lifecycle is simulated here frame by frame.
+"""
+
+from __future__ import annotations
+
+import pytest
+import torch
+
+from vllm_omni.model_executor.models.minimax_music3.chunking import chunk_windows
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    AR_CHUNK_FRAMES,
+    AR_CHUNK_HOP_FRAMES,
+    NUM_CODEBOOKS,
+)
+from vllm_omni.model_executor.models.minimax_music3.frame_state import (
+    FrameBuffer,
+    RequestARState,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# One scalar per frame instead of 32768 dims: the bookkeeping is what is under
+# test, and a frame's value doubles as its identity.
+_WIDTH = 4
+
+
+def _frame(index: int) -> torch.Tensor:
+    return torch.full((1, _WIDTH), float(index))
+
+
+def _identities(chunk: torch.Tensor) -> list[int]:
+    return [int(row[0].item()) for row in chunk]
+
+
+class TestFrameBuffer:
+    def test_slice_returns_the_requested_frames(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        assert _identities(buffer.slice(3, 7)) == [3, 4, 5, 6]
+
+    def test_slice_is_relative_to_the_live_base(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(4)
+        assert buffer.base_frame == 4
+        assert _identities(buffer.slice(4, 8)) == [4, 5, 6, 7]
+
+    def test_slicing_a_discarded_range_raises(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(5)
+        with pytest.raises(RuntimeError, match="unavailable"):
+            buffer.slice(3, 6)
+
+    def test_slicing_past_the_end_raises(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(4):
+            buffer.append(_frame(index))
+        with pytest.raises(RuntimeError, match="unavailable"):
+            buffer.slice(2, 9)
+
+    def test_discard_is_clamped_at_both_ends(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(5):
+            buffer.append(_frame(index))
+        buffer.discard_before(-3)
+        assert buffer.base_frame == 0
+        buffer.discard_before(99)
+        assert buffer.base_frame == 5
+        assert buffer.end_frame == 5
+
+    def test_peak_tracks_the_high_water_mark(self) -> None:
+        buffer = FrameBuffer()
+        for index in range(10):
+            buffer.append(_frame(index))
+        buffer.discard_before(8)
+        assert buffer.peak_frames == 10
+
+
+class TestWindowLifecycle:
+    """Drive the state exactly as the talker does, one frame per step."""
+
+    @staticmethod
+    def _run(total_frames: int) -> tuple[list[list[int]], int]:
+        """Generate ``total_frames`` frames, returning the emitted spans.
+
+        Returns:
+            The frame identities of every emitted span, and the peak buffer
+            occupancy reached along the way.
+        """
+        state = RequestARState(seed=1, max_frames=total_frames)
+        spans: list[list[int]] = []
+        for index in range(total_frames):
+            state.frames.append(_frame(index))
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+            state.generated_frames += 1
+            window = state.ready_window()
+            if window is not None:
+                spans.append(_identities(state.take_window(window, is_final=False)))
+
+        # Terminal flush: everything not yet emitted, as one contiguous span.
+        start = state.emitted_windows * AR_CHUNK_HOP_FRAMES
+        end = state.frames.end_frame
+        if end > start:
+            spans.append(_identities(state.frames.slice(start, end)))
+        return spans, state.frames.peak_frames
+
+    def test_short_song_is_one_span(self) -> None:
+        spans, _ = self._run(50)
+        assert spans == [list(range(50))]
+
+    def test_every_frame_is_delivered_exactly_once_in_order(self) -> None:
+        """Spans overlap by design, but their union must be the whole song.
+
+        Reconstruct by taking each span's non-overlapping remainder, which is
+        what the acoustic stage does after cropping.
+        """
+        for total in (250, 437, 600, 901):
+            spans, _ = self._run(total)
+            rebuilt: list[int] = []
+            for span in spans:
+                rebuilt.extend(frame for frame in span if frame >= len(rebuilt))
+            assert rebuilt == list(range(total)), f"gap or duplicate at {total} frames"
+
+    def test_windows_advance_by_the_hop(self) -> None:
+        spans, _ = self._run(600)
+        for index, span in enumerate(spans[:-1]):
+            assert span[0] == index * AR_CHUNK_HOP_FRAMES
+            assert len(span) == AR_CHUNK_FRAMES
+
+    def test_mid_stream_windows_are_never_the_final_frame(self) -> None:
+        """A window is held back until a further frame proves it is not last."""
+        total = 600
+        spans, _ = self._run(total)
+        for span in spans[:-1]:
+            assert span[-1] < total - 1
+
+    def test_terminal_span_covers_the_tail(self) -> None:
+        total = 437
+        spans, _ = self._run(total)
+        assert spans[-1][-1] == total - 1
+
+    def test_terminal_span_is_bounded(self) -> None:
+        """The tail is at most one full window plus one hop."""
+        for total in (250, 301, 437, 600, 901, 1000):
+            spans, _ = self._run(total)
+            assert len(spans[-1]) <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
+
+    def test_buffer_stays_bounded_over_a_long_song(self) -> None:
+        """Peak occupancy must not grow with song length; that is the point
+        of evicting retired frames."""
+        _, peak_short = self._run(600)
+        _, peak_long = self._run(3000)
+        assert peak_long == peak_short
+        assert peak_long <= AR_CHUNK_FRAMES + AR_CHUNK_HOP_FRAMES
+
+    def test_emitted_window_count_matches_the_window_plan(self) -> None:
+        """What the stream emitted must reconcile with chunk_windows()."""
+        for total in (250, 437, 600):
+            spans, _ = self._run(total)
+            planned = chunk_windows(total)
+            # Every window is accounted for: the pre-terminal spans are one
+            # window each, and the terminal span carries whatever remains.
+            assert len(spans) <= len(planned)
+            assert spans[0][0] == planned[0].start
+
+
+class TestSamplingPositions:
+    def test_each_frame_reserves_one_slot_per_codebook(self) -> None:
+        """Codebook streams must never collide, or codes repeat across frames."""
+        state = RequestARState(seed=1, max_frames=10)
+        positions = [state.advance_position() for _ in range(5)]
+        assert positions == [0, 8, 16, 24, 32]
+
+    def test_draw_slots_are_globally_unique(self) -> None:
+        state = RequestARState(seed=1, max_frames=64)
+        slots: set[int] = set()
+        for _ in range(64):
+            base = state.advance_position()
+            for codebook in range(NUM_CODEBOOKS):
+                slot = base + codebook
+                assert slot not in slots
+                slots.add(slot)
+        assert len(slots) == 64 * NUM_CODEBOOKS
+
+
+class TestCodeRetention:
+    def test_codes_span_matches_the_frame_span(self) -> None:
+        state = RequestARState(seed=1, max_frames=10)
+        for index in range(10):
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+        codes = state.take_codes(2, 6)
+        assert codes is not None
+        assert codes.shape == (4, NUM_CODEBOOKS)
+        assert [int(row[0]) for row in codes] == [2, 3, 4, 5]
+
+    def test_out_of_range_span_is_dropped_not_raised(self) -> None:
+        """Codes are diagnostic; the audio path must not fail without them."""
+        state = RequestARState(seed=1, max_frames=10)
+        for index in range(4):
+            state.codes.append(torch.full((1, NUM_CODEBOOKS), index))
+        assert state.take_codes(2, 99) is None
+        assert state.take_codes(0, 0) is None

tests/model_executor/models/test_minimax_music3_prompt.py

@@ -0,0 +1,203 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniMax Music 3 prompt assembly and window arithmetic.
+
+The prompt transformations are not cosmetic: they decide token positions, and
+token positions decide the audio. These tests pin the documented behaviours
+that a well-meaning cleanup would otherwise "fix".
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.chunking import (
+    chunk_windows,
+    crop_sample_bounds,
+    overlap_mel_length,
+)
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    AR_CHUNK_FRAMES,
+    AR_CHUNK_HOP_FRAMES,
+    BLEND_MEL_FRAMES,
+    DAV_HOP_SAMPLES,
+    HOP_MEL_FRAMES,
+)
+from vllm_omni.model_executor.models.minimax_music3.prompt import (
+    SPECIAL_TOKEN_IDS,
+    build_cfg_null_token_ids,
+    build_prompt,
+    clean_caption,
+    normalize_lyrics,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestNormalizeLyrics:
+    def test_prepends_start_tag(self) -> None:
+        assert normalize_lyrics("[Verse]\nhello").startswith("[start]\n")
+
+    def test_lowercases_structure_tags(self) -> None:
+        assert "[verse]" in normalize_lyrics("[Verse]\nhello")
+        assert "[Verse]" not in normalize_lyrics("[Verse]\nhello")
+
+    def test_text_after_a_leading_tag_is_dropped(self) -> None:
+        """Documented contract: a tag must sit on its own line.
+
+        Lyrics written next to a tag are silently lost, which is why the
+        serving docs tell callers to put tags on their own line.
+        """
+        out = normalize_lyrics("[Verse] Walking down the street")
+        assert "Walking" not in out
+        assert "[verse]" in out
+
+    def test_text_on_its_own_line_survives(self) -> None:
+        out = normalize_lyrics("[Verse]\nWalking down the street")
+        assert "Walking down the street" in out
+
+    def test_bracket_space_becomes_newline(self) -> None:
+        assert normalize_lyrics("a [chorus]") == "[start]\na\n[chorus]"
+
+
+class TestCleanCaption:
+    def test_special_token_form_is_rewritten_not_passed_through(self) -> None:
+        """A caption must not be able to inject a real special token."""
+        assert clean_caption("<|bpm 92|>") == "bpm is 92"
+        assert "<|" not in clean_caption("<|bpm 92|>")
+
+    def test_lone_tag_keeps_its_inner_text(self) -> None:
+        assert clean_caption("<|instrumental|>") == "instrumental"
+
+    def test_markdown_is_stripped(self) -> None:
+        assert clean_caption("## Heading") == "Heading"
+        assert clean_caption("**bold**") == "bold"
+        assert clean_caption("- item") == "item"
+
+    def test_blank_lines_collapse_at_two_not_three(self) -> None:
+        """The source collapses runs of two or more, not three or more.
+
+        This shifts blank-line token positions and therefore the audio, so it
+        is pinned deliberately.
+        """
+        assert clean_caption("a\n\nb") == "a\nb"
+        assert clean_caption("a\n\n\n\nb") == "a\nb"
+
+
+class TestBuildPrompt:
+    def test_framing_tokens_are_present_and_ordered(self) -> None:
+        prompt = build_prompt("a bright pop song", "[Verse]\nhello")
+        assert prompt.startswith("<|im_start|><|caption_start|>")
+        assert prompt.endswith("<|im_end|><|audio_start|>")
+        assert prompt.index("<|caption_end|>") < prompt.index("<|lyrics_start|>")
+
+
+class TestBuildCfgNullTokenIds:
+    @staticmethod
+    def _prompt_ids(body_len: int) -> list[int]:
+        return (
+            [SPECIAL_TOKEN_IDS["<|im_start|>"]]
+            + list(range(1000, 1000 + body_len))
+            + [SPECIAL_TOKEN_IDS["<|im_end|>"], SPECIAL_TOKEN_IDS["<|audio_start|>"]]
+        )
+
+    def test_length_is_preserved(self) -> None:
+        """Equal length is what lets the two rows decode in lockstep."""
+        ids = self._prompt_ids(12)
+        assert len(build_cfg_null_token_ids(ids)) == len(ids)
+
+    def test_body_is_replaced_and_framing_survives(self) -> None:
+        ids = self._prompt_ids(12)
+        null = build_cfg_null_token_ids(ids)
+        assert null[0] == SPECIAL_TOKEN_IDS["<|im_start|>"]
+        assert null[-2] == SPECIAL_TOKEN_IDS["<|im_end|>"]
+        assert null[-1] == SPECIAL_TOKEN_IDS["<|audio_start|>"]
+        assert set(null[1:-2]) == {SPECIAL_TOKEN_IDS["<|audio_cfg|>"]}
+
+    def test_input_is_not_mutated(self) -> None:
+        ids = self._prompt_ids(6)
+        before = list(ids)
+        build_cfg_null_token_ids(ids)
+        assert ids == before
+
+    def test_minimal_prompt_with_empty_body(self) -> None:
+        ids = [
+            SPECIAL_TOKEN_IDS["<|im_start|>"],
+            SPECIAL_TOKEN_IDS["<|audio_cfg|>"],
+            SPECIAL_TOKEN_IDS["<|im_end|>"],
+            SPECIAL_TOKEN_IDS["<|audio_start|>"],
+        ]
+        assert build_cfg_null_token_ids(ids) == ids
+
+    def test_too_short_is_rejected(self) -> None:
+        with pytest.raises(ValueError, match="too short"):
+            build_cfg_null_token_ids([1, 2, 3])
+
+    def test_wrong_framing_is_rejected(self) -> None:
+        with pytest.raises(ValueError, match="framing"):
+            build_cfg_null_token_ids([9, 9, 9, 9, 9])
+
+
+class TestChunkWindows:
+    def test_no_frames_yields_nothing(self) -> None:
+        assert chunk_windows(0) == []
+
+    def test_short_song_is_one_final_window(self) -> None:
+        windows = chunk_windows(50)
+        assert len(windows) == 1
+        assert (windows[0].start, windows[0].end) == (0, 50)
+        assert windows[0].is_first and windows[0].is_last
+
+    def test_exactly_one_window(self) -> None:
+        windows = chunk_windows(AR_CHUNK_FRAMES)
+        assert len(windows) == 1
+        assert windows[0].is_last
+
+    def test_windows_advance_by_the_hop(self) -> None:
+        windows = chunk_windows(450)
+        starts = [w.start for w in windows]
+        assert starts == list(range(0, starts[-1] + 1, AR_CHUNK_HOP_FRAMES))
+        assert windows[-1].is_last
+        assert not any(w.is_last for w in windows[:-1])
+
+    def test_windows_cover_every_frame(self) -> None:
+        frames = 437
+        windows = chunk_windows(frames)
+        covered: set[int] = set()
+        for window in windows:
+            covered.update(range(window.start, window.end))
+        assert covered == set(range(frames))
+
+    def test_final_window_is_clamped_to_the_song(self) -> None:
+        windows = chunk_windows(437)
+        assert windows[-1].end == 437
+
+    def test_only_the_first_window_is_first(self) -> None:
+        windows = chunk_windows(600)
+        assert windows[0].is_first
+        assert not any(w.is_first for w in windows[1:])
+
+
+class TestCropBounds:
+    def test_single_window_song_is_not_cropped(self) -> None:
+        window = chunk_windows(50)[0]
+        assert crop_sample_bounds(window) == (0, 0)
+
+    def test_first_window_keeps_its_head(self) -> None:
+        left, right = crop_sample_bounds(chunk_windows(600)[0])
+        assert left == 0
+        assert right == (HOP_MEL_FRAMES - BLEND_MEL_FRAMES) * DAV_HOP_SAMPLES
+
+    def test_last_window_keeps_its_tail(self) -> None:
+        left, right = crop_sample_bounds(chunk_windows(600)[-1])
+        assert left == BLEND_MEL_FRAMES * DAV_HOP_SAMPLES
+        assert right == 0
+
+    def test_middle_window_is_cropped_on_both_sides(self) -> None:
+        windows = chunk_windows(900)
+        middle = windows[len(windows) // 2]
+        left, right = crop_sample_bounds(middle)
+        assert left > 0 and right > 0
+
+    def test_overlap_is_half_a_window(self) -> None:
+        assert overlap_mel_length() == HOP_MEL_FRAMES // 2

tests/model_executor/models/test_minimax_music3_row_binding.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Regression tests for MiniMax Music 3 row-to-request binding.
+
+vLLM compacts and reorders the persistent batch whenever a request finishes or
+is preempted, so a batch row number is only meaningful within the step that
+produced it. Every failure pinned here was a real defect: state keyed by row
+survived into a step where the row belonged to someone else, or state keyed by
+one spelling of the request id was released under another and leaked.
+
+These exercise the pure helpers and the binding arithmetic directly, so they
+need no GPU and no engine.
+"""
+
+from __future__ import annotations
+
+import pytest
+
+from vllm_omni.model_executor.models.minimax_music3.constants import (
+    DEFAULT_MAX_AUDIO_FRAMES,
+)
+from vllm_omni.model_executor.models.minimax_music3.talker import (
+    _CFG_UNCOND_SUFFIX,
+    _first_int,
+    _pair_id_from_request,
+    _request_key,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+class TestRequestKey:
+    def test_engine_list_form_is_unwrapped(self) -> None:
+        """The engine injects global_request_id as a single-element list.
+
+        Keying state on the list's ``str()`` would produce "['req-1']", which
+        never matches the bare id that ``on_requests_finished`` delivers, so
+        every finished song would leak its frame buffer.
+        """
+        assert _request_key(["req-1"]) == "req-1"
+
+    def test_bare_string_is_unchanged(self) -> None:
+        assert _request_key("req-1") == "req-1"
+
+    def test_list_and_bare_forms_agree(self) -> None:
+        """Allocation and release must land on the same key."""
+        assert _request_key(["req-1"]) == _request_key("req-1")
+
+    def test_empty_and_missing_are_stable(self) -> None:
+        assert _request_key(None) == ""
+        assert _request_key([]) == ""
+
+    def test_tuple_form(self) -> None:
+        assert _request_key(("req-9",)) == "req-9"
+
+
+class TestPairIdDerivation:
+    def test_conditioned_request_is_its_own_pair(self) -> None:
+        assert _pair_id_from_request("req-1") == "req-1"
+
+    def test_companion_maps_to_its_partner(self) -> None:
+        assert _pair_id_from_request(f"req-1{_CFG_UNCOND_SUFFIX}") == "req-1"
+
+    def test_pair_members_agree_on_the_pair_id(self) -> None:
+        cond = "req-42"
+        uncond = f"req-42{_CFG_UNCOND_SUFFIX}"
+        assert _pair_id_from_request(cond) == _pair_id_from_request(uncond)
+
+    def test_a_request_ending_in_a_similar_string_is_not_confused(self) -> None:
+        assert _pair_id_from_request("req-cfg_uncond") == "req-cfg_uncond"
+
+
+class TestFirstInt:
+    def test_serving_layer_list_wrapping(self) -> None:
+        """tts_params values arrive wrapped in a single-element list."""
+        assert _first_int([750], DEFAULT_MAX_AUDIO_FRAMES) == 750
+
+    def test_bare_value(self) -> None:
+        assert _first_int(750, DEFAULT_MAX_AUDIO_FRAMES) == 750
+
+    def test_missing_falls_back(self) -> None:
+        assert _first_int(None, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int([], DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+
+    def test_nonsense_falls_back_rather_than_raising(self) -> None:
+        assert _first_int("abc", DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int(0, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+        assert _first_int(-5, DEFAULT_MAX_AUDIO_FRAMES) == DEFAULT_MAX_AUDIO_FRAMES
+
+
+def _feedback_plan(
+    request_ids: list[str],
+    computed: list[int],
+    scheduled: list[int],
+    prompt_tokens: dict[str, int],
+    has_codes: set[str],
+) -> list[int]:
+    """Mirror of the talker's feedback row selection.
+
+    Returns the flat token offsets whose embedding gets replaced by an audio
+    frame. Kept as a standalone function so the arithmetic can be checked
+    without constructing a model.
+    """
+    offsets: list[int] = []
+    running = 0
+    for count in scheduled:
+        offsets.append(running)
+        running += count
+
+    rows: list[int] = []
+    for row, request_id in enumerate(request_ids):
+        if request_id not in has_codes:
+            continue
+        if scheduled[row] != 1 or computed[row] < prompt_tokens.get(request_id, 0):
+            continue
+        rows.append(offsets[row])
+    return rows
+
+
+class TestFeedbackRowSelection:
+    """The guard that decides which token embeddings become audio frames."""
+
+    def test_decode_only_step_replaces_every_row(self) -> None:
+        ids = ["a", "a__cfg_uncond", "b", "b__cfg_uncond"]
+        rows = _feedback_plan(
+            ids,
+            computed=[120, 120, 90, 90],
+            scheduled=[1, 1, 1, 1],
+            prompt_tokens={i: 50 for i in ids},
+            has_codes=set(ids),
+        )
+        assert rows == [0, 1, 2, 3]
+
+    def test_prefill_row_is_never_overwritten(self) -> None:
+        """A prompt token must keep its real embedding.
+
+        The earlier guard compared the flat token count against a row budget,
+        so a short prompt slipped through and had its first embeddings
+        replaced with a previous request's audio frame.
+        """
+        rows = _feedback_plan(
+            ["a", "b"],
+            computed=[0, 300],
+            scheduled=[8, 1],
+            prompt_tokens={"a": 8, "b": 40},
+            has_codes={"b"},
+        )
+        # Row 0 occupies flat tokens 0..7 and must be untouched; row 1's single
+        # decode token sits at flat offset 8.
+        assert rows == [8]
+
+    def test_short_prefill_below_the_row_budget_is_still_safe(self) -> None:
+        """The exact shape that defeated the old token-count guard."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[0],
+            scheduled=[10],
+            prompt_tokens={"a": 10},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_chunked_prefill_midway_is_not_fed_back(self) -> None:
+        rows = _feedback_plan(
+            ["a"],
+            computed=[2048],
+            scheduled=[2048],
+            prompt_tokens={"a": 5000},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_final_prefill_chunk_is_not_fed_back(self) -> None:
+        """The step that completes a prompt still carries prompt tokens."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[4096],
+            scheduled=[904],
+            prompt_tokens={"a": 5000},
+            has_codes={"a"},
+        )
+        assert rows == []
+
+    def test_mixed_batch_offsets_follow_the_flat_layout(self) -> None:
+        """Tokens are a flat concatenation, so offsets are a running sum."""
+        ids = ["p", "d1", "d2"]
+        rows = _feedback_plan(
+            ids,
+            computed=[0, 500, 700],
+            scheduled=[64, 1, 1],
+            prompt_tokens={"p": 64, "d1": 40, "d2": 40},
+            has_codes={"d1", "d2"},
+        )
+        assert rows == [64, 65]
+
+    def test_request_without_codes_yet_is_skipped(self) -> None:
+        """The first decode step after prefill has no previous frame."""
+        rows = _feedback_plan(
+            ["a"],
+            computed=[50],
+            scheduled=[1],
+            prompt_tokens={"a": 50},
+            has_codes=set(),
+        )
+        assert rows == []
+
+    def test_reordered_batch_binds_by_request_not_position(self) -> None:
+        """The same two requests in swapped rows get swapped offsets.
+
+        This is the compaction case: holding codes in a row-indexed buffer
+        would have handed each request the other's frame.
+        """
+        prompt_tokens = {"a": 50, "b": 50}
+        first = _feedback_plan(["a", "b"], [100, 200], [1, 1], prompt_tokens, {"a"})
+        swapped = _feedback_plan(["b", "a"], [200, 100], [1, 1], prompt_tokens, {"a"})
+        assert first == [0]
+        assert swapped == [1]