[Pooling] Report actual input token usage for scoring APIs (#56573)
Signed-off-by: Taneem Ibrahim <taneem.ibrahim@gmail.com>
Taneem Ibrahim committed
Sep 13, 2026 at 08:20 UTC
2e9f7bb45475c7e461079284a73ec9c53761c6b5
2 files changed
+17
-15
tests/models/language/pooling/test_max_tokens_per_doc.py
+8
-8
@@ -80,10 +80,10 @@ RERANK_CONFIGS = [
80
],
81
# This model has produced both prompt-token totals in CI/local cache;
82
# keep truncation checks exact while tolerating the boundary delta.
83
- without_truncated_prompt_tokens=(285, 286),
84
- with_max_tokens_per_query_prompt_tokens=(155, 156),
85
- with_max_tokens_per_doc_prompt_tokens=155,
86
- with_max_tokens_per_query_and_doc_prompt_tokens=25,
83
+ without_truncated_prompt_tokens=(284, 285),
84
+ with_max_tokens_per_query_prompt_tokens=(154, 155),
85
+ with_max_tokens_per_doc_prompt_tokens=154,
86
+ with_max_tokens_per_query_and_doc_prompt_tokens=24,
87
),
88
# 4. late-interaction
89
TestConfig(
@@ -94,10 +94,10 @@ RERANK_CONFIGS = [
94
"512",
95
"--trust-remote-code",
96
],
97
- without_truncated_prompt_tokens=285,
98
- with_max_tokens_per_query_prompt_tokens=155,
99
- with_max_tokens_per_doc_prompt_tokens=155,
100
- with_max_tokens_per_query_and_doc_prompt_tokens=25,
97
+ without_truncated_prompt_tokens=284,
98
+ with_max_tokens_per_query_prompt_tokens=154,
99
+ with_max_tokens_per_doc_prompt_tokens=154,
100
+ with_max_tokens_per_query_and_doc_prompt_tokens=24,
101
),
102
# 5. jinaai/jina-reranker-v3
103
TestConfig(
vllm/entrypoints/pooling/scoring/serving.py
+9
-7
@@ -80,6 +80,11 @@ class ServingScores(PoolingServing):
80
request_id = ctx.request_id
81
created_time = ctx.created_time
82
model_name = ctx.model_name
83
+ assert ctx.engine_inputs is not None
84
+ num_prompt_tokens = sum(
85
+ self._extract_prompt_len(engine_input["prompts"])
86
+ for engine_input in ctx.engine_inputs
87
+ )
88
89
if isinstance(ctx.request, ScoreRequest):
90
return self._request_output_to_score_response(
@@ -87,6 +92,7 @@ class ServingScores(PoolingServing):
92
request_id,
93
created_time,
94
model_name,
95
+ num_prompt_tokens,
96
)
97
elif isinstance(ctx.request, RerankRequest):
98
return self._request_output_to_rerank_response(
@@ -95,6 +101,7 @@ class ServingScores(PoolingServing):
101
model_name,
102
ctx.request.documents,
103
ctx.request.top_n if ctx.request.top_n > 0 else len(final_res_batch),
104
+ num_prompt_tokens,
105
)
106
else:
107
raise ValueError(f"Invalid {self.request_id_prefix} request type")
@@ -105,9 +112,9 @@ class ServingScores(PoolingServing):
112
request_id: str,
113
created_time: int,
114
model_name: str,
115
+ num_prompt_tokens: int,
116
) -> JSONResponse:
117
items: list[ScoreResponseData] = []
110
- num_prompt_tokens = 0
118
119
for idx, final_res in enumerate(final_res_batch):
120
classify_res = ScoringRequestOutput.from_base(final_res)
@@ -116,10 +123,7 @@ class ServingScores(PoolingServing):
123
index=idx,
124
score=classify_res.outputs.score,
125
)
119
- prompt_token_ids = final_res.prompt_token_ids
120
-
126
items.append(item)
122
- num_prompt_tokens += len(prompt_token_ids)
127
128
usage = UsageInfo(
129
prompt_tokens=num_prompt_tokens,
@@ -143,12 +147,12 @@ class ServingScores(PoolingServing):
147
model_name: str,
148
documents: ScoreInput | list[ScoreInput],
149
top_n: int,
150
+ num_prompt_tokens: int,
151
) -> JSONResponse:
152
if not isinstance(documents, list):
153
documents = [documents]
154
155
results: list[RerankResult] = []
151
- num_prompt_tokens = 0
156
for idx, final_res in enumerate(final_res_batch):
157
classify_res = ScoringRequestOutput.from_base(final_res)
158
@@ -166,8 +170,6 @@ class ServingScores(PoolingServing):
170
relevance_score=classify_res.outputs.score,
171
)
172
results.append(result)
169
- prompt_token_ids = final_res.prompt_token_ids
170
- num_prompt_tokens += len(prompt_token_ids)
173
174
# sort by relevance, then return the top n if set
175
results.sort(key=lambda x: x.relevance_score, reverse=True)