sglang 0.4.1.post5__py3-none-any.whl → 0.4.1.post7__py3-none-any.whl

sglang/srt/layers/sampler.py

@@ -1,11 +1,12 @@
 import logging
-from typing import List
+from typing import Dict, List
 
 import torch
 from torch import nn
 
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.schedule_batch import global_server_args_dict
+from sglang.srt.sampling.custom_logit_processor import CustomLogitProcessor
 from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
 from sglang.srt.utils import crash_on_warnings, is_flashinfer_available
 
@@ -35,6 +36,10 @@ class Sampler(nn.Module):
     ):
         logits = logits_output.next_token_logits
 
+        # Apply the custom logit processors if registered in the sampling info.
+        if sampling_info.has_custom_logit_processor:
+            self._apply_custom_logit_processor(logits, sampling_info)
+
         if self.use_nan_detectioin and torch.any(torch.isnan(logits)):
             logger.warning("Detected errors during sampling! NaN in the logits.")
             logits = torch.where(
@@ -121,6 +126,39 @@ class Sampler(nn.Module):
 
         return batch_next_token_ids
 
+    def _apply_custom_logit_processor(
+        self, logits: torch.Tensor, sampling_batch_info: SamplingBatchInfo
+    ):
+        """Apply custom logit processors to the logits.
+        This function will modify the logits in-place."""
+
+        assert logits.shape[0] == len(sampling_batch_info), (
+            f"The batch size of logits ({logits.shape[0]}) does not match the batch size of "
+            f"sampling_batch_info ({len(sampling_batch_info)})"
+        )
+
+        for _, (
+            processor,
+            batch_mask,
+        ) in sampling_batch_info.custom_logit_processor.items():
+            # Get the batch indices that need to be processed
+            batch_indices = batch_mask.nonzero(as_tuple=True)[0]
+
+            assert batch_mask.shape[0] == len(sampling_batch_info), (
+                f"The number of batch mask ({batch_mask.shape[0]}) does not match the number of "
+                f"sampling_batch_info ({len(sampling_batch_info)})"
+            )
+
+            # Apply the processor to the logits
+            logits[batch_mask] = processor(
+                logits[batch_mask],
+                [sampling_batch_info.custom_params[i] for i in batch_indices],
+            )
+
+            logger.debug(
+                f"Custom logit processor {processor.__class__.__name__} is applied."
+            )
+
 
 def top_k_top_p_min_p_sampling_from_probs_torch(
     probs: torch.Tensor,

sglang/srt/layers/vocab_parallel_embedding.py

@@ -6,13 +6,13 @@ from typing import List, Optional, Sequence, Tuple
 import torch
 import torch.nn.functional as F
 from torch.nn.parameter import Parameter, UninitializedParameter
-from vllm.distributed import (
+
+from sglang.srt.distributed import (
     divide,
     get_tensor_model_parallel_rank,
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_reduce,
 )
-
 from sglang.srt.layers.parameter import BasevLLMParameter
 from sglang.srt.layers.quantization.base_config import (
     QuantizationConfig,
@@ -220,6 +220,7 @@ class VocabParallelEmbedding(torch.nn.Module):
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
         enable_tp: bool = True,
+        use_presharded_weights: bool = False,
     ):
         super().__init__()
         self.quant_config = quant_config
@@ -236,6 +237,12 @@ class VocabParallelEmbedding(torch.nn.Module):
         self.padding_size = padding_size
         self.org_vocab_size = org_num_embeddings or num_embeddings
         num_added_embeddings = num_embeddings - self.org_vocab_size
+        self.use_presharded_weights = use_presharded_weights
+        if use_presharded_weights:
+            assert (
+                num_added_embeddings == 0
+            ), "Lora is not supported with presharded weights."
+
         self.org_vocab_size_padded = pad_vocab_size(
             self.org_vocab_size, self.padding_size
         )
@@ -447,10 +454,14 @@ class VocabParallelEmbedding(torch.nn.Module):
             start_idx = start_idx // packed_factor
             shard_size = shard_size // packed_factor
         else:
-            assert loaded_weight.shape[output_dim] == self.org_vocab_size
+            assert loaded_weight.shape[output_dim] == (
+                self.org_vocab_size
+                // (self.tp_size if self.use_presharded_weights else 1)
+            )
 
         # Copy the data.
-        loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
+        if not self.use_presharded_weights:
+            loaded_weight = loaded_weight.narrow(output_dim, start_idx, shard_size)
         param[: loaded_weight.shape[0]].data.copy_(loaded_weight)
         param[loaded_weight.shape[0] :].data.fill_(0)
 
@@ -514,6 +525,7 @@ class ParallelLMHead(VocabParallelEmbedding):
         padding_size: int = DEFAULT_VOCAB_PADDING_SIZE,
         quant_config: Optional[QuantizationConfig] = None,
         prefix: str = "",
+        use_presharded_weights: bool = False,
     ):
         super().__init__(
             num_embeddings,
@@ -523,6 +535,7 @@ class ParallelLMHead(VocabParallelEmbedding):
             padding_size,
             quant_config,
             prefix,
+            use_presharded_weights=use_presharded_weights,
         )
         self.quant_config = quant_config
         if bias:

sglang/srt/lora/lora.py

@@ -19,18 +19,11 @@
 # https://github.com/vllm-project/vllm/blob/4abf6336ec65c270343eb895e7b18786e9274176/vllm/lora/layers.py
 
 
-import json
-import os
 import re
-from typing import Any, Dict, List, Optional, Tuple
 
-import safetensors.torch
 import torch
 from torch import nn
-from vllm.model_executor.layers.vocab_parallel_embedding import (
-    ParallelLMHead,
-    VocabParallelEmbedding,
-)
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
 
 from sglang.srt.layers.linear import (
     ColumnParallelLinear,
@@ -38,7 +31,6 @@ from sglang.srt.layers.linear import (
     QKVParallelLinear,
     RowParallelLinear,
 )
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_loader.loader import DefaultModelLoader
 
 

sglang/srt/managers/configure_logging.py

@@ -0,0 +1,46 @@
+"""
+Copyright 2023-2025 SGLang Team
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+"""
+Configure the logging settings of a server.
+
+Usage:
+python3 -m sglang.srt.managers.configure_logging --url http://localhost:30000
+"""
+
+import argparse
+
+import requests
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--url", type=str, default="http://localhost:30000")
+    parser.add_argument("--log-requests", action="store_true")
+    parser.add_argument(
+        "--dump-requests-folder", type=str, default="/tmp/sglang_request_dump"
+    )
+    parser.add_argument("--dump-requests-threshold", type=int, default=1000)
+    args = parser.parse_args()
+
+    response = requests.post(
+        args.url + "/configure_logging",
+        json={
+            "log_requests": args.log_requests,
+            "log_requests_level": 1,  # Log full requests
+            "dump_requests_folder": args.dump_requests_folder,
+            "dump_requests_threshold": args.dump_requests_threshold,
+        },
+    )
+    assert response.status_code == 200

sglang/srt/managers/data_parallel_controller.py

@@ -23,6 +23,7 @@ import psutil
 import setproctitle
 import zmq
 
+from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 from sglang.srt.managers.io_struct import (
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
@@ -55,6 +56,7 @@ class DataParallelController:
 
     def __init__(self, server_args, port_args) -> None:
         # Parse args
+        self.max_total_num_tokens = None
         self.server_args = server_args
         self.port_args = port_args
         self.load_balance_method = LoadBalanceMethod.from_str(
@@ -63,9 +65,10 @@ class DataParallelController:
 
         # Init inter-process communication
         self.context = zmq.Context(1 + server_args.dp_size)
-        self.recv_from_tokenizer = get_zmq_socket(
-            self.context, zmq.PULL, port_args.scheduler_input_ipc_name
-        )
+        if server_args.node_rank == 0:
+            self.recv_from_tokenizer = get_zmq_socket(
+                self.context, zmq.PULL, port_args.scheduler_input_ipc_name, False
+            )
 
         # Dispatch method
         self.round_robin_counter = 0
@@ -75,33 +78,50 @@ class DataParallelController:
         }
         self.dispatching = dispatch_lookup[self.load_balance_method]
 
-        # Start data parallel workers
-        base_gpu_id = 0
+        # Launch data parallel workers
+        self.scheduler_procs = []
         self.workers = [None] * server_args.dp_size
 
+        if not server_args.enable_dp_attention:
+            dp_port_args = self.launch_dp_schedulers(server_args, port_args)
+        else:
+            dp_port_args = self.launch_dp_attention_schedulers(server_args, port_args)
+
+        # Only node rank 0 runs the real data parallel controller that dispatches the requests.
+        if server_args.node_rank == 0:
+            for dp_rank in range(server_args.dp_size):
+                self.workers[dp_rank] = get_zmq_socket(
+                    self.context,
+                    zmq.PUSH,
+                    dp_port_args[dp_rank].scheduler_input_ipc_name,
+                    True,
+                )
+
+        self.max_req_input_len = None
+
+    def launch_dp_schedulers(self, server_args, port_args):
+        base_gpu_id = 0
+
         threads = []
         sockets = []
+        dp_port_args = []
         for dp_rank in range(server_args.dp_size):
             tmp_port_args = PortArgs.init_new(server_args)
             tmp_port_args.tokenizer_ipc_name = port_args.tokenizer_ipc_name
             tmp_port_args.detokenizer_ipc_name = port_args.detokenizer_ipc_name
+            dp_port_args.append(tmp_port_args)
 
-            if server_args.enable_dp_attention:
-                # Data parallelism resues the tensor parallelism group,
-                # so all dp ranks should use the same nccl port.
-                tmp_port_args.nccl_port = port_args.nccl_port
-            else:
-                # This port is checked free in PortArgs.init_new.
-                # We hold it first so that the next dp worker gets a different port
-                sockets.append(bind_port(tmp_port_args.nccl_port))
+            # This port is checked free in PortArgs.init_new.
+            # We hold it first so that the next dp worker gets a different port
+            sockets.append(bind_port(tmp_port_args.nccl_port))
 
             # Create a thread for each worker
             thread = threading.Thread(
-                target=self.launch_worker_func,
+                target=self.launch_tensor_parallel_group,
                 args=(server_args, tmp_port_args, base_gpu_id, dp_rank),
             )
             threads.append(thread)
-            base_gpu_id += 1 if server_args.enable_dp_attention else server_args.tp_size
+            base_gpu_id += server_args.tp_size
 
         # Free all sockets before starting the threads to launch TP workers
         for sock in sockets:
@@ -113,26 +133,14 @@ class DataParallelController:
         for thread in threads:
             thread.join()
 
-    def launch_worker_func(
-        self,
-        server_args: ServerArgs,
-        port_args: PortArgs,
-        base_gpu_id: int,
-        dp_rank: int,
-    ):
-        logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
+        return dp_port_args
 
-        launch_func_ = (
-            self.launch_tensor_parallel_process
-            if server_args.enable_dp_attention
-            else self.launch_tensor_parallel_group
-        )
-        self.workers[dp_rank] = launch_func_(
-            server_args,
-            port_args,
-            base_gpu_id,
-            dp_rank,
-        )
+    def launch_dp_attention_schedulers(self, server_args, port_args):
+        self.launch_tensor_parallel_group(server_args, port_args, 0, None)
+        dp_port_args = []
+        for dp_rank in range(server_args.dp_size):
+            dp_port_args.append(PortArgs.init_new(server_args, dp_rank))
+        return dp_port_args
 
     def launch_tensor_parallel_group(
         self,
@@ -141,8 +149,10 @@ class DataParallelController:
         base_gpu_id: int,
         dp_rank: int,
     ):
+        if not server_args.enable_dp_attention:
+            logger.info(f"Launch DP{dp_rank} starting at GPU #{base_gpu_id}.")
+
         # Launch tensor parallel scheduler processes
-        scheduler_procs = []
         scheduler_pipe_readers = []
         tp_size_per_node = server_args.tp_size // server_args.nnodes
         tp_rank_range = range(
@@ -150,52 +160,39 @@ class DataParallelController:
             tp_size_per_node * (server_args.node_rank + 1),
         )
         for tp_rank in tp_rank_range:
+            rank_port_args = port_args
+
+            if server_args.enable_dp_attention:
+                # dp attention has different sharding logic
+                _, _, dp_rank = compute_dp_attention_world_info(
+                    server_args.enable_dp_attention,
+                    tp_rank,
+                    server_args.tp_size,
+                    server_args.dp_size,
+                )
+                # compute zmq ports for this dp rank
+                rank_port_args = PortArgs.init_new(server_args, dp_rank)
+                # Data parallelism resues the tensor parallelism group,
+                # so all dp ranks should use the same nccl port.
+                rank_port_args.nccl_port = port_args.nccl_port
+
             reader, writer = mp.Pipe(duplex=False)
             gpu_id = server_args.base_gpu_id + base_gpu_id + tp_rank % tp_size_per_node
             proc = mp.Process(
                 target=run_scheduler_process,
-                args=(server_args, port_args, gpu_id, tp_rank, dp_rank, writer),
+                args=(server_args, rank_port_args, gpu_id, tp_rank, dp_rank, writer),
             )
             proc.start()
-            scheduler_procs.append(proc)
+            self.scheduler_procs.append(proc)
             scheduler_pipe_readers.append(reader)
 
-        send_to = get_zmq_socket(
-            self.context, zmq.PUSH, port_args.scheduler_input_ipc_name
-        )
-
-        # Wait for model to finish loading and get max token nums
+        # Wait for model to finish loading
         scheduler_info = []
         for i in range(len(scheduler_pipe_readers)):
             scheduler_info.append(scheduler_pipe_readers[i].recv())
 
         self.max_total_num_tokens = scheduler_info[0]["max_total_num_tokens"]
-
-        return send_to
-
-    def launch_tensor_parallel_process(
-        self,
-        server_args: ServerArgs,
-        port_args: PortArgs,
-        base_gpu_id: int,
-        dp_rank: int,
-    ):
-        reader, writer = mp.Pipe(duplex=False)
-        gpu_id = base_gpu_id
-        tp_rank = dp_rank
-        proc = mp.Process(
-            target=run_scheduler_process,
-            args=(server_args, port_args, gpu_id, tp_rank, dp_rank, writer),
-        )
-        proc.start()
-        send_to = get_zmq_socket(
-            self.context, zmq.PUSH, port_args.scheduler_input_ipc_name
-        )
-
-        scheduler_info = reader.recv()
-        self.max_total_num_tokens = scheduler_info["max_total_num_tokens"]
-
-        return send_to
+        self.max_req_input_len = scheduler_info[0]["max_req_input_len"]
 
     def round_robin_scheduler(self, req):
         self.workers[self.round_robin_counter].send_pyobj(req)
@@ -221,8 +218,8 @@ class DataParallelController:
                 ):
                     self.dispatching(recv_req)
                 else:
-                    # Send other control messages to all workers
-                    for worker in self.workers:
+                    # Send other control messages to first worker of tp group
+                    for worker in self.workers[:: self.server_args.tp_size]:
                         worker.send_pyobj(recv_req)
 
 
@@ -238,9 +235,19 @@ def run_data_parallel_controller_process(
     try:
         controller = DataParallelController(server_args, port_args)
         pipe_writer.send(
-            {"status": "ready", "max_total_num_tokens": controller.max_total_num_tokens}
+            {
+                "status": "ready",
+                "max_total_num_tokens": controller.max_total_num_tokens,
+                "max_req_input_len": controller.max_req_input_len,
+            }
         )
-        controller.event_loop()
+        if server_args.node_rank == 0:
+            controller.event_loop()
+        for proc in controller.scheduler_procs:
+            proc.join()
+            logger.error(
+                f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}"
+            )
     except Exception:
         traceback = get_exception_traceback()
         logger.error(f"DataParallelController hit an exception: {traceback}")

sglang/srt/managers/detokenizer_manager.py

@@ -15,6 +15,7 @@
 
 import dataclasses
 import logging
+import os
 import signal
 from collections import OrderedDict
 from typing import Dict, List, Union
@@ -35,6 +36,12 @@ from sglang.utils import find_printable_text, get_exception_traceback
 
 logger = logging.getLogger(__name__)
 
+# Maximum number of request states that detokenizer can hold. When exceeded,
+# oldest request states will be evicted. Default: 65536 (1<<16).
+# For more details, see: https://github.com/sgl-project/sglang/issues/2812
+# Use power of 2 values for better memory allocation.
+DETOKENIZER_MAX_STATES = int(os.environ.get("SGLANG_DETOKENIZER_MAX_STATES", 1 << 16))
+
 
 @dataclasses.dataclass
 class DecodeStatus:
@@ -58,10 +65,10 @@ class DetokenizerManager:
         # Init inter-process communication
         context = zmq.Context(2)
         self.recv_from_scheduler = get_zmq_socket(
-            context, zmq.PULL, port_args.detokenizer_ipc_name
+            context, zmq.PULL, port_args.detokenizer_ipc_name, True
         )
         self.send_to_tokenizer = get_zmq_socket(
-            context, zmq.PUSH, port_args.tokenizer_ipc_name
+            context, zmq.PUSH, port_args.tokenizer_ipc_name, False
         )
 
         if server_args.skip_tokenizer_init:
@@ -71,9 +78,10 @@ class DetokenizerManager:
                 server_args.tokenizer_path,
                 tokenizer_mode=server_args.tokenizer_mode,
                 trust_remote_code=server_args.trust_remote_code,
+                revision=server_args.revision,
             )
 
-        self.decode_status = LimitedCapacityDict()
+        self.decode_status = LimitedCapacityDict(capacity=DETOKENIZER_MAX_STATES)
 
     def trim_matched_stop(
         self, output: Union[str, List[int]], finished_reason: Dict, no_stop_trim: bool
@@ -155,7 +163,17 @@ class DetokenizerManager:
             # Incremental decoding
             output_strs = []
             for i in range(bs):
-                s = self.decode_status[recv_obj.rids[i]]
+                try:
+                    s = self.decode_status[recv_obj.rids[i]]
+                except KeyError:
+                    raise RuntimeError(
+                        f"Decode status not found for request {recv_obj.rids[i]}. "
+                        "It may be due to the request being evicted from the decode status due to memory pressure. "
+                        "Please increase the maximum number of requests by setting "
+                        "the SGLANG_DETOKENIZER_MAX_STATES environment variable to a bigger value than the default value. "
+                        f"The current value is {DETOKENIZER_MAX_STATES}. "
+                        "For more details, see: https://github.com/sgl-project/sglang/issues/2812"
+                    )
                 new_text = read_texts[i][len(surr_texts[i]) :]
                 if recv_obj.finished_reasons[i] is None:
                     # Streaming chunk: update the decode status
@@ -181,8 +199,6 @@ class DetokenizerManager:
                     finished_reasons=recv_obj.finished_reasons,
                     output_strs=output_strs,
                     prompt_tokens=recv_obj.prompt_tokens,
-                    origin_input_ids=recv_obj.origin_input_ids,
-                    output_ids=recv_obj.output_ids,
                     completion_tokens=recv_obj.completion_tokens,
                     cached_tokens=recv_obj.cached_tokens,
                     input_token_logprobs_val=recv_obj.input_token_logprobs_val,
@@ -193,13 +209,12 @@ class DetokenizerManager:
                     input_top_logprobs_idx=recv_obj.input_top_logprobs_idx,
                     output_top_logprobs_val=recv_obj.output_top_logprobs_val,
                     output_top_logprobs_idx=recv_obj.output_top_logprobs_idx,
-                    normalized_prompt_logprob=recv_obj.normalized_prompt_logprob,
                 )
             )
 
 
 class LimitedCapacityDict(OrderedDict):
-    def __init__(self, capacity=1 << 15, *args, **kwargs):
+    def __init__(self, capacity: int, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.capacity = capacity