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

sglang/srt/flush_cache.py

@@ -1,4 +1,6 @@
 """
+Flush the KV cache.
+
 Usage:
 python3 -m sglang.srt.flush_cache --url http://localhost:30000
 """

sglang/srt/hf_transformers_utils.py

@@ -1,10 +1,10 @@
 """Utilities for Huggingface Transformers."""
 
+import functools
 import json
 import os
 import warnings
-import functools
-from typing import Optional, Union, AbstractSet, Collection, Literal
+from typing import AbstractSet, Collection, Literal, Optional, Union
 
 from huggingface_hub import snapshot_download
 from transformers import (
@@ -88,6 +88,9 @@ def get_tokenizer(
     if tokenizer_name.endswith(".json"):
         return TiktokenTokenizer(tokenizer_name)
 
+    if tokenizer_name.endswith(".model"):
+        return SentencePieceTokenizer(tokenizer_name)
+
     """Gets a tokenizer for the given model name via Huggingface."""
     if is_multimodal_model(tokenizer_name):
         processor = get_processor(
@@ -179,6 +182,8 @@ def get_processor(
 class TiktokenTokenizer:
     def __init__(self, tokenizer_path):
         import tiktoken
+        from jinja2 import Template
+
         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
@@ -190,7 +195,8 @@ class TiktokenTokenizer:
             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"]
+            bytes(item["bytes"]).decode(): item["token"]
+            for item in tok_dict["special_tokens"]
         }
         assert tok_dict["word_split"] == "V1"
 
@@ -202,7 +208,10 @@ class TiktokenTokenizer:
         }
         if "default_allowed_special" in tok_dict:
             default_allowed_special = set(
-                [bytes(bytes_list).decode() for bytes_list in tok_dict["default_allowed_special"]]
+                [
+                    bytes(bytes_list).decode()
+                    for bytes_list in tok_dict["default_allowed_special"]
+                ]
             )
         else:
             default_allowed_special = None
@@ -211,25 +220,35 @@ class TiktokenTokenizer:
 
         tokenizer = tiktoken.Encoding(**kwargs)
         tokenizer._default_allowed_special = default_allowed_special or set()
+        tokenizer._default_allowed_special |= {"<|separator|>"}
 
         def encode_patched(
             self,
             text: str,
             *,
-            allowed_special: Union[Literal["all"], AbstractSet[str]] = set(),  # noqa: B006
+            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
+                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
+        self.chat_template = Template(
+            "{% for message in messages %}{% if message['role'] == 'user' %}{{ 'Human: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'system' %}{{ 'System: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'assistant' %}{{ 'Assistant: '  + message['content'] + '<|separator|>\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
+        )
 
     def encode(self, x, add_special_tokens=False):
         return self.tokenizer.encode(x)
@@ -237,10 +256,50 @@ class TiktokenTokenizer:
     def decode(self, x):
         return self.tokenizer.decode(x)
 
-    def batch_decode(self, batch, skip_special_tokens=True,  spaces_between_special_tokens=False):
+    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")
+    def apply_chat_template(self, messages, tokenize, add_generation_prompt):
+        ret = self.chat_template.render(
+            messages=messages, add_generation_prompt=add_generation_prompt
+        )
+        return self.encode(ret) if tokenize else ret
+
+
+class SentencePieceTokenizer:
+    def __init__(self, tokenizer_path):
+        import sentencepiece as spm
+        from jinja2 import Template
+
+        tokenizer = spm.SentencePieceProcessor(model_file=tokenizer_path)
+
+        # Convert to HF interface
+        self.tokenizer = tokenizer
+        self.eos_token_id = tokenizer.eos_id()
+        self.vocab_size = tokenizer.vocab_size()
+        self.chat_template = Template(
+            "{% for message in messages %}{% if message['role'] == 'user' %}{{ 'Human: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'system' %}{{ 'System: ' + message['content'].strip() + '<|separator|>\n\n' }}{% elif message['role'] == 'assistant' %}{{ 'Assistant: '  + message['content'] + '<|separator|>\n\n' }}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
+        )
+
+    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)
+
+    def apply_chat_template(self, messages, tokenize, add_generation_prompt):
+        ret = self.chat_template.render(
+            messages=messages, add_generation_prompt=add_generation_prompt
+        )
+        return self.encode(ret) if tokenize else ret

sglang/srt/layers/extend_attention.py

@@ -191,6 +191,7 @@ def extend_attention_fwd(
     b_seq_len_extend,
     max_len_in_batch,
     max_len_extend,
+    sm_scale=None,
     logit_cap=-1,
 ):
     """
@@ -213,7 +214,7 @@ def extend_attention_fwd(
     else:
         BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)
 
-    sm_scale = 1.0 / (Lq**0.5)
+    sm_scale = 1.0 / (Lq**0.5) if sm_scale is None else sm_scale
     batch_size, head_num = b_seq_len.shape[0], q_extend.shape[1]
     kv_group_num = q_extend.shape[1] // k_extend.shape[1]