sglang 0.2.9.post1__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):
+# 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:
+## 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,17 +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 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
@@ -47,25 +61,50 @@ from sglang.srt.utils import suppress_other_loggers
 
 @dataclasses.dataclass
 class BenchArgs:
-    batch_size: int = 1
-    input_len: int = 1024
-    output_len: int = 4
+    run_name: str = "before"
+    batch_size: Tuple[int] = (1,)
+    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("--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("--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, 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):
-        attrs = [attr.name for attr in dataclasses.fields(cls)]
-        return cls(**{attr: getattr(args, attr) for attr in attrs})
+        # use the default value's type to case the args into correct types.
+        attrs = [(attr.name, type(attr.default)) for attr in dataclasses.fields(cls)]
+        return cls(
+            **{attr: attr_type(getattr(args, attr)) for attr, attr_type in attrs}
+        )
 
 
 def load_model(server_args, tp_rank):
@@ -93,7 +132,7 @@ def load_model(server_args, tp_rank):
     return model_runner, tokenizer
 
 
-def prepare_inputs(bench_args, tokenizer):
+def prepare_inputs_for_correctness_test(bench_args, tokenizer):
     prompts = [
         "The capital of France is",
         "The capital of the United Kindom is",
@@ -119,7 +158,9 @@ def prepare_inputs(bench_args, tokenizer):
     return input_ids, reqs
 
 
-def prepare_extend_inputs(bench_args, input_ids, reqs, model_runner):
+def prepare_extend_inputs_for_correctness_test(
+    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 :]
@@ -129,8 +170,8 @@ def prepare_extend_inputs(bench_args, input_ids, reqs, model_runner):
     return reqs
 
 
-def prepare_synthetic_inputs(bench_args, tokenizer):
-    input_ids = np.ones((bench_args.batch_size, bench_args.input_len), dtype=np.int32)
+def prepare_synthetic_inputs_for_latency_test(batch_size, input_len):
+    input_ids = np.ones((batch_size, input_len), dtype=np.int32)
     sampling_params = SamplingParams(
         temperature=0,
         max_new_tokens=BenchArgs.output_len,
@@ -148,7 +189,7 @@ def prepare_synthetic_inputs(bench_args, tokenizer):
 
 
 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,
@@ -179,7 +220,7 @@ def correctness_test(
     model_runner, tokenizer = load_model(server_args, tp_rank)
 
     # Prepare inputs
-    input_ids, reqs = prepare_inputs(bench_args, tokenizer)
+    input_ids, reqs = prepare_inputs_for_correctness_test(bench_args, tokenizer)
 
     if bench_args.cut_len > 0:
         # Prefill
@@ -187,7 +228,9 @@ def correctness_test(
         rank_print("prefill logits (first half)", next_token_logits)
 
     # Prepare extend inputs
-    reqs = prepare_extend_inputs(bench_args, input_ids, reqs, model_runner)
+    reqs = prepare_extend_inputs_for_correctness_test(
+        bench_args, input_ids, reqs, model_runner
+    )
 
     # Extend
     next_token_ids, next_token_logits, batch = extend(reqs, model_runner)
@@ -205,6 +248,74 @@ def correctness_test(
         rank_print(tokenizer.decode(output_ids[i]))
 
 
+@torch.inference_mode()
+def latency_test_run_once(
+    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": run_name,
+        "batch_size": batch_size,
+        "input_len": input_len,
+        "output_len": output_len,
+    }
+
+    tot_latency = 0
+
+    # Prefill
+    torch.cuda.synchronize()
+    tic = time.time()
+    next_token_ids, _, batch = extend(reqs, model_runner)
+    torch.cuda.synchronize()
+    prefill_latency = time.time() - tic
+    tot_latency += prefill_latency
+    throughput = input_len * batch_size / prefill_latency
+    rank_print(
+        f"Prefill. latency: {prefill_latency:6.5f} s, throughput: {throughput:9.2f} token/s"
+    )
+    measurement_results["prefill_latency"] = prefill_latency
+    measurement_results["prefill_throughput"] = throughput
+
+    # 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 = 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 = 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"
+    )
+    measurement_results["avg_decode_latency"] = avg_decode_latency
+    measurement_results["avg_decode_throughput"] = avg_decode_throughput
+
+    throughput = (input_len + output_len) * batch_size / tot_latency
+    rank_print(
+        f"Total. latency: {tot_latency:6.3f} s, throughput: {throughput:9.2f} token/s"
+    )
+    measurement_results["total_latency"] = tot_latency
+    measurement_results["total_throughput"] = throughput
+    return measurement_results
+
+
 def latency_test(
     server_args,
     bench_args,
@@ -214,75 +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)}"
-    )
 
-    # Prepare inputs
-    reqs = prepare_synthetic_inputs(bench_args, tokenizer)
+    # Prepare inputs for warm up
+    reqs = prepare_synthetic_inputs_for_latency_test(
+        bench_args.batch_size[0], bench_args.input_len[0]
+    )
 
-    def clear():
-        model_runner.req_to_token_pool.clear()
-        model_runner.token_to_kv_pool.clear()
+    # Warm up
+    latency_test_run_once(
+        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
+    )
 
-    @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"
+    # Run the sweep
+    result_list = []
+    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)
 
-        # 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"
-        )
+    # Write results in jsonlines format on rank 0.
+    if tp_rank == 0 and bench_args.result_filename:
+        import jsonlines
 
-        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"
-        )
+        with jsonlines.open(bench_args.result_filename, "a") as f:
+            f.write_all(result_list)
 
-    # Warm up
-    run_once(4)
-    clear()
 
-    # Run again
-    run_once(bench_args.output_len)
+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)
@@ -310,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

@@ -13,6 +13,8 @@ import torch
 PACKAGE_LIST = [
     "sglang",
     "flashinfer",
+    "triton",
+    "transformers",
     "requests",
     "tqdm",
     "numpy",
@@ -72,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():
@@ -117,6 +135,7 @@ def _get_cuda_driver_version():
     """
     Get CUDA driver version.
     """
+    versions = set()
     try:
         output = subprocess.check_output(
             [
@@ -125,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(