[Model] Stop MiniCPM-o Thinker at TTS boundaries (#7463)

Signed-off-by: Zhu Shihao <zhushihao@bytedance.com>
Co-authored-by: Zhu Shihao <zhushihao@bytedance.com>

docs/configuration/stage_configs.md

@@ -38,7 +38,7 @@ Common `StagePipelineConfig` fields include:
 | `custom_process_input_func` | Processor applied to this stage's incoming payload. |
 | `custom_process_next_stage_input_func` | Processor used for full-payload handoff to the next stage. |
 | `async_chunk_process_next_stage_input_func` | Processor used for async chunk handoff. |
-| `sampling_constraints` | Model-owned sampling constraints that deploy defaults cannot override. |
+| `sampling_constraints` | Model-owned sampling constraints that deploy defaults cannot override. Scalar values replace deploy defaults; required `stop_token_ids` extend and deduplicate them. |
 
 To add or change topology, define and register a new pipeline variant. Use
 deploy YAML only for runtime placement, resource sizing, connectors, and other
@@ -97,7 +97,7 @@ Each entry under `stages:` accepts any `StageDeployConfig` field directly (no ne
 | `devices` | str \| null | optional | `null` | Device list assigned to this stage. The number of device ids must equal this stage's local world size (`tensor_parallel_size` × local data-parallel size × `pipeline_parallel_size`, or `num_replicas` × that product for a replica pool); a mismatch fails early — see the note below. |
 | `output_connectors` | dict \| null | optional | `null` | Keyed by `to_stage_<n>`; values are names registered under top-level `connectors:`. |
 | `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_<n>`; values are names registered under top-level `connectors:`. |
-| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Deep-merged with pipeline `sampling_constraints` (pipeline wins). |
+| `default_sampling_params` | dict \| null | optional | `null` | Baseline sampling params. Merged with pipeline `sampling_constraints`; scalar constraints win, while required `stop_token_ids` are appended and deduplicated. |
 | `engine_extras` | dict | optional | `{}` | Catch-all for engine fields not listed above; deep-merged across overlays and forwarded to the stage engine. |
 
 **Note:** a stage's `devices` count must equal its local world size (`tensor_parallel_size` × `data_parallel_size_local` × `pipeline_parallel_size`, falling back to global `data_parallel_size` when the local size is unset), or `num_replicas` × that product for a replica pool. A mismatch fails early and names the offending stage. A top-level `--tensor-parallel-size` is broadcast to every stage, so it can make a single-GPU stage violate this contract ([issue #5003](gh-issue:5003)); fix that case with `--stage-overrides` (set `tensor_parallel_size` and `devices` together per stage) or set TP only on the multi-GPU stage.

tests/config/test_config_factory.py

@@ -1332,6 +1332,10 @@ class TestDeployConfigLoading:
         assert deploy.duplex_session.max_sessions == max_sessions
         assert [stage.session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
         assert [stage.to_omegaconf().session_mode for stage in stages] == ["duplex", "duplex", "duplex"]
+        assert stages[0].yaml_extras["default_sampling_params"]["stop_token_ids"] == [
+            151704,
+            151645,
+        ]
 
     def test_load_minicpmo_default_deploy_config(self):
         deploy_path = Path(get_deploy_config_path("minicpmo_4_5.yaml"))
@@ -3212,7 +3216,7 @@ class TestSentinelDefaultPrecedence:
 
 
 class TestSamplingConstraintsPrecedence:
-    """Test that pipeline sampling_constraints override deploy defaults."""
+    """Test scalar constraint precedence and additive required stop tokens."""
 
     def test_constraints_win(self):
         deploy_path = Path(get_deploy_config_path("qwen3_omni_moe.yaml"))
@@ -3232,6 +3236,47 @@ class TestSamplingConstraintsPrecedence:
         # Pipeline says stop_token_ids=[2150] for talker
         assert stages[1].yaml_extras["default_sampling_params"]["stop_token_ids"] == [2150]
 
+    def test_required_stop_tokens_extend_yaml_defaults(self, tmp_path):
+        pipeline = PipelineConfig(
+            model_type="required_stop_tokens",
+            stages=(
+                StagePipelineConfig(
+                    stage_id=0,
+                    model_stage="thinker",
+                    final_output=True,
+                    sampling_constraints={
+                        "detokenize": True,
+                        "stop_token_ids": [151704, 151645],
+                    },
+                ),
+            ),
+        )
+        deploy_path = tmp_path / "required-stop-tokens.yaml"
+        deploy_path.write_text(
+            """
+stages:
+  - stage_id: 0
+    default_sampling_params:
+      max_tokens: 7
+      stop_token_ids: [100, 151704]
+""",
+            encoding="utf-8",
+        )
+
+        legacy_stage = merge_pipeline_deploy(pipeline, load_deploy_config(deploy_path))[0]
+        structured_stage = VllmOmniConfig.from_pipeline_config(
+            pipeline,
+            deploy_config_path=str(deploy_path),
+        ).stage_by_id(0)
+        expected = {
+            "max_tokens": 7,
+            "detokenize": True,
+            "stop_token_ids": [100, 151704, 151645],
+        }
+
+        assert legacy_stage.yaml_extras["default_sampling_params"] == expected
+        assert structured_stage.model_config.default_sampling_params == expected
+
 
 class TestPipelineConfigResolvers:
     @pytest.mark.parametrize("resolver", [obj for obj in OMNI_PIPELINES.values() if callable(obj)])

tests/entrypoints/test_omni_entrypoints.py

@@ -200,20 +200,28 @@ def _make_base():
     return obj
 
 
-def test_resolve_sampling_params_list_preserves_stage_constraints():
+def test_resolve_sampling_params_list_merges_required_stop_tokens():
     base = _make_base()
     base.engine.num_stages = 1
-    base.default_sampling_params_list = [SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[42])]
+    required_stop_ids = [151704, 151645]
+    base.default_sampling_params_list = [
+        SamplingParams(max_tokens=1000, detokenize=False, stop_token_ids=[99, *required_stop_ids])
+    ]
     base.engine.stage_configs = [
         StageConfig(
             stage_id=0,
             model_stage="dummy-model",
-            sampling_constraints={"detokenize": False, "stop_token_ids": [42]},
+            sampling_constraints={"detokenize": False, "stop_token_ids": required_stop_ids},
         ).to_omegaconf()
     ]
     base.sampling_constraints_list = base._get_sampling_constraints_list(base.engine.stage_configs)
-    assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": [42]}]
-    caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[7])
+    assert base.sampling_constraints_list == [{"detokenize": False, "stop_token_ids": required_stop_ids}]
+
+    resolved_defaults = base.resolve_sampling_params_list(None)
+
+    assert resolved_defaults[0].stop_token_ids == [99, 151704, 151645]
+
+    caller_params = SamplingParams(seed=1234, max_tokens=7, detokenize=True, stop_token_ids=[100, 151704])
 
     resolved = base.resolve_sampling_params_list(caller_params)
 
@@ -221,10 +229,10 @@ def test_resolve_sampling_params_list_preserves_stage_constraints():
     assert resolved[0].seed == 1234
     assert resolved[0].max_tokens == 7
     assert resolved[0].detokenize is False
-    assert resolved[0].stop_token_ids == [42]
-    assert 42 in resolved[0]._all_stop_token_ids
+    assert resolved[0].stop_token_ids == [100, 151704, 151645]
+    assert {100, 151704, 151645}.issubset(resolved[0]._all_stop_token_ids)
     assert caller_params.detokenize is True
-    assert caller_params.stop_token_ids == [7]
+    assert caller_params.stop_token_ids == [100, 151704]
 
 
 @pytest.mark.parametrize("use_defaults", [False, True])

tests/model_executor/models/minicpmo_4_5/test_llm2tts.py

@@ -9,7 +9,7 @@ Covers ``vllm_omni.model_executor.stage_input_processors.minicpmo_4_5_omni.llm2t
   - both inputs missing -> raises
   - structured model_intermediate_buffer carries the thinker ids and text
   - scheduler prompt tokens follow the selected TTS region or output tokens
-  - MiniCPM-o 4.5 TTS region detection on 151703 / 151704 tokens
+  - MiniCPM-o 4.5 TTS region detection on 151703 / 151704 / 151645 tokens
   - plain chat without TTS markers conditions on the generated assistant span
   - prompt arg is normalized to a list and ``multi_modal_data`` is gated by
     ``requires_multimodal_data``
@@ -260,13 +260,14 @@ class TestTtsRegionDetection:
         )
         return result[0]["model_intermediate_buffer"], hidden
 
