Stop at max_output_tokens, not one past it

mstar/model/orpheus/orpheus_model.py

@@ -387,7 +387,8 @@ def get_initial_forward_pass_args(
     # -------------------------------------------------------------------
 
     def get_loop_caps(self, **model_kwargs) -> dict[str, int]:
-        return {"decode_loop": self.get_max_output_tokens(**model_kwargs)}
+        # the prefill samples the first of max_output_tokens, so the loop runs one fewer
+        return {"decode_loop": max(1, self.get_max_output_tokens(**model_kwargs) - 1)}
 
     def kv_labels(self) -> dict[str, dict[str, KVLabel]]:
         """The LLM's prompt is its token ids. Its generated audio codes never

mstar/model/qwen3_omni/qwen3_omni_model.py

@@ -1435,9 +1435,11 @@ def _kv_layout(self, plan: list[PromptPart], tensors: NameToTensorList, audio_ou
         }
 
     def get_loop_caps(self, **model_kwargs) -> dict[str, int]:
+        # the Thinker's last prefill emits its first text token and
+        # talker_last_prefill its first codec frame, so each loop runs one fewer
         return {
-            "thinker_decode_loop": self.get_max_output_tokens(**model_kwargs),
-            "talker_decode_loop": max(1, self._get_max_talker_output_tokens(**model_kwargs)),
+            "thinker_decode_loop": max(1, self.get_max_output_tokens(**model_kwargs) - 1),
+            "talker_decode_loop": max(1, self._get_max_talker_output_tokens(**model_kwargs) - 1),
         }
 
     def kv_labels(self) -> dict[str, dict[str, KVLabel]]:

mstar/model/qwen3_omni/submodules.py

@@ -1706,22 +1706,17 @@ def check_stop(
         request_info: CurrentForwardPassInfo,
         outputs: dict[str, list[torch.Tensor]],
     ) -> set[str] | StopCheck:
-        """Stop on codec_eos or the frame limit. The codec_eos frame stays out
-        of the Code2Wav stream: it is not a codebook index, and only this host
-        copy tells it from audio."""
+        """Stop on codec_eos. The frame budget is the loop's cap
+        (``Qwen3OmniModel.get_loop_caps``), which the graph runtime enforces.
+        The codec_eos frame stays out of the Code2Wav stream: it is not a
+        codebook index, and only this host copy tells it from audio."""
         if "layer0_codes" not in outputs:
             return set()
         token = outputs["layer0_codes"][0].item()
         eos_token_id = self.config.talker.codec_eos_token_id
-        max_tokens = request_info.step_metadata.get(
-            "talker_max_tokens", request_info.max_tokens
-        )
-        is_eos = eos_token_id is not None and eos_token_id == token
-        stops = set()
-        if is_eos or \
-                (request_info.dynamic_loop_iter_counts.get("talker_decode_loop", 0) + 1 >= max_tokens):
-            stops = {"talker_decode_loop"}
-        return StopCheck(stops, withheld=("codec_tokens",)) if is_eos else stops
+        if eos_token_id is None or eos_token_id != token:
+            return set()
+        return StopCheck({"talker_decode_loop"}, withheld=("codec_tokens",))
 
     def cleanup_request(self, request_id: str) -> None:
         """Remove per-request state when a request completes."""

test/modular/test_model_kv_layouts.py

@@ -262,6 +262,34 @@ def test_canary_decode_cap_leaves_the_prefills_token_in_the_budget(model_name, m
     assert model.get_loop_caps(max_output_tokens=max_tokens)["decode_loop"] == max(1, max_tokens - 1)
 
 
+def _first_token_at_prefill(model_name):
+    """(model, decode loop) for a model whose prefill samples the reply's first token."""
+    if model_name == "higgs_audio":
+        return _higgs(), "decode_loop"
+    if model_name == "whisper":
+        return _whisper(), "decode_loop"
+    from mstar.model.orpheus.orpheus_model import OrpheusModel
+    from mstar.model.qwen3_5.qwen3_5_model import Qwen3_5DenseModel
+    from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel
+
+    cls, loop = {
+        "orpheus": (OrpheusModel, "decode_loop"),
+        "qwen3_5": (Qwen3_5DenseModel, "decode_loop"),
+        "qwen3_omni": (Qwen3OmniModel, "thinker_decode_loop"),
+    }[model_name]
+    return object.__new__(cls), loop
+
+
+@pytest.mark.parametrize("max_tokens", [2, 3, 64])
+@pytest.mark.parametrize("model_name", ["higgs_audio", "orpheus", "qwen3_5", "qwen3_omni", "whisper"])
+def test_the_prefills_token_and_the_decode_cap_add_up_to_max_output_tokens(model_name, max_tokens):
+    """The decode loop's cap alone ends it (``TokenStop``), one token a step,
+    after the prefill's."""
+    model, loop = _first_token_at_prefill(model_name)
+
+    assert 1 + model.get_loop_caps(max_output_tokens=max_tokens)[loop] == max_tokens
+
+
 @pytest.mark.parametrize("variant", ["chatterbox", "turbo"])
 def test_a_reference_clip_is_laid_out_at_the_tokens_the_tokenizer_makes_of_it(variant):
     from mstar.model.chatterbox.components.audio_frontend import resample, trim_reference
@@ -446,19 +474,38 @@ def _omni_submodule():
     return SimpleNamespace(config=SimpleNamespace(im_end_token_id=0, talker=SimpleNamespace(codec_eos_token_id=0)))
 
 
-@pytest.mark.parametrize("talker_max_tokens", [0, 1, 64])
-def test_qwen3_omni_caps_the_talkers_decode_where_its_stop_check_ends_it(talker_max_tokens):
-    pytest.importorskip("flashinfer")
+@pytest.mark.parametrize("talker_max_tokens", [2, 3, 64])
+def test_qwen3_omni_talker_cap_alone_gives_the_frame_budget(talker_max_tokens):
+    """talker_last_prefill streams the first frame and each decode iteration
+    one more. The loop's cap ends decode, and the stop check looks only for
+    codec_eos."""
     from mstar.model.qwen3_omni.qwen3_omni_model import Qwen3OmniModel
     from mstar.model.qwen3_omni.submodules import TalkerSubmodule
 
     cap = Qwen3OmniModel.get_loop_caps(object.__new__(Qwen3OmniModel), talker_max_output_tokens=talker_max_tokens)
 
-    assert cap["talker_decode_loop"] == _stops_after(
+    assert 1 + cap["talker_decode_loop"] == talker_max_tokens
+    assert _stops_after(
         lambda *args: TalkerSubmodule.check_stop(_omni_submodule(), *args), "talker_decode_loop", "talker_decode",
         2048, {"layer0_codes": [torch.tensor([NOT_EOS])]}, max_tokens=2048,
         step_metadata={"talker_max_tokens": talker_max_tokens},
-    ), "the cap and the stop check end decode on different iterations"
+    ) == 2048, "the stop check ended decode on a frame that is not codec_eos"
+
+
+def test_qwen3_omni_talker_last_prefill_stops_no_loop():
+    """The last prefill's frame comes before talker_decode_loop exists, so a
+    stop there would name a loop its walk doesn't have."""
+    from mstar.model.qwen3_omni.submodules import TalkerSubmodule
+
+    info = ForwardPassInfoWrapper(
+        info=CurrentForwardPassInfo(
+            request_id="r", graph_walk="talker_last_prefill", fwd_index=0, random_seed=0,
+            max_tokens=2048, step_metadata={"talker_max_tokens": 2},
+        ),
+        metadata=InputMetadata(),
+    )
+
+    assert not TalkerSubmodule.check_stop(_omni_submodule(), "r", info, {"layer0_codes": [torch.tensor([NOT_EOS])]})
 
 
 def _tiny_qwen3_tts_talker():