sglang 0.2.13__py3-none-any.whl → 0.2.14.post1__py3-none-any.whl

sglang/srt/layers/logits_processor.py

@@ -55,6 +55,9 @@ class LogitsMetadata:
     extend_start_loc: Optional[torch.Tensor] = None
     top_logprobs_nums: Optional[List[int]] = None
 
+    extend_seq_lens_cpu: List[int] = None
+    logprob_start_lens_cpu: List[int] = None
+
     @classmethod
     def from_input_metadata(cls, input_metadata: InputMetadata):
         return cls(
@@ -63,22 +66,30 @@ class LogitsMetadata:
             extend_start_loc=input_metadata.extend_start_loc,
             return_logprob=input_metadata.return_logprob,
             top_logprobs_nums=input_metadata.top_logprobs_nums,
+            extend_seq_lens_cpu=input_metadata.extend_seq_lens_cpu,
+            logprob_start_lens_cpu=input_metadata.logprob_start_lens_cpu,
         )
 
 
 class LogitsProcessor(nn.Module):
-    def __init__(self, config):
+    def __init__(self, config, skip_all_gather: bool = False):
         super().__init__()
         self.config = config
-        self.tp_size = get_tensor_model_parallel_world_size()
+        self.do_tensor_parallel_all_gather = (
+            not skip_all_gather and get_tensor_model_parallel_world_size() > 1
+        )
 
     def _get_normalized_prompt_logprobs(
-        self, input_token_logprobs, logits_metadata: LogitsMetadata
+        self,
+        input_token_logprobs: torch.Tensor,
+        cum_start_len0: torch.Tensor,
+        cum_start_len1: torch.Tensor,
+        logits_metadata: LogitsMetadata,
     ):
         logprobs_cumsum = torch.cumsum(input_token_logprobs, dim=0, dtype=torch.float32)
 
-        start = logits_metadata.extend_start_loc.clone()
-        end = start + logits_metadata.extend_seq_lens - 2
+        start = logits_metadata.extend_start_loc.clone() - cum_start_len0
+        end = start + logits_metadata.extend_seq_lens - 2 - cum_start_len1
         start.clamp_(min=0, max=input_token_logprobs.shape[0] - 1)
         end.clamp_(min=0, max=input_token_logprobs.shape[0] - 1)
         sum_logp = (
@@ -91,7 +102,7 @@ class LogitsProcessor(nn.Module):
         return normalized_prompt_logprobs
 
     @staticmethod
-    def get_top_logprobs(all_logprobs, logits_metadata: LogitsMetadata):
+    def get_top_logprobs(all_logprobs: torch.Tensor, logits_metadata: LogitsMetadata):
         if logits_metadata.forward_mode == ForwardMode.DECODE:
             output_top_logprobs = []
             max_k = max(logits_metadata.top_logprobs_nums)
@@ -105,7 +116,7 @@ class LogitsProcessor(nn.Module):
             # TODO: vectorize the code below
             input_top_logprobs, output_top_logprobs = [], []
             pt = 0
-            extend_seq_lens_cpu = logits_metadata.extend_seq_lens.tolist()
+            extend_seq_lens_cpu = logits_metadata.extend_seq_lens_cpu
 
             max_k = max(logits_metadata.top_logprobs_nums)
             ret = all_logprobs.topk(max_k, dim=1)
@@ -113,26 +124,30 @@ class LogitsProcessor(nn.Module):
             indices = ret.indices.tolist()
 
             for i, extend_seq_len in enumerate(extend_seq_lens_cpu):
+                start_len = logits_metadata.logprob_start_lens_cpu[i]
+                pruned_len = extend_seq_len - start_len
+
                 if extend_seq_len == 0:
                     input_top_logprobs.append([])
                     output_top_logprobs.append([])
                     continue
+
                 k = logits_metadata.top_logprobs_nums[i]
                 input_top_logprobs.append(
                     [
                         list(zip(values[pt + j][:k], indices[pt + j][:k]))
-                        for j in range(extend_seq_len - 1)
+                        for j in range(pruned_len - 1)
                     ]
                 )
                 output_top_logprobs.append(
                     list(
                         zip(
-                            values[pt + extend_seq_len - 1][:k],
-                            indices[pt + extend_seq_len - 1][:k],
+                            values[pt + pruned_len - 1][:k],
+                            indices[pt + pruned_len - 1][:k],
                         )
                     )
                 )
-                pt += extend_seq_len
+                pt += pruned_len
 
             return input_top_logprobs, output_top_logprobs
 
@@ -159,13 +174,13 @@ class LogitsProcessor(nn.Module):
             last_hidden = hidden_states[last_index]
 
         last_logits = torch.matmul(last_hidden, weight.T)
-        if self.tp_size > 1:
+        if self.do_tensor_parallel_all_gather:
             last_logits = tensor_model_parallel_all_gather(last_logits)
         last_logits = last_logits[:, : self.config.vocab_size].float()
 
         if hasattr(self.config, "final_logit_softcapping"):
             last_logits.div_(self.config.final_logit_softcapping)
-            last_logits = torch.tanh(last_logits)
+            torch.tanh(last_logits, out=last_logits)
             last_logits.mul_(self.config.final_logit_softcapping)
 
         # Return only last_logits if logprob is not requested
@@ -203,14 +218,30 @@ class LogitsProcessor(nn.Module):
                     output_top_logprobs=output_top_logprobs,
                 )
             else:
