sglang 0.5.3rc2__py3-none-any.whl → 0.5.4.post1__py3-none-any.whl

sglang/srt/sampling/custom_logit_processor.py

@@ -1,18 +1,22 @@
 import json
 from abc import ABC, abstractmethod
 from functools import lru_cache
-from typing import Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 import dill
+import orjson
 import torch
 
+if TYPE_CHECKING:
+    from sglang.srt.managers.schedule_batch import Req
+
 
 @lru_cache(maxsize=None)
 def _cache_from_str(json_str: str):
     """Deserialize a json string to a Callable object.
     This function is cached to avoid redundant deserialization.
     """
-    data = json.loads(json_str)
+    data = orjson.loads(json_str)
     return dill.loads(bytes.fromhex(data["callable"]))
 
 
@@ -51,3 +55,74 @@ class DisallowedTokensLogitsProcessor(CustomLogitProcessor):
         ), f"{custom_param_list=}"
         logits[..., disallowed_token_ids] = -float("inf")
         return logits
+
+
+class ThinkingBudgetLogitProcessor(CustomLogitProcessor):
+    """A logit processor that controls the length of thinking."""
+
+    THINKING_START_TOKEN_ID: int
+    THINKING_END_TOKEN_ID: int
+    NEW_LINE_TOKEN_ID: int
+
+    def __call__(self, logits, custom_param_list: list[dict[str, Any]]):
+        if custom_param_list is None or not custom_param_list:
+            return logits
+        for i, param_dict in enumerate(custom_param_list):
+            if param_dict is None:
+                continue
+
+            thinking_budget: int | None = param_dict.get("thinking_budget")
+
+            # Skip if thinking_budget is unset, or not an integer, or negative
+            if (
+                thinking_budget is None
+                or not isinstance(thinking_budget, int)
+                or thinking_budget < 0
+            ):
+                continue
+            req: Req = param_dict.get("__req__")
+            cur_ids: list[int] = [*req.origin_input_ids, *req.output_ids]
+
+            # Check if out of thinking stage
+            if (
+                self.THINKING_START_TOKEN_ID not in cur_ids
+                or self.THINKING_END_TOKEN_ID in cur_ids
+            ):
+                continue
+
+            # Find the index of the thinking start token
+            start_index = cur_ids.index(self.THINKING_START_TOKEN_ID)
+
+            # Count the number of tokens after the thinking start token
+            num_tokens_after_start = len(cur_ids) - start_index - 1
+
+            if num_tokens_after_start < thinking_budget:
+                continue
+
+            # Ensure new line token before thinking end token
+            if not req.output_ids or req.output_ids[-1] != self.NEW_LINE_TOKEN_ID:
+                logits[i, :] = -float("inf")
+                logits[i, self.NEW_LINE_TOKEN_ID] = 0.0
+                continue
+
+            # Assign highest probability to the thinking end token
+            logits[i, :] = -float("inf")
+            logits[i, self.THINKING_END_TOKEN_ID] = 0.0
+
+        return logits
+
+
+class Qwen3ThinkingBudgetLogitProcessor(ThinkingBudgetLogitProcessor):
+    """A logit processor that controls the length of thinking for Qwen3 models."""
+
+    THINKING_START_TOKEN_ID: int = 151667
+    THINKING_END_TOKEN_ID: int = 151668
+    NEW_LINE_TOKEN_ID: int = 198
+
+
+class DeepSeekR1ThinkingBudgetLogitProcessor(ThinkingBudgetLogitProcessor):
+    """A logit processor that controls the length of thinking for DeepSeek-R1 models."""
+
+    THINKING_START_TOKEN_ID: int = 128798
+    THINKING_END_TOKEN_ID: int = 128799
+    NEW_LINE_TOKEN_ID: int = 201

sglang/srt/sampling/sampling_batch_info.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import dataclasses
 import logging
-import threading
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
 
 import torch
@@ -10,6 +9,7 @@ import torch
 import sglang.srt.sampling.penaltylib as penaltylib
 from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
 from sglang.srt.sampling.sampling_params import TOP_K_ALL
+from sglang.srt.server_args import get_global_server_args
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -44,12 +44,9 @@ class SamplingBatchInfo:
     vocab_mask: Optional[torch.Tensor] = None
     apply_mask_func: Optional[Callable[[torch.Tensor, torch.Tensor], None]] = None
 
-    # An event used for overlap schedule
-    sampling_info_done: Optional[threading.Event] = None
-
     # Penalizer
     penalizer_orchestrator: Optional[penaltylib.BatchedPenalizerOrchestrator] = None
-    linear_penalty: torch.Tensor = None
+    acc_linear_penalties: torch.Tensor = None  # Used in the overlap mode
 
     # Whether any request has custom logit processor
     has_custom_logit_processor: bool = False
@@ -69,16 +66,10 @@ class SamplingBatchInfo:
     # Handle logit bias
     logit_bias: Optional[torch.Tensor] = None
 
-    @classmethod
-    def _get_global_server_args_dict(cls):
-        from sglang.srt.managers.schedule_batch import global_server_args_dict
-
-        return global_server_args_dict
-
     @classmethod
     def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
-        global_server_args_dict = cls._get_global_server_args_dict()
-        enable_deterministic = global_server_args_dict["enable_deterministic_inference"]
+        global_server_args = get_global_server_args()
+        enable_deterministic = global_server_args.enable_deterministic_inference
 
         reqs = batch.reqs
         device = batch.device
@@ -115,10 +106,9 @@ class SamplingBatchInfo:
                         logit_bias[i, int(key)] = value
 
         # Check if any request has custom logit processor
-        has_custom_logit_processor = global_server_args_dict[
-            "enable_custom_logit_processor"
-        ] and any(  # check the flag first.
-            r.custom_logit_processor for r in reqs
+        has_custom_logit_processor = (
+            global_server_args.enable_custom_logit_processor
+            and any(r.custom_logit_processor for r in reqs)  # check the flag first.
         )  # then check the requests.
 
         if has_custom_logit_processor:
@@ -217,19 +207,19 @@ class SamplingBatchInfo:
 
     def update_penalties(self):
         if self.penalizer_orchestrator.is_required:
-            self.linear_penalty = torch.zeros(
+            self.acc_linear_penalties = torch.zeros(
                 (len(self.temperatures), self.vocab_size),
                 dtype=torch.float32,
                 device=self.temperatures.device,
             )
-            self.penalizer_orchestrator.apply(self.linear_penalty)
+            self.penalizer_orchestrator.apply(self.acc_linear_penalties)
         else:
-            self.linear_penalty = None
+            self.acc_linear_penalties = None
 
     def apply_logits_bias(self, logits: torch.Tensor):
-        if self.linear_penalty is not None:
+        if self.acc_linear_penalties is not None:
             # Used in the overlap mode
-            logits.add_(self.linear_penalty)
+            logits.add_(self.acc_linear_penalties)
 
         if self.penalizer_orchestrator and self.penalizer_orchestrator.is_required:
             # Used in the non-overlap mode
@@ -370,6 +360,11 @@ class SamplingBatchInfo:
         self.need_top_k_sampling |= other.need_top_k_sampling
         self.need_min_p_sampling |= other.need_min_p_sampling
 
+    def copy_for_forward(self):
+        # Accumulate the penalty into a pre-allocated buffer to get rid of the dependency of `penalizer_orchestrator` later
+        self.update_penalties()
+        return dataclasses.replace(self, penalizer_orchestrator=None)
+
 
 def merge_bias_tensor(
     lhs: Optional[torch.Tensor],

sglang/srt/sampling/sampling_params.py

@@ -13,13 +13,15 @@
 # ==============================================================================
 """Sampling parameters for text generation."""
 
+import logging
+import sre_parse
 from typing import Any, Dict, List, Optional, Union
 
-from sglang.srt.utils import get_bool_env_var
-
 _SAMPLING_EPS = 1e-6
 TOP_K_ALL = 1 << 30
 
+logger = logging.getLogger(__name__)
+
 
 class SamplingParams:
     """
@@ -35,6 +37,7 @@ class SamplingParams:
         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,
@@ -63,6 +66,7 @@ class SamplingParams:
             self.stop_token_ids = set(stop_token_ids)
         else:
             self.stop_token_ids = None
+        self.stop_regex_strs = stop_regex
         self.temperature = temperature
         self.top_p = top_p
         self.top_k = top_k
@@ -170,3 +174,67 @@ class SamplingParams:
                 else:
                     stop_str_max_len = max(stop_str_max_len, len(stop_str))
             self.stop_str_max_len = stop_str_max_len
+
+        # Process stop regex strings
+        if self.stop_regex_strs is None:
+            self.stop_regex_strs = []
+            self.stop_regex_max_len = 0
+        else:
+            if isinstance(self.stop_regex_strs, str):
+                self.stop_regex_strs = [self.stop_regex_strs]
+
+            stop_regex_max_len = 0
+            for stop_regex in self.stop_regex_strs:
+                stop_regex_max_len = max(
+                    stop_regex_max_len, get_max_seq_length(stop_regex)
+                )
+
+            self.stop_regex_max_len = stop_regex_max_len
+
+
+# This function gets a strict upperbound on the maximum number of tokens that would need
+# to be buffered to match the input regex string
+# NOTE: in the worst case, one character that needs to be buffered corresponds to one
+# token
+def get_max_seq_length(regex_str: str):
+    return _max_length_from_subpattern(sre_parse.parse(regex_str))
+
+
+MAX_LEN = 2**30
+
+
+def _max_length_from_subpattern(subpattern: sre_parse.SubPattern):
+    total = 0
+    for token, value in subpattern:
+        if token in {
+            sre_parse.LITERAL,  # `value` is any one character
+            sre_parse.IN,  # Any character within `value`
+            sre_parse.ANY,  # "."
+        }:
+            total += 1
+        elif token == sre_parse.SUBPATTERN:
+            # EG: (a\d+) ->
+            # [(SUBPATTERN,
+            #   (1, 0, 0, [(LITERAL, 97),
+            #              (MAX_REPEAT, (1, MAXREPEAT, [(IN, [(CATEGORY, CATEGORY_DIGIT)])]))]))]
+            _, _, _, inner_subpattern = value
+            total += _max_length_from_subpattern(inner_subpattern)
+        elif token == sre_parse.BRANCH:
+            _, branches = value
+            total += max(_max_length_from_subpattern(branch) for branch in branches)
+        elif token in {sre_parse.MAX_REPEAT, sre_parse.MIN_REPEAT}:
+            _, max_num_repeat, inner_subpattern = value
+            if max_num_repeat == sre_parse.MAXREPEAT:
+                total += MAX_LEN
+            else:
+                total += max_num_repeat * _max_length_from_subpattern(inner_subpattern)
+        elif token == sre_parse.AT:
+            # These are zero-width assertions like ^, $, and \b that don't add to the max
+            # length
+            total += 0
+        else:
+            logger.warning(f"Got unhandled regex token: {token}")
+
+            total += MAX_LEN
+
+    return total