[Pooling] Report actual input token usage for scoring APIs (#56573)

Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>

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(

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)