sglang 0.5.3__py3-none-any.whl → 0.5.3.post1__py3-none-any.whl

sglang/srt/managers/tokenizer_manager.py

@@ -182,6 +182,8 @@ class TokenizerManager(TokenizerCommunicatorMixin):
             if speculative_algorithm.is_none()
             else server_args.speculative_num_draft_tokens
         )
+        # Initialize delimiter text for multi-item scoring (will be set after tokenizer is loaded)
+        self.multi_item_delimiter_text = None
 
         if self.model_config.is_multimodal:
             import_processors("sglang.srt.multimodal.processors")
@@ -223,6 +225,7 @@ class TokenizerManager(TokenizerCommunicatorMixin):
                 self.processor = _processor
                 self.tokenizer = get_tokenizer_from_processor(self.processor)
                 os.environ["TOKENIZERS_PARALLELISM"] = "false"
+                self._initialize_multi_item_delimiter_text()
         else:
             self.mm_processor = self.processor = None
 
@@ -235,6 +238,7 @@ class TokenizerManager(TokenizerCommunicatorMixin):
                     trust_remote_code=server_args.trust_remote_code,
                     revision=server_args.revision,
                 )
+                self._initialize_multi_item_delimiter_text()
         # Initialize async dynamic batch tokenizer if enabled (common for both multimodal and non-multimodal)
         if (
             server_args.enable_dynamic_batch_tokenizer
@@ -1678,6 +1682,201 @@ class TokenizerManager(TokenizerCommunicatorMixin):
             if len(self.model_update_tmp) == self.server_args.dp_size:
                 self.model_update_result.set_result(self.model_update_tmp)
 
+    def _initialize_multi_item_delimiter_text(self):
+        """Initialize multi-item delimiter text from token ID after tokenizer is loaded."""
+        if (
+            hasattr(self.server_args, "multi_item_scoring_delimiter")
+            and self.server_args.multi_item_scoring_delimiter is not None
+            and self.tokenizer is not None
+        ):
+            try:
+                self.multi_item_delimiter_text = self.tokenizer.decode(
+                    [self.server_args.multi_item_scoring_delimiter],
+                    skip_special_tokens=False,
+                )
+            except Exception as e:
+                logger.warning(
+                    f"Failed to decode delimiter token {self.server_args.multi_item_scoring_delimiter}: {e}"
+                )
+                self.multi_item_delimiter_text = None
+
+    def _build_multi_item_token_sequence(
+        self, query: List[int], items: List[List[int]], delimiter_token_id: int
+    ) -> List[int]:
+        """
+        Build a single token sequence for multi-item scoring.
+        Format: query<delimiter>item1<delimiter>item2<delimiter>item3<delimiter>
+
+        Args:
+            query: Query token IDs
+            items: List of item token ID sequences
+            delimiter_token_id: Token ID to use as delimiter
+
+        Returns:
+            Combined token sequence
+        """
+        combined_sequence = query[:]  # Start with query
+
+        for item in items:
+            combined_sequence.append(delimiter_token_id)  # Add delimiter
+            combined_sequence.extend(item)  # Add item tokens
+
+        # Add final delimiter after the last item for logprob extraction
+        combined_sequence.append(delimiter_token_id)
+
+        return combined_sequence
+
+    def _extract_logprobs_for_tokens(
+        self, logprobs_data: List, label_token_ids: List[int]
+    ) -> Dict[int, float]:
+        """
+        Extract logprobs for specified token IDs from logprobs data.
+
+        Args:
+            logprobs_data: List of (logprob, token_id, text) tuples
+            label_token_ids: Token IDs to extract logprobs for
+
+        Returns:
+            Dictionary mapping token_id to logprob
+        """
+        logprobs = {}
+        if logprobs_data:
+            for logprob, token_id, _ in logprobs_data:
+                if token_id in label_token_ids:
+                    logprobs[token_id] = logprob
+        return logprobs
+
+    def _convert_logprobs_to_scores(
+        self,
+        logprobs: Dict[int, float],
+        label_token_ids: List[int],
+        apply_softmax: bool,
+    ) -> List[float]:
+        """
+        Convert logprobs dictionary to ordered score list.
+
+        Args:
+            logprobs: Dictionary mapping token_id to logprob
+            label_token_ids: Token IDs in desired order
+            apply_softmax: Whether to apply softmax normalization
+
+        Returns:
+            List of scores in the same order as label_token_ids
+        """
+        score_list = [
+            logprobs.get(token_id, float("-inf")) for token_id in label_token_ids
+        ]
+
+        if apply_softmax:
+            score_list = torch.softmax(torch.tensor(score_list), dim=0).tolist()
+        else:
+            # Convert logprobs to probabilities if not using softmax
+            score_list = [
+                math.exp(x) if x != float("-inf") else 0.0 for x in score_list
+            ]
+
+        return score_list
+
+    def _process_multi_item_scoring_results(
+        self,
+        results: Any,
+        items: List,
+        label_token_ids: List[int],
+        apply_softmax: bool,
+        batch_request=None,
+    ) -> List[List[float]]:
+        """
+        Process results from multi-item scoring request.
+        Extracts logprobs at delimiter positions from input_token_ids_logprobs.
+
+        Args:
+            results: Results from generate_request
+            items: List of items being scored
+            label_token_ids: Token IDs to extract scores for
+            apply_softmax: Whether to apply softmax normalization
+            batch_request: The original batch request containing input sequence
+
+        Returns:
+            List of score lists, one for each item
+        """
+        single_result = results[0] if isinstance(results, list) else results
+
+        # For multi-item scoring, logprobs are in input_token_ids_logprobs
+        input_logprobs = single_result["meta_info"].get("input_token_ids_logprobs", [])
+
+        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."
+            )
+
+        scores = []
+        num_items = len(items) if isinstance(items, list) else 1
+
+        # Check if we have the expected number of logprobs
+        expected_logprobs_count = num_items + 1
+        if len(input_logprobs) != expected_logprobs_count:
+            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>')}"
+            )
+
+        # Skip the first delimiter (between query and first item) and process remaining delimiter positions
+        # We want to exclude the first one since it represents the boundary between query and first item, not an item boundary
+        start_idx = 1 if len(input_logprobs) > 1 else 0
+
+        # Process logprobs for each item position (excluding first delimiter)
+        for item_idx in range(num_items):
+            logprob_idx = start_idx + item_idx
+            item_logprobs_data = input_logprobs[logprob_idx]
+            logprobs = self._extract_logprobs_for_tokens(
+                item_logprobs_data, label_token_ids
+            )
+            score_list = self._convert_logprobs_to_scores(
+                logprobs, label_token_ids, apply_softmax
+            )
+            scores.append(score_list)
+
+        return scores
+
+    def _process_single_item_scoring_results(
+        self, results: Any, label_token_ids: List[int], apply_softmax: bool
+    ) -> List[List[float]]:
+        """
+        Process results from single-item scoring request.
+        Single-item scoring results are stored in output_token_ids_logprobs.
+
+        Args:
+            results: Results from generate_request
+            label_token_ids: Token IDs to extract scores for
+            apply_softmax: Whether to apply softmax normalization
+
+        Returns:
+            List of score lists, one for each result
+        """
+        scores = []
+
+        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", [])
+
+            if not output_logprobs or len(output_logprobs) == 0:
+                raise RuntimeError(
+                    f"output_logprobs is empty for request {result['meta_info'].get('id', '<unknown>')}."
+                )
+
+            # Extract logprobs for the first (and only) position
+            logprobs = self._extract_logprobs_for_tokens(
+                output_logprobs[0], label_token_ids
+            )
+            score_list = self._convert_logprobs_to_scores(
+                logprobs, label_token_ids, apply_softmax
+            )
+            scores.append(score_list)
+
+        return scores
+
     async def score_request(
         self,
         query: Optional[Union[str, List[int]]] = None,
@@ -1688,7 +1887,29 @@ class TokenizerManager(TokenizerCommunicatorMixin):
         request: Optional[Any] = None,
     ) -> List[List[float]]:
         """
-        See Engine.score() for more details.
+        Score the probability of specified token IDs appearing after the given (query + item) pair.
+
+        This method supports two scoring approaches:
+        1. Single-Item scoring (default): Process each query+item pair independently
+        2. Multi-Item scoring: When multi_item_scoring_delimiter is set, combine query and
+           multiple items into a single sequence using delimiter for efficient processing.
+           Note: item_first parameter is ignored in multi-item scoring mode since it uses
+           a fixed format: query<delimiter>item1<delimiter>item2<delimiter>item3<delimiter>
+
+           Multi-item scoring works with both text and pre-tokenized inputs:
+           - Text: query<delimiter_text>item1<delimiter_text>item2<delimiter_text>item3<delimiter_text>
+           - Tokens: query<delimiter_token_id>item1<delimiter_token_id>item2<delimiter_token_id>item3<delimiter_token_id>
+
+        Args:
+            query: The query text or pre-tokenized query token IDs
+            items: The item text(s) or pre-tokenized item token IDs
+            label_token_ids: List of token IDs to compute probabilities for
+            apply_softmax: Whether to normalize probabilities using softmax
+            item_first: If True, prepend items to query. Ignored for multi-item scoring.
+            request: Optional FastAPI request object
+
+        Returns:
+            List of lists containing probabilities for each item and each label token
         """
         if label_token_ids is None:
             raise ValueError("label_token_ids must be provided")
@@ -1701,9 +1922,17 @@ class TokenizerManager(TokenizerCommunicatorMixin):
                         f"Token ID {token_id} is out of vocabulary (vocab size: {vocab_size})"
                     )
 
+        # Check if multi-item scoring is enabled by presence of delimiter
+        use_multi_item_scoring = (
+            self.server_args.multi_item_scoring_delimiter is not None
+            and self.multi_item_delimiter_text is not None
+        )
+
         batch_request = GenerateReqInput(
             token_ids_logprob=label_token_ids,
             return_logprob=True,
+            # Set logprob_start_len=0 for multi-item scoring since we want logprobs at all delimiter positions
+            logprob_start_len=0 if use_multi_item_scoring else -1,
             stream=False,
             sampling_params={"max_new_tokens": 0},
         )
@@ -1715,12 +1944,23 @@ class TokenizerManager(TokenizerCommunicatorMixin):
         ):
             # Both query and items are text
             items_list = [items] if isinstance(items, str) else items
-            if item_first:
-                prompts = [f"{item}{query}" for item in items_list]
-            else:
-                prompts = [f"{query}{item}" for item in items_list]
 
-            batch_request.text = prompts
+            if use_multi_item_scoring:
+                # Multi-item scoring: create single prompt with delimiter text
+                # Always use format: query<delimiter>item1<delimiter>item2<delimiter>item3<delimiter>
+                # (item_first is ignored for multi-item scoring)
+                delimiter = self.multi_item_delimiter_text
+                combined_items = delimiter.join(items_list)
+                # Add final delimiter after the last item for logprob extraction
+                single_prompt = f"{query}{delimiter}{combined_items}{delimiter}"
+                batch_request.text = [single_prompt]
+            else:
+                # Single-item scoring: create separate prompts for each item
+                if item_first:
+                    prompts = [f"{item}{query}" for item in items_list]
+                else:
+                    prompts = [f"{query}{item}" for item in items_list]
+                batch_request.text = prompts
 
         elif (
             isinstance(query, list)
@@ -1729,61 +1969,38 @@ class TokenizerManager(TokenizerCommunicatorMixin):
             and isinstance(items[0], list)
         ):
             # Both query and items are token IDs
