[Score API][18132] return token usage in Score API response (#18381)

This commit is contained in:
Guy Stone
2026-03-03 16:45:35 -05:00
committed by GitHub
parent b0f26698f5
commit f749802402
7 changed files with 106 additions and 39 deletions

View File

@@ -67,6 +67,7 @@ from sglang.srt.managers.multi_tokenizer_mixin import MultiTokenizerRouter
from sglang.srt.managers.scheduler import run_scheduler_process
from sglang.srt.managers.template_manager import TemplateManager
from sglang.srt.managers.tokenizer_manager import TokenizerManager
from sglang.srt.managers.tokenizer_manager_multiitem_mixin import ScoreResult
from sglang.srt.model_loader.remote_instance_weight_loader_utils import (
parse_remote_instance_transfer_engine_info_from_scheduler_infos,
)
@@ -771,7 +772,7 @@ class Engine(EngineBase):
label_token_ids: Optional[List[int]] = None,
apply_softmax: bool = False,
item_first: bool = False,
) -> List[List[float]]:
) -> ScoreResult:
"""
Score the probability of specified token IDs appearing after the given (query + item) pair. For example:
query = "<|user|>Is the following city the capital of France? "
@@ -796,8 +797,9 @@ class Engine(EngineBase):
item_first: If True, prepend items to query. Otherwise append items to query.
Returns:
List of dictionaries mapping token IDs to their probabilities for each item.
Each dictionary in the list corresponds to one item input.
ScoreResult with:
scores: List of lists containing probabilities for each item and each label token
prompt_tokens: The number of prompt tokens processed.
Raises:
ValueError: If query is not provided, or if items is not provided,
@@ -821,7 +823,7 @@ class Engine(EngineBase):
label_token_ids: Optional[List[int]] = None,
apply_softmax: bool = False,
item_first: bool = False,
) -> List[List[float]]:
) -> ScoreResult:
"""
Asynchronous version of score method.

View File