-    def test_4_5_markers_detected(self) -> None:
+    @pytest.mark.parametrize("tts_end_id", [151704, 151645])
+    def test_4_5_markers_detected(self, tts_end_id: int) -> None:
         # prompt:        [10, 11]
-        # output:        [151703, 30, 31, 151704, 40]
-        # full sequence: [10, 11, 151703, 30, 31, 151704, 40]
+        # output:        [151703, 30, 31, end, 40]
+        # full sequence: [10, 11, 151703, 30, 31, end, 40]
         #                  0   1     2    3   4     5     6
-        # 4.5 BOS at idx 2 -> slice starts at 3; EOS at idx 5 -> slice ends at 5.
-        buffer, hidden = self._run([10, 11], [151703, 30, 31, 151704, 40])
+        # 4.5 BOS at idx 2 -> slice starts at 3; end at idx 5 -> slice ends at 5.
+        buffer, hidden = self._run([10, 11], [151703, 30, 31, tts_end_id, 40])
         assert buffer["ids"]["tts"] == [30, 31]
         assert torch.equal(torch.tensor(buffer["hidden_states"]["tts"]), hidden[3:5])
 

tests/model_executor/models/minicpmo_4_5/test_pipeline.py

@@ -91,6 +91,10 @@ class TestPipelineTopology:
         assert thinker.final_output_type == "text"
         assert thinker.owns_tokenizer is True
         assert thinker.requires_multimodal_data is True