-            if item_first:
-                input_ids_list = [item + query for item in items]
+            if use_multi_item_scoring:
+                # Multi-item scoring: concatenate with delimiter token ID
+                # Format: query<delimiter_token_id>item1<delimiter_token_id>item2<delimiter_token_id>item3<delimiter_token_id>
+                delimiter_token_id = self.server_args.multi_item_scoring_delimiter
+                combined_input_ids = self._build_multi_item_token_sequence(
+                    query, items, delimiter_token_id
+                )
+                batch_request.input_ids = [combined_input_ids]
             else:
-                input_ids_list = [query + item for item in items]
-
-            batch_request.input_ids = input_ids_list
+                # Single-item scoring: process each item separately
+                if item_first:
+                    input_ids_list = [item + query for item in items]
+                else:
+                    input_ids_list = [query + item for item in items]
+                batch_request.input_ids = input_ids_list
         else:
             raise ValueError(
                 "Invalid combination of query/items types for score_request."
             )
 
         results = await self.generate_request(batch_request, request).__anext__()
-        scores = []
-
-        for result in results:
-            # Get logprobs for each token
-            logprobs = {}
-
-            # For scoring requests, we read from output_token_ids_logprobs since we want
-            # the logprobs for specific tokens mentioned in the label_token_ids at
-            # the next position after the last token in the prompt
-            output_logprobs = result["meta_info"].get("output_token_ids_logprobs", [])
-
-            # Check if output_logprobs is properly populated
-            if (
-                output_logprobs is None
-                or not output_logprobs
-                or len(output_logprobs) == 0
-            ):
-                raise RuntimeError(
-                    f"output_logprobs is empty for request {result['meta_info'].get('id', '<unknown>')}. "
-                    "This indicates token_ids_logprobs were not computed properly for the scoring request."
-                )
-
-            for logprob, token_id, _ in output_logprobs[0]:
-                if token_id in label_token_ids:
-                    logprobs[token_id] = logprob
 
-            # Get scores in order of label_token_ids
-            score_list = [
-                logprobs.get(token_id, float("-inf")) for token_id in label_token_ids
-            ]
-
-            # Apply softmax to logprobs if needed
-            if apply_softmax:
-                score_list = torch.softmax(torch.tensor(score_list), dim=0).tolist()
-            else:
-                # Convert logprobs to probabilities if not using softmax
-                score_list = [
-                    math.exp(x) if x != float("-inf") else 0.0 for x in score_list
-                ]
-
-            scores.append(score_list)
-
-        return scores
+        if use_multi_item_scoring:
+            # Multi-item scoring: extract scores from input_token_ids_logprobs
+            return self._process_multi_item_scoring_results(
+                results, items, label_token_ids, apply_softmax, batch_request
+            )
+        else:
+            # Single-item scoring: process each result separately
+            return self._process_single_item_scoring_results(
+                results, label_token_ids, apply_softmax
+            )
 
     async def watch_load_thread(self):
         # Only for dp_controller when dp_size > 1

sglang/srt/managers/tp_worker.py

@@ -15,14 +15,12 @@
 from __future__ import annotations
 
 import logging
-import threading
-from typing import TYPE_CHECKING, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.distributed import get_pp_group, get_world_group
-from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.io_struct import (
     DestroyWeightsUpdateGroupReqInput,
     GetWeightsByNameReqInput,
@@ -36,13 +34,10 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightsFromTensorReqInput,
 )
 from sglang.srt.managers.schedule_batch import ModelWorkerBatch, global_server_args_dict
+from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
-from sglang.srt.model_executor.forward_batch_info import (
-    ForwardBatch,
-    ForwardBatchOutput,
-    PPProxyTensors,
-)
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
 from sglang.srt.model_executor.model_runner import ModelRunner
 from sglang.srt.server_args import ServerArgs
 from sglang.srt.utils import MultiprocessingSerializer, broadcast_pyobj, set_random_seed
