Convert SamplingParams to msgspec Struct (#29198)

Co-authored-by: Rain Jiang <96632942+rainj-me@users.noreply.github.com>

python/sglang/srt/sampling/sampling_params.py

@@ -15,7 +15,9 @@
 
 import logging
 import math
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Set, Union
+
+import msgspec
 
 # sre_parse is deprecated in Python 3.11+, use re._parser instead
 try:
@@ -59,7 +61,7 @@ def raise_if_tokenizer_required(
         )
 
 
-class SamplingParams:
+class SamplingParams(msgspec.Struct, kw_only=True, omit_defaults=True):
     """
     The sampling parameters.
 
@@ -68,77 +70,89 @@ class SamplingParams:
     for the documentation.
     """
 
-    def __init__(
-        self,
-        max_new_tokens: int = 128,
-        stop: Optional[Union[str, List[str]]] = None,
-        stop_token_ids: Optional[List[int]] = None,
-        stop_regex: Optional[Union[str, List[str]]] = None,
-        temperature: float = 1.0,
-        top_p: float = 1.0,
-        top_k: int = -1,
-        min_p: float = 0.0,
-        frequency_penalty: float = 0.0,
-        presence_penalty: float = 0.0,
-        repetition_penalty: float = 1.0,
-        min_new_tokens: int = 0,
-        n: int = 1,
-        json_schema: Optional[str] = None,
-        regex: Optional[str] = None,
-        ebnf: Optional[str] = None,
-        structural_tag: Optional[str] = None,
-        ignore_eos: bool = False,
-        skip_special_tokens: bool = True,
-        spaces_between_special_tokens: bool = True,
-        no_stop_trim: bool = False,
-        custom_params: Optional[Dict[str, Any]] = None,
-        stream_interval: Optional[int] = None,
-        logit_bias: Optional[Dict[str, float]] = None,
-        sampling_seed: Optional[int] = None,
-    ) -> None:
+    # --- API parameters (set by callers) ---
+    max_new_tokens: Optional[int] = 128
+    stop: Optional[Union[str, List[str]]] = (
+        None  # API input alias, copied to stop_strs then cleared in normalize()
+    )
+    stop_token_ids: Optional[Set[int]] = None
+    stop_regex: Optional[Union[str, List[str]]] = (
+        None  # API input alias, copied to stop_regex_strs then cleared in normalize()
+    )
+    temperature: float = 1.0
+    top_p: float = 1.0
+    top_k: int = TOP_K_ALL
+    min_p: float = 0.0
+    frequency_penalty: float = 0.0
+    presence_penalty: float = 0.0
+    repetition_penalty: float = 1.0
+    min_new_tokens: int = 0
+    n: int = 1
+    json_schema: Optional[str] = None
+    regex: Optional[str] = None
+    ebnf: Optional[str] = None
+    structural_tag: Optional[str] = None
+    ignore_eos: bool = False
+    skip_special_tokens: bool = True
+    spaces_between_special_tokens: bool = True
+    no_stop_trim: bool = False
+    custom_params: Optional[Dict[str, Any]] = None
+    stream_interval: Optional[int] = None
+    logit_bias: Optional[Dict[str, float]] = None
+    sampling_seed: Optional[int] = None
+
+    # --- Internal fields (populated by __post_init__ or normalize(), not API-facing) ---
+    stop_strs: Optional[Union[str, List[str]]] = None  # from stop
+    stop_regex_strs: Optional[Union[str, List[str]]] = None  # from stop_regex
+    stop_str_max_len: int = 0  # set by normalize()
+    stop_regex_max_len: int = 0  # set by normalize()
+    is_normalized: bool = False  # set by normalize()
+
+    def __post_init__(self):
         # For non-optional params, treat None as "use default" so that callers
         # (e.g. /generate) can pass null without crashing verify().
-        self.max_new_tokens = max_new_tokens
-        self.stop_strs = stop
-        if stop_token_ids:
-            filtered = {int(t) for t in stop_token_ids if t is not None}
+
+        # msgspec calls __post_init__ after deserialization. Once normalize()
+        # has populated tokenizer-derived fields, avoid resetting them.
+        if self.is_normalized:
+            return
+
+        self.stop_strs = self.stop
+        if self.stop_token_ids:
+            filtered = {int(t) for t in self.stop_token_ids if t is not None}
             self.stop_token_ids = filtered or None
         else:
             self.stop_token_ids = None
-        self.stop_regex_strs = stop_regex
-        self.temperature = temperature if temperature is not None else 1.0
-        self.top_p = top_p if top_p is not None else 1.0
-        self.top_k = top_k if top_k is not None else -1
-        self.min_p = min_p if min_p is not None else 0.0
+        self.stop_regex_strs = self.stop_regex
+        self.temperature = self.temperature if self.temperature is not None else 1.0
+        self.top_p = self.top_p if self.top_p is not None else 1.0
+        self.top_k = self.top_k if self.top_k is not None else -1
+        self.min_p = self.min_p if self.min_p is not None else 0.0
         self.frequency_penalty = (
-            frequency_penalty if frequency_penalty is not None else 0.0
+            self.frequency_penalty if self.frequency_penalty is not None else 0.0
         )
         self.presence_penalty = (
-            presence_penalty if presence_penalty is not None else 0.0
+            self.presence_penalty if self.presence_penalty is not None else 0.0
         )
         self.repetition_penalty = (
-            repetition_penalty if repetition_penalty is not None else 1.0
+            self.repetition_penalty if self.repetition_penalty is not None else 1.0
         )
-        self.min_new_tokens = min_new_tokens if min_new_tokens is not None else 0
-        self.regex = regex
-        self.n = n if n is not None else 1
-        self.json_schema = json_schema
-        self.ebnf = ebnf
-        self.structural_tag = structural_tag
-        self.ignore_eos = ignore_eos if ignore_eos is not None else False
+        self.min_new_tokens = (
+            self.min_new_tokens if self.min_new_tokens is not None else 0
+        )
+        self.n = self.n if self.n is not None else 1
+        self.ignore_eos = self.ignore_eos if self.ignore_eos is not None else False
         self.skip_special_tokens = (
-            skip_special_tokens if skip_special_tokens is not None else True
+            self.skip_special_tokens if self.skip_special_tokens is not None else True
         )
         self.spaces_between_special_tokens = (
-            spaces_between_special_tokens
-            if spaces_between_special_tokens is not None
+            self.spaces_between_special_tokens
+            if self.spaces_between_special_tokens is not None
             else True
         )
-        self.no_stop_trim = no_stop_trim if no_stop_trim is not None else False
-        self.custom_params = custom_params
-        self.stream_interval = stream_interval
-        self.logit_bias = logit_bias
-        self.sampling_seed = sampling_seed
+        self.no_stop_trim = (
+            self.no_stop_trim if self.no_stop_trim is not None else False
+        )
 
         # Process some special cases
         if 0 <= self.temperature < _SAMPLING_EPS:
@@ -245,6 +259,11 @@ class SamplingParams:
             tokenizer, self.stop_strs, self.stop_regex_strs, self.min_new_tokens
         )
 
+        # Clear API input aliases so omit_defaults=True drops them from the wire.
+        self.stop = None
+        self.stop_regex = None
+        self.is_normalized = True
+
 
 # This function gets a strict upperbound on the maximum number of tokens that would need
 # to be buffered to match the input regex string

scripts/ci/utils/compute_partitions.py

@@ -38,7 +38,7 @@ _TARGET_BACKENDS = {HWBackend.CUDA, HWBackend.CPU}
 # base-a is the critical-path entry gate; pin its fanout to smoke-coverage
 # defaults instead of est_time. max_parallel = size (no throttle).
 _BASE_A_OVERRIDES = {
-    "base-a-test-cpu": 4,
+    "base-a-test-cpu": 8,
     "base-a-test-1-gpu-small": 1,
 }
 

test/registered/unit/sampling/test_sampling_params.py

@@ -6,9 +6,12 @@ register_cpu_ci(est_time=7, suite="base-a-test-cpu")
 register_cpu_ci(est_time=7, suite="base-c-test-cpu")
 register_xpu_ci(est_time=10, suite="stage-a-test-1-gpu-xpu")
 
+import copy
 import unittest
 from unittest.mock import MagicMock
 
+import msgspec
+
 from sglang.srt.sampling.sampling_params import (
     MAX_LEN,
     TOP_K_ALL,
@@ -363,6 +366,84 @@ class TestSamplingParamsNormalize(CustomTestCase):
         self.assertEqual(sp.stop_regex_max_len, 3)
 
 
+class TestSamplingParamsMsgspecStruct(CustomTestCase):
+
+    def test_copy_remains_mutable_and_independent(self):
+        sp = SamplingParams(max_new_tokens=8, custom_params={"a": 1})
+
+        copied = copy.copy(sp)
+        copied.max_new_tokens = 16
+        copied.custom_params = {"b": 2}
+
+        self.assertEqual(sp.max_new_tokens, 8)
+        self.assertEqual(sp.custom_params, {"a": 1})
+        self.assertEqual(copied.max_new_tokens, 16)
+        self.assertEqual(copied.custom_params, {"b": 2})
+
+    def test_none_values_still_use_constructor_defaults(self):
+        sp = SamplingParams(
+            temperature=None,
+            top_p=None,
+            top_k=None,
+            min_p=None,
+            frequency_penalty=None,
+            presence_penalty=None,
+            repetition_penalty=None,
+            min_new_tokens=None,
+            n=None,
+            ignore_eos=None,
+            skip_special_tokens=None,
+            spaces_between_special_tokens=None,
+            no_stop_trim=None,
+        )
+
+        self.assertEqual(sp.temperature, 1.0)
+        self.assertEqual(sp.top_p, 1.0)
+        self.assertEqual(sp.top_k, TOP_K_ALL)
+        self.assertEqual(sp.min_p, 0.0)
+        self.assertEqual(sp.frequency_penalty, 0.0)
+        self.assertEqual(sp.presence_penalty, 0.0)
+        self.assertEqual(sp.repetition_penalty, 1.0)
+        self.assertEqual(sp.min_new_tokens, 0)
+        self.assertEqual(sp.n, 1)
+        self.assertFalse(sp.ignore_eos)
+        self.assertTrue(sp.skip_special_tokens)
+        self.assertTrue(sp.spaces_between_special_tokens)
+        self.assertFalse(sp.no_stop_trim)
+
+    def test_msgpack_omits_default_fields(self):
+        encoded = msgspec.msgpack.encode(SamplingParams())
+
+        self.assertEqual(msgspec.msgpack.decode(encoded), {})
+
+    def test_msgpack_round_trip_preserves_normalized_state(self):
+        tokenizer = MagicMock()
+        tokenizer.encode.side_effect = lambda s, add_special_tokens=False: {
+            "hello": [101, 102],
+            "world": [201],
+        }[s]
+        sp = SamplingParams(
+            stop=["hello", "world"],
+            stop_regex=r"[a-z]{3}",
+            stop_token_ids=[1, 2],
+            temperature=0.5,
+        )
+        sp.normalize(tokenizer)
+
+        encoder = msgspec.msgpack.Encoder()
+        decoder = msgspec.msgpack.Decoder(SamplingParams)
+        rebuilt = decoder.decode(encoder.encode(sp))
+
+        self.assertIsInstance(rebuilt, SamplingParams)
+        self.assertTrue(rebuilt.is_normalized)
+        self.assertEqual(rebuilt.stop_strs, ["hello", "world"])
+        self.assertEqual(rebuilt.stop_str_max_len, 2)
+        self.assertEqual(rebuilt.stop_regex_strs, [r"[a-z]{3}"])
+        self.assertEqual(rebuilt.stop_regex_max_len, 3)
+        self.assertEqual(rebuilt.stop_token_ids, {1, 2})
+        self.assertEqual(rebuilt.temperature, 0.5)
+
+
 class TestRegexMaxLength(CustomTestCase):
 
     def test_literal_string(self):