@@ -376,13 +376,13 @@ class OpenAIServingRerank(OpenAIServingBase):
for doc in request.documents
]
probs = await self.tokenizer_manager.score_prompts(
result = await self.tokenizer_manager.score_prompts(
prompts,
label_token_ids=[self._yes_token_id, self._no_token_id],
apply_softmax=False,
request=raw_request,
)
scores = [_qwen3_rerank_score(p[0], p[1]) for p in probs]
scores = [_qwen3_rerank_score(s[0], s[1]) for s in result.scores]
except ValueError as e:
return self.create_error_response(str(e))
except Exception as e:

View File

@@ -7,6 +7,7 @@ from sglang.srt.entrypoints.openai.protocol import (
ErrorResponse,
ScoringRequest,
ScoringResponse,
UsageInfo,
)
from sglang.srt.entrypoints.openai.serving_base import OpenAIServingBase
@@ -42,7 +43,7 @@ class OpenAIServingScore(OpenAIServingBase):
"""Handle the scoring request"""
try:
# Use tokenizer_manager's score_request method directly
scores = await self.tokenizer_manager.score_request(
result = await self.tokenizer_manager.score_request(
query=request.query,
items=request.items,
label_token_ids=request.label_token_ids,
@@ -51,10 +52,13 @@ class OpenAIServingScore(OpenAIServingBase):
request=raw_request,
)
# Create response with just the scores, without usage info
response = ScoringResponse(
scores=scores,
scores=result.scores,
model=request.model,
usage=UsageInfo(
prompt_tokens=result.prompt_tokens,
total_tokens=result.prompt_tokens,
),
)
return response

View File

@@ -1,5 +1,6 @@
import logging
import math
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Union
from sglang.srt.managers.io_struct import GenerateReqInput
@@ -7,6 +8,12 @@ from sglang.srt.managers.io_struct import GenerateReqInput
logger = logging.getLogger(__name__)
@dataclass(frozen=True, slots=True)
class ScoreResult:
scores: List[List[float]]
prompt_tokens: int
class TokenizerManagerMultiItemMixin:
async def score_prompts(
self,
@@ -14,7 +21,7 @@ class TokenizerManagerMultiItemMixin:
label_token_ids: List[int],
apply_softmax: bool = False,
request: Optional[Any] = None,
) -> List[List[float]]:
) -> ScoreResult:
"""
Score probabilities of specified token IDs after each *full prompt*.
@@ -29,7 +36,9 @@ class TokenizerManagerMultiItemMixin:
request: Optional FastAPI request object.
Returns:
List of score lists, one for each prompt, each in the order of label_token_ids.
ScoreResult with:
scores: List of score lists, one for each prompt, each in the order of label_token_ids.
prompt_tokens: The number of prompt tokens processed.
"""
# Text prompts
if isinstance(prompts, str) or (
@@ -108,7 +117,7 @@ class TokenizerManagerMultiItemMixin:
label_token_ids: List[int],
apply_softmax: bool,
batch_request=None,
) -> List[List[float]]:
) -> ScoreResult:
"""
Process results from multi-item scoring request.
Extracts logprobs at delimiter positions from input_token_ids_logprobs.
@@ -121,17 +130,22 @@ class TokenizerManagerMultiItemMixin:
batch_request: The original batch request containing input sequence
Returns:
List of score lists, one for each item
ScoreResult with:
scores: List of score lists, one for each prompt, each in the order of label_token_ids.
prompt_tokens: The number of prompt tokens processed.
"""
single_result = results[0] if isinstance(results, list) else results
result = results[0] if isinstance(results, list) else results
meta_info = result.get("meta_info", {})
# For multi-item scoring, logprobs are in input_token_ids_logprobs
input_logprobs = single_result["meta_info"].get("input_token_ids_logprobs", [])
input_logprobs = meta_info.get("input_token_ids_logprobs", [])
prompt_tokens = meta_info.get("prompt_tokens", 0)
request_id = meta_info.get("id", "<unknown>")
if not input_logprobs:
raise RuntimeError(
f"input_token_ids_logprobs is empty for multi-item scoring request {single_result['meta_info'].get('id', '<unknown>')}. "
"This indicates token_ids_logprobs were not computed properly for Mutil Item Scoring."
f"input_token_ids_logprobs is empty for multi-item scoring request {request_id}. "
"This indicates token_ids_logprobs were not computed properly for Multi-Item Scoring."
)
scores = []
@@ -143,7 +157,7 @@ class TokenizerManagerMultiItemMixin:
raise RuntimeError(
f"Expected {expected_logprobs_count} input_token_ids_logprobs for multi-item scoring "
f"with {num_items} items, but got {len(input_logprobs)}. "
f"Request ID: {single_result['meta_info'].get('id', '<unknown>')}"
f"Request ID: {request_id}"
)
# Skip the first delimiter (between query and first item) and process remaining delimiter positions
@@ -162,11 +176,11 @@ class TokenizerManagerMultiItemMixin:
)
scores.append(score_list)
return scores
return ScoreResult(scores=scores, prompt_tokens=prompt_tokens)
def _process_single_item_scoring_results(
self, results: Any, label_token_ids: List[int], apply_softmax: bool
) -> List[List[float]]:
) -> ScoreResult:
"""
Process results from single-item scoring request.
Single-item scoring results are stored in output_token_ids_logprobs.
@@ -177,13 +191,17 @@ class TokenizerManagerMultiItemMixin:
apply_softmax: Whether to apply softmax normalization
Returns:
List of score lists, one for each result
ScoreResult with:
scores: List of score lists, one for each prompt, each in the order of label_token_ids.
prompt_tokens: The number of prompt tokens processed.
"""
scores = []
prompt_tokens = 0
for result in results:
# For single-item scoring, logprobs are in output_token_ids_logprobs
output_logprobs = result["meta_info"].get("output_token_ids_logprobs", [])
prompt_tokens += result["meta_info"].get("prompt_tokens", 0)
if not output_logprobs or len(output_logprobs) == 0:
raise RuntimeError(
@@ -199,7 +217,7 @@ class TokenizerManagerMultiItemMixin:
)
scores.append(score_list)
return scores
return ScoreResult(scores=scores, prompt_tokens=prompt_tokens)
async def score_request(
self,
@@ -209,7 +227,7 @@ class TokenizerManagerMultiItemMixin:
apply_softmax: bool = False,
item_first: bool = False,
request: Optional[Any] = None,
) -> List[List[float]]:
) -> ScoreResult:
"""
Score the probability of specified token IDs appearing after the given (query + item) pair.
@@ -233,11 +251,18 @@ class TokenizerManagerMultiItemMixin:
request: Optional FastAPI request object
Returns:
List of lists containing probabilities for each item and each label token
ScoreResult with:
scores: List of score lists, one for each prompt, each in the order of label_token_ids.
prompt_tokens: The number of prompt tokens processed.
"""
if label_token_ids is None:
raise ValueError("label_token_ids must be provided")
if items is None:
raise ValueError("items must be provided")
if not items:
return ScoreResult(scores=[], prompt_tokens=0)
if self.tokenizer is not None:
vocab_size = self.tokenizer.vocab_size
for token_id in label_token_ids: