sglang 0.2.10__py3-none-any.whl → 0.2.11__py3-none-any.whl

sglang/__init__.py

@@ -22,6 +22,11 @@ from sglang.api import (
     user_end,
     video,
 )
+from sglang.lang.choices import (
+    greedy_token_selection,
+    token_length_normalized,
+    unconditional_likelihood_normalized,
+)
 
 # SGLang DSL APIs
 __all__ = [
@@ -45,6 +50,9 @@ __all__ = [
     "user_begin",
     "user_end",
     "video",
+    "greedy_token_selection",
+    "token_length_normalized",
+    "unconditional_likelihood_normalized",
 ]
 
 # Global Configurations

sglang/api.py

@@ -6,6 +6,7 @@ from typing import Callable, List, Optional, Union
 
 from sglang.global_config import global_config
 from sglang.lang.backend.base_backend import BaseBackend
+from sglang.lang.choices import ChoicesSamplingMethod, token_length_normalized
 from sglang.lang.ir import (
     SglExpr,
     SglExprList,
@@ -73,12 +74,18 @@ def gen(
     return_text_in_logprobs: Optional[bool] = None,
     dtype: Optional[type] = None,
     choices: Optional[List[str]] = None,
+    choices_method: Optional[ChoicesSamplingMethod] = None,
     regex: Optional[str] = None,
 ):
     """Call the model to generate. See the meaning of the arguments in docs/en/sampling_params.md"""
 
     if choices:
-        return SglSelect(name, choices, 0.0 if temperature is None else temperature)
+        return SglSelect(
+            name,
+            choices,
+            0.0 if temperature is None else temperature,
+            token_length_normalized if choices_method is None else choices_method,
+        )
 
     # check regex is valid
     if regex is not None:
@@ -186,9 +193,10 @@ def select(
     name: Optional[str] = None,
     choices: Optional[List[str]] = None,
     temperature: float = 0.0,
+    choices_method: ChoicesSamplingMethod = token_length_normalized,
 ):
     assert choices is not None
-    return SglSelect(name, choices, temperature)
+    return SglSelect(name, choices, temperature, choices_method)
 
 
 def _role_common(name: str, expr: Optional[SglExpr] = None):

sglang/bench_latency.py

@@ -1,13 +1,21 @@
 """
 Benchmark the latency of a given model. It accepts arguments similar to those of launch_server.py.
 
-# Usage (latency test) with dummy weights:
+# Usage (latency test)
+## with dummy weights:
 python -m sglang.bench_latency --model-path meta-llama/Meta-Llama-3-8B-Instruct --load-format dummy
+## sweep through multiple data points and store (append) the results in a jsonl file:
+python -m sglang.bench_latency --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --result-filename out.jsonl
+## do some changes, and store the results under a different run_name:
+python -m sglang.bench_latency --model-path meta-llama/Meta-Llama-3-8B-Instruct --batch 1 12 14 --input-len 256 512 --output-len 32 256 --result-filename out.jsonl --run-name after
+## plot the results in series of lines:
+python -m sglang.bench_latency --result-filename out.jsonl --graph-sql="select run_name, batch_size, prefill_throughput from results"
+
 
 # Usage (correctness test):
 python -m sglang.bench_latency --model-path TinyLlama/TinyLlama-1.1B-Chat-v0.4 --correct
 
-### Reference output (of the correctness test above, can be gpu dependent):
+## Reference output (of the correctness test above, can be gpu dependent):
 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]],
@@ -28,19 +36,23 @@ I'm going to the park
 
 import argparse
 import dataclasses
+import itertools
 import logging
 import multiprocessing
+import os
+import sqlite3
 import time
 from typing import Tuple
 
-import jsonlines
 import numpy as np
+import pandas as pd
 import torch
 import torch.distributed as dist
 
 from sglang.srt.hf_transformers_utils import get_tokenizer
-from sglang.srt.managers.schedule_batch import Batch, ForwardMode, Req
+from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.model_config import ModelConfig
+from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.sampling_params import SamplingParams
 from sglang.srt.server_args import ServerArgs
@@ -49,26 +61,42 @@ from sglang.srt.utils import suppress_other_loggers
 
 @dataclasses.dataclass
 class BenchArgs:
+    run_name: str = "before"
     batch_size: Tuple[int] = (1,)
-    input_len: int = 1024
-    output_len: int = 4
+    input_len: Tuple[int] = (1024,)
+    output_len: Tuple[int] = (4,)
     result_filename: str = ""
     correctness_test: bool = False
     # This is only used for correctness test
     cut_len: int = 4
+    # Plotting args
+    graph_sql: str = (
+        "select run_name, batch_size, prefill_throughput from results where run_name='before'"
+    )
+    graph_filename: str = "out.png"
 
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser):
+        parser.add_argument("--run-name", type=str, default=BenchArgs.run_name)
         parser.add_argument(
             "--batch-size", type=int, nargs="+", 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(
+            "--input-len", type=int, nargs="+", default=BenchArgs.input_len
+        )
+        parser.add_argument(
+            "--output-len", type=int, nargs="+", default=BenchArgs.output_len
+        )
         parser.add_argument(
             "--result-filename", type=str, default=BenchArgs.result_filename
         )
         parser.add_argument("--correctness-test", action="store_true")
         parser.add_argument("--cut-len", type=int, default=BenchArgs.cut_len)
+        # graphing
+        parser.add_argument("--graph-sql", type=str, default=BenchArgs.graph_sql)
+        parser.add_argument(
+            "--graph-filename", type=str, default=BenchArgs.graph_filename
+        )
 
     @classmethod
     def from_cli_args(cls, args: argparse.Namespace):
@@ -161,7 +189,7 @@ def prepare_synthetic_inputs_for_latency_test(batch_size, input_len):
 
 
 def extend(reqs, model_runner):
-    batch = Batch.init_new(
+    batch = ScheduleBatch.init_new(
         reqs=reqs,
         req_to_token_pool=model_runner.req_to_token_pool,
         token_to_kv_pool=model_runner.token_to_kv_pool,
@@ -222,15 +250,21 @@ def correctness_test(
 
 @torch.inference_mode()
 def latency_test_run_once(
-    model_runner, rank_print, reqs, batch_size, input_len, output_len
+    run_name, model_runner, rank_print, reqs, batch_size, input_len, output_len
 ):
+    max_batch_size = model_runner.max_total_num_tokens // (input_len + output_len)
+    if batch_size > max_batch_size:
+        rank_print(
+            f"skipping ({batch_size}, {input_len}, {output_len}) due to max batch size limit"
+        )
+        return
 
     # Clear the pools.
     model_runner.req_to_token_pool.clear()
     model_runner.token_to_kv_pool.clear()
 
     measurement_results = {
-        "run_name": "before",
+        "run_name": run_name,
         "batch_size": batch_size,
         "input_len": input_len,
         "output_len": output_len,
@@ -291,49 +325,119 @@ def latency_test(
 
     # Load the model
     model_runner, tokenizer = load_model(server_args, tp_rank)
-    rank_print(
-        f"max_batch_size={model_runner.max_total_num_tokens // (bench_args.input_len + bench_args.output_len)}"
-    )
 
-    # To make this PR easier to review, for now, only do the first element in batch_size tuple.
-    bench_args.batch_size = bench_args.batch_size[0]
-
-    # Prepare inputs
+    # Prepare inputs for warm up
     reqs = prepare_synthetic_inputs_for_latency_test(
-        bench_args.batch_size, bench_args.input_len
+        bench_args.batch_size[0], bench_args.input_len[0]
     )
 
     # Warm up
     latency_test_run_once(
-        model_runner, rank_print, reqs, bench_args.batch_size, bench_args.input_len, 4
+        bench_args.run_name,
+        model_runner,
+        rank_print,
+        reqs,
+        bench_args.batch_size[0],
+        bench_args.input_len[0],
+        4,  # shorter decoding to speed up the warmup
     )
 
-    # Run again
+    # Run the sweep
     result_list = []
-    result_list.append(
-        latency_test_run_once(
-            model_runner,
-            rank_print,
-            reqs,
-            bench_args.batch_size,
-            bench_args.input_len,
-            bench_args.output_len,
+    for bs, il, ol in itertools.product(
+        bench_args.batch_size, bench_args.input_len, bench_args.output_len
+    ):
+        req = prepare_synthetic_inputs_for_latency_test(bs, il)
+        ret = latency_test_run_once(
+            bench_args.run_name, model_runner, rank_print, reqs, bs, il, ol
         )
-    )
+        if ret is not None:
+            result_list.append(ret)
+
+    # Write results in jsonlines format on rank 0.
+    if tp_rank == 0 and bench_args.result_filename:
+        import jsonlines
 
-    # Write results in jsonlines format.
-    if bench_args.result_filename:
         with jsonlines.open(bench_args.result_filename, "a") as f:
             f.write_all(result_list)
 
 
+def plot_latency_test(
+    server_args,
+    bench_args,
+    tp_rank,
+):
+    assert tp_rank == 0
+
+    # read the jsonl file and put in sqlite
+    df = pd.read_json(bench_args.result_filename, lines=True)
+    conn = sqlite3.connect(":memory:")
+    cur = conn.cursor()
+
+    # get the columns and their types
+    column_names = list(df.iloc[0].keys())
+    type_dict = {
+        str: "TEXT",
+        np.int64: "INTEGER",
+        np.float64: "FLOAT",
+    }
+    column_types = [type_dict[type(i)] for i in list(df.iloc[0])]
+
+    # create the table
+    cur.execute(
+        f"""
+        CREATE TABLE IF NOT EXISTS results (
+            {", ".join([f"{name} {type}" for name, type in zip(column_names, column_types)])}
+        )
+    """
+    )
+    conn.commit()
+
+    # write the results to DB
+    df.to_sql("results", conn, if_exists="replace", index=False)
+    conn.commit()
+
+    # read it back using sql
+    df = pd.read_sql_query(bench_args.graph_sql, conn)
+    conn.close()
+
+    # plot it and save to a file
+    import matplotlib.pyplot as plt
+
+    assert (
+        len(df.columns) == 3
+    ), f"The sql should have fetched <series, x, y> columns, not {df.columns}"
+    for label in df[df.columns[0]].unique():
+        q = f"{df.columns[0]}=='{label}'"
+        series = df.query(q)
+        plt.plot(series[df.columns[1]], series[df.columns[2]], label=q, marker="o")
+    plt.xlabel(df.columns[1])
+    plt.ylabel(df.columns[2])
+    plt.legend()
+    plt.savefig(bench_args.graph_filename, dpi=300)
+
+    # if in kitty, just dump it to the terminal
+    if os.environ["TERM"] == "xterm-kitty":
+        os.system(
+            f"kitty icat --use-window-size 1,1,600,600 {bench_args.graph_filename}"
+        )
+
+
 def main(server_args, bench_args):
-    print(bench_args)
 
-    if bench_args.correctness_test:
-        work_func = correctness_test
+    if server_args.model_path:
+        if bench_args.correctness_test:
+            work_func = correctness_test
+        else:
+            work_func = latency_test
+    elif os.path.isfile(bench_args.result_filename):
+        assert bench_args.graph_filename, "please provide a filename for the graph"
+        work_func = plot_latency_test
     else:
-        work_func = latency_test
+        raise ValueError(
+            "Provide --model-path for running the tests or "
+            "provide --result-filename for plotting the results"
+        )
 
     if server_args.tp_size == 1:
         work_func(server_args, bench_args, 0)
@@ -361,6 +465,11 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     ServerArgs.add_cli_args(parser)
     BenchArgs.add_cli_args(parser)
+    # For this script, model-path is not required
+    assert (
+        parser._actions[1].option_strings[0] == "--model-path"
+    ), "options changed, this code need to be updated"
+    parser._actions[1].required = False
     args = parser.parse_args()
 
     server_args = ServerArgs.from_cli_args(args)

sglang/check_env.py

@@ -14,6 +14,7 @@ PACKAGE_LIST = [
     "sglang",
     "flashinfer",
     "triton",
+    "transformers",
     "requests",
     "tqdm",
     "numpy",
@@ -73,10 +74,26 @@ def _get_gpu_info():
     Get information about available GPUs.
     """
     devices = defaultdict(list)
+    capabilities = defaultdict(list)
     for k in range(torch.cuda.device_count()):
         devices[torch.cuda.get_device_name(k)].append(str(k))
+        capability = torch.cuda.get_device_capability(k)
+        capabilities[f"{capability[0]}.{capability[1]}"].append(str(k))
 
-    return {f"GPU {','.join(device_ids)}": name for name, device_ids in devices.items()}
+    gpu_info = {}
+    for name, device_ids in devices.items():
+        gpu_info[f"GPU {','.join(device_ids)}"] = name
+
+    if len(capabilities) == 1:
+        # All GPUs have the same compute capability
+        cap, gpu_ids = list(capabilities.items())[0]
+        gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap
+    else:
+        # GPUs have different compute capabilities
+        for cap, gpu_ids in capabilities.items():
+            gpu_info[f"GPU {','.join(gpu_ids)} Compute Capability"] = cap
+
+    return gpu_info
 
 
 def _get_cuda_version_info():
@@ -118,6 +135,7 @@ def _get_cuda_driver_version():
     """
     Get CUDA driver version.
     """
+    versions = set()
     try:
         output = subprocess.check_output(
             [
@@ -126,7 +144,11 @@ def _get_cuda_driver_version():
                 "--format=csv,noheader,nounits",
             ]
         )
-        return {"CUDA Driver Version": output.decode().strip()}
+        versions = set(output.decode().strip().split("\n"))
+        if len(versions) == 1:
+            return {"CUDA Driver Version": versions.pop()}
+        else:
+            return {"CUDA Driver Versions": ", ".join(sorted(versions))}
     except subprocess.SubprocessError:
         return {"CUDA Driver Version": "Not Available"}
 

sglang/global_config.py

@@ -19,7 +19,6 @@ class GlobalConfig:
         self.init_new_token_ratio = 0.7
         self.base_min_new_token_ratio = 0.1
         self.new_token_ratio_decay = 0.001
-        self.new_token_ratio_recovery = 0.05
 
         # Runtime constants: The threshold (number of tokens) to trigger layer-wise cuda sync.
         # This can improve the speed for large batch sizes during prefill.

sglang/lang/backend/base_backend.py

@@ -1,6 +1,7 @@
 from typing import Callable, List, Optional, Union
 
 from sglang.lang.chat_template import get_chat_template
+from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod
 from sglang.lang.interpreter import StreamExecutor
 from sglang.lang.ir import SglSamplingParams
 
@@ -64,7 +65,8 @@ class BaseBackend:
         s: StreamExecutor,
         choices: List[str],
         temperature: float,
-    ):
+        choices_method: Optional[ChoicesSamplingMethod] = None,
+    ) -> ChoicesDecision:
         raise NotImplementedError()
 
     def concatenate_and_append(self, src_rids: List[str], dst_rid: str):

sglang/lang/backend/openai.py

@@ -8,6 +8,7 @@ import numpy as np
 
 from sglang.lang.backend.base_backend import BaseBackend
 from sglang.lang.chat_template import ChatTemplate, get_chat_template_by_model_path
+from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod
 from sglang.lang.interpreter import StreamExecutor
 from sglang.lang.ir import SglSamplingParams
 
@@ -296,7 +297,9 @@ class OpenAI(BaseBackend):
         s: StreamExecutor,
         choices: List[str],
         temperature: float,
-    ):
+        choices_method: ChoicesSamplingMethod,
+    ) -> ChoicesDecision:
+        """Note: `choices_method` is not used by the OpenAI backend."""
         if self.is_chat_model:
             raise NotImplementedError(
                 "select/choices is not supported for chat models. "
@@ -354,8 +357,10 @@ class OpenAI(BaseBackend):
 
             prompt_tokens.append(ret_token)
 
-        decision = choices[np.argmax(scores)]
-        return decision, scores, None, None
+        return ChoicesDecision(
+            decision=choices[np.argmax(scores)],
+            meta_info={"scores": scores},
+        )
 
 
 def openai_completion(

sglang/lang/backend/runtime_endpoint.py

@@ -1,17 +1,21 @@
 import json
 from typing import List, Optional
 
-import numpy as np
-
 from sglang.global_config import global_config
 from sglang.lang.backend.base_backend import BaseBackend
 from sglang.lang.chat_template import get_chat_template_by_model_path
+from sglang.lang.choices import (
+    ChoicesDecision,
+    ChoicesSamplingMethod,
+    token_length_normalized,
+)
 from sglang.lang.interpreter import StreamExecutor
 from sglang.lang.ir import SglSamplingParams
 from sglang.utils import http_request
 
 
 class RuntimeEndpoint(BaseBackend):
+
     def __init__(
         self,
         base_url: str,
@@ -43,7 +47,7 @@ class RuntimeEndpoint(BaseBackend):
     def flush_cache(self):
         res = http_request(
             self.base_url + "/flush_cache",
-            auth_token=self.auth_token,
+            api_key=self.api_key,
             verify=self.verify,
         )
         self._assert_success(res)
@@ -51,7 +55,7 @@ class RuntimeEndpoint(BaseBackend):
     def get_server_args(self):
         res = http_request(
             self.base_url + "/get_server_args",
-            auth_token=self.auth_token,
+            api_key=self.api_key,
             verify=self.verify,
         )
         self._assert_success(res)
@@ -208,20 +212,14 @@ class RuntimeEndpoint(BaseBackend):
         s: StreamExecutor,
         choices: List[str],
         temperature: float,
-    ):
+        choices_method: ChoicesSamplingMethod,
+    ) -> ChoicesDecision:
         assert temperature <= 1e-5
 
         # Cache common prefix
         data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
-        self._add_images(s, data)
-        res = http_request(
-            self.base_url + "/generate",
-            json=data,
-            api_key=self.api_key,
-            verify=self.verify,
-        )
-        self._assert_success(res)
-        prompt_len = res.json()["meta_info"]["prompt_tokens"]
+        obj = self._generate_http_request(s, data)
+        prompt_len = obj["meta_info"]["prompt_tokens"]
 
         # Compute logprob
         data = {
@@ -230,27 +228,35 @@ class RuntimeEndpoint(BaseBackend):
             "return_logprob": True,
             "logprob_start_len": max(prompt_len - 2, 0),
         }
-        self._add_images(s, data)
-        res = http_request(
-            self.base_url + "/generate",
-            json=data,
-            api_key=self.api_key,
-            verify=self.verify,
-        )
-        self._assert_success(res)
-        obj = res.json()
+        obj = self._generate_http_request(s, data)
+
         normalized_prompt_logprobs = [
             r["meta_info"]["normalized_prompt_logprob"] for r in obj
         ]
-        decision = choices[np.argmax(normalized_prompt_logprobs)]
         input_token_logprobs = [r["meta_info"]["input_token_logprobs"] for r in obj]
         output_token_logprobs = [r["meta_info"]["output_token_logprobs"] for r in obj]
 
-        return (
-            decision,
-            normalized_prompt_logprobs,
-            input_token_logprobs,
-            output_token_logprobs,
+        # Compute unconditional logprobs if required
+        if choices_method.requires_unconditional_logprobs:
+            input_ids = [[el[1] for el in subl] for subl in input_token_logprobs]
+            data = {
+                "input_ids": input_ids,
+                "sampling_params": {"max_new_tokens": 0},
+                "return_logprob": True,
+            }
+            obj = self._generate_http_request(s, data)
+            unconditional_token_logprobs = [
+                r["meta_info"]["input_token_logprobs"] for r in obj
+            ]
+        else:
+            unconditional_token_logprobs = None
+
+        return choices_method(
+            choices=choices,
+            normalized_prompt_logprobs=normalized_prompt_logprobs,
+            input_token_logprobs=input_token_logprobs,
+            output_token_logprobs=output_token_logprobs,
+            unconditional_token_logprobs=unconditional_token_logprobs,
         )
 
     def concatenate_and_append(self, src_rids: List[str], dst_rid: str):
@@ -262,6 +268,17 @@ class RuntimeEndpoint(BaseBackend):
         )
         self._assert_success(res)
 
+    def _generate_http_request(self, s: StreamExecutor, data):
+        self._add_images(s, data)
+        res = http_request(
+            self.base_url + "/generate",
+            json=data,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+        return res.json()
+
     def _add_images(self, s: StreamExecutor, data):
         if s.images_:
             assert len(s.images_) == 1, "Only support one image."