From bcf298c28c449dfc3b5b45f07a6e04b63c147a05 Mon Sep 17 00:00:00 2001 From: Vladislav Nosivskoy Date: Thu, 18 Jun 2026 02:02:30 +0300 Subject: [PATCH] [HiCache & Bench] add cache hit breakdown in bench_serving (#22053) Signed-off-by: Vladislav Nosivskoy Co-authored-by: Zhiqiang Xie --- python/sglang/bench_serving.py | 115 +++++++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index a11adfa22..4e56b1857 100644 --- a/python/sglang/bench_serving.py +++ b/python/sglang/bench_serving.py @@ -106,6 +106,8 @@ class RequestFuncOutput: error: str = "" output_len: int = 0 start_time: float = 0.0 + cached_tokens: int = 0 + cached_tokens_details: Optional[Dict[str, Any]] = None @staticmethod def init_new(request_func_input: RequestFuncInput): @@ -229,6 +231,19 @@ async def async_request_trt_llm( return output +def _extract_cache_from_sglext(data, output): + """Extract cache hit details from sglext in OAI-compatible responses.""" + sglext = data.get("sglext") or {} + details = sglext.get("cached_tokens_details") + if details: + output.cached_tokens = ( + (details.get("device") or 0) + + (details.get("host") or 0) + + (details.get("storage") or 0) + ) + output.cached_tokens_details = details + + # set ignore_eos True by default async def async_request_openai_completions( request_func_input: RequestFuncInput, @@ -302,6 +317,9 @@ async def async_request_openai_completions( else: data = json.loads(chunk) + if getattr(args, "cache_report", False): + _extract_cache_from_sglext(data, output) + # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated @@ -455,6 +473,8 @@ async def async_request_openai_chat_completions( output.output_len = response_json.get("usage", {}).get( "completion_tokens", output_len ) + if getattr(args, "cache_report", False): + _extract_cache_from_sglext(response_json, output) else: # Streaming response async for chunk_bytes in response.content: @@ -474,6 +494,9 @@ async def async_request_openai_chat_completions( "completion_tokens", output_len ) + if getattr(args, "cache_report", False): + _extract_cache_from_sglext(data, output) + choices = data.get("choices") or [] if not choices: continue @@ -675,6 +698,13 @@ async def async_request_sglang_generate( # NOTE: Some completion API might have a last # usage summary response without a token so we # want to check a token was generated + if getattr(args, "cache_report", False): + _meta = data.get("meta_info") or {} + output.cached_tokens = _meta.get("cached_tokens", 0) + output.cached_tokens_details = _meta.get( + "cached_tokens_details" + ) + if "text" in data and data["text"]: timestamp = time.perf_counter() generated_text = data["text"] @@ -1608,6 +1638,58 @@ async def benchmark( print("{:<40} {:<10.2f}".format("P95 ITL (ms):", metrics.p95_itl_ms)) print("{:<40} {:<10.2f}".format("P99 ITL (ms):", metrics.p99_itl_ms)) print("{:<40} {:<10.2f}".format("Max ITL (ms):", metrics.max_itl_ms)) + if args.cache_report: + total_prompt_tokens = 0 + total_cached = 0 + total_device = total_host = total_storage = 0 + storage_backend_name = None + has_details = False + for o in outputs: + if not o.success: + continue + total_prompt_tokens += o.prompt_len + total_cached += o.cached_tokens + if o.cached_tokens_details: + has_details = True + total_device += o.cached_tokens_details.get("device") or 0 + total_host += o.cached_tokens_details.get("host") or 0 + s = o.cached_tokens_details.get("storage") or 0 + if s: + total_storage += s + storage_backend_name = o.cached_tokens_details.get( + "storage_backend" + ) + hit_rate = ( + total_cached / total_prompt_tokens * 100 if total_prompt_tokens > 0 else 0.0 + ) + + print("{s:{c}^{n}}".format(s="Cache Hit Details", n=50, c="-")) + print("{:<40} {:<10}".format("Total prompt tokens:", total_prompt_tokens)) + print("{:<40} {:<10}".format("Total cached tokens:", total_cached)) + if has_details and total_cached > 0: + print("{:<40} {:<10}".format(" Device:", total_device)) + print("{:<40} {:<10}".format(" Host:", total_host)) + if total_storage > 0: + label = ( + f" Storage ({storage_backend_name}):" + if storage_backend_name + else " Storage:" + ) + print("{:<40} {:<10}".format(label, total_storage)) + print("{:<40} {:.1f}%".format("Cache hit rate:", hit_rate)) + if has_details and total_cached > 0: + device_pct = total_device / total_cached * 100 + host_pct = total_host / total_cached * 100 + print("{:<40} {:.1f}%".format(" Device:", device_pct)) + print("{:<40} {:.1f}%".format(" Host:", host_pct)) + if total_storage > 0: + storage_pct = total_storage / total_cached * 100 + label = ( + f" Storage ({storage_backend_name}):" + if storage_backend_name + else " Storage:" + ) + print("{:<40} {:.1f}%".format(label, storage_pct)) print("=" * 50) resp = requests.get(base_url + "/server_info", headers=get_auth_headers()) @@ -1672,6 +1754,17 @@ async def benchmark( "max_output_tokens_per_s": metrics.max_output_tokens_per_s, "max_concurrent_requests": metrics.max_concurrent_requests, } + + if args.cache_report: + result["cache_report"] = { + "total_prompt_tokens": total_prompt_tokens, + "total_cached_tokens": total_cached, + "cache_hit_rate_pct": round(hit_rate, 2), + "device_cached_tokens": total_device if has_details else None, + "host_cached_tokens": total_host if has_details else None, + "storage_cached_tokens": (total_storage if total_storage > 0 else None), + "storage_backend": storage_backend_name, + } else: print(f"Error running benchmark for request rate: {request_rate}") print("-" * 30) @@ -1703,6 +1796,12 @@ async def benchmark( "errors": [output.error for output in outputs], } + if args.cache_report: + result_details["cached_tokens"] = [o.cached_tokens for o in outputs] + result_details["cached_tokens_details"] = [ + o.cached_tokens_details for o in outputs + ] + # Append results to a JSONL file with open(output_file_name, "a") as file: if args.output_details: @@ -1778,6 +1877,9 @@ def run_benchmark(args_: argparse.Namespace): if not hasattr(args, "served_model_name"): args.served_model_name = None + if not hasattr(args, "cache_report"): + args.cache_report = False + if getattr(args, "print_requests", False): assert args.backend == "sglang-oai-chat" # only support this now @@ -1792,6 +1894,13 @@ def run_benchmark(args_: argparse.Namespace): if args.extra_request_body: extra_request_body = json.loads(args.extra_request_body) + if args.cache_report: + sglang_backends = ("sglang", "sglang-native", "sglang-oai", "sglang-oai-chat") + if args.backend not in sglang_backends: + print("WARNING: --cache-report is only supported with sglang backends.") + elif args.backend in ("sglang-oai", "sglang-oai-chat"): + extra_request_body["return_cached_tokens_details"] = True + # Inject bootstrap fields for fake decode benchmarking if getattr(args, "fake_prefill", False): extra_request_body["bootstrap_host"] = FAKE_BOOTSTRAP_HOST @@ -2259,6 +2368,12 @@ if __name__ == "__main__": action="store_true", help="Return routed experts.", ) + parser.add_argument( + "--cache-report", + action="store_true", + help="Collect and display cache hit statistics after the benchmark. " + "Supported with sglang backends (native, oai, oai-chat).", + ) parser.add_argument("--seed", type=int, default=1, help="The random seed.") parser.add_argument( "--disable-ignore-eos",