sglang 0.1.12__py3-none-any.whl → 0.1.14__py3-none-any.whl

sglang/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.12"
+__version__ = "0.1.14"
 
 from sglang.api import *
 from sglang.global_config import global_config

sglang/api.py

@@ -44,6 +44,20 @@ def set_default_backend(backend: BaseBackend):
     global_config.default_backend = backend
 
 
+def flush_cache(backend: BaseBackend = None):
+    backend = backend or global_config.default_backend
+    if backend is None:
+        return False
+    return backend.flush_cache()
+
+
+def get_server_args(backend: BaseBackend = None):
+    backend = backend or global_config.default_backend
+    if backend is None:
+        return None
+    return backend.get_server_args()
+
+
 def gen(
     name: Optional[str] = None,
     max_tokens: Optional[int] = None,

sglang/backend/anthropic.py

@@ -30,13 +30,17 @@ class Anthropic(BaseBackend):
         s: StreamExecutor,
         sampling_params: SglSamplingParams,
     ):
-        prompt = s.text_
-        ret = anthropic.Anthropic().completions.create(
+        if s.messages_:
+            messages = s.messages_
+        else:
+            messages = [{"role": "user", "content": s.text_}]
+
+        ret = anthropic.Anthropic().messages.create(
             model=self.model_name,
-            prompt=prompt,
+            messages=messages,
             **sampling_params.to_anthropic_kwargs(),
         )
-        comp = ret.completion
+        comp = ret.content[0].text
 
         return comp, {}
 
@@ -45,13 +49,15 @@ class Anthropic(BaseBackend):
         s: StreamExecutor,
         sampling_params: SglSamplingParams,
     ):
-        prompt = s.text_
-        generator = anthropic.Anthropic().completions.create(
+        if s.messages_:
+            messages = s.messages_
+        else:
+            messages = [{"role": "user", "content": s.text_}]
+
+        with anthropic.Anthropic().messages.stream(
             model=self.model_name,
-            prompt=prompt,
-            stream=True,
+            messages=messages,
             **sampling_params.to_anthropic_kwargs(),
-        )
-
-        for ret in generator:
-            yield ret.completion, {}
+        ) as stream:
+            for text in stream.text_stream:
+                yield text, {}

sglang/backend/base_backend.py

@@ -72,3 +72,9 @@ class BaseBackend:
 
     def shutdown(self):
         pass
+
+    def flush_cache(self):
+        pass
+
+    def get_server_args(self):
+        pass

sglang/backend/openai.py

@@ -4,7 +4,7 @@ from typing import Callable, List, Optional, Union
 
 import numpy as np
 from sglang.backend.base_backend import BaseBackend
-from sglang.lang.chat_template import get_chat_template
+from sglang.lang.chat_template import ChatTemplate, get_chat_template_by_model_path
 from sglang.lang.interpreter import StreamExecutor
 from sglang.lang.ir import SglSamplingParams
 
@@ -41,23 +41,45 @@ INSTRUCT_MODEL_NAMES = [
 
 
 class OpenAI(BaseBackend):
-    def __init__(self, model_name, *args, **kwargs):
+    def __init__(
+        self,
+        model_name: str,
+        is_chat_model: Optional[bool] = None,
+        chat_template: Optional[ChatTemplate] = None,
+        is_azure: bool = False,
+        *args,
+        **kwargs,
+    ):
         super().__init__()
 
         if isinstance(openai, Exception):
             raise openai
 
-        self.client = openai.OpenAI(*args, **kwargs)
+        if is_azure:
+            self.client = openai.AzureOpenAI(*args, **kwargs)
+        else:
+            self.client = openai.OpenAI(*args, **kwargs)
+
         self.model_name = model_name
-        self.tokenizer = tiktoken.encoding_for_model(model_name)
+        try:
+            self.tokenizer = tiktoken.encoding_for_model(model_name)
+        except KeyError:
+            self.tokenizer = tiktoken.get_encoding("cl100k_base")
         self.logit_bias_int = create_logit_bias_int(self.tokenizer)
 
-        if model_name in INSTRUCT_MODEL_NAMES:
-            self.is_chat_model = False
+        self.chat_template = chat_template or get_chat_template_by_model_path(
+            model_name
+        )
+
+        if is_chat_model is not None:
+            self.is_chat_model = is_chat_model
         else:
-            self.is_chat_model = True
+            if model_name in INSTRUCT_MODEL_NAMES:
+                self.is_chat_model = False
+            else:
+                self.is_chat_model = True
 
-        self.chat_template = get_chat_template("default")
+        self.chat_begin_str = self.chat_template.role_prefix_and_suffix["assistant"][0]
 
     def get_chat_template(self):
         return self.chat_template
@@ -69,7 +91,7 @@ class OpenAI(BaseBackend):
     ):
         if sampling_params.dtype is None:
             if self.is_chat_model:
