Skip to content

fix(serve): emit OPTION_SPECULATIVE_CONFIG so vLLM enables speculative decoding #6143

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZealSV wants to merge 1 commit into
base: master
Choose a base branch
from
+106 −9
Open

fix(serve): emit OPTION_SPECULATIVE_CONFIG so vLLM enables speculative decoding #6143

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 48 additions & 5 deletions sagemaker-serve/src/sagemaker/serve/model_builder_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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"},
)
Expand Down Expand Up @@ -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"},
Expand Down Expand Up @@ -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},
Expand Down
62 changes: 58 additions & 4 deletions sagemaker-serve/tests/unit/test_model_builder_utils_optimization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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)."""
Expand Down
Loading