sglang 0.1.16__py3-none-any.whl → 0.1.17__py3-none-any.whl

sglang/lang/interpreter.py

@@ -6,6 +6,7 @@ import multiprocessing
 import queue
 import threading
 import uuid
+import warnings
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
 from typing import Any, Callable, Dict, List, Optional, Union
@@ -30,7 +31,11 @@ from sglang.lang.ir import (
     SglVarScopeEnd,
     SglVideo,
 )
-from sglang.utils import encode_image_base64, encode_video_base64, get_exception_traceback
+from sglang.utils import (
+    encode_image_base64,
+    encode_video_base64,
+    get_exception_traceback,
+)
 
 
 def run_internal(state, program, func_args, func_kwargs, sync):
@@ -61,7 +66,7 @@ def run_program(
         default_sampling_para,
         chat_template=None,
         stream=stream,
-        api_num_spec_tokens=program.api_num_spec_tokens,
+        num_api_spec_tokens=program.num_api_spec_tokens,
     )
     state = ProgramState(stream_executor)
 
@@ -173,7 +178,7 @@ class StreamExecutor:
         default_sampling_para,
         chat_template,
         stream,
-        api_num_spec_tokens=None,
+        num_api_spec_tokens=None,
         use_thread=True,
     ):
         self.sid = uuid.uuid4().hex
@@ -181,20 +186,16 @@ class StreamExecutor:
         self.arguments: Dict[str, Any] = arguments
         self.default_sampling_para = default_sampling_para
         self.stream = stream
-        self.api_num_spec_tokens = api_num_spec_tokens
 
         self.variables = {}  # Dict[name: str -> value: str]
         self.variable_event = {}  # Dict[name: str -> event: threading.Event]
         self.meta_info = {}  # Dict[name: str -> info: str]
         self.is_finished = False
-        self.error = None
+        self.error_ = None
 
         # For completion
         self.text_ = ""  # The full text
 
-        # For speculative execution
-        self.speculated_text = ""
-
         # For chat
         self.messages_ = []  # The messages in the OpenAI API format
         self.chat_template = chat_template or self.backend.get_chat_template()
@@ -208,6 +209,10 @@ class StreamExecutor:
         # For fork/join
         self.fork_start_text_pos = None
 
+        # For speculative execution
+        self.num_api_spec_tokens = num_api_spec_tokens
+        self.speculated_text = ""
+
         # Worker thread
         self.use_thread = use_thread
         if self.use_thread:
@@ -286,6 +291,8 @@ class StreamExecutor:
             exes[i].fork_start_text_pos = len(self.text_)
             exes[i].images_ = list(self.images_)
 
+            # TODO(ying): handle API speculative execution
+
         return exes
 
     def text(self):
@@ -296,6 +303,10 @@ class StreamExecutor:
         self.sync()
         return self.messages_
 
+    def error(self):
+        self.sync()
+        return self.error_
+
     def end(self):
         if self.use_thread:
             if self.worker.is_alive():
@@ -314,7 +325,7 @@ class StreamExecutor:
             try:
                 self._execute(expr)
             except Exception as e:
-                # print(f"Error in stream_executor: {get_exception_traceback()}")
+                warnings.warn(f"Error in stream_executor: {get_exception_traceback()}")
                 error = e
                 break
             self.queue.task_done()
@@ -334,7 +345,7 @@ class StreamExecutor:
             if self.stream_var_event:
                 for name in self.stream_var_event:
                     self.stream_var_event[name].set()
-            self.error = error
+            self.error_ = error
 
         if self.stream_text_event:
             self.stream_text_event.set()
@@ -383,12 +394,23 @@ class StreamExecutor:
         else:
             raise ValueError(f"Unknown type: {type(other)}")
 
-    def _execute_fill(self, value: str):
+    def _execute_fill(self, value: str, prefix=False):
         value = str(value)
+
+        if (
+            self.cur_role == "assistant"
+            and self.num_api_spec_tokens is not None
+            and self.backend.is_chat_model
+            and not prefix
+        ):
+            self.backend.spec_fill(value)
+            return
+
         if self.speculated_text.startswith(value):
             self.speculated_text = self.speculated_text[len(value) :]
         else:
             self.speculated_text = ""
+
         self.text_ += value
 
     def _execute_image(self, expr: SglImage):
@@ -413,65 +435,80 @@ class StreamExecutor:
         # if global_config.eager_fill_image:
         #     self.backend.fill_image(self)
 
+    def _spec_gen(self, sampling_params):
+        stop = sampling_params.stop
+        max_new_tokens = sampling_params.max_new_tokens
+        meta_info = {}
+
+        def regen():
+            nonlocal meta_info
+
+            sampling_params.max_new_tokens = max(
+                sampling_params.max_new_tokens, self.num_api_spec_tokens
+            )
+            sampling_params.stop = None
+            self.speculated_text, meta_info = self.backend.generate(
+                self, sampling_params=sampling_params
+            )
+
+        def find_stop():
+            if isinstance(stop, str):
+                return self.speculated_text.find(stop)
+            elif isinstance(stop, (tuple, list)):
+                pos = -1
+                for stop_str in stop:
+                    stop_pos = self.speculated_text.find(stop_str)
+                    if stop_pos != -1 and (pos == -1 or stop_pos < pos):
+                        pos = stop_pos
+                return pos
+            else:
+                raise Exception("Wrong type of stop in sampling parameters.")
+
+        if stop is None:
+            if len(self.speculated_text) < max_new_tokens:
+                regen()
+            comp = self.speculated_text[:max_new_tokens]
+            self.speculated_text = self.speculated_text[max_new_tokens:]
+        elif isinstance(stop, (str, list, tuple)):
+            if self.speculated_text == "":
+                regen()
+            stop_pos = find_stop()
+            if stop_pos == -1:
+                stop_pos = min(
+                    sampling_params.max_new_tokens,
+                    len(self.speculated_text),
+                )
+            comp = self.speculated_text[:stop_pos]
+            self.speculated_text = self.speculated_text[stop_pos:]
+        else:
+            raise ValueError("Wrong type of stop in sampling parameters.")
+
+        return comp, meta_info
+
     def _execute_gen(self, expr: SglGen):
         sampling_params = self._resolve_sampling_params(expr.sampling_params)
         name = expr.name
 
         if not self.stream:
-            if self.api_num_spec_tokens is not None:
-                stop = sampling_params.stop
-                max_new_tokens = sampling_params.max_new_tokens
-                meta_info = {}
-
-                def regen():
-                    sampling_params.max_new_tokens = max(
-                        sampling_params.max_new_tokens, self.api_num_spec_tokens
-                    )
-                    sampling_params.stop = None
-                    self.speculated_text, meta_info = self.backend.generate(
-                        self, sampling_params=sampling_params
-                    )
-
-                def find_stop():
-                    if isinstance(stop, str):
-                        return self.speculated_text.find(stop), len(stop)
-                    elif isinstance(stop, (tuple, list)):
-                        pos = -1
-                        stop_len = 0
-                        for stop_str in stop:
-                            stop_pos = self.speculated_text.find(stop_str)
-                            if stop_pos != -1 and (pos == -1 or stop_pos < pos):
-                                pos = stop_pos
-                                stop_len = len(stop_str)
-                        return pos, stop_len
-                    else:
-                        raise Exception("Wrong type of stop in sampling parameters.")
-
-                if stop is None:
-                    if len(self.speculated_text) < max_new_tokens:
-                        regen()
-                    comp = self.speculated_text[:max_new_tokens]
-                    self.speculated_text = self.speculated_text[max_new_tokens:]
-                elif isinstance(stop, (str, list, tuple)):
-                    if self.speculated_text == "":
-                        regen()
-                    stop_pos, stop_len = find_stop()
-                    if stop_pos == -1:
-                        stop_pos, stop_len = (
-                            min(
-                                sampling_params.max_new_tokens,
-                                len(self.speculated_text),
-                            ),
-                            0,
-                        )
-                    comp = self.speculated_text[:stop_pos]
-                    self.speculated_text = self.speculated_text[stop_pos:]
-                else:
-                    raise ValueError("Wrong type of stop in sampling parameters.")
-            else:
+            if self.num_api_spec_tokens is None:
                 comp, meta_info = self.backend.generate(
-                    self, sampling_params=sampling_params
+                    self,
+                    sampling_params=sampling_params,
                 )
+            else:
+                if self.backend.is_chat_model:
+                    # Speculative execution on models with only chat interface.
+                    # Store the calls into a temporary list.
+                    # They will be lazily executed later.
+                    comp, meta_info = self.backend.generate(
+                        self,
+                        sampling_params=sampling_params,
+                        spec_var_name=name,
+                    )
+                    return
+
+                else: # Speculative execution on models with completion interface
+                    comp, meta_info = self._spec_gen(sampling_params)
 
             self.text_ += comp
 