+        assert thinker.sampling_constraints == {
+            "detokenize": True,
+            "stop_token_ids": [151704, 151645],
+        }
 
     def test_talker_stage(self, pipeline: PipelineConfig) -> None:
         talker = pipeline.get_stage(1)

vllm_omni/config/omni_config.py

@@ -47,6 +47,7 @@ from vllm_omni.config.stage_config import (
     _select_processor_funcs,
     build_stage_runtime_overrides,
     load_deploy_config,
+    merge_sampling_constraints,
     normalize_pipeline_cli_overrides,
     reconcile_diffusion_attention_overrides,
 )
@@ -1347,10 +1348,10 @@ def _stage_sampling_params(
     stage_deploy: StageDeployConfig | None,
     topology: StagePipelineConfig,
 ) -> dict[str, Any] | None:
-    sampling: dict[str, Any] = {}
-    if stage_deploy is not None and stage_deploy.default_sampling_params:
-        sampling.update(_copy_value(stage_deploy.default_sampling_params))
-    sampling.update(_copy_value(topology.sampling_constraints))
+    sampling = merge_sampling_constraints(
+        _copy_value(stage_deploy.default_sampling_params) if stage_deploy is not None else None,
+        _copy_value(topology.sampling_constraints),
+    )
     return sampling or None
 
 

vllm_omni/config/stage_config.py

@@ -970,16 +970,29 @@ def _build_engine_args(
     return engine_args
 
 
+def merge_sampling_constraints(
+    sampling_params: Mapping[str, Any] | None,
+    constraints: Mapping[str, Any],
+) -> dict[str, Any]:
+    """Apply scalar constraints and extend stop tokens without mutating inputs."""
+    resolved_constraints = dict(constraints)
+    if "stop_token_ids" in resolved_constraints:
+        caller_stop_ids = (sampling_params or {}).get("stop_token_ids") or []
+        required_stop_ids = resolved_constraints["stop_token_ids"] or []
+        resolved_constraints["stop_token_ids"] = list(dict.fromkeys([*caller_stop_ids, *required_stop_ids]))
+    return {**(sampling_params or {}), **resolved_constraints}
+
+
 def _build_extras(
     ps: StagePipelineConfig,
     ds: StageDeployConfig | None,
 ) -> dict[str, Any]:
     """Assemble ``yaml_extras`` (sampling + connectors + pipeline extras)."""
     extras: dict[str, Any] = {}
-    sampling: dict[str, Any] = {}
-    if ds is not None and ds.default_sampling_params:
-        sampling.update(ds.default_sampling_params)
-    sampling.update(ps.sampling_constraints)
+    sampling = merge_sampling_constraints(
+        ds.default_sampling_params if ds is not None else None,
+        ps.sampling_constraints,
+    )
     if sampling:
         extras["default_sampling_params"] = sampling
     if ds is not None and ds.default_pooling_params:

vllm_omni/entrypoints/omni_base.py

@@ -18,6 +18,7 @@ from vllm.transformers_utils.repo_utils import file_or_path_exists
 from vllm.transformers_utils.runai_utils import is_runai_obj_uri
 from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 
+from vllm_omni.config.stage_config import merge_sampling_constraints
 from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
 from vllm_omni.engine.messages import (
     EngineQueueMessage,
@@ -372,12 +373,12 @@ class OmniBase(PDDisaggregationMixin):
 
     @staticmethod
     def _apply_sampling_constraints(params: Any, constraints: Mapping[str, Any]) -> Any:
-        """Rebuild params with pipeline-required settings without mutating caller input."""
+        """Apply pipeline requirements, merging required stops with caller stops."""
         if not constraints:
             return params
         if isinstance(params, Mapping):
-            return {**params, **constraints}
-        if is_dataclass(params):
+            values = dict(params)
+        elif is_dataclass(params):
             values = {field.name: getattr(params, field.name) for field in fields(params) if field.init}
         elif struct_fields := getattr(params, "__struct_fields__", None):
             values = {
@@ -387,7 +388,11 @@ class OmniBase(PDDisaggregationMixin):
             }
         else:
             raise TypeError(f"Expected a mapping, dataclass, or msgspec struct, got {type(params).__name__}")
-        return type(params)(**{**values, **constraints})
+
+        resolved = merge_sampling_constraints(values, constraints)
+        if isinstance(params, Mapping):
+            return resolved
+        return type(params)(**resolved)
 
     def _record_request_failure_once(self, request_id: str, reason: str) -> None:
         req_state = self.request_states.get(request_id)

vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py

@@ -53,7 +53,13 @@ MINICPMO_4_5_PIPELINE = PipelineConfig(
             owns_tokenizer=True,
             requires_multimodal_data=True,
             engine_output_type="latent",
-            sampling_constraints={"detokenize": True},
+            sampling_constraints={
+                "detokenize": True,
+                # The llm2tts bridge discards this boundary and every row
+                # after it, so stop Stage 0 as soon as either valid
+                # MiniCPM-o 4.5 turn terminator is sampled.
+                "stop_token_ids": [151704, 151645],
+            },
         ),
         StagePipelineConfig(
             stage_id=1,