From 68f02749817edf7166718c3bb08cd7f3940765bd Mon Sep 17 00:00:00 2001 From: joonhyung-lee Date: Thu, 20 Aug 2026 21:51:43 +0900 Subject: [PATCH] feat(bench): add --dp-rank-roundrobin for explicit DP-rank routing Add an opt-in `--dp-rank-roundrobin N` flag to benchmark_serving.py that pins request i to data-parallel rank i % N via the `X-data-parallel-rank` HTTP header, bypassing the server-side DP balancer. Under injection=inf, vLLM's internal balancer can dispatch requests unevenly across DP ranks, leaving some engines idle while others queue, which depresses output TPS on high-DP configs. This flag lets a benchmark measure the ideal round-robin routing baseline as an A/B against the default balancer, with no server/SDK change. - RequestFuncInput gains an `extra_headers` field, merged into the request headers by the OpenAI completions and chat-completions request funcs. - benchmark() assigns the header per request when the flag is set. --- utils/bench_serving/backend_request_func.py | 5 +++++ utils/bench_serving/benchmark_serving.py | 23 ++++++++++++++++++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/utils/bench_serving/backend_request_func.py b/utils/bench_serving/backend_request_func.py index 32331a398f..14e4202446 100644 --- a/utils/bench_serving/backend_request_func.py +++ b/utils/bench_serving/backend_request_func.py @@ -30,6 +30,7 @@ class RequestFuncInput: extra_body: Optional[dict] = None multi_modal_content: Optional[dict] = None ignore_eos: bool = False + extra_headers: Optional[dict] = None @dataclass @@ -262,6 +263,8 @@ async def async_request_openai_completions( headers = { "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}" } + if request_func_input.extra_headers: + headers.update(request_func_input.extra_headers) output = RequestFuncOutput() output.prompt_len = request_func_input.prompt_len @@ -368,6 +371,8 @@ async def async_request_openai_chat_completions( "Content-Type": "application/json", "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", } + if request_func_input.extra_headers: + headers.update(request_func_input.extra_headers) output = RequestFuncOutput() output.prompt_len = request_func_input.prompt_len diff --git a/utils/bench_serving/benchmark_serving.py b/utils/bench_serving/benchmark_serving.py index 70334ea167..97c1692b60 100644 --- a/utils/bench_serving/benchmark_serving.py +++ b/utils/bench_serving/benchmark_serving.py @@ -342,6 +342,7 @@ async def benchmark( goodput_config_dict: Dict[str, float], max_concurrency: Optional[int], lora_modules: Optional[List[str]], + dp_rank_roundrobin: Optional[int] = None, ): if backend in ASYNC_REQUEST_FUNCS: request_func = ASYNC_REQUEST_FUNCS[backend] @@ -439,6 +440,7 @@ async def limited_request_func(request_func_input, pbar): benchmark_start_time = time.perf_counter() tasks: List[asyncio.Task] = [] + req_idx = 0 async for request in get_request(input_requests, request_rate, burstiness): prompt, prompt_len, output_len, mm_content = request req_model_id, req_model_name = model_id, model_name @@ -446,6 +448,15 @@ async def limited_request_func(request_func_input, pbar): req_lora_module = next(lora_modules) req_model_id, req_model_name = req_lora_module, req_lora_module + # Explicit round-robin over DP ranks: pin request i to rank i % N via + # vLLM's X-data-parallel-rank header, bypassing the server-side balancer. + extra_headers = None + if dp_rank_roundrobin: + extra_headers = { + "X-data-parallel-rank": str(req_idx % dp_rank_roundrobin) + } + req_idx += 1 + request_func_input = RequestFuncInput(model=req_model_id, model_name=req_model_name, prompt=prompt, @@ -455,7 +466,8 @@ async def limited_request_func(request_func_input, pbar): logprobs=logprobs, best_of=best_of, multi_modal_content=mm_content, - ignore_eos=ignore_eos) + ignore_eos=ignore_eos, + extra_headers=extra_headers) tasks.append( asyncio.create_task( limited_request_func(request_func_input=request_func_input, @@ -698,6 +710,7 @@ def main(args: argparse.Namespace): goodput_config_dict=goodput_config_dict, max_concurrency=args.max_concurrency, lora_modules=args.lora_modules, + dp_rank_roundrobin=args.dp_rank_roundrobin, )) # Save config and results to json @@ -869,6 +882,14 @@ def main(args: argparse.Namespace): "results in a more uniform arrival of requests.", ) parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--dp-rank-roundrobin", + type=int, + default=None, + help="If set to N, pin request i to data-parallel rank i %% N via the " + "X-data-parallel-rank header, bypassing the server-side DP balancer " + "(explicit round-robin routing). Use N = data_parallel_size.", + ) parser.add_argument( "--trust-remote-code", action="store_true",