@@ -479,6 +516,9 @@ class StreamExecutor:
             self.meta_info[name] = meta_info
             self.variable_event[name].set()
         else:
+            assert (
+                self.num_api_spec_tokens is None
+            ), "stream is not supported with api speculative execution"
             generator = self.backend.generate_stream(
                 self, sampling_params=sampling_params
             )
@@ -534,10 +574,19 @@ class StreamExecutor:
 
         prefix, _ = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_)
 
-        self._execute_fill(prefix)
+        self._execute_fill(prefix, prefix=True)
         self.cur_role_begin_pos = len(self.text_)
 
     def _execute_role_end(self, expr: SglRoleEnd):
+        if (
+            self.cur_role == "assistant"
+            and self.num_api_spec_tokens is not None
+            and self.backend.is_chat_model
+        ):
+            # Execute the stored lazy generation calls
+            self.backend.role_end_generate(self)
+        self.cur_role = None
+
         new_text = self.text_[self.cur_role_begin_pos :].lstrip()
 
         _, suffix = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_)
@@ -564,8 +613,6 @@ class StreamExecutor:
             # OpenAI chat API format
             self.messages_.append({"role": expr.role, "content": new_text})
 
-        self.cur_role = None
-
     def _execute_var_scope_begin(self, expr: SglVarScopeBegin):
         self.variables[expr.name] = int(len(self.text_))
 
@@ -709,7 +756,7 @@ class ProgramState:
         return self.stream_executor.sync()
 
     def error(self):
-        return self.stream_executor.error
+        return self.stream_executor.error()
 
     def text_iter(self, var_name: Optional[str] = None):
         if self.stream_executor.stream:

sglang/lang/ir.py

@@ -81,6 +81,21 @@ class SglSamplingParams:
             "top_p": self.top_p,
             "top_k": self.top_k,
         }
+        
+    def to_litellm_kwargs(self):
+        if self.regex is not None:
+            warnings.warn(
+                "Regular expression is not supported in the LiteLLM backend."
+            )
+        return {
+            "max_tokens": self.max_new_tokens,
+            "stop": self.stop or None,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "frequency_penalty": self.frequency_penalty,
+            "presence_penalty": self.presence_penalty,
+        }
 
     def to_srt_kwargs(self):
         return {
@@ -97,9 +112,9 @@ class SglSamplingParams:
 
 
 class SglFunction:
-    def __init__(self, func, api_num_spec_tokens=None, bind_arguments=None):
+    def __init__(self, func, num_api_spec_tokens=None, bind_arguments=None):
         self.func = func
-        self.api_num_spec_tokens = api_num_spec_tokens
+        self.num_api_spec_tokens = num_api_spec_tokens
         self.bind_arguments = bind_arguments or {}
         self.pin_prefix_rid = None
 

sglang/srt/constrained/fsm_cache.py

@@ -6,6 +6,9 @@ class FSMCache(BaseCache):
     def __init__(self, tokenizer_path, tokenizer_args_dict, enable=True):
         super().__init__(enable=enable)
 
+        if tokenizer_path.endswith(".json"):
+            return
+
         from importlib.metadata import version
 
         if version("outlines") >= "0.0.35":

sglang/srt/flush_cache.py

@@ -13,4 +13,4 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     response = requests.get(args.url + "/flush_cache")
-    assert response.status_code == 200
+    assert response.status_code == 200

sglang/srt/hf_transformers_utils.py

@@ -3,7 +3,8 @@
 import json
 import os
 import warnings
-from typing import List, Optional, Tuple, Union
+import functools
+from typing import Optional, Union, AbstractSet, Collection, Literal
 
 from huggingface_hub import snapshot_download
 from transformers import (
@@ -84,6 +85,9 @@ def get_tokenizer(
     tokenizer_revision: Optional[str] = None,
     **kwargs,
 ) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
+    if tokenizer_name.endswith(".json"):
+        return TiktokenTokenizer(tokenizer_name)
+
     """Gets a tokenizer for the given model name via Huggingface."""
     if is_multimodal_model(tokenizer_name):
         processor = get_processor(
@@ -170,3 +174,73 @@ def get_processor(
         **kwargs,
     )
     return processor
+
+
+class TiktokenTokenizer:
+    def __init__(self, tokenizer_path):
+        import tiktoken
+        PAT_STR_B = r"""(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+
+        # Read JSON
+        name = "tmp-json"
+        with open(tokenizer_path, "rb") as fin:
+            tok_dict = json.load(fin)
+
+        mergeable_ranks = {
+            bytes(item["bytes"]): item["token"] for item in tok_dict["regular_tokens"]
+        }
+        special_tokens = {
+            bytes(item["bytes"]).decode(): item["token"] for item in tok_dict["special_tokens"]
+        }
+        assert tok_dict["word_split"] == "V1"
+
+        kwargs = {
+            "name": name,
+            "pat_str": tok_dict.get("pat_str", PAT_STR_B),
+            "mergeable_ranks": mergeable_ranks,
+            "special_tokens": special_tokens,
+        }
+        if "default_allowed_special" in tok_dict:
+            default_allowed_special = set(
+                [bytes(bytes_list).decode() for bytes_list in tok_dict["default_allowed_special"]]
+            )
+        else:
+            default_allowed_special = None
+        if "vocab_size" in tok_dict:
+            kwargs["explicit_n_vocab"] = tok_dict["vocab_size"]
+
+        tokenizer = tiktoken.Encoding(**kwargs)
+        tokenizer._default_allowed_special = default_allowed_special or set()
+
+        def encode_patched(
+            self,
+            text: str,
+            *,
+            allowed_special: Union[Literal["all"], AbstractSet[str]] = set(),  # noqa: B006
+            disallowed_special: Union[Literal["all"], Collection[str]] = "all",
+        ) -> list[int]:
+            if isinstance(allowed_special, set):
+                allowed_special |= self._default_allowed_special
+            return tiktoken.Encoding.encode(
+                self, text, allowed_special=allowed_special, disallowed_special=disallowed_special
+            )
+        tokenizer.encode = functools.partial(encode_patched, tokenizer)
+
+        # Convert to HF interface
+        self.tokenizer = tokenizer
+        self.eos_token_id = tokenizer._special_tokens["<|eos|>"]
+        self.vocab_size = tokenizer.n_vocab
+
+    def encode(self, x, add_special_tokens=False):
+        return self.tokenizer.encode(x)
+
+    def decode(self, x):
+        return self.tokenizer.decode(x)
+
+    def batch_decode(self, batch, skip_special_tokens=True,  spaces_between_special_tokens=False):
+        if isinstance(batch[0], int):
+            batch = [[x] for x in batch]
+        return self.tokenizer.decode_batch(batch)
+
+    def convert_ids_to_tokens(self, index):
+        return self.tokenizer.decode_single_token_bytes(index).decode("utf-8", errors="ignore")

sglang/srt/layers/extend_attention.py

@@ -8,6 +8,12 @@ from sglang.srt.utils import wrap_kernel_launcher
 CUDA_CAPABILITY = torch.cuda.get_device_capability()
 
 
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
 @triton.jit
 def _fwd_kernel(
     Q_Extend,
@@ -39,6 +45,7 @@ def _fwd_kernel(
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_M: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    logit_cap: tl.constexpr,
 ):
     cur_seq = tl.program_id(0)
     cur_head = tl.program_id(1)
@@ -90,6 +97,10 @@ def _fwd_kernel(
         qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
         qk += tl.dot(q, k)
         qk *= sm_scale
+
+        if logit_cap > 0:
+            qk = logit_cap * tanh(qk / logit_cap)
+
         qk = tl.where(mask_m[:, None] & mask_n[None, :], qk, float("-inf"))
 
         n_e_max = tl.maximum(tl.max(qk, 1), e_max)
@@ -126,6 +137,10 @@ def _fwd_kernel(
         qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
         qk += tl.dot(q, k)
         qk *= sm_scale
+
+        if logit_cap > 0:
+            qk = logit_cap * tanh(qk / logit_cap)
+
         mask_causual = (cur_block_m * BLOCK_M + offs_m[:, None]) >= (
             start_n + offs_n[None, :]
         )
@@ -176,6 +191,7 @@ def extend_attention_fwd(
     b_seq_len_extend,
     max_len_in_batch,
     max_len_extend,
+    logit_cap=-1,
 ):
     """
     q_extend, k_extend, v_extend, o_extend: contiguous tensors
@@ -271,6 +287,7 @@ def extend_attention_fwd(
         BLOCK_N=BLOCK_N,
         num_warps=num_warps,
         num_stages=num_stages,
+        logit_cap=logit_cap,
     )
     cached_kernel = wrap_kernel_launcher(_fwd_kernel)