@@ -236,9 +231,8 @@ class TpModelWorker:
     def forward_batch_generation(
         self,
         model_worker_batch: ModelWorkerBatch,
-        launch_done: Optional[threading.Event] = None,
         is_verify: bool = False,
-    ) -> ForwardBatchOutput:
+    ) -> GenerationBatchResult:
         # update the consumer index of hicache to the running batch
         self.set_hicache_consumer(model_worker_batch.hicache_consumer_index)
 
@@ -256,32 +250,49 @@ class TpModelWorker:
             logits_output, can_run_cuda_graph = self.model_runner.forward(
                 forward_batch, pp_proxy_tensors=pp_proxy_tensors
             )
-            if launch_done is not None:
-                launch_done.set()
-
-            skip_sample = is_verify or model_worker_batch.is_prefill_only
-            next_token_ids = None
-
-            if not skip_sample:
-                next_token_ids = self.model_runner.sample(logits_output, forward_batch)
-            elif model_worker_batch.return_logprob and not is_verify:
-                # NOTE: Compute logprobs without full sampling
-                self.model_runner.compute_logprobs_only(
-                    logits_output, model_worker_batch
-                )
-
-            return ForwardBatchOutput(
+            batch_result = GenerationBatchResult(
                 logits_output=logits_output,
-                next_token_ids=next_token_ids,
                 can_run_cuda_graph=can_run_cuda_graph,
             )
+
+            if is_verify:
+                # Skip sampling and return logits for target forward
+                return batch_result
+
+            if model_worker_batch.delay_sample_launch:
+                batch_result.delay_sample_launch = True
+                batch_result.forward_batch = forward_batch
+                return batch_result
+
+            if model_worker_batch.is_prefill_only:
+                # For prefill-only requests, create dummy token IDs on CPU
+                # The size should match the batch size (number of sequences), not total tokens
+                batch_result.next_token_ids = torch.zeros(
+                    len(model_worker_batch.seq_lens),
+                    dtype=torch.long,
+                    device=model_worker_batch.input_ids.device,
+                )
+                if (
+                    model_worker_batch.return_logprob
+                    and logits_output.next_token_logits is not None
+                ):
+                    # NOTE: Compute logprobs without full sampling
+                    self.model_runner.compute_logprobs_only(
+                        logits_output, model_worker_batch
+                    )
+            else:
+                batch_result.next_token_ids = self.model_runner.sample(
+                    logits_output, forward_batch
+                )
+
+            return batch_result
         else:
             pp_proxy_tensors, can_run_cuda_graph = self.model_runner.forward(
                 forward_batch,
                 pp_proxy_tensors=pp_proxy_tensors,
             )
-            return ForwardBatchOutput(
-                pp_proxy_tensors=pp_proxy_tensors,
+            return GenerationBatchResult(
+                pp_hidden_states_proxy_tensors=pp_proxy_tensors,
                 can_run_cuda_graph=can_run_cuda_graph,
             )
 

sglang/srt/mem_cache/allocator.py

@@ -274,10 +274,15 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
         self.full_to_swa_index_mapping[free_index] = 0
 
     def backup_state(self):
-        raise NotImplementedError
+        return [
+            self.full_attn_allocator.backup_state(),
+            self.swa_attn_allocator.backup_state(),
+        ]
 
     def restore_state(self, state):
-        raise NotImplementedError
+        assert len(state) == 2
+        self.full_attn_allocator.restore_state(state[0])
+        self.swa_attn_allocator.restore_state(state[1])
 
     def clear(self):
         self.swa_attn_allocator.clear()

sglang/srt/mem_cache/chunk_cache.py

@@ -60,7 +60,7 @@ class ChunkCache(BasePrefixCache):
         ]
 
         # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
-        req.prefix_indices = kv_indices
+        req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True)
 
     def evict(self, num_tokens: int):
         pass