diff --git a/rust/sglang-server/src/message/sampling.rs b/rust/sglang-server/src/message/sampling.rs index 2498d565d..bc0d8c50b 100644 --- a/rust/sglang-server/src/message/sampling.rs +++ b/rust/sglang-server/src/message/sampling.rs @@ -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), + Object(BTreeMap), +} + +#[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>, #[serde(default)] pub sampling_seed: Option, - /// 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, + pub custom_params: Option>, // 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::(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`. diff --git a/test/registered/unit/managers/test_io_struct.py b/test/registered/unit/managers/test_io_struct.py index 759054a48..a44abaa8d 100644 --- a/test/registered/unit/managers/test_io_struct.py +++ b/test/registered/unit/managers/test_io_struct.py @@ -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=[ diff --git a/test/registered/unit/sampling/test_sampling_params.py b/test/registered/unit/sampling/test_sampling_params.py index 8f4b89e55..8f897000d 100644 --- a/test/registered/unit/sampling/test_sampling_params.py +++ b/test/registered/unit/sampling/test_sampling_params.py @@ -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)