diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 8bf1ab524..96ed322df 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -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. diff --git a/python/sglang/srt/entrypoints/openai/serving_rerank.py b/python/sglang/srt/entrypoints/openai/serving_rerank.py index 0c33c0e81..853d44da0 100644 --- a/python/sglang/srt/entrypoints/openai/serving_rerank.py +++ b/python/sglang/srt/entrypoints/openai/serving_rerank.py @@ -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: diff --git a/python/sglang/srt/entrypoints/openai/serving_score.py b/python/sglang/srt/entrypoints/openai/serving_score.py index 19f788ad8..e9fb5f8c0 100644 --- a/python/sglang/srt/entrypoints/openai/serving_score.py +++ b/python/sglang/srt/entrypoints/openai/serving_score.py @@ -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 diff --git a/python/sglang/srt/managers/tokenizer_manager_multiitem_mixin.py b/python/sglang/srt/managers/tokenizer_manager_multiitem_mixin.py index 2ab5dd11c..bbc685d3f 100644 --- a/python/sglang/srt/managers/tokenizer_manager_multiitem_mixin.py +++ b/python/sglang/srt/managers/tokenizer_manager_multiitem_mixin.py @@ -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", "") if not input_logprobs: raise RuntimeError( - f"input_token_ids_logprobs is empty for multi-item scoring request {single_result['meta_info'].get('id', '')}. " - "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', '')}" + 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: diff --git a/test/registered/core/test_score_api.py b/test/registered/core/test_score_api.py index 465d9d233..4110337ee 100644 --- a/test/registered/core/test_score_api.py +++ b/test/registered/core/test_score_api.py @@ -164,7 +164,7 @@ class TestScoreAPI(CustomTestCase): label_token_ids=label_token_ids, apply_softmax=True, item_first=case["item_first"], - ) + ).scores # Get scores from HuggingFace using the same parameters hf_scores = self.compute_hf_scores( @@ -193,7 +193,7 @@ class TestScoreAPI(CustomTestCase): items=texts, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores self.assertEqual( len(scores), @@ -245,7 +245,7 @@ class TestScoreAPI(CustomTestCase): items=items, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores # Verify we got results self.assertEqual(len(scores), len(items)) @@ -306,15 +306,18 @@ class TestScoreAPI(CustomTestCase): label_token_ids = [9454, 2753] # "Yes" and "No" tokens # Get scores using SGLang - scores = self.engine.score( + result = self.engine.score( query=query, items=items, label_token_ids=label_token_ids, apply_softmax=True, ) + scores = result.scores + prompt_tokens = result.prompt_tokens # Verify we get the expected number of scores self.assertEqual(len(scores), len(items), "Should get one score list per item") + self.assertGreater(prompt_tokens, 0, "Should have positive prompt_tokens") # Verify each score list has the correct length for i, score_list in enumerate(scores): @@ -348,14 +351,14 @@ class TestScoreAPI(CustomTestCase): items=items, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores scores2 = self.engine.score( query=query, items=items, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores # Results should be identical (deterministic) self.assertEqual(len(scores1), len(scores2), "Should get same number of items") @@ -391,7 +394,7 @@ class TestScoreAPI(CustomTestCase): items=items, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores self.assertEqual( len(scores), len(items), f"Should get {len(items)} score lists" @@ -411,14 +414,19 @@ class TestScoreAPI(CustomTestCase): items = [] label_token_ids = [1, 2] - scores = self.engine.score( + result = self.engine.score( query=query, items=items, label_token_ids=label_token_ids, apply_softmax=True, ) + scores = result.scores + prompt_tokens = result.prompt_tokens self.assertEqual(len(scores), 0, "Should return empty list for empty items") + self.assertEqual( + prompt_tokens, 0, "Should return 0 prompt_tokens for empty items" + ) def test_multi_item_scoring_single_item(self): """Test multi-item scoring with single item (should work like regular scoring).""" @@ -426,18 +434,21 @@ class TestScoreAPI(CustomTestCase): items = ["Paris"] label_token_ids = [1, 2, 3] - scores = self.engine.score( + result = self.engine.score( query=query, items=items, label_token_ids=label_token_ids, apply_softmax=True, ) + scores = result.scores + prompt_tokens = result.prompt_tokens self.assertEqual(len(scores), 1, "Should get one score list") self.assertEqual( len(scores[0]), len(label_token_ids), "Should have correct number of scores" ) self.assertAlmostEqual(sum(scores[0]), 1.0, places=6) + self.assertGreater(prompt_tokens, 0, "Should have positive prompt_tokens") def test_multi_item_scoring_different_queries(self): """Test multi-item scoring with different types of queries.""" @@ -459,7 +470,7 @@ class TestScoreAPI(CustomTestCase): items=items, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores self.assertEqual( len(scores), @@ -490,7 +501,7 @@ class TestScoreAPI(CustomTestCase): items=items, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores self.assertEqual(len(scores), len(items)) @@ -513,7 +524,7 @@ class TestScoreAPI(CustomTestCase): items=items, label_token_ids=label_token_ids, apply_softmax=False, # No softmax - ) + ).scores self.assertEqual(len(scores), len(items)) @@ -537,7 +548,7 @@ class TestScoreAPI(CustomTestCase): items=items, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores self.assertEqual(len(scores), len(items), "Should handle large batches") @@ -556,7 +567,7 @@ class TestScoreAPI(CustomTestCase): items=items, label_token_ids=label_token_ids, apply_softmax=True, - ) + ).scores self.assertEqual(len(scores), len(items)) diff --git a/test/registered/openai_server/basic/test_openai_server.py b/test/registered/openai_server/basic/test_openai_server.py index 7d348e579..17fc78764 100644 --- a/test/registered/openai_server/basic/test_openai_server.py +++ b/test/registered/openai_server/basic/test_openai_server.py @@ -1038,6 +1038,18 @@ class TestOpenAIV1Score(CustomTestCase): msg=f"Score {i} probabilities should sum to 1", ) + # Verify usage + self.assertIn("usage", response, "Response should have a 'usage' field") + self.assertGreater(response["usage"]["prompt_tokens"], 0) + self.assertEqual( + response["usage"]["prompt_tokens"], response["usage"]["total_tokens"] + ) + self.assertEqual( + response["usage"]["completion_tokens"], + 0, + "completion_tokens should be 0 for /v1/score", + ) + def test_score_token_input(self): """Test scoring with token IDs input""" query = "The capital of France is" @@ -1088,6 +1100,18 @@ class TestOpenAIV1Score(CustomTestCase): msg=f"Score {i} probabilities should sum to 1", ) + # Verify usage + self.assertIn("usage", response, "Response should have a 'usage' field") + self.assertGreater(response["usage"]["prompt_tokens"], 0) + self.assertEqual( + response["usage"]["prompt_tokens"], response["usage"]["total_tokens"] + ) + self.assertEqual( + response["usage"]["completion_tokens"], + 0, + "completion_tokens should be 0 for /v1/score", + ) + def test_score_error_handling(self): """Test error handling for invalid inputs""" query = "The capital of France is" diff --git a/test/registered/openai_server/basic/test_serving_rerank.py b/test/registered/openai_server/basic/test_serving_rerank.py index 55eff93a7..4ed4599eb 100644 --- a/test/registered/openai_server/basic/test_serving_rerank.py +++ b/test/registered/openai_server/basic/test_serving_rerank.py @@ -3,6 +3,7 @@ import unittest from unittest.mock import Mock from sglang.srt.entrypoints.openai.protocol import V1RerankReqInput +from sglang.srt.managers.tokenizer_manager_multiitem_mixin import ScoreResult from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci # Keep consistent with other openai_server/basic unit tests. @@ -163,7 +164,7 @@ class TestOpenAIServingRerankUnit(unittest.TestCase): # Return [p_yes, p_no] for each prompt assert len(prompts) == 2 assert label_token_ids and len(label_token_ids) == 2 - return [[0.9, 0.1], [0.2, 0.8]] + return ScoreResult(scores=[[0.9, 0.1], [0.2, 0.8]], prompt_tokens=42) handler = OpenAIServingRerank(_TM()) req = V1RerankReqInput(query="q", documents=["d1", "d2"], return_documents=True)