API Perf: Replace pydantic per-element validation with C loop validation (#26355)

benchmark/io/bench_input_ids_validator.py

@@ -0,0 +1,67 @@
+"""Microbenchmark: cost of validating `input_ids` during FastAPI body binding.
+
+Compares two validators on a single `input_ids` field:
+- GenerateReqInputPydanticValidator: default pydantic walk (per-element type check)
+- GenerateReqInputCustomValidator: C-loop validator (validate_optional_list_i64_1d_2d)
+
+Usage:
+    python benchmark/io/bench_input_ids_validator.py
+"""
+
+import time
+from dataclasses import dataclass
+from typing import Annotated, List, Optional, Union
+
+from pydantic import PlainValidator, TypeAdapter
+
+from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d
+
+
+@dataclass
+class GenerateReqInputPydanticValidator:
+    """Default pydantic — walks every element of input_ids to type-check."""
+
+    input_ids: Optional[Union[List[List[int]], List[int]]] = None
+
+
+@dataclass
+class GenerateReqInputCustomValidator:
+    """C-loop validator via PlainValidator."""
+
+    input_ids: Annotated[
+        Optional[Union[List[List[int]], List[int]]],
+        PlainValidator(validate_optional_list_i64_1d_2d),
+    ] = None
+
+
+_ta_pydantic = TypeAdapter(GenerateReqInputPydanticValidator)
+_ta_custom = TypeAdapter(GenerateReqInputCustomValidator)
+
+
+def _time(fn, n_iter=30):
+    t0 = time.perf_counter()
+    for _ in range(n_iter):
+        fn()
+    return (time.perf_counter() - t0) * 1000 / n_iter
+
+
+def main():
+    print(
+        f"{'n_tokens':>9s} | {'default pydantic (ms)':>22s} | "
+        f"{'rigid i64 validator (ms)':>26s}"
+    )
+    print("-" * 65)
+
+    for n in [1_000, 10_000, 100_000, 1_000_000]:
+        d = {"input_ids": list(range(1, n + 1))}
+
+        p1 = _time(lambda: _ta_pydantic.validate_python(d))
+        p2 = _time(lambda: _ta_custom.validate_python(d))
+
+        print(f"{n:>9d} | {p1:>22.3f} | {p2:>26.3f}")
+
+    print("\nLegend: mean over 30 iters, in ms.")
+
+
+if __name__ == "__main__":
+    main()

python/sglang/srt/managers/io_struct.py

@@ -25,9 +25,10 @@ from array import array
 from collections import Counter
 from dataclasses import dataclass, field
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union
+from typing import TYPE_CHECKING, Annotated, Any, Dict, List, Literal, Optional, Union
 
 import torch
+from pydantic import PlainValidator
 
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.managers.embed_types import PositionalEmbeds
@@ -40,6 +41,7 @@ from sglang.srt.observability.req_time_stats import (
 )
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.utils import ImageData, VideoData
+from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d
 
 # Handle serialization of Image for pydantic
 if TYPE_CHECKING:
@@ -136,8 +138,13 @@ MultimodalDataInputFormat = Union[
 class GenerateReqInput(BaseReq):
     # The input prompt. It can be a single prompt or a batch of prompts.
     text: Optional[Union[List[str], str]] = None
-    # The token ids for text; one can specify either text or input_ids
+    # The token ids for text.
-    input_ids: Optional[Union[List[List[int]], List[int]]] = None
+    #
+    # Use C-loop validator to replace Pydantic per-element type check for efficiency.
+    input_ids: Annotated[
+        Optional[Union[List[List[int]], List[int]]],
+        PlainValidator(validate_optional_list_i64_1d_2d),
+    ] = None
     # The embeddings for input_ids; one can specify either text or input_ids or input_embeds.
     input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None
     # The image input. It can be an image instance, file name, URL, or base64 encoded string.

python/sglang/srt/utils/field_validators.py

@@ -0,0 +1,81 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Lightweight, reusable validators for hot-path API fields.
+
+These are intended to be paired with ``pydantic.PlainValidator`` on
+dataclass fields whose JSON shape is large or homogeneously typed, where
+pydantic's default per-element walk has been measured to dominate
+request latency.
+
+Usage::
+
+    from typing import Annotated, List, Optional, Union
+    from pydantic import PlainValidator
+    from sglang.srt.utils.field_validators import validate_optional_list_i64_1d_2d
+
+    @dataclass
+    class MyReq:
+        input_ids: Annotated[
+            Optional[Union[List[List[int]], List[int]]],
+            PlainValidator(validate_optional_list_i64_1d_2d),
+        ] = None
+"""
+
+from __future__ import annotations
+
+from array import array
+from typing import Any
+
+
+def validate_list_i64_1d(v: Any) -> list[int]:
+    """Validates type: list[int]"""
+    if v is None:
+        raise ValueError("must not be None")
+    if not isinstance(v, list):
+        raise ValueError(f"must be list; got {type(v).__name__}")
+    if not v:
+        return v
+    if not isinstance(v[0], int):
+        raise ValueError(f"elements must be int; got {type(v[0]).__name__}")
+    try:
+        array("q", v)
+    except (TypeError, OverflowError) as e:
+        raise ValueError(f"contains non-int64 element: {e}") from None
+    return v
+
+
+def validate_optional_list_i64_1d_2d(
+    v: Any,
+) -> list[int] | list[list[int]] | None:
+    """Validates type: list[int] | list[list[int]] | None"""
+    if v is None:
+        # Accept None
+        return v
+    if not isinstance(v, list):
+        raise ValueError(f"must be list or null; got {type(v).__name__}")
+    if not v:
+        # Accept empty list
+        return v
+    if isinstance(v[0], int):
+        # Accept list[int]
+        return validate_list_i64_1d(v)
+    if isinstance(v[0], list):
+        # Accept list[list[int]]
+        for i, row in enumerate(v):
+            try:
+                validate_list_i64_1d(row)
+            except ValueError as e:
+                raise ValueError(f"row {i}: {e}") from None
+        return v
+    raise ValueError(f"elements must be int or list; got {type(v[0]).__name__}")

test/registered/unit/utils/test_field_validators.py

@@ -0,0 +1,101 @@
+import unittest
+
+from sglang.srt.utils.field_validators import (
+    validate_list_i64_1d,
+    validate_optional_list_i64_1d_2d,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+class TestValidateListI64_1d(CustomTestCase):
+    """`validate_list_i64_1d` rejects anything that array('q', v) can't accept."""
+
+    def test_accept_int_list(self):
+        v = [1, 2, 3]
+        self.assertIs(validate_list_i64_1d(v), v)
+
+    def test_accept_empty_list(self):
+        v = []
+        self.assertIs(validate_list_i64_1d(v), v)
+
+    def test_accept_int64_boundaries(self):
+        for v in ([-(2**63)], [2**63 - 1], [0]):
+            self.assertIs(validate_list_i64_1d(v), v)
+
+    def test_reject_none(self):
+        with self.assertRaisesRegex(ValueError, "must not be None"):
+            validate_list_i64_1d(None)
+
+    def test_reject_non_list(self):
+        for v in ((1, 2, 3), "abc", {1: 2}, 42, 3.14):
+            with self.subTest(v=v):
+                with self.assertRaisesRegex(ValueError, "must be list"):
+                    validate_list_i64_1d(v)
+
+    def test_reject_non_int_first_element(self):
+        for v in ([1.5, 2, 3], ["a", "b"], [None, 1]):
+            with self.subTest(v=v):
+                with self.assertRaisesRegex(ValueError, "elements must be int"):
+                    validate_list_i64_1d(v)
+
+    def test_reject_overflow_int64(self):
+        # 2**63 overflows signed int64.
+        with self.assertRaisesRegex(ValueError, "non-int64 element"):
+            validate_list_i64_1d([0, 2**63])
+
+    def test_reject_non_int_later_element(self):
+        # First-element fast path passes but C loop rejects later float.
+        with self.assertRaisesRegex(ValueError, "non-int64 element"):
+            validate_list_i64_1d([1, 2, 3.5])
+
+
+class TestValidateOptionalListI64_1d_2d(CustomTestCase):
+    """`validate_optional_list_i64_1d_2d` accepts None | [] | list[int] | list[list[int]]."""
+
+    def test_accept_none(self):
+        self.assertIsNone(validate_optional_list_i64_1d_2d(None))
+
+    def test_accept_empty_list(self):
+        v = []
+        self.assertIs(validate_optional_list_i64_1d_2d(v), v)
+
+    def test_accept_1d_int_list(self):
+        v = [1, 2, 3]
+        self.assertIs(validate_optional_list_i64_1d_2d(v), v)
+
+    def test_accept_2d_int_list(self):
+        v = [[1, 2], [3, 4, 5]]
+        self.assertIs(validate_optional_list_i64_1d_2d(v), v)
+
+    def test_accept_2d_with_empty_row(self):
+        v = [[], [1, 2]]
+        self.assertIs(validate_optional_list_i64_1d_2d(v), v)
+
+    def test_reject_non_list_top_level(self):
+        for v in ((1, 2), "abc", 42, 3.14, {1: 2}):
+            with self.subTest(v=v):
+                with self.assertRaisesRegex(ValueError, "must be list or null"):
+                    validate_optional_list_i64_1d_2d(v)
+
+    def test_reject_mixed_first_element_type(self):
+        with self.assertRaisesRegex(ValueError, "elements must be int or list"):
+            validate_optional_list_i64_1d_2d([1.5, 2.5])
+
+    def test_reject_overflow_in_1d(self):
+        with self.assertRaisesRegex(ValueError, "non-int64 element"):
+            validate_optional_list_i64_1d_2d([0, 2**63])
+
+    def test_reject_bad_row_in_2d_reports_index(self):
+        with self.assertRaisesRegex(ValueError, "row 1:"):
+            validate_optional_list_i64_1d_2d([[1, 2], [3, "x"]])
+
+    def test_reject_overflow_in_2d_reports_index(self):
+        with self.assertRaisesRegex(ValueError, "row 0:.*non-int64"):
+            validate_optional_list_i64_1d_2d([[2**63], [1, 2]])
+
+
+if __name__ == "__main__":
+    unittest.main()