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

sglang/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.17"
+__version__ = "0.1.18"
 
 # SGL API Components
 from sglang.api import (
@@ -24,10 +24,10 @@ from sglang.api import (
 
 # SGL Backends
 from sglang.backend.anthropic import Anthropic
+from sglang.backend.litellm import LiteLLM
 from sglang.backend.openai import OpenAI
 from sglang.backend.runtime_endpoint import RuntimeEndpoint
 from sglang.backend.vertexai import VertexAI
-from sglang.backend.litellm import LiteLLM
 
 # Global Configurations
 from sglang.global_config import global_config

sglang/api.py

@@ -1,4 +1,4 @@
-"""Some Public API Definitions"""
+"""Public APIs of the language."""
 
 import os
 import re
@@ -43,14 +43,14 @@ def set_default_backend(backend: BaseBackend):
     global_config.default_backend = backend
 
 
-def flush_cache(backend: BaseBackend = None):
+def flush_cache(backend: Optional[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):
+def get_server_args(backend: Optional[BaseBackend] = None):
     backend = backend or global_config.default_backend
     if backend is None:
         return None
@@ -158,7 +158,7 @@ def video(path: str, num_frames: int):
 
 def select(
     name: Optional[str] = None,
-    choices: List[str] = None,
+    choices: Optional[List[str]] = None,
     temperature: float = 0.0,
 ):
     assert choices is not None

sglang/backend/litellm.py

@@ -13,7 +13,6 @@ except ImportError as e:
 
 
 class LiteLLM(BaseBackend):
-
     def __init__(
         self,
         model_name,
@@ -33,7 +32,8 @@ class LiteLLM(BaseBackend):
         self.model_name = model_name
 
         self.chat_template = chat_template or get_chat_template_by_model_path(
-            model_name)
+            model_name
+        )
 
         self.client_params = {
             "api_key": api_key,

sglang/backend/openai.py

@@ -1,7 +1,7 @@
+import dataclasses
 import logging
 import time
 import warnings
-import dataclasses
 from typing import Callable, List, Optional, Union
 
 import numpy as np
@@ -105,14 +105,16 @@ class OpenAI(BaseBackend):
     def get_chat_template(self):
         return self.chat_template
 
-    def _prepare_spec_execution(self, sampling_params: SglSamplingParams,
-                                num_api_spec_tokens: int, spec_var_name: str):
+    def _prepare_spec_execution(
+        self,
+        sampling_params: SglSamplingParams,
+        num_api_spec_tokens: int,
+        spec_var_name: str,
+    ):
         if "max_tokens" not in self.spec_kwargs:
             self.spec_kwargs["max_tokens"] = num_api_spec_tokens
         else:
-            assert (
-                self.spec_kwargs["max_tokens"] == num_api_spec_tokens
-            )
+            assert self.spec_kwargs["max_tokens"] == num_api_spec_tokens
 
         params = sampling_params.to_openai_kwargs()
         for key, value in params.items():
@@ -151,8 +153,9 @@ class OpenAI(BaseBackend):
                         )
                     prompt = s.messages_
                 else:
-                    return self._prepare_spec_execution(sampling_params,
-                        s.num_api_spec_tokens, spec_var_name)
+                    return self._prepare_spec_execution(
+                        sampling_params, s.num_api_spec_tokens, spec_var_name
+                    )
             else:
                 prompt = s.text_
 
@@ -325,7 +328,7 @@ class OpenAI(BaseBackend):
             ret_str = ret.choices[0].text
             ret_token = self.tokenizer.encode(ret_str)[0]
             self.token_usage.prompt_tokens += ret.usage.prompt_tokens
-            self.token_usage.completion_tokens= ret.usage.completion_tokens
+            self.token_usage.completion_tokens = ret.usage.completion_tokens
 
             # TODO:
             # 1. return logits as the scores
@@ -355,7 +358,9 @@ class OpenAI(BaseBackend):
         return decision, scores, None, None
 
 
-def openai_completion(client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs):
+def openai_completion(
+    client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs
+):
     for attempt in range(retries):
         try:
             if is_chat:
@@ -385,15 +390,19 @@ def openai_completion(client, token_usage, is_chat=None, retries=3, prompt=None,
     return comp
 
 
-def openai_completion_stream(client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs):
+def openai_completion_stream(
+    client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs
+):
     for attempt in range(retries):
         try:
             if is_chat:
                 if "stop" in kwargs and kwargs["stop"] is None:
                     kwargs.pop("stop")
                 generator = client.chat.completions.create(
-                    messages=prompt, stream=True, stream_options={"include_usage": True},
-                    **kwargs
+                    messages=prompt,
+                    stream=True,
+                    stream_options={"include_usage": True},
+                    **kwargs,
                 )
                 for ret in generator:
                     if len(ret.choices) == 0:
@@ -405,8 +414,10 @@ def openai_completion_stream(client, token_usage, is_chat=None, retries=3, promp
                     yield content or "", {}
             else:
                 generator = client.completions.create(
-                    prompt=prompt, stream=True, stream_options={"include_usage": True},
-                    **kwargs
+                    prompt=prompt,
+                    stream=True,
+                    stream_options={"include_usage": True},
+                    **kwargs,
                 )
                 for ret in generator:
                     if len(ret.choices) == 0:

sglang/bench_latency.py

@@ -0,0 +1,299 @@
+"""
+Benchmark the latency of a given model. It accepts arguments similar to those of launch_server.py.
+
+# Usage (latency test):
+python -m sglang.bench_latency --model-path meta-llama/Meta-Llama-3-8B-Instruct --load-format dummy
+
+# Usage (correctness test):
+python -m sglang.bench_latency --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct
+
+### Reference output:
+prefill logits (first half) tensor([[-10.0312,  -9.5000,   0.8936,  ...,  -4.9414,  -3.2402,  -3.3633],
+        [-10.0312,  -9.5000,   0.8936,  ...,  -4.9414,  -3.2402,  -3.3633],
+        [ -9.1875, -10.2500,   2.7109,  ...,  -4.3359,  -4.0664,  -4.1328]],
+       device='cuda:0', dtype=torch.float16)
+prefill logits (final) tensor([[-8.3203, -7.1211,  3.3379,  ..., -4.9570, -4.1328, -3.4141],
+        [-8.9062, -9.0156,  4.1445,  ..., -4.9922, -4.4961, -4.0742],
+        [-9.6328, -9.0547,  4.0117,  ..., -5.3047, -4.7148, -4.4609]],
+       device='cuda:0', dtype=torch.float16)
+<s> The capital of France is.
+The capital of the United States is Washington, D.C.
+
+<s> The capital of the United Kindom is.
+The capital of the United Kingdom is London.
+The capital of the
+<s> Today is a sunny day and I like go for a walk in the park.
+I'm going to the park
+"""
+
+import argparse
+import dataclasses
+import logging
+import multiprocessing
+import time
+
+import numpy as np
+import torch
+import torch.distributed as dist
+
+from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.managers.controller.infer_batch import Batch, ForwardMode, Req
+from sglang.srt.managers.controller.model_runner import ModelRunner
+from sglang.srt.model_config import ModelConfig
+from sglang.srt.sampling_params import SamplingParams
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import suppress_other_loggers
+
+
+@dataclasses.dataclass
+class BenchArgs:
+    batch_size: int = 1
+    input_len: int = 1024
+    output_len: int = 4
+    correctness_test: bool = False
+    # This is only used for correctness test
+    cut_len: int = 4
+
+    @staticmethod
+    def add_cli_args(parser: argparse.ArgumentParser):
+        parser.add_argument("--batch-size", type=int, default=BenchArgs.batch_size)
+        parser.add_argument("--input-len", type=int, default=BenchArgs.input_len)
+        parser.add_argument("--output-len", type=int, default=BenchArgs.output_len)
+        parser.add_argument("--correctness-test", action="store_true")
+        parser.add_argument("--cut-len", type=int, default=BenchArgs.cut_len)
+
+    @classmethod
+    def from_cli_args(cls, args: argparse.Namespace):
+        attrs = [attr.name for attr in dataclasses.fields(cls)]
+        return cls(**{attr: getattr(args, attr) for attr in attrs})
+
+
+def load_model(server_args, tp_rank):
+    suppress_other_loggers()
+
+    model_config = ModelConfig(path=server_args.model_path)
+    model_runner = ModelRunner(
+        model_config=model_config,
+        mem_fraction_static=server_args.mem_fraction_static,
+        gpu_id=tp_rank,
+        tp_rank=tp_rank,
+        tp_size=server_args.tp_size,
+        nccl_port=28888,
+        server_args=server_args,
+    )
+    print(f"max_total_num_tokens={model_runner.max_total_num_tokens}")
+    tokenizer = get_tokenizer(
+        server_args.tokenizer_path,
+        tokenizer_mode=server_args.tokenizer_mode,
+        trust_remote_code=server_args.trust_remote_code,
+    )
+    if server_args.tp_size > 1:
+        dist.barrier()
+    return model_runner, tokenizer
+
+
+def prepare_inputs(bench_args, tokenizer):
+    prompts = [
+        "The capital of France is",
+        "The capital of the United Kindom is",
+        "Today is a sunny day and I like",
+    ]
+    input_ids = [tokenizer.encode(p) for p in prompts]
+    sampling_params = SamplingParams(
+        temperature=0,
+        max_new_tokens=BenchArgs.output_len,
+    )
+
+    reqs = []
+    for i in range(len(prompts)):
+        assert len(input_ids[i]) > bench_args.cut_len
+
+        tmp_input_ids = input_ids[i][:bench_args.cut_len]
+        req = Req(rid=i, origin_input_text=prompts[i], origin_input_ids=tmp_input_ids)
+        req.prefix_indices = []
+        req.sampling_params = sampling_params
+        req.input_ids = req.origin_input_ids
+        reqs.append(req)
+
+    return input_ids, reqs
+
+
+def prepare_extend_inputs(bench_args, input_ids, reqs, model_runner):
+    for i in range(len(reqs)):
+        req = reqs[i]
+        req.input_ids += input_ids[i][bench_args.cut_len:]
+        req.prefix_indices = model_runner.req_to_token_pool.req_to_token[
+            i, :bench_args.cut_len
+        ]
+    return reqs
+
+
+def prepare_synthetic_inputs(bench_args, tokenizer):
+    input_ids = np.ones((bench_args.batch_size, bench_args.input_len), dtype=np.int32)
+    sampling_params = SamplingParams(
+        temperature=0,
+        max_new_tokens=BenchArgs.output_len,
+    )
+
+    reqs = []
+    for i in range(len(input_ids)):
+        req = Req(rid=i, origin_input_text="", origin_input_ids=list(input_ids[i]))
+        req.prefix_indices = []
+        req.sampling_params = sampling_params
+        req.input_ids = req.origin_input_ids
+        reqs.append(req)
+
+    return reqs
+
+
+def extend(reqs, model_runner):
+    batch = Batch.init_new(
+        reqs=reqs,
+        req_to_token_pool=model_runner.req_to_token_pool,
+        token_to_kv_pool=model_runner.token_to_kv_pool,
+        tree_cache=None)
+    batch.prepare_for_extend(model_runner.model_config.vocab_size, None)
+    output = model_runner.forward(batch, ForwardMode.EXTEND)
+    next_token_ids, _ = batch.sample(output.next_token_logits)
+    return next_token_ids, output.next_token_logits, batch
+
+
+def decode(input_token_ids, batch, model_runner):
+    batch.prepare_for_decode(input_token_ids.cpu().numpy())
+    output = model_runner.forward(batch, ForwardMode.DECODE)
+    next_token_ids, _ = batch.sample(output.next_token_logits)
+    return next_token_ids, output.next_token_logits
+
+
+def correctness_test(
+    server_args,
+    bench_args,
+    tp_rank,
+):
+    rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
+
+    # Load the model
+    model_runner, tokenizer = load_model(server_args, tp_rank)
+
+    # Prepare inputs
+    input_ids, reqs = prepare_inputs(bench_args, tokenizer)
+
+    # Prefill
+    next_token_ids, next_token_logits, batch = extend(reqs, model_runner)
+    rank_print("prefill logits (first half)", next_token_logits)
+
+    # Prepare extend inputs
+    reqs = prepare_extend_inputs(bench_args, input_ids, reqs, model_runner)
+
+    # Extend
+    next_token_ids, next_token_logits, batch = extend(reqs, model_runner)
+    rank_print("prefill logits (final)", next_token_logits)
+
+    # Decode
+    output_ids = [list(req.input_ids) for req in reqs]
+    for _ in range(bench_args.output_len):
+        next_token_ids, _ = decode(next_token_ids, batch, model_runner)
+        for i in range(len(reqs)):
+            output_ids[i].append(next_token_ids[i])
+
+    # Print
+    for i in range(len(reqs)):
+        print(tokenizer.decode(output_ids[i]))
+
+
+def latency_test(
+    server_args,
+    bench_args,
+    tp_rank,
+):
+    rank_print = print if tp_rank == 0 else lambda *args, **kwargs: None
+
+    # Load the model
+    model_runner, tokenizer = load_model(server_args, tp_rank)
+    print(f"max_batch_size={model_runner.max_total_num_tokens // (bench_args.input_len + bench_args.output_len)}")
+
+    # Prepare inputs
+    reqs = prepare_synthetic_inputs(bench_args, tokenizer)
+
+    def clear():
+        model_runner.req_to_token_pool.clear()
+        model_runner.token_to_kv_pool.clear()
+
+    @torch.inference_mode()
+    def run_once(output_len):
+        # Prefill
+        torch.cuda.synchronize()
+        tot_latency = 0
+        tic = time.time()
+        next_token_ids, _, batch = extend(reqs, model_runner)
+        torch.cuda.synchronize()
+        prefill_latency = time.time() - tic
+        tot_latency += prefill_latency
+        throughput = bench_args.input_len * bench_args.batch_size / prefill_latency
+        rank_print(f"Prefill. latency: {prefill_latency:6.5f} s, throughput: {throughput:9.2f} token/s")
+
+        # Decode
+        for i in range(output_len):
+            torch.cuda.synchronize()
+            tic = time.time()
+            next_token_ids, _ = decode(next_token_ids, batch, model_runner)
+            torch.cuda.synchronize()
+            latency = time.time() - tic
+            tot_latency += latency
+            throughput = bench_args.batch_size / latency
+            if i < 5: rank_print(f"Decode.  latency: {latency:6.5f} s, throughput: {throughput:9.2f} token/s")
+        avg_decode_latency = (tot_latency - prefill_latency) / output_len
+        avg_decode_throughput = bench_args.batch_size / avg_decode_latency
+        rank_print(f"Decode.  avg latency: {avg_decode_latency:6.5f} s, avg throughput: {avg_decode_throughput:9.2f} token/s")
+        
+        throughput = (bench_args.input_len + bench_args.output_len) * bench_args.batch_size / tot_latency
+        rank_print(f"Total. latency: {tot_latency:6.3f} s, throughput: {throughput:9.2f} token/s")
+
+    # Warm up
+    run_once(4)
+    clear()
+
+    # Run again
+    run_once(bench_args.output_len)
+
+
+def main(server_args, bench_args):
+    print(bench_args)
+
+    if bench_args.correctness_test:
+        work_func = correctness_test
+    else:
+        work_func = latency_test
+
+    workers = []
+    for tp_rank in range(server_args.tp_size):
+        proc = multiprocessing.Process(
+            target=work_func,
+            args=(
+                server_args,
+                bench_args,
+                tp_rank,
+            ),
+        )
+        proc.start()
+        workers.append(proc)
+
+    for proc in workers:
+        proc.join()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    ServerArgs.add_cli_args(parser)
+    BenchArgs.add_cli_args(parser)
+    args = parser.parse_args()
+
+    server_args = ServerArgs.from_cli_args(args)
+    bench_args = BenchArgs.from_cli_args(args)
+
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format="%(message)s",
+    )
+
+    main(server_args, bench_args)

sglang/global_config.py

@@ -27,7 +27,7 @@ class GlobalConfig:
 
         # Request dependency time due to network delay
         self.request_dependency_delay = 0.02
-        self.wait_for_new_request_delay = 0.0004
+        self.wait_for_new_request_delay = 0.0006
 
         # New generation token ratio estimation
         self.base_new_token_ratio = 0.4
@@ -35,5 +35,8 @@ class GlobalConfig:
         self.new_token_ratio_decay = 0.0001
         self.new_token_ratio_recovery = 0.05
 
+        # The threshold (number of tokens) to trigger layer-wise cuda sync.
+        # This can improve the speed for large batch sizes during prefill.
+        self.layer_sync_threshold = 8192
 
 global_config = GlobalConfig()

sglang/lang/compiler.py

@@ -4,7 +4,7 @@ from queue import Queue
 from typing import List, Union
 
 from sglang.global_config import global_config
-from sglang.lang.interpreter import ProgramState, StreamExecutor, pin_program
+from sglang.lang.interpreter import ProgramState, StreamExecutor, cache_program
 from sglang.lang.ir import (
     SglArgument,
     SglConstantText,
@@ -184,7 +184,7 @@ class CompiledFunction:
 
         # Extract prefix by tracing and cache it
         if len(batch_kwargs) > 1:
-            pin_program(self.function, backend)
+            cache_program(self.function, backend)
 
         # Run all programs
         if num_threads == "auto":

sglang/lang/interpreter.py

@@ -507,7 +507,7 @@ class StreamExecutor:
                     )
                     return
 
-                else: # Speculative execution on models with completion interface
+                else:  # Speculative execution on models with completion interface
                     comp, meta_info = self._spec_gen(sampling_params)
 
             self.text_ += comp

sglang/lang/ir.py

@@ -81,12 +81,10 @@ 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."
-            )
+            warnings.warn("Regular expression is not supported in the LiteLLM backend.")
         return {
             "max_tokens": self.max_new_tokens,
             "stop": self.stop or None,
@@ -122,6 +120,7 @@ class SglFunction:
         argspec = inspect.getfullargspec(func)
         assert argspec.args[0] == "s", 'The first argument must be "s"'
         self.arg_names = argspec.args[1:]
+        self.arg_defaults = argspec.defaults if argspec.defaults is not None else []
 
     def bind(self, **kwargs):
         assert all(key in self.arg_names for key in kwargs)
@@ -180,7 +179,18 @@ class SglFunction:
         assert isinstance(batch_kwargs, (list, tuple))
         if len(batch_kwargs) == 0:
             return []
-        assert isinstance(batch_kwargs[0], dict)
+        if not isinstance(batch_kwargs[0], dict):
+            num_programs = len(batch_kwargs)
+            # change the list of argument values to dict of arg_name -> arg_value
+            batch_kwargs = [
+                {self.arg_names[i]: v for i, v in enumerate(arg_values)}
+                for arg_values in batch_kwargs
+                if isinstance(arg_values, (list, tuple)) and
+                   len(self.arg_names) - len(self.arg_defaults) <= len(arg_values) <= len(self.arg_names)
+            ]
+            # Ensure to raise an exception if the number of arguments mismatch
+            if len(batch_kwargs) != num_programs:
+                raise Exception("Given arguments mismatch the SGL function signature")
 
         default_sampling_para = SglSamplingParams(
             max_new_tokens=max_new_tokens,

sglang/launch_server.py

@@ -1,6 +1,9 @@
+"""Launch the inference server."""
+
 import argparse
 
-from sglang.srt.server import ServerArgs, launch_server
+from sglang.srt.server import launch_server
+from sglang.srt.server_args import ServerArgs
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()

sglang/launch_server_llavavid.py

@@ -1,10 +1,11 @@
+"""Launch the inference server for Llava-video model."""
+
 import argparse
 import multiprocessing as mp
 
 from sglang.srt.server import ServerArgs, launch_server
 
 if __name__ == "__main__":
-
     model_overide_args = {}
 
     model_overide_args["mm_spatial_pool_stride"] = 2

sglang/srt/constrained/__init__.py

@@ -1,13 +1,19 @@
 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.regex import FSMInfo, make_deterministic_fsm
-from outlines.models.transformers import TransformerTokenizer
 from pydantic import BaseModel
 
+try:
+    from outlines.caching import cache as disk_cache
+    from outlines.fsm.guide import RegexGuide
+    from outlines.caching import disable_cache
+    from outlines.fsm.guide import RegexGuide
+    from outlines.fsm.regex import FSMInfo, make_byte_level_fsm, make_deterministic_fsm
+    from outlines.models.transformers import TransformerTokenizer
+except ImportError as e:
+    print(f'\nError: {e}. Please install a new version of outlines by `pip install "outlines>=0.0.44"`\n')
+    raise
+
 try:
     from outlines.fsm.json_schema import build_regex_from_object
 except ImportError:
@@ -28,11 +34,12 @@ except ImportError:
 
 
 __all__ = [
-    "RegexFSM",
+    "RegexGuide",
     "FSMInfo",
     "make_deterministic_fsm",
     "build_regex_from_object",
     "TransformerTokenizer",
     "disk_cache",
     "disable_cache",
+    "make_byte_level_fsm",
 ]

sglang/srt/constrained/fsm_cache.py

@@ -1,4 +1,6 @@
-from sglang.srt.constrained import RegexFSM, TransformerTokenizer
+"""Cache for the compressed finite state machine."""
+
+from sglang.srt.constrained import RegexGuide, TransformerTokenizer
 from sglang.srt.constrained.base_cache import BaseCache
 
 
@@ -6,7 +8,8 @@ class FSMCache(BaseCache):
     def __init__(self, tokenizer_path, tokenizer_args_dict, enable=True):
         super().__init__(enable=enable)
 
-        if tokenizer_path.endswith(".json"):
+        if tokenizer_path.endswith(".json") or tokenizer_path.endswith(".model"):
+            # Do not support TiktokenTokenizer or SentencePieceTokenizer
             return
 
         from importlib.metadata import version
@@ -25,4 +28,4 @@ class FSMCache(BaseCache):
             )
 
     def init_value(self, regex):
-        return RegexFSM(regex, self.outlines_tokenizer)
+        return RegexGuide(regex, self.outlines_tokenizer)