minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Rust] Keep sampling and scheduler wire schemas in sync (#37222)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-08-31 12:31:07 -07:00
committed by GitHub
parent 48098b5f23
commit 2138494272
3 changed files with 144 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-server/src/message/sampling.rs
+93 -13
View File
@@ -28,6 +28,32 @@ const MAX_STOP_REGEX_LEN: usize = 256;
/// (`re._MAXCACHE`), so past that every pattern recompiles on every decode step.
const MAX_STOP_REGEX_COUNT: usize = 32;
/// JSON values accepted by Python's `CustomParamValue`: a scalar, a list of
/// scalars, or a string-keyed object whose values are scalars.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum CustomParamValue {
Null(()),
Bool(bool),
Signed(i64),
Unsigned(u64),
Float(f64),
String(String),
List(Vec<JsonScalar>),
Object(BTreeMap<String, JsonScalar>),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum JsonScalar {
Null(()),
Bool(bool),
Signed(i64),
Unsigned(u64),
Float(f64),
String(String),
}
/// One module per field default, each exposing the two hooks serde needs under
/// one name: `default` (key absent) and `deserialize` (key present — including
/// an explicit `null`, which Python's `__post_init__` maps back to the default:
@@ -180,10 +206,10 @@ pub struct SamplingParams {
pub logit_bias: Option<BTreeMap<String, f64>>,
#[serde(default)]
pub sampling_seed: Option<i64>,
/// Opaque JSON object forwarded to a custom logit processor. Python types it
/// as `Dict[str, JsonScalar | list | dict]`; it is never inspected here.
/// JSON object forwarded to a custom logit processor. Its values match
/// Python's `CustomParamValue` exactly.
#[serde(default)]
pub custom_params: Option<serde_json::Value>,
pub custom_params: Option<BTreeMap<String, CustomParamValue>>,
// Normalized internal fields.
//
@@ -329,6 +355,16 @@ impl SamplingParams {
if self.top_k == -1 {
self.top_k = TOP_K_ALL; // -1 disables top_k → whole vocabulary
}
for constraint in [
&mut self.json_schema,
&mut self.regex,
&mut self.ebnf,
&mut self.structural_tag,
] {
if constraint.as_deref() == Some("") {
*constraint = None;
}
}
}
/// Python `normalize(tokenizer)`: size the stop match windows, reject
@@ -402,6 +438,13 @@ impl SamplingParams {
/// Python `verify(vocab_size)` — the same ranges, messages and mutual
/// exclusions, plus the rust-server `n == 1` restriction.
fn verify(&self, vocab_size: u64) -> Result<(), Error> {
if let Some(beam_width) = self.beam_width
&& beam_width < 1
{
return Err(bad(format!(
"beam_width must be at least 1, got {beam_width}."
)));
}
if !self.temperature.is_finite() || self.temperature < 0.0 {
return Err(bad(format!(
"temperature must be a non-negative finite number, got {}",
@@ -477,13 +520,18 @@ impl SamplingParams {
}
}
// Grammars are mutually exclusive.
let grammars = [&self.json_schema, &self.regex, &self.ebnf]
.iter()
.filter(|g| g.is_some())
.count();
let grammars = [
&self.json_schema,
&self.regex,
&self.ebnf,
&self.structural_tag,
]
.iter()
.filter(|g| g.is_some())
.count();
if grammars > 1 {
return Err(bad(
"Only one of regex, json_schema, or ebnf can be set".into()
"Only one of json_schema, regex, ebnf, or structural_tag can be set".into(),
));
}
// Not a Python restriction: the rust from_scheduler maps one rid to one response,
@@ -497,11 +545,6 @@ impl SamplingParams {
)));
}
if let Some(beam_width) = self.beam_width {
if beam_width < 1 {
return Err(bad(format!(
"beam_width must be at least 1, got {beam_width}."
)));
}
// Also not a Python restriction: beam search returns its candidates
// in `meta_info.beam_results`, which from_scheduler does not carry.
if beam_width > 1 {
@@ -912,6 +955,43 @@ mod tests {
assert_eq!(once.top_k, twice.top_k);
}
#[test]
fn empty_grammar_constraints_are_unset() {
let sp = norm(r#"{"json_schema":"","regex":"","ebnf":"","structural_tag":""}"#);
assert!(sp.json_schema.is_none());
assert!(sp.regex.is_none());
assert!(sp.ebnf.is_none());
assert!(sp.structural_tag.is_none());
}
#[test]
fn structural_tag_is_mutually_exclusive_with_other_grammars() {
for field in ["json_schema", "regex", "ebnf"] {
let json = format!(r#"{{"structural_tag":"tag","{field}":"other"}}"#);
assert!(norm_err(&json).to_string().contains("Only one of"));
}
}
#[test]
fn custom_params_matches_python_shape() {
assert!(
norm(r#"{"custom_params":{"null":null,"bool":true,"int":1,"float":1.5,"str":"x","list":[1,"x",null],"object":{"x":1}}}"#)
.custom_params
.is_some()
);
for json in [
r#"{"custom_params":[]}"#,
r#"{"custom_params":{"nested_list":[[1]]}}"#,
r#"{"custom_params":{"nested_object":{"x":{"y":1}}}}"#,
] {
assert!(
serde_json::from_str::<SamplingParams>(json).is_err(),
"{json} must not parse"
);
}
}
/// `skip_tokenizer_init` has no tokenizer, so the text-matching stop features
/// and `min_new_tokens` (needs eos_token_id) are 400s, not silent no-ops.
/// Mirrors Python `raise_if_tokenizer_required`.
test/registered/unit/managers/test_io_struct.py
+29
View File
@@ -1,7 +1,9 @@
import copy
import re
import unittest
import weakref
from array import array
from pathlib import Path
import msgspec
import numpy as np
@@ -41,6 +43,33 @@ register_cpu_ci(est_time=6, suite="base-c-test-cpu")
class TestTokenizedReqInputMsgpack(unittest.TestCase):
def test_rust_tokenized_generate_schema_stays_in_lockstep(self):
"""Compare the Rust wire declaration with the imported Python schema."""
rust_path = (
Path(__file__).resolve().parents[4]
/ "rust/sglang-server/src/message/io_struct.rs"
)
source = rust_path.read_text()
start = source.index("pub(super) TokenizedGenerateReqInput<'a> {")
end = source.index("\n }\n}", start)
rust_fields = (
"rid",
"http_worker_ipc",
*re.findall(r"^\s*([a-z][a-z0-9_]*):", source[start:end], re.MULTILINE),
)
python_fields = TokenizedGenerateReqInput.__struct_fields__
self.assertEqual(python_fields[: len(rust_fields)], rust_fields)
self.assertTrue(
all(
default is not msgspec.NODEFAULT
for default in TokenizedGenerateReqInput.__struct_defaults__[
len(rust_fields) :
]
),
"Rust may omit only a defaulted suffix of the Python wire schema",
)
def _make_mm_inputs(self, device="cpu"):
return MultimodalProcessorOutput(
mm_items=[
test/registered/unit/sampling/test_sampling_params.py
+22 -6
View File
@@ -7,7 +7,9 @@ register_cpu_ci(est_time=8, suite="base-c-test-cpu")
register_xpu_ci(est_time=10, suite="stage-a-test-1-gpu-xpu")
import copy
import re
import unittest
from pathlib import Path
from unittest.mock import MagicMock
import msgspec
@@ -25,7 +27,6 @@ from sglang.test.test_utils import CustomTestCase
class TestSamplingParamsInit(CustomTestCase):
def test_zero_temperature_becomes_greedy(self):
"""Test greedy conversion when temperature is 0."""
sp = SamplingParams(temperature=0.0)
@@ -88,7 +89,6 @@ class TestSamplingParamsInit(CustomTestCase):
class TestSamplingParamsVerify(CustomTestCase):
VOCAB_SIZE = 32000
GRAMMAR_VALUES = {
"json_schema": '{"type":"object"}',
@@ -314,13 +314,12 @@ class TestSamplingParamsVerify(CustomTestCase):
class TestSamplingParamsNormalize(CustomTestCase):
def _mock_tokenizer(self, encode_map=None):
"""Create a mock tokenizer that returns predetermined token lists."""
tokenizer = MagicMock()
if encode_map:
tokenizer.encode.side_effect = (
lambda s, add_special_tokens=False: encode_map.get(s, [1])
tokenizer.encode.side_effect = lambda s, add_special_tokens=False: (
encode_map.get(s, [1])
)
else:
tokenizer.encode.return_value = [1] # Default: 1 token
@@ -428,6 +427,24 @@ class TestSamplingParamsNormalize(CustomTestCase):
class TestSamplingParamsMsgspecStruct(CustomTestCase):
def test_rust_sampling_schema_stays_in_lockstep(self):
"""Compare Rust fields with the imported Python wire schema."""
rust_path = (
Path(__file__).resolve().parents[4]
/ "rust/sglang-server/src/message/sampling.rs"
)
source = rust_path.read_text()
start = source.index("pub struct SamplingParams {")
end = source.index("\n}\n\n/// The `/generate`", start)
rust_fields = tuple(
re.findall(
r"^\s*pub ([a-z][a-z0-9_]*):",
source[start:end],
re.MULTILINE,
)
)
self.assertEqual(SamplingParams.__struct_fields__, rust_fields)
def test_copy_remains_mutable_and_independent(self):
sp = SamplingParams(max_new_tokens=8, custom_params={"a": 1})
@@ -501,7 +518,6 @@ class TestSamplingParamsMsgspecStruct(CustomTestCase):
class TestRegexMaxLength(CustomTestCase):
def test_literal_string(self):
"""Test that plain string 'abc' gives max length 3."""
self.assertEqual(get_max_seq_length("abc"), 3)