-                all_logits = torch.matmul(hidden_states, weight.T)
-                if self.tp_size > 1:
+                pt, states, pruned_input_ids = 0, [], []
+                for i, extend_len in enumerate(logits_metadata.extend_seq_lens_cpu):
+                    start_len = logits_metadata.logprob_start_lens_cpu[i]
+                    states.append(hidden_states[pt + start_len : pt + extend_len])
+                    pruned_input_ids.append(input_ids[pt + start_len : pt + extend_len])
+                    pt += extend_len
+
+                states = torch.cat(states, dim=0)
+                pruned_input_ids = torch.cat(pruned_input_ids, dim=0)
+
+                cum_start_len1 = torch.tensor(
+                    logits_metadata.logprob_start_lens_cpu, device="cuda"
+                ).cumsum(0)
+                cum_start_len0 = torch.zeros_like(cum_start_len1)
+                cum_start_len0[1:] = cum_start_len1[:-1]
+
+                all_logits = torch.matmul(states, weight.T)
+                if self.do_tensor_parallel_all_gather:
                     all_logits = tensor_model_parallel_all_gather(all_logits)
                 all_logits = all_logits[:, : self.config.vocab_size].float()
 
                 if hasattr(self.config, "final_logit_softcapping"):
                     all_logits.div_(self.config.final_logit_softcapping)
-                    all_logits = torch.tanh(all_logits)
+                    torch.tanh(all_logits, out=all_logits)
                     all_logits.mul_(self.config.final_logit_softcapping)
 
                 all_logprobs = all_logits
@@ -228,19 +259,25 @@ class LogitsProcessor(nn.Module):
                 else:
                     input_top_logprobs = output_top_logprobs = None
 
-                last_logprobs = all_logprobs[last_index]
+                last_logprobs = all_logprobs[last_index - cum_start_len1]
 
                 # Compute the logprobs and normalized logprobs for the prefill tokens.
                 # Note that we pad a zero at the end of each sequence for easy computation.
                 input_token_logprobs = all_logprobs[
                     torch.arange(all_logprobs.shape[0], device="cuda"),
-                    torch.cat([input_ids[1:], torch.tensor([0], device="cuda")]),
+                    torch.cat([pruned_input_ids[1:], torch.tensor([0], device="cuda")]),
                 ]
 
                 normalized_prompt_logprobs = self._get_normalized_prompt_logprobs(
-                    input_token_logprobs, logits_metadata
+                    input_token_logprobs,
+                    cum_start_len0,
+                    cum_start_len1,
+                    logits_metadata,
                 )
 
+                # Remove the last token logprob for the prefill tokens.
+                input_token_logprobs = input_token_logprobs[:-1]
+
                 return LogitProcessorOutput(
                     next_token_logits=last_logits,
                     next_token_logprobs=last_logprobs,

sglang/srt/layers/radix_attention.py

@@ -203,7 +203,6 @@ class RadixAttention(nn.Module):
             return self.decode_forward(q, k, v, input_metadata)
 
     def store_kv_cache(self, cache_k, cache_v, input_metadata: InputMetadata):
-        k_cache = input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id)
-        v_cache = input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id)
-        k_cache[input_metadata.out_cache_loc] = cache_k
-        v_cache[input_metadata.out_cache_loc] = cache_v
+        input_metadata.token_to_kv_pool.set_kv_buffer(
+            self.layer_id, input_metadata.out_cache_loc, cache_k, cache_v
+        )

sglang/srt/layers/sampler.py

@@ -0,0 +1,101 @@
+import logging
+
+import torch
+from flashinfer.sampling import (
+    min_p_sampling_from_probs,
+    top_k_renorm_prob,
+    top_k_top_p_sampling_from_probs,
+    top_p_renorm_prob,
+)
+from vllm.model_executor.custom_op import CustomOp
+
+# TODO: move this dict to another place
+from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
+
+logger = logging.getLogger(__name__)
+
+
+class Sampler(CustomOp):
+    def __init__(self):
+        super().__init__()
+
+    def forward_cuda(self, logits: torch.Tensor, sampling_info: SamplingBatchInfo):
+        # Post process logits
+        logits = logits.contiguous()
+        logits.div_(sampling_info.temperatures)
+        if sampling_info.logit_bias is not None:
+            logits.add_(sampling_info.logit_bias)
+
+        if sampling_info.vocab_mask is not None:
+            logits = logits.masked_fill(~sampling_info.vocab_mask, float("-inf"))
+
+        logits = sampling_info.penalizer_orchestrator.apply(logits)
+
+        probs = torch.softmax(logits, dim=-1)
+
+        if not global_server_args_dict["disable_flashinfer_sampling"]:
+            max_top_k_round, batch_size = 32, probs.shape[0]
+            uniform_samples = torch.rand(
+                (max_top_k_round, batch_size), device=probs.device
+            )
+            if sampling_info.min_ps.any():
+                probs = top_k_renorm_prob(probs, sampling_info.top_ks)
+                probs = top_p_renorm_prob(probs, sampling_info.top_ps)
+                batch_next_token_ids, success = min_p_sampling_from_probs(
+                    probs, uniform_samples, sampling_info.min_ps
+                )
+            else:
+                batch_next_token_ids, success = top_k_top_p_sampling_from_probs(
+                    probs, uniform_samples, sampling_info.top_ks, sampling_info.top_ps
+                )
+        else:
+            # Here we provide a slower fallback implementation.
+            batch_next_token_ids, success = top_k_top_p_min_p_sampling_from_probs_torch(
+                probs, sampling_info.top_ks, sampling_info.top_ps, sampling_info.min_ps
+            )
+
+        if not torch.all(success):
+            logging.warning("Sampling failed, fallback to top_k=1 strategy")
+            probs = probs.masked_fill(torch.isnan(probs), 0.0)
+            argmax_ids = torch.argmax(probs, dim=-1)
+            batch_next_token_ids = torch.where(
+                success, batch_next_token_ids, argmax_ids
+            )
+
+        return batch_next_token_ids
+
+    def forward_native():
+        raise NotImplementedError("Native forward is not implemented yet.")
+
+
+def top_k_top_p_min_p_sampling_from_probs_torch(
+    probs: torch.Tensor,
+    top_ks: torch.Tensor,
+    top_ps: torch.Tensor,
+    min_ps: torch.Tensor,
+):
+    """A top-k, top-p and min-p sampling implementation with native pytorch operations."""
+    probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    min_p_thresholds = probs_sort[:, 0] * min_ps
+    probs_sort[(probs_sum - probs_sort) > top_ps.view(-1, 1)] = 0.0
+    probs_sort[
+        torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1)
+        >= top_ks.view(-1, 1)
+    ] = 0.0
+    probs_sort[probs_sort < min_p_thresholds.view(-1, 1)] = 0.0
+    probs_sort.div_(probs_sort.max(dim=-1, keepdim=True)[0])
+    try:
+        sampled_index = torch.multinomial(probs_sort, num_samples=1)
+    except RuntimeError as e:
+        logger.warning(f"Sampling error: {e}")
+        batch_next_token_ids = torch.zeros(
+            (probs_sort.shape[0],), dtype=torch.int32, device=probs.device
+        )
+        success = torch.zeros(probs.shape[0], dtype=torch.bool, device=probs.device)
+        return batch_next_token_ids, success
+
+    batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(-1)
+    success = torch.ones(probs.shape[0], dtype=torch.bool, device=probs.device)
+    return batch_next_token_ids, success

sglang/srt/managers/controller_multi.py

@@ -21,7 +21,6 @@ Each data parallel worker can manage multiple tensor parallel workers.
 import dataclasses
 import logging
 import multiprocessing
-import os
 from enum import Enum, auto
 
 import numpy as np
@@ -36,7 +35,7 @@ from sglang.srt.managers.io_struct import (
     TokenizedGenerateReqInput,
 )
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import kill_parent_process
+from sglang.srt.utils import configure_logger, kill_parent_process
 from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
@@ -194,10 +193,7 @@ def start_controller_process(
 ):
     """Start a controller process."""
 
-    logging.basicConfig(
-        level=getattr(logging, server_args.log_level.upper()),
-        format="%(message)s",
-    )
+    configure_logger(server_args)
 
     try:
         controller = ControllerMulti(server_args, port_args, model_overide_args)
@@ -212,6 +208,4 @@ def start_controller_process(
     except Exception:
         logger.error("Exception in ControllerMulti:\n" + get_exception_traceback())
     finally:
-        for w in controller.workers:
-            os.kill(w.proc.pid, 9)
         kill_parent_process()

sglang/srt/managers/controller_single.py

@@ -17,7 +17,6 @@ limitations under the License.
 
 import logging
 import multiprocessing
-import os
 from typing import List
 
 import zmq
@@ -28,7 +27,7 @@ from sglang.srt.managers.tp_worker import (
     launch_tp_servers,
 )
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import kill_parent_process
+from sglang.srt.utils import configure_logger, kill_parent_process
 from sglang.utils import get_exception_traceback
 
 logger = logging.getLogger(__name__)
@@ -53,7 +52,7 @@ class ControllerSingle:
         self.dp_worker_id = dp_worker_id
         self.mp_queue = mp_queue
 
-        # Init communication
+        # Init inter-process communication
         context = zmq.Context(2)
 
         if not self.is_dp_worker:
@@ -134,11 +133,11 @@ def start_controller_process(
     queue: multiprocessing.connection.Connection = None,
 ):
     """Start a controller process."""
-
-    logging.basicConfig(
-        level=getattr(logging, server_args.log_level.upper()),
-        format="%(message)s",
-    )
+    if is_data_parallel_worker:
+        logger_prefix = f" DP{dp_worker_id} TP0"
+    else:
+        logger_prefix = " TP0"
+    configure_logger(server_args, prefix=logger_prefix)
 
     if not is_data_parallel_worker:
         tp_size_local = server_args.tp_size // server_args.nnodes
@@ -167,6 +166,4 @@ def start_controller_process(
     except Exception:
         logger.error("Exception in ControllerSingle:\n" + get_exception_traceback())
     finally:
-        for t in controller.tp_procs:
-            os.kill(t.pid, 9)
         kill_parent_process()

sglang/srt/managers/detokenizer_manager.py

@@ -17,7 +17,6 @@ limitations under the License.
 
 import asyncio
 import dataclasses
-import inspect
 from typing import List
 
 import uvloop
@@ -29,6 +28,7 @@ from sglang.srt.managers.io_struct import (
     BatchEmbeddingOut,
     BatchStrOut,
     BatchTokenIDOut,
+    UpdateWeightReqOutput,
 )
 from sglang.srt.managers.schedule_batch import FINISH_MATCHED_STR
 from sglang.srt.server_args import PortArgs, ServerArgs
@@ -39,6 +39,8 @@ asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 @dataclasses.dataclass
 class DecodeStatus:
+    """Store the status of incremental decoding."""
+
     vid: int
     decoded_text: str
     decode_ids: List[int]
@@ -47,11 +49,14 @@ class DecodeStatus:
 
 
 class DetokenizerManager:
+    """DetokenizerManager is a process that detokenizes the token ids."""
+
     def __init__(
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
     ):
+        # Init inter-process communication
         context = zmq.asyncio.Context(2)
         self.recv_from_router = context.socket(zmq.PULL)
         self.recv_from_router.bind(f"tcp://127.0.0.1:{port_args.detokenizer_port}")
@@ -71,10 +76,13 @@ class DetokenizerManager:
         self.decode_status = {}
 
     async def handle_loop(self):
+        """The event loop that handles requests"""
+
         while True:
-            recv_obj: BatchTokenIDOut = await self.recv_from_router.recv_pyobj()
+            recv_obj = await self.recv_from_router.recv_pyobj()
 
             if isinstance(recv_obj, BatchEmbeddingOut):
+                # If it is embedding model, no detokenization is needed.
                 self.send_to_tokenizer.send_pyobj(
                     BatchEmbeddingOut(
                         rids=recv_obj.rids,
@@ -84,15 +92,18 @@ class DetokenizerManager:
                     )
                 )
                 continue
+            elif isinstance(recv_obj, UpdateWeightReqOutput):
+                # If it is a weight update request, no detokenization is needed.
+                self.send_to_tokenizer.send_pyobj(recv_obj)
+                continue
+            elif self.tokenizer is None:
+                # If the tokenizer is skipped, no detokenization is needed
+                self.send_to_tokenizer.send_pyobj(recv_obj)
+                continue
 
             assert isinstance(recv_obj, BatchTokenIDOut)
             bs = len(recv_obj.rids)
 
-            if self.tokenizer is None:
-                # Send BatchTokenIDOut if no tokenizer init'ed.
-                self.send_to_tokenizer.send_pyobj(recv_obj)
-                continue
-
             # Initialize decode status
             read_ids, surr_ids = [], []
             for i in range(bs):
@@ -126,8 +137,7 @@ class DetokenizerManager:
                 spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
             )
 
-            # Trim stop str
-            # TODO(lmzheng): handle the case where multiple stop strs are hit
+            # Incremental decoding
             output_strs = []
             for i in range(bs):
                 s = self.decode_status[recv_obj.rids[i]]
@@ -144,6 +154,7 @@ class DetokenizerManager:
 
                 output_strs.append(s.decoded_text + new_text)
 
+                # Trim stop str. TODO(lmzheng): handle the case where multiple stop strs are hit
                 if isinstance(recv_obj.finished_reason[i], FINISH_MATCHED_STR):
                     pos = output_strs[i].find(recv_obj.finished_reason[i].matched)
                     if pos != -1:

sglang/srt/managers/io_struct.py

@@ -22,10 +22,8 @@ import uuid
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
-import torch
-
 from sglang.srt.managers.schedule_batch import BaseFinishReason
-from sglang.srt.sampling_params import SamplingParams
+from sglang.srt.sampling.sampling_params import SamplingParams
 
 
 @dataclass
@@ -43,9 +41,9 @@ class GenerateReqInput:
     rid: Optional[Union[List[str], str]] = None
     # Whether to return logprobs.
     return_logprob: Optional[Union[List[bool], bool]] = None
-    # The start location of the prompt for return_logprob.
+    # If return logprobs, the start location in the prompt for returning logprobs.
     logprob_start_len: Optional[Union[List[int], int]] = None
-    # The number of top logprobs to return.
+    # If return logprobs, the number of top logprobs to return at each position.
     top_logprobs_num: Optional[Union[List[int], int]] = None
     # Whether to detokenize tokens in text in the returned logprobs.
     return_text_in_logprobs: bool = False
@@ -77,7 +75,7 @@ class GenerateReqInput:
             if self.return_logprob is None:
                 self.return_logprob = False
             if self.logprob_start_len is None:
-                self.logprob_start_len = 0
+                self.logprob_start_len = -1
             if self.top_logprobs_num is None:
                 self.top_logprobs_num = 0
         else:
@@ -143,7 +141,7 @@ class GenerateReqInput:
                 self.return_logprob = [self.return_logprob] * num
 
             if self.logprob_start_len is None:
-                self.logprob_start_len = [0] * num
+                self.logprob_start_len = [-1] * num
             elif not isinstance(self.logprob_start_len, list):
                 self.logprob_start_len = [self.logprob_start_len] * num
 
@@ -155,16 +153,27 @@ class GenerateReqInput:
 
 @dataclass
 class TokenizedGenerateReqInput:
+    # The request id
     rid: str
+    # The input text
     input_text: str
+    # The input token ids
     input_ids: List[int]
+    # The pixel values for input images
     pixel_values: List[float]
+    # The hash of input images
     image_hash: int
+    # The image size
     image_size: List[int]
+    # The sampling parameters
     sampling_params: SamplingParams
+    # Whether to return the logprobs
     return_logprob: bool
+    # If return logprobs, the start location in the prompt for returning logprobs.
     logprob_start_len: int
+    # If return logprobs, the number of top logprobs to return at each position.
     top_logprobs_num: int
+    # Whether to stream output
     stream: bool
 
 
@@ -215,15 +224,21 @@ class EmbeddingReqInput:
 
 @dataclass
 class TokenizedEmbeddingReqInput:
+    # The request id
     rid: str
+    # The input text
     input_text: str
+    # The input token ids
     input_ids: List[int]
+    # Dummy sampling params for compatibility
     sampling_params: SamplingParams
 
 
 @dataclass
 class BatchTokenIDOut:
+    # The request id
     rids: List[str]
+    # The version id to sync decode status with in detokenizer_manager
     vids: List[int]
     decoded_texts: List[str]
     decode_ids: List[int]
@@ -236,17 +251,25 @@ class BatchTokenIDOut:
 
 @dataclass
 class BatchStrOut:
+    # The request id
     rids: List[str]
+    # The output decoded strings
     output_strs: List[str]
+    # The meta info
     meta_info: List[Dict]
+    # The finish reason
     finished_reason: List[BaseFinishReason]
 
 
 @dataclass
 class BatchEmbeddingOut:
+    # The request id
     rids: List[str]
+    # The output embedding
     embeddings: List[List[float]]
+    # The meta info
     meta_info: List[Dict]
+    # The finish reason
     finished_reason: List[BaseFinishReason]
 
 
@@ -256,10 +279,20 @@ class FlushCacheReq:
 
 
 @dataclass
-class AbortReq:
-    rid: str
+class UpdateWeightReqInput:
+    # The model path with the new weights
+    model_path: str
+    # The format to load the weights
+    load_format: Optional[str] = None
 
 
 @dataclass
-class DetokenizeReqInput:
-    input_ids: List[int]
+class UpdateWeightReqOutput:
+    success: bool
+    message: str
+
+
+@dataclass
+class AbortReq:
+    # The request id
+    rid: str

sglang/srt/managers/policy_scheduler.py

@@ -111,11 +111,14 @@ class PrefillAdder:
         rem_total_tokens: int,
         rem_input_tokens: int,
         rem_chunk_tokens: Optional[int],
+        mixed_with_decode_tokens: int = 0,
     ):
         self.tree_cache = tree_cache
-        self.rem_total_tokens = rem_total_tokens
-        self.rem_input_tokens = rem_input_tokens
+        self.rem_total_tokens = rem_total_tokens - mixed_with_decode_tokens
+        self.rem_input_tokens = rem_input_tokens - mixed_with_decode_tokens
         self.rem_chunk_tokens = rem_chunk_tokens
+        if self.rem_chunk_tokens is not None:
+            self.rem_chunk_tokens -= mixed_with_decode_tokens
 
         self.can_run_list = []
         self.new_inflight_req = None