From d11fa3e2ee49f114c14d5c638952600ca151dbec Mon Sep 17 00:00:00 2001 From: ZealSV Date: Fri, 31 Jul 2026 15:58:02 -0700 Subject: [PATCH] fix(serve): emit OPTION_SPECULATIVE_CONFIG so vLLM enables speculative decoding The speculative-decoding path set only OPTION_SPECULATIVE_DRAFT_MODEL, the env var the legacy lmi-dist rolling batcher read. The current LMI container runs the vLLM V1 engine, which reads OPTION_SPECULATIVE_CONFIG (a JSON blob) and ignores OPTION_SPECULATIVE_DRAFT_MODEL. As a result a model requested with speculative_decoding_config deploys and serves, but vLLM leaves speculative decoding off (speculative_config=None) -- the draft model is staged but never used. Add _set_speculative_draft_model_env, which emits BOTH vars: the legacy path var (back-compat with older containers) and OPTION_SPECULATIVE_CONFIG {"method":"draft_model","model":,"num_speculative_tokens":N}. Wire it into the custom, JumpStart, and deployment-config SD paths. num_speculative_tokens defaults to 5 and is overridable via NumSpeculativeTokens in the config. Unit: 6 SD tests pass (3 new asserting OPTION_SPECULATIVE_CONFIG + override). --- .../sagemaker/serve/model_builder_utils.py | 53 ++++++++++++++-- .../test_model_builder_utils_optimization.py | 62 +++++++++++++++++-- 2 files changed, 106 insertions(+), 9 deletions(-) diff --git a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py index e58ea4d7ad..e22abf40df 100644 --- a/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py +++ b/sagemaker-serve/src/sagemaker/serve/model_builder_utils.py @@ -138,6 +138,7 @@ def build(self): from sagemaker.serve.model_server.triton.config_template import CONFIG_TEMPLATE SPECULATIVE_DRAFT_MODEL = "/opt/ml/additional-model-data-sources" +DEFAULT_NUM_SPECULATIVE_TOKENS = 5 _DJL_MODEL_BUILDER_ENTRY_POINT = "inference.py" _NO_JS_MODEL_EX = "HuggingFace JumpStart Model ID not detected. Building for HuggingFace Model ID." _JS_SCOPE = "inference" @@ -1838,6 +1839,46 @@ def _generate_channel_name(self, additional_model_data_sources: Optional[List[Di return channel_name + def _set_speculative_draft_model_env( + self, + draft_model_path: str, + speculative_decoding_config: Optional[Dict] = None, + ) -> None: + """Set the environment variables that enable speculative decoding. + + Emits two variables for the LMI container: + + - ``OPTION_SPECULATIVE_DRAFT_MODEL``: the draft-model path, read by the + legacy ``lmi-dist`` rolling batcher. + - ``OPTION_SPECULATIVE_CONFIG``: a JSON blob read by the current vLLM + engine, which does not consume ``OPTION_SPECULATIVE_DRAFT_MODEL``. + Without this, the model deploys but vLLM leaves speculative decoding + disabled (``speculative_config=None``). + + Args: + draft_model_path (str): Draft-model path or identifier passed to the container. + speculative_decoding_config (Optional[Dict]): The speculative decoding config; read + for ``NumSpeculativeTokens`` when present. + """ + num_speculative_tokens = DEFAULT_NUM_SPECULATIVE_TOKENS + if speculative_decoding_config: + num_speculative_tokens = speculative_decoding_config.get( + "NumSpeculativeTokens", num_speculative_tokens + ) + + self.env_vars.update( + { + "OPTION_SPECULATIVE_DRAFT_MODEL": draft_model_path, + "OPTION_SPECULATIVE_CONFIG": json.dumps( + { + "method": "draft_model", + "model": draft_model_path, + "num_speculative_tokens": num_speculative_tokens, + } + ), + } + ) + def _generate_additional_model_data_sources( self, model_source: str, @@ -1963,7 +2004,9 @@ def _custom_speculative_decoding( else: speculative_draft_model = additional_model_source - self.env_vars.update({"OPTION_SPECULATIVE_DRAFT_MODEL": speculative_draft_model}) + self._set_speculative_draft_model_env( + speculative_draft_model, speculative_decoding_config + ) self.add_tags( {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "custom"}, ) @@ -2034,8 +2077,8 @@ def _jumpstart_speculative_decoding( accept_eula, ) - self.env_vars.update( - {"OPTION_SPECULATIVE_DRAFT_MODEL": f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/"} + self._set_speculative_draft_model_env( + f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/", speculative_decoding_config ) self.add_tags( {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": "jumpstart"}, @@ -2267,8 +2310,8 @@ def _set_additional_model_source( "to `Auto` instead." ) - self.env_vars.update( - {"OPTION_SPECULATIVE_DRAFT_MODEL": f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/"} + self._set_speculative_draft_model_env( + f"{SPECULATIVE_DRAFT_MODEL}/{channel_name}/", speculative_decoding_config ) self.add_tags( {"Key": Tag.SPECULATIVE_DRAFT_MODEL_PROVIDER, "Value": model_provider}, diff --git a/sagemaker-serve/tests/unit/test_model_builder_utils_optimization.py b/sagemaker-serve/tests/unit/test_model_builder_utils_optimization.py index e3e37dc9a5..b4e8362218 100644 --- a/sagemaker-serve/tests/unit/test_model_builder_utils_optimization.py +++ b/sagemaker-serve/tests/unit/test_model_builder_utils_optimization.py @@ -3,11 +3,15 @@ Targets uncovered optimization and deployment config functionality. """ +import json import unittest from unittest.mock import Mock, patch, MagicMock import tempfile -from sagemaker.serve.model_builder_utils import _ModelBuilderUtils +from sagemaker.serve.model_builder_utils import ( + _ModelBuilderUtils, + DEFAULT_NUM_SPECULATIVE_TOKENS, +) from sagemaker.core.enums import Tag @@ -122,14 +126,64 @@ def test_custom_speculative_decoding_with_eula(self): utils.additional_model_data_sources = [] utils.env_vars = {} utils._tags = [] - + config = {"ModelSource": "s3://bucket/draft-model", "AcceptEula": True} - + utils._custom_speculative_decoding(config, False) - + self.assertEqual(len(utils.additional_model_data_sources), 1) self.assertIn("ModelAccessConfig", utils.additional_model_data_sources[0]["S3DataSource"]) + def test_custom_speculative_decoding_emits_speculative_config(self): + """Custom SD must emit OPTION_SPECULATIVE_CONFIG (read by the vLLM engine), + not only the legacy OPTION_SPECULATIVE_DRAFT_MODEL.""" + utils = _ModelBuilderUtils() + utils.additional_model_data_sources = [] + utils.env_vars = {} + utils._tags = [] + + utils._custom_speculative_decoding({"ModelSource": "/local/path/to/model"}, False) + + self.assertIn("OPTION_SPECULATIVE_CONFIG", utils.env_vars) + spec = json.loads(utils.env_vars["OPTION_SPECULATIVE_CONFIG"]) + self.assertEqual(spec["method"], "draft_model") + self.assertEqual(spec["model"], "/local/path/to/model") + self.assertEqual(spec["num_speculative_tokens"], DEFAULT_NUM_SPECULATIVE_TOKENS) + + def test_custom_speculative_decoding_num_tokens_override(self): + """NumSpeculativeTokens in the config overrides the default.""" + utils = _ModelBuilderUtils() + utils.additional_model_data_sources = [] + utils.env_vars = {} + utils._tags = [] + + utils._custom_speculative_decoding( + {"ModelSource": "/local/path/to/model", "NumSpeculativeTokens": 3}, False + ) + + spec = json.loads(utils.env_vars["OPTION_SPECULATIVE_CONFIG"]) + self.assertEqual(spec["num_speculative_tokens"], 3) + + +class TestSetSpeculativeDraftModelEnv(unittest.TestCase): + """Test _set_speculative_draft_model_env helper directly.""" + + def test_emits_both_env_vars(self): + """Helper emits the legacy path var AND the vLLM JSON config.""" + utils = _ModelBuilderUtils() + utils.env_vars = {} + + utils._set_speculative_draft_model_env("/opt/ml/additional-model-data-sources/draft_model") + + self.assertEqual( + utils.env_vars["OPTION_SPECULATIVE_DRAFT_MODEL"], + "/opt/ml/additional-model-data-sources/draft_model", + ) + spec = json.loads(utils.env_vars["OPTION_SPECULATIVE_CONFIG"]) + self.assertEqual(spec["method"], "draft_model") + self.assertEqual(spec["model"], "/opt/ml/additional-model-data-sources/draft_model") + self.assertEqual(spec["num_speculative_tokens"], DEFAULT_NUM_SPECULATIVE_TOKENS) + class TestJumpStartSpeculativeDecoding(unittest.TestCase): """Test _jumpstart_speculative_decoding method - skipped (requires ModelBuilder context)."""