[Bugfix] Measure complete pooling responses (#56760)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Sep 15, 2026 at 08:21 UTC
64856080b05054fc62754f37234f7b483b753445
1 file changed
+15
-1
vllm/benchmarks/lib/endpoint_request_func.py
+15
-1
@@ -606,6 +606,7 @@ async def _run_pooling_request(
606
headers: dict[str, Any],
607
pbar: tqdm | None = None,
608
num_input_sequences: int = 1,
609
+ prompt_len: int = 0,
610
) -> RequestFuncOutput:
611
output = RequestFuncOutput(num_input_sequences=num_input_sequences)
612
st = time.perf_counter()
@@ -613,11 +614,19 @@ async def _run_pooling_request(
614
try:
615
async with session.post(url=api_url, headers=headers, json=payload) as response:
616
if response.status == 200:
617
+ encoding_format = payload.get("encoding_format", "float")
618
+ if encoding_format in ("bytes", "bytes_only"):
619
+ async for _ in response.content.iter_any():
620
+ pass
621
+ else:
622
+ await response.read()
623
output.ttft = output.latency = time.perf_counter() - st
624
618
- if payload.get("encoding_format", "float") == "bytes":
625
+ if encoding_format == "bytes":
626
metadata = json.loads(response.headers["metadata"])
627
usage = metadata.get("usage", {})
628
+ elif encoding_format == "bytes_only":
629
+ usage = {"prompt_tokens": prompt_len}
630
else:
631
data = await response.json()
632
usage = data.get("usage", {})
@@ -672,6 +681,7 @@ async def async_request_openai_embeddings(
681
headers=headers,
682
pbar=pbar,
683
num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
684
+ prompt_len=request_func_input.prompt_len,
685
)
686
687
@@ -709,6 +719,7 @@ async def async_request_vllm_rerank(
719
headers=headers,
720
pbar=pbar,
721
num_input_sequences=len(request_func_input.prompt) - 1,
722
+ prompt_len=request_func_input.prompt_len,
723
)
724
725
@@ -743,6 +754,7 @@ async def async_request_openai_embeddings_chat(
754
payload=payload,
755
headers=headers,
756
pbar=pbar,
757
+ prompt_len=request_func_input.prompt_len,
758
)
759
760
@@ -847,6 +859,7 @@ async def async_request_infinity_embeddings(
859
headers=headers,
860
pbar=pbar,
861
num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
862
+ prompt_len=request_func_input.prompt_len,
863
)
864
865
@@ -896,6 +909,7 @@ async def async_request_vllm_pooling(
909
headers=headers,
910
pbar=pbar,
911
num_input_sequences=_get_num_input_sequences(request_func_input.prompt),
912
+ prompt_len=request_func_input.prompt_len,
913
)
914
915