From 2e9f7bb45475c7e461079284a73ec9c53761c6b5 Mon Sep 17 00:00:00 2001 From: Taneem Ibrahim Date: Sun, 13 Sep 2026 08:20:21 -0500 Subject: [PATCH] [Pooling] Report actual input token usage for scoring APIs (#56573) Signed-off-by: Taneem Ibrahim --- .../language/pooling/test_max_tokens_per_doc.py | 16 ++++++++-------- vllm/entrypoints/pooling/scoring/serving.py | 16 +++++++++------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/tests/models/language/pooling/test_max_tokens_per_doc.py b/tests/models/language/pooling/test_max_tokens_per_doc.py index b6d6ba1454..d99aa49909 100644 --- a/tests/models/language/pooling/test_max_tokens_per_doc.py +++ b/tests/models/language/pooling/test_max_tokens_per_doc.py @@ -80,10 +80,10 @@ RERANK_CONFIGS = [ ], # This model has produced both prompt-token totals in CI/local cache; # keep truncation checks exact while tolerating the boundary delta. - without_truncated_prompt_tokens=(285, 286), - with_max_tokens_per_query_prompt_tokens=(155, 156), - with_max_tokens_per_doc_prompt_tokens=155, - with_max_tokens_per_query_and_doc_prompt_tokens=25, + without_truncated_prompt_tokens=(284, 285), + with_max_tokens_per_query_prompt_tokens=(154, 155), + with_max_tokens_per_doc_prompt_tokens=154, + with_max_tokens_per_query_and_doc_prompt_tokens=24, ), # 4. late-interaction TestConfig( @@ -94,10 +94,10 @@ RERANK_CONFIGS = [ "512", "--trust-remote-code", ], - without_truncated_prompt_tokens=285, - with_max_tokens_per_query_prompt_tokens=155, - with_max_tokens_per_doc_prompt_tokens=155, - with_max_tokens_per_query_and_doc_prompt_tokens=25, + without_truncated_prompt_tokens=284, + with_max_tokens_per_query_prompt_tokens=154, + with_max_tokens_per_doc_prompt_tokens=154, + with_max_tokens_per_query_and_doc_prompt_tokens=24, ), # 5. jinaai/jina-reranker-v3 TestConfig( diff --git a/vllm/entrypoints/pooling/scoring/serving.py b/vllm/entrypoints/pooling/scoring/serving.py index b32a7a5730..674522b7eb 100644 --- a/vllm/entrypoints/pooling/scoring/serving.py +++ b/vllm/entrypoints/pooling/scoring/serving.py @@ -80,6 +80,11 @@ class ServingScores(PoolingServing): request_id = ctx.request_id created_time = ctx.created_time model_name = ctx.model_name + assert ctx.engine_inputs is not None + num_prompt_tokens = sum( + self._extract_prompt_len(engine_input["prompts"]) + for engine_input in ctx.engine_inputs + ) if isinstance(ctx.request, ScoreRequest): return self._request_output_to_score_response( @@ -87,6 +92,7 @@ class ServingScores(PoolingServing): request_id, created_time, model_name, + num_prompt_tokens, ) elif isinstance(ctx.request, RerankRequest): return self._request_output_to_rerank_response( @@ -95,6 +101,7 @@ class ServingScores(PoolingServing): model_name, ctx.request.documents, ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch), + num_prompt_tokens, ) else: raise ValueError(f"Invalid {self.request_id_prefix} request type") @@ -105,9 +112,9 @@ class ServingScores(PoolingServing): request_id: str, created_time: int, model_name: str, + num_prompt_tokens: int, ) -> JSONResponse: items: list[ScoreResponseData] = [] - num_prompt_tokens = 0 for idx, final_res in enumerate(final_res_batch): classify_res = ScoringRequestOutput.from_base(final_res) @@ -116,10 +123,7 @@ class ServingScores(PoolingServing): index=idx, score=classify_res.outputs.score, ) - prompt_token_ids = final_res.prompt_token_ids - items.append(item) - num_prompt_tokens += len(prompt_token_ids) usage = UsageInfo( prompt_tokens=num_prompt_tokens, @@ -143,12 +147,12 @@ class ServingScores(PoolingServing): model_name: str, documents: ScoreInput | list[ScoreInput], top_n: int, + num_prompt_tokens: int, ) -> JSONResponse: if not isinstance(documents, list): documents = [documents] results: list[RerankResult] = [] - num_prompt_tokens = 0 for idx, final_res in enumerate(final_res_batch): classify_res = ScoringRequestOutput.from_base(final_res) @@ -166,8 +170,6 @@ class ServingScores(PoolingServing): relevance_score=classify_res.outputs.score, ) results.append(result) - prompt_token_ids = final_res.prompt_token_ids - num_prompt_tokens += len(prompt_token_ids) # sort by relevance, then return the top n if set results.sort(key=lambda x: x.relevance_score, reverse=True)