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

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

@@ -26,7 +26,17 @@ class GlobalConfig:
         self.concate_and_append_mode = "no_adjust"
 
         # Request dependency time due to network delay
-        self.request_dependency_time = 0.03
-
+        self.request_dependency_delay = 0.02
+        self.wait_for_new_request_delay = 0.0006
+
+        # New generation token ratio estimation
+        self.base_new_token_ratio = 0.4
+        self.base_min_new_token_ratio = 0.2
+        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

@@ -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: