[Benchmark] Add sglang-embedding backend to bench_serving (#20017)

Co-authored-by: Satyam Kumar <satyamk@linkedin.com>

python/sglang/bench_serving.py

@@ -47,6 +47,8 @@ from sglang.benchmark.utils import (
 
 _ROUTING_KEY_HEADER = "X-SMG-Routing-Key"
 
+_EMBEDDING_UNSUPPORTED_DATASETS = {"image", "mmmu", "mooncake"}
+
 TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) and (
     shutil.which("gnuplot") is not None
 )
@@ -699,6 +701,56 @@ async def async_request_sglang_generate(
     return output
 
 
+async def async_request_openai_embeddings(
+    request_func_input: RequestFuncInput,
+    pbar: Optional[tqdm] = None,
+) -> RequestFuncOutput:
+    api_url = request_func_input.api_url
+
+    async with _create_bench_client_session() as session:
+        payload = {
+            "input": request_func_input.prompt,
+            "model": request_func_input.model,
+        }
+
+        if request_func_input.lora_name:
+            payload["model"] = request_func_input.lora_name
+            payload["lora_path"] = request_func_input.lora_name
+
+        payload.update(request_func_input.extra_request_body)
+
+        headers = get_request_headers()
+        if request_func_input.routing_key:
+            headers[_ROUTING_KEY_HEADER] = request_func_input.routing_key
+
+        output = RequestFuncOutput.init_new(request_func_input)
+
+        st = time.perf_counter()
+        output.start_time = st
+        try:
+            async with session.post(
+                url=api_url, json=payload, headers=headers
+            ) as response:
+                if response.status == 200:
+                    await response.json()
+                    output.latency = time.perf_counter() - st
+                    output.success = True
+                    output.output_len = 0
+                else:
+                    output.error = (
+                        (response.reason or "") + ": " + (await response.text())
+                    )
+                    output.success = False
+        except Exception:
+            output.success = False
+            exc_info = sys.exc_info()
+            output.error = "".join(traceback.format_exception(*exc_info))
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
 async def async_request_gserver(
     request_func_input: RequestFuncInput,
     pbar: Optional[tqdm] = None,
@@ -808,6 +860,7 @@ ASYNC_REQUEST_FUNCS = {
     "sglang-native": async_request_sglang_generate,
     "sglang-oai": async_request_openai_completions,
     "sglang-oai-chat": async_request_openai_chat_completions,
+    "sglang-embedding": async_request_openai_embeddings,
     "vllm": async_request_openai_completions,
     "vllm-chat": async_request_openai_chat_completions,
     "lmdeploy": async_request_openai_completions,
@@ -1403,10 +1456,13 @@ async def benchmark(
                 "Total input vision tokens:", metrics.total_input_vision
             )
         )
+    is_embedding = backend == "sglang-embedding"
+    if not is_embedding:
         print("{:<40} {:<10}".format("Total generated tokens:", metrics.total_output))
         print(
             "{:<40} {:<10}".format(
-            "Total generated tokens (retokenized):", metrics.total_output_retokenized
+                "Total generated tokens (retokenized):",
+                metrics.total_output_retokenized,
             )
         )
     print(
@@ -1419,6 +1475,7 @@ async def benchmark(
             "Input token throughput (tok/s):", metrics.input_throughput
         )
     )
+    if not is_embedding:
         print(
             "{:<40} {:<10.2f}".format(
                 "Output token throughput (tok/s):", metrics.output_throughput
@@ -1426,7 +1483,8 @@ async def benchmark(
         )
         print(
             "{:<40} {:<10.2f}".format(
-            "Peak output token throughput (tok/s):", metrics.max_output_tokens_per_s
+                "Peak output token throughput (tok/s):",
+                metrics.max_output_tokens_per_s,
             )
         )
     print(
@@ -1434,6 +1492,7 @@ async def benchmark(
             "Peak concurrent requests:", metrics.max_concurrent_requests
         )
     )
+    if not is_embedding:
         print(
             "{:<40} {:<10.2f}".format(
                 "Total token throughput (tok/s):", metrics.total_throughput
@@ -1457,12 +1516,15 @@ async def benchmark(
     print(
         "{:<40} {:<10.2f}".format("P99 E2E Latency (ms):", metrics.p99_e2e_latency_ms)
     )
+    if not is_embedding:
         print("{s:{c}^{n}}".format(s="Time to First Token", n=50, c="-"))
         print("{:<40} {:<10.2f}".format("Mean TTFT (ms):", metrics.mean_ttft_ms))
         print("{:<40} {:<10.2f}".format("Median TTFT (ms):", metrics.median_ttft_ms))
         print("{:<40} {:<10.2f}".format("P99 TTFT (ms):", metrics.p99_ttft_ms))
         print(
-        "{s:{c}^{n}}".format(s="Time per Output Token (excl. 1st token)", n=50, c="-")
+            "{s:{c}^{n}}".format(
+                s="Time per Output Token (excl. 1st token)", n=50, c="-"
+            )
         )
         print("{:<40} {:<10.2f}".format("Mean TPOT (ms):", metrics.mean_tpot_ms))
         print("{:<40} {:<10.2f}".format("Median TPOT (ms):", metrics.median_tpot_ms))
@@ -1670,7 +1732,13 @@ def run_benchmark(args_: argparse.Namespace):
         else f"http://{args.host}:{args.port}/v1/models"
     )
 
-    if args.backend in ["sglang", "sglang-native"]:
+    if args.backend == "sglang-embedding":
+        api_url = (
+            f"{args.base_url}/v1/embeddings"
+            if args.base_url
+            else f"http://{args.host}:{args.port}/v1/embeddings"
+        )
+    elif args.backend in ["sglang", "sglang-native"]:
         api_url = (
             f"{args.base_url}/generate"
             if args.base_url
@@ -1739,12 +1807,19 @@ def run_benchmark(args_: argparse.Namespace):
         print("No model specified or found. Please provide a model using `--model`.")
         sys.exit(1)
 
-    if not check_chat_template(args.model):
+    if args.backend != "sglang-embedding" and not check_chat_template(args.model):
         print(
             "\nWARNING It is recommended to use the `Chat` or `Instruct` model for benchmarking.\n"
             "Because when the tokenizer counts the output tokens, if there is gibberish, it might count incorrectly.\n"
         )
 
+    if (
+        args.backend == "sglang-embedding"
+        and args.dataset_name in _EMBEDDING_UNSUPPORTED_DATASETS
+    ):
+        print(f"{args.dataset_name} dataset is unsupported for embeddings benchmark")
+        sys.exit(1)
+
     if args.dataset_name in ["image", "mmmu"]:
         args.apply_chat_template = True
         assert (

python/sglang/benchmark/datasets/common.py

@@ -54,6 +54,9 @@ class BaseDataset(ABC):
 
 
 def compute_random_lens(full_len: int, range_ratio: float, num: int) -> List[int]:
+    # full_len=0 is valid for embedding benchmarks where no output tokens are generated
+    if full_len <= 0:
+        return [0] * num
     return np.random.randint(
         max(int(full_len * range_ratio), 1),
         full_len + 1,