-                if not s.text_.endswith("ASSISTANT:"):
+                if not s.text_.endswith(self.chat_begin_str):
                     raise RuntimeError(
                         "This use case is not supported. "
                         "For OpenAI chat models, sgl.gen must be right after sgl.assistant"
@@ -122,7 +144,11 @@ class OpenAI(BaseBackend):
     ):
         if sampling_params.dtype is None:
             if self.is_chat_model:
-                assert s.text_.endswith("ASSISTANT:")
+                if not s.text_.endswith(self.chat_begin_str):
+                    raise RuntimeError(
+                        "This use case is not supported. "
+                        "For OpenAI chat models, sgl.gen must be right after sgl.assistant"
+                    )
                 prompt = s.messages_
             else:
                 prompt = s.text_
@@ -137,7 +163,7 @@ class OpenAI(BaseBackend):
             )
             return generator
         else:
-            raise ValueError(f"Unknown dtype: {dtype}")
+            raise ValueError(f"Unknown dtype: {sampling_params.dtype}")
 
     def select(
         self,
@@ -241,7 +267,10 @@ def openai_completion_stream(client, retries=3, is_chat=None, prompt=None, **kwa
                     messages=prompt, stream=True, **kwargs
                 )
                 for ret in generator:
-                    content = ret.choices[0].delta.content
+                    try:
+                        content = ret.choices[0].delta.content
+                    except IndexError:
+                        content = None
                     yield content or "", {}
             else:
                 generator = client.completions.create(

sglang/backend/runtime_endpoint.py

@@ -12,15 +12,26 @@ from sglang.utils import encode_image_base64, find_printable_text, http_request
 
 
 class RuntimeEndpoint(BaseBackend):
-    def __init__(self, base_url, auth_token=None):
+    def __init__(
+        self,
+        base_url: str,
+        auth_token: Optional[str] = None,
+        api_key: Optional[str] = None,
+        verify: Optional[str] = None,
+    ):
         super().__init__()
         self.support_concate_and_append = True
 
         self.base_url = base_url
         self.auth_token = auth_token
+        self.api_key = api_key
+        self.verify = verify
 
         res = http_request(
-            self.base_url + "/get_model_info", auth_token=self.auth_token
+            self.base_url + "/get_model_info",
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         assert res.status_code == 200
         self.model_info = res.json()
@@ -32,6 +43,22 @@ class RuntimeEndpoint(BaseBackend):
     def get_model_name(self):
         return self.model_info["model_path"]
 
+    def flush_cache(self):
+        res = http_request(
+            self.base_url + "/flush_cache",
+            auth_token=self.auth_token,
+            verify=self.verify,
+        )
+        return res.status_code == 200
+
+    def get_server_args(self):
+        res = http_request(
+            self.base_url + "/get_server_args",
+            auth_token=self.auth_token,
+            verify=self.verify,
+        )
+        return res.json()
+
     def get_chat_template(self):
         return self.chat_template
 
@@ -40,6 +67,8 @@ class RuntimeEndpoint(BaseBackend):
             self.base_url + "/generate",
             json={"text": prefix_str, "sampling_params": {"max_new_tokens": 0}},
             auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         assert res.status_code == 200
 
@@ -48,6 +77,8 @@ class RuntimeEndpoint(BaseBackend):
             self.base_url + "/generate",
             json={"text": s.text_, "sampling_params": {"max_new_tokens": 0}},
             auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         assert res.status_code == 200
 
@@ -55,7 +86,11 @@ class RuntimeEndpoint(BaseBackend):
         data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
         self._add_images(s, data)
         res = http_request(
-            self.base_url + "/generate", json=data, auth_token=self.auth_token
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         assert res.status_code == 200
 
@@ -87,7 +122,11 @@ class RuntimeEndpoint(BaseBackend):
         self._add_images(s, data)
 
         res = http_request(
-            self.base_url + "/generate", json=data, auth_token=self.auth_token
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         obj = res.json()
         comp = obj["text"]
@@ -126,6 +165,8 @@ class RuntimeEndpoint(BaseBackend):
             json=data,
             stream=True,
             auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         pos = 0
 
@@ -157,7 +198,11 @@ class RuntimeEndpoint(BaseBackend):
         data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
         self._add_images(s, data)
         res = http_request(
-            self.base_url + "/generate", json=data, auth_token=self.auth_token
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         assert res.status_code == 200
         prompt_len = res.json()["meta_info"]["prompt_tokens"]
@@ -171,7 +216,11 @@ class RuntimeEndpoint(BaseBackend):
         }
         self._add_images(s, data)
         res = http_request(
-            self.base_url + "/generate", json=data, auth_token=self.auth_token
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         assert res.status_code == 200
         obj = res.json()
@@ -188,6 +237,8 @@ class RuntimeEndpoint(BaseBackend):
             self.base_url + "/concate_and_append_request",
             json={"src_rids": src_rids, "dst_rid": dst_rid},
             auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
         )
         assert res.status_code == 200
 

sglang/lang/chat_template.py

@@ -12,42 +12,43 @@ class ChatTemplateStyle(Enum):
 class ChatTemplate:
     name: str
     default_system_prompt: str
-    role_prefix_and_suffix: Dict[str, Tuple[str]]
+    role_prefix_and_suffix: Dict[str, Tuple[str, str]]
     stop_str: List[str] = ()
     image_token: str = "<image>"
     style: ChatTemplateStyle = ChatTemplateStyle.PLAIN
 
-    def get_prefix_and_suffix(self, role, hist_messages):
-        if self.style == ChatTemplateStyle.PLAIN:
-            return self.role_prefix_and_suffix[role]
-        elif self.style == ChatTemplateStyle.LLAMA2:
-            if len(hist_messages) == 0 and role == "system":
-                return (
-                    self.role_prefix_and_suffix["user"][0]
-                    + self.role_prefix_and_suffix["system"][0],
-                    self.role_prefix_and_suffix["system"][1],
+    def get_prefix_and_suffix(
+        self, role: str, hist_messages: List[Dict]
+    ) -> Tuple[str, str]:
+        prefix, suffix = self.role_prefix_and_suffix.get(role, ("", ""))
+
+        if self.style == ChatTemplateStyle.LLAMA2:
+            if role == "system" and not hist_messages:
+                user_prefix, _ = self.role_prefix_and_suffix.get("user", ("", ""))
+                system_prefix, system_suffix = self.role_prefix_and_suffix.get(
+                    "system", ("", "")
                 )
+                return (user_prefix + system_prefix, system_suffix)
             elif (
-                len(hist_messages) == 1
-                and role == "user"
+                role == "user"
+                and len(hist_messages) == 1
                 and hist_messages[0]["content"] is not None
             ):
-                return ("", self.role_prefix_and_suffix["user"][1])
-            return self.role_prefix_and_suffix[role]
-        else:
-            raise ValueError(f"Invalid style: {self.style}")
+                return ("", suffix)
+
+        return prefix, suffix
 
-    def get_prompt(self, messages):
+    def get_prompt(self, messages: List[Dict]) -> str:
         prompt = ""
-        for i in range(len(messages)):
-            role, content = messages[i]["role"], messages[i]["content"]
+        for i, message in enumerate(messages):
+            role, content = message["role"], message["content"]
             if role == "system" and content is None:
                 content = self.default_system_prompt
                 if content is None:
                     continue
 
             prefix, suffix = self.get_prefix_and_suffix(role, messages[:i])
-            prompt += prefix + content + suffix
+            prompt += f"{prefix}{content}{suffix}"
         return prompt
 
 
@@ -106,9 +107,9 @@ register_chat_template(
         name="chatml",
         default_system_prompt=None,
         role_prefix_and_suffix={
-            "system": ("<|im_start|>system\n", "\n<|im_end|>\n"),
-            "user": ("<|im_start|>user\n", "\n<|im_end|>\n"),
-            "assistant": ("<|im_start|>assistant\n", "\n<|im_end|>\n"),
+            "system": ("<|im_start|>system\n", "<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"),
         },
         style=ChatTemplateStyle.PLAIN,
         stop_str=("<|im_end|>",),
@@ -121,9 +122,9 @@ register_chat_template(
         name="chatml-llava",
         default_system_prompt="Answer the questions.",
         role_prefix_and_suffix={
-            "system": ("<|im_start|>system\n", "\n<|im_end|>\n"),
-            "user": ("<|im_start|>user\n", "\n<|im_end|>\n"),
-            "assistant": ("<|im_start|>assistant\n", "\n<|im_end|>\n"),
+            "system": ("<|im_start|>system\n", "<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"),
         },
         style=ChatTemplateStyle.PLAIN,
         stop_str=("<|im_end|>",),
@@ -178,6 +179,19 @@ register_chat_template(
     )
 )
 
+register_chat_template(
+    ChatTemplate(
+        name="gemma-it",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("", ""),
+            "user": ("<start_of_turn>user\n", "<end_of_turn>\n"),
+            "assistant": ("<start_of_turn>model\n", "<end_of_turn>\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+    )
+)
+
 
 @register_chat_template_matching_function
 def match_vicuna(model_path: str):
@@ -218,6 +232,13 @@ def match_chat_yi(model_path: str):
         return get_chat_template("yi")
 
 
+@register_chat_template_matching_function
+def match_gemma_it(model_path: str):
+    model_path = model_path.lower()
+    if "gemma" in model_path and "it" in model_path:
+        return get_chat_template("gemma-it")
+
+
 if __name__ == "__main__":
     messages = [
         {"role": "system", "content": None},  # None means default

sglang/lang/interpreter.py

@@ -245,6 +245,9 @@ class StreamExecutor:
             self.variable_event[name].wait()
         return self.variables[name]
 
+    def set_var(self, name, value):
+        self.variables[name] = value
+
     def get_meta_info(self, name):
         if name in self.variable_event:
             self.variable_event[name].wait()
@@ -583,6 +586,10 @@ class StreamExecutor:
         if self.chat_template.stop_str:
             if not clone:
                 clone = self.default_sampling_para.clone()
+            if clone.stop == ():
+                clone.stop = []
+            elif isinstance(clone.stop, str):
+                clone.stop = [clone.stop]
             clone.stop += self.chat_template.stop_str
 
         return clone or self.default_sampling_para
@@ -679,7 +686,7 @@ class ProgramState:
             if var_name is None:
                 yield self.text()
             else:
-                yield self.get_var(name)
+                yield self.get_var(var_name)
 
     async def text_async_iter(
         self, var_name: Optional[str] = None, return_meta_data: bool = False
@@ -717,11 +724,14 @@ class ProgramState:
             if var_name is None:
                 yield self.text()
             else:
-                yield self.get_var(name)
+                yield self.get_var(var_name)
 
     def get_var(self, name):
         return self.stream_executor.get_var(name)
 
+    def set_var(self, name, value):
+        return self.stream_executor.set_var(name, value)
+
     def get_meta_info(self, name):
         return self.stream_executor.get_meta_info(name)
 
@@ -732,6 +742,9 @@ class ProgramState:
     def __getitem__(self, name):
         return self.get_var(name)
 
+    def __setitem__(self, name, value):
+        self.set_var(name, value)
+
     def __del__(self):
         self.stream_executor.end()
 

sglang/lang/ir.py

@@ -73,7 +73,7 @@ class SglSamplingParams:
                 "Regular expression is not supported in the Anthropic backend."
             )
         return {
-            "max_tokens_to_sample": self.max_new_tokens,
+            "max_tokens": self.max_new_tokens,
             "stop_sequences": (
                 self.stop if isinstance(self.stop, (list, tuple)) else [self.stop]
             ),

sglang/srt/constrained/__init__.py

@@ -1,9 +1,31 @@
+import json
+from typing import Dict, Optional, Union
+
 from outlines.caching import cache as disk_cache
 from outlines.caching import disable_cache
 from outlines.fsm.fsm import RegexFSM
-from outlines.fsm.json_schema import build_regex_from_object
 from outlines.fsm.regex import FSMInfo, make_deterministic_fsm
 from outlines.models.transformers import TransformerTokenizer
+from pydantic import BaseModel
+
+try:
+    from outlines.fsm.json_schema import build_regex_from_object
+except ImportError:
+    # Since outlines 0.0.32, build_regex_from_object is replaced by build_regex_from_schema,
+    # which only accepts string schema as input.
+    from outlines.fsm.json_schema import build_regex_from_schema
+
+    def build_regex_from_object(
+        object: Union[str, BaseModel, Dict], whitespace_pattern: Optional[str] = None
+    ):
+        if isinstance(object, type(BaseModel)):
+            schema = json.dumps(object.model_json_schema())
+        elif isinstance(object, Dict):
+            schema = json.dumps(object)
+        else:
+            schema = object
+        return build_regex_from_schema(schema, whitespace_pattern)
+
 
 __all__ = [
     "RegexFSM",

sglang/srt/constrained/fsm_cache.py

@@ -5,9 +5,20 @@ from sglang.srt.constrained.base_cache import BaseCache
 class FSMCache(BaseCache):
     def __init__(self, tokenizer_path, tokenizer_args_dict, enable=True):
         super().__init__(enable=enable)
-        self.outlines_tokenizer = TransformerTokenizer(
-            tokenizer_path, **tokenizer_args_dict
-        )
+
+        from importlib.metadata import version
+        if version("outlines") >= "0.0.35":
+            from transformers import AutoTokenizer
+
+            tokenizer_args_dict.setdefault("padding_side", "left")
+            tokenizer = AutoTokenizer.from_pretrained(
+                tokenizer_path, **tokenizer_args_dict
+            )
+            self.outlines_tokenizer = TransformerTokenizer(tokenizer)
+        else:
+            self.outlines_tokenizer = TransformerTokenizer(
+                tokenizer_path, **tokenizer_args_dict
+            )
 
     def init_value(self, regex):
         return RegexFSM(regex, self.outlines_tokenizer)

sglang/srt/layers/context_flashattention_nopad.py

@@ -129,7 +129,7 @@ def context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len):
 
     Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
     assert Lq == Lk and Lk == Lv
-    assert Lk in {16, 32, 64, 128}
+    assert Lk in {16, 32, 64, 128, 256}
 
     sm_scale = 1.0 / (Lq**0.5)
     batch, head = b_seq_len.shape[0], q.shape[1]

sglang/srt/layers/extend_attention.py

@@ -181,19 +181,20 @@ def extend_attention_fwd(
 
     k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
     """
-    if CUDA_CAPABILITY[0] >= 8:
-        BLOCK_M, BLOCK_N = 128, 128
-    else:
-        BLOCK_M, BLOCK_N = 64, 64
-
     Lq, Lk, Lv, Lo = (
         q_extend.shape[-1],
         k_extend.shape[-1],
         v_extend.shape[-1],
         o_extend.shape[-1],
     )
+
     assert Lq == Lk and Lk == Lv and Lv == Lo
-    assert Lq in {16, 32, 64, 128}
+    assert Lq in {16, 32, 64, 128, 256}
+
+    if CUDA_CAPABILITY[0] >= 8:
+        BLOCK_M, BLOCK_N = (128, 128) if Lq <= 128 else (64, 64)
+    else:
+        BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)
 
     sm_scale = 1.0 / (Lq**0.5)
     batch_size, head_num = b_seq_len.shape[0], q_extend.shape[1]

sglang/srt/layers/radix_attention.py

@@ -1,15 +1,9 @@
-from typing import List
-
 import torch
 from sglang.srt.layers.context_flashattention_nopad import context_attention_fwd
 from sglang.srt.layers.extend_attention import extend_attention_fwd
 from sglang.srt.layers.token_attention import token_attention_fwd
 from sglang.srt.managers.router.model_runner import ForwardMode, InputMetadata
 from torch import nn
-from vllm.model_executor.parallel_utils.parallel_state import (
-    get_tensor_model_parallel_rank,
-    get_tensor_model_parallel_world_size,
-)
 
 
 class RadixAttention(nn.Module):
@@ -21,11 +15,9 @@ class RadixAttention(nn.Module):
         self.head_dim = head_dim
         self.layer_id = layer_id
 
-        from sglang.srt.managers.router.model_runner import global_model_mode
-
-        self.use_flashinfer = "flashinfer" in global_model_mode
+        from sglang.srt.managers.router.model_runner import global_server_args_dict
 
-        if self.use_flashinfer:
+        if global_server_args_dict.get("enable_flashinfer", False):
             self.prefill_forward = self.prefill_forward_flashinfer
             self.extend_forward = self.prefill_forward_flashinfer
             self.decode_forward = self.decode_forward_flashinfer