diff --git a/docs/configuration/stage_configs.md b/docs/configuration/stage_configs.md index 855390739..07beef2de 100644 --- a/docs/configuration/stage_configs.md +++ b/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_`; values are names registered under top-level `connectors:`. | | `input_connectors` | dict \| null | optional | `null` | Keyed by `from_stage_`; 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. diff --git a/tests/config/test_config_factory.py b/tests/config/test_config_factory.py index af898ae50..82c1fb32e 100644 --- a/tests/config/test_config_factory.py +++ b/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)]) diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 6fe04fdb7..be7a0ffac 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/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]) diff --git a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py b/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py index 3fd4867cb..d93755fec 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_llm2tts.py +++ b/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]) diff --git a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py b/tests/model_executor/models/minicpmo_4_5/test_pipeline.py index f8fdfdcd2..e12010ce4 100644 --- a/tests/model_executor/models/minicpmo_4_5/test_pipeline.py +++ b/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) diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index b58887099..0b8723c0e 100644 --- a/vllm_omni/config/omni_config.py +++ b/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 diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 063a95b0e..9c9186c16 100644 --- a/vllm_omni/config/stage_config.py +++ b/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: diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 899ff116d..ead0e6f1d 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/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) diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py b/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py index 0ca61a346..e48328d0b 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/pipeline.py +++ b/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,