Make bench_one_batch_server compatible for more backends (#18512)

python/sglang/test/bench_one_batch_server_internal.py

@@ -109,6 +109,7 @@ class BenchArgs:
     append_to_github_summary: bool = True
     seed: int = 42
     cache_hit_rate: float = 0.0
+    backend: str = "sglang"
     server_args_for_metrics: Optional[List[str]] = None
 
     @staticmethod
@@ -193,6 +194,13 @@ class BenchArgs:
             help="Cache hit rate for benchmarking (0.0-1.0). "
             "0.0 means no cache hits (flush all), 0.4 means 40%% of input tokens are cached.",
         )
+        parser.add_argument(
+            "--backend",
+            type=str,
+            default=BenchArgs.backend,
+            choices=["sglang", "vllm"],
+            help="Backend server type (sglang or vllm).",
+        )
         parser.add_argument(
             "--server-args-for-metrics",
             type=str,
@@ -289,8 +297,10 @@ def _warmup_cache(
     cache_hit_rate: float,
     dataset_name: str = "random",
     image_data: Optional[List] = None,
+    backend: str = "sglang",
+    model_name: Optional[str] = None,
 ):
-    """Warm up the cache by sending prefix tokens to populate the radix cache.
+    """Warm up the cache by sending prefix tokens to populate the radix/prefix cache.
 
     Args:
         url: Server URL
@@ -299,6 +309,8 @@ def _warmup_cache(
         cache_hit_rate: Fraction of input tokens to cache (0.0-1.0)
         dataset_name: Name of the dataset (used to determine if image data should be included)
         image_data: Optional image data for VLM models
+        backend: Backend server type ("sglang" or "vllm")
+        model_name: Model name (required for vllm backend)
     """
     cached_token_len = int(input_len * cache_hit_rate)
     if cached_token_len <= 0:
@@ -310,6 +322,17 @@ def _warmup_cache(
     )
     # Create prefix input_ids for cache warming
     cache_warmup_input_ids = [ids[:cached_token_len] for ids in input_ids]
+
+    if backend == "vllm":
+        cache_warmup_payload = {
+            "model": model_name,
+            "prompt": cache_warmup_input_ids,
+            "max_tokens": 1,
+            "temperature": 0.0,
+            "stream": False,
+        }
+        gen_url = url + "/v1/completions"
+    else:
         cache_warmup_payload = {
             "input_ids": cache_warmup_input_ids,
             "sampling_params": {
@@ -322,9 +345,10 @@ def _warmup_cache(
         if dataset_name == "mmmu" and image_data is not None:
             # include image data in cache warmup
             cache_warmup_payload["image_data"] = image_data
+        gen_url = url + "/generate"
 
     warmup_response = requests.post(
-        url + "/generate",
+        gen_url,
         json=cache_warmup_payload,
         timeout=DEFAULT_TIMEOUT,
     )
@@ -353,7 +377,14 @@ def run_one_case(
     dataset_path: str = BenchArgs.dataset_path,
     parallel_batch: bool = False,
     cache_hit_rate: float = BenchArgs.cache_hit_rate,
+    backend: str = "sglang",
+    model_name: Optional[str] = None,
 ):
+    if backend == "vllm":
+        # You need to have export VLLM_SERVER_DEV_MODE=1 in your environment to use this endpoint.
+        response = requests.post(url + "/reset_prefix_cache", timeout=DEFAULT_TIMEOUT)
+        response.raise_for_status()
+    else:
         response = requests.post(url + "/flush_cache", timeout=DEFAULT_TIMEOUT)
         response.raise_for_status()
 
@@ -378,6 +409,32 @@ def run_one_case(
             return_text=False,
         )
 
+    # Extract input_ids from requests
+    if dataset_name == "mmmu":
+        input_ids = []
+        # for vlms, tokenizer is an instance of AutoProcessor
+        tokenizer = tokenizer.tokenizer
+        for input_req in input_requests:
+            input_ids += [tokenizer.encode(input_req.prompt)]
+        image_data = [req.image_data for req in input_requests]
+    else:
+        input_ids = [req.prompt for req in input_requests]
+        image_data = None
+
+    # Build payload based on backend
+    if backend == "vllm":
+        payload = {
+            "model": model_name,
+            "prompt": input_ids,
+            "max_tokens": output_len,
+            "temperature": temperature,
+            "stream": True,
+            "ignore_eos": True,
+        }
+        if return_logprob:
+            payload["logprobs"] = 1
+        gen_url = url + "/v1/completions"
+    else:
         # Load sampling parameters
         use_structured_outputs = False
         if use_structured_outputs:
@@ -404,19 +461,10 @@ def run_one_case(
             "stream": True,
             **({"parallel_batch": parallel_batch} if parallel_batch else {}),
         }
-    if dataset_name == "mmmu":
-        # vlm
-        input_ids = []
-        # for vlms, tokenizer is an instance of AutoProcessor
-        tokenizer = tokenizer.tokenizer
-        for input_req in input_requests:
-            input_ids += [tokenizer.encode(input_req.prompt)]
-        payload["image_data"] = [req.image_data for req in input_requests]
-
-    else:
-        input_ids = [req.prompt for req in input_requests]
-
         payload["input_ids"] = input_ids
+        if image_data is not None:
+            payload["image_data"] = image_data
+        gen_url = url + "/generate"
 
     # Warm up cache if cache_hit_rate > 0.0
     if cache_hit_rate > 0.0:
@@ -426,7 +474,9 @@ def run_one_case(
             input_len=input_len,
             cache_hit_rate=cache_hit_rate,
             dataset_name=dataset_name,
-            image_data=payload.get("image_data"),
+            image_data=image_data,
+            backend=backend,
+            model_name=model_name,
         )
 
     # Turn on profiler
@@ -447,7 +497,7 @@ def run_one_case(
     # Run the request
     tic = time.perf_counter()
     response = requests.post(
-        url + "/generate",
+        gen_url,
         json=payload,
         stream=True,
         timeout=DEFAULT_TIMEOUT,
@@ -456,6 +506,25 @@ def run_one_case(
 
     # Get the TTFT of the last request in the batch
     last_ttft = 0.0
+    if backend == "vllm":
+        # Parse OpenAI-compatible streaming format from vLLM
+        first_token_indices = set()
+        for chunk in response.iter_lines(decode_unicode=False):
+            chunk = chunk.decode("utf-8")
+            if chunk and chunk.startswith("data:"):
+                data_str = chunk[5:].strip()
+                if data_str == "[DONE]":
+                    break
+                data = json.loads(data_str)
+                if "error" in data:
+                    raise RuntimeError(f"Request has failed. {data}.")
+                for choice in data.get("choices", []):
+                    idx = choice["index"]
+                    if idx not in first_token_indices:
+                        first_token_indices.add(idx)
+                        if len(first_token_indices) == batch_size:
+                            last_ttft = time.perf_counter() - tic
+    else:
         for chunk in response.iter_lines(decode_unicode=False):
             chunk = chunk.decode("utf-8")
             if chunk and chunk.startswith("data:"):
@@ -478,6 +547,11 @@ def run_one_case(
     output_throughput = batch_size * output_len / (latency - last_ttft)
     overall_throughput = batch_size * (input_len + output_len) / latency
 
+    if backend == "vllm":
+        # vLLM does not expose these metrics via API
+        last_gen_throughput = -1
+        acc_length = -1
+    else:
         response = requests.get(url + "/get_server_info", timeout=DEFAULT_TIMEOUT)
         response.raise_for_status()
         server_info = response.json()
@@ -638,7 +712,30 @@ def run_benchmark_internal(
     else:
         proc, base_url = launch_server_process(launch_server_func, server_args)
 
-    # Get tokenizer
+    # Get tokenizer and server info
+    if bench_args.backend == "vllm":
+        # For vLLM, get model name from /v1/models endpoint
+        print(f"Connecting to vLLM server at {base_url}...")
+        response = requests.get(base_url + "/v1/models", timeout=DEFAULT_TIMEOUT)
+        response.raise_for_status()
+        model_list = response.json().get("data", [])
+        if not model_list:
+            raise RuntimeError("No models found on vLLM server via /v1/models")
+        model_name = model_list[0]["id"]
+        print(f"Found model: {model_name}")
+        print(f"Loading tokenizer for {model_name}...")
+        if bench_args.dataset_name == "mmmu":
+            tokenizer = get_processor(model_name)
+        else:
+            tokenizer = get_tokenizer(model_name)
+        print("Tokenizer loaded.")
+
+        server_info = {"model_name": model_name}
+        # vLLM does not expose token capacity or max running requests via API
+        skip_token_capacity_threshold = float("inf")
+        skip_max_running_requests_threshold = float("inf")
+    else:
+        model_name = None
         response = requests.get(base_url + "/get_server_info", timeout=DEFAULT_TIMEOUT)
         response.raise_for_status()
         server_info = response.json()
@@ -688,6 +785,8 @@ def run_benchmark_internal(
                 dataset_name=bench_args.dataset_name,
                 dataset_path=bench_args.dataset_path,
                 parallel_batch=bench_args.parallel_batch,
+                backend=bench_args.backend,
+                model_name=model_name,
             )
         print("=" * 8 + " Warmup End   " + "=" * 8 + "\n")
 
@@ -721,6 +820,8 @@ def run_benchmark_internal(
                     dataset_path=bench_args.dataset_path,
                     parallel_batch=bench_args.parallel_batch,
                     cache_hit_rate=bench_args.cache_hit_rate,
+                    backend=bench_args.backend,
+                    model_name=model_name,
                 )
             )
 
@@ -761,6 +862,8 @@ def run_benchmark_internal(
                             profile_by_stage=bench_args.profile_by_stage,
                             profile_prefix=profile_prefix,
                             profile_output_dir=bench_args.profile_output_dir,
+                            backend=bench_args.backend,
+                            model_name=model_name,
                         )
                     )
             except Exception as e: