sglang 0.3.2__py3-none-any.whl → 0.3.3__py3-none-any.whl

sglang/srt/sampling/sampling_batch_info.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, List
 import torch
 
 import sglang.srt.sampling.penaltylib as penaltylib
+from sglang.srt.constrained import RegexGuide
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ScheduleBatch
@@ -13,22 +14,24 @@ if TYPE_CHECKING:
 
 @dataclasses.dataclass
 class SamplingBatchInfo:
-    # Basic Info
-    vocab_size: int
-
     # Batched sampling params
-    temperatures: torch.Tensor = None
-    top_ps: torch.Tensor = None
-    top_ks: torch.Tensor = None
-    min_ps: torch.Tensor = None
+    temperatures: torch.Tensor
+    top_ps: torch.Tensor
+    top_ks: torch.Tensor
+    min_ps: torch.Tensor
 
     # Dispatch in CUDA graph
-    need_min_p_sampling: bool = False
+    need_min_p_sampling: bool
 
     # Bias Tensors
+    vocab_size: int
     logit_bias: torch.Tensor = None
     vocab_mask: torch.Tensor = None
 
+    # FSM states
+    regex_fsms: List[RegexGuide] = None
+    regex_fsm_states: List[int] = None
+
     # Penalizer
     penalizer_orchestrator: penaltylib.BatchedPenalizerOrchestrator = None
     linear_penalties: torch.Tensor = None
@@ -37,24 +40,30 @@ class SamplingBatchInfo:
     @classmethod
     def from_schedule_batch(cls, batch: ScheduleBatch, vocab_size: int):
         reqs = batch.reqs
-        ret = cls(vocab_size=vocab_size)
-
-        with torch.device("cuda"):
-            ret.temperatures = torch.tensor(
+        with batch.input_ids.device:
+            temperatures = torch.tensor(
                 [r.sampling_params.temperature for r in reqs],
                 dtype=torch.float,
             ).view(-1, 1)
-            ret.top_ps = torch.tensor(
+            top_ps = torch.tensor(
                 [r.sampling_params.top_p for r in reqs], dtype=torch.float
             )
-            ret.top_ks = torch.tensor(
+            top_ks = torch.tensor(
                 [r.sampling_params.top_k for r in reqs], dtype=torch.int
             )
-            ret.min_ps = torch.tensor(
+            min_ps = torch.tensor(
                 [r.sampling_params.min_p for r in reqs], dtype=torch.float
             )
 
-        ret.need_min_p_sampling = any(r.sampling_params.min_p > 0 for r in reqs)
+        ret = cls(
+            temperatures=temperatures,
+            top_ps=top_ps,
+            top_ks=top_ks,
+            min_ps=min_ps,
+            need_min_p_sampling=any(r.sampling_params.min_p > 0 for r in reqs),
+            vocab_size=vocab_size,
+        )
+        # TODO (lianmin): `need_min_p_sampling` needs to be updated in filter and merge.
 
         # Each penalizers will do nothing if they evaluate themselves as not required by looking at
         # the sampling_params of the requests (See {_is_required()} of each penalizers). So this
@@ -102,24 +111,24 @@ class SamplingBatchInfo:
                         )
                     self.linear_penalties = penalizer.apply(self.linear_penalties)
 
-    def update_regex_vocab_mask(self, batch: ScheduleBatch):
-        has_regex = any(req.regex_fsm is not None for req in batch.reqs)
+    def update_regex_vocab_mask(self):
+        has_regex = self.regex_fsms and any(regex_fsm for regex_fsm in self.regex_fsms)
 
         # Reset the vocab mask
         self.vocab_mask = None
 
         if has_regex:
             self.vocab_mask = torch.zeros(
-                batch.batch_size(), self.vocab_size, dtype=torch.bool, device="cuda"
+                len(self.temperatures), self.vocab_size, dtype=torch.bool, device="cuda"
             )
-            for i, req in enumerate(batch.reqs):
-                if req.regex_fsm is not None:
+            for i, regex_fsm in enumerate(self.regex_fsms):
+                if regex_fsm is not None:
                     self.vocab_mask[i].fill_(1)
                     self.vocab_mask[i][
-                        req.regex_fsm.get_next_instruction(req.regex_fsm_state).tokens
+                        regex_fsm.get_next_instruction(self.regex_fsm_states[i]).tokens
                     ] = 0
 
-    def filter(self, unfinished_indices: List[int], new_indices: torch.Tensor):
+    def filter_batch(self, unfinished_indices: List[int], new_indices: torch.Tensor):
         self.penalizer_orchestrator.filter(unfinished_indices, new_indices)
 
         for item in [
@@ -129,9 +138,9 @@ class SamplingBatchInfo:
             "min_ps",
             "logit_bias",
         ]:
-            self_val = getattr(self, item, None)
-            if self_val is not None:  # logit_bias can be None
-                setattr(self, item, self_val[new_indices])
+            value = getattr(self, item, None)
+            if value is not None:  # logit_bias can be None
+                setattr(self, item, value[new_indices])
 
     @staticmethod
     def merge_bias_tensor(
@@ -153,7 +162,7 @@ class SamplingBatchInfo:
 
         return None
 
-    def merge(self, other: "SamplingBatchInfo"):
+    def merge_batch(self, other: "SamplingBatchInfo"):
         self.penalizer_orchestrator.merge(other.penalizer_orchestrator)
 
         for item in [

sglang/srt/sampling/sampling_params.py

@@ -26,7 +26,7 @@ class SamplingParams:
         max_new_tokens: int = 128,
         min_new_tokens: int = 0,
         stop: Optional[Union[str, List[str]]] = None,
-        stop_token_ids: Optional[List[int]] = [],
+        stop_token_ids: Optional[List[int]] = None,
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = -1,
@@ -49,6 +49,8 @@ class SamplingParams:
         self.presence_penalty = presence_penalty
         self.repetition_penalty = repetition_penalty
         self.stop_strs = stop
+        if stop_token_ids is None:
+            stop_token_ids = []
         self.stop_token_ids = {*stop_token_ids}
         self.max_new_tokens = max_new_tokens
         self.min_new_tokens = min_new_tokens

sglang/srt/server.py

@@ -19,11 +19,13 @@ SRT = SGLang Runtime.
 """
 
 import asyncio
+import atexit
 import dataclasses
 import json
 import logging
 import multiprocessing as mp
 import os
+import random
 import threading
 import time
 from http import HTTPStatus
@@ -41,21 +43,15 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, Response, StreamingResponse
 
 from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
-from sglang.srt.constrained import disable_cache
 from sglang.srt.hf_transformers_utils import get_tokenizer
-from sglang.srt.managers.controller_multi import (
-    start_controller_process as start_controller_process_multi,
-)
-from sglang.srt.managers.controller_single import launch_tp_servers
-from sglang.srt.managers.controller_single import (
-    start_controller_process as start_controller_process_single,
-)
-from sglang.srt.managers.detokenizer_manager import start_detokenizer_process
+from sglang.srt.managers.detokenizer_manager import run_detokenizer_process
 from sglang.srt.managers.io_struct import (
     EmbeddingReqInput,
     GenerateReqInput,
+    RewardReqInput,
     UpdateWeightReqInput,
 )
+from sglang.srt.managers.scheduler import run_scheduler_process
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
 from sglang.srt.openai_api.adapter import (
     load_chat_template_for_openai_api,
@@ -74,15 +70,12 @@ from sglang.srt.openai_api.protocol import ModelCard, ModelList
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import (
     add_api_key_middleware,
-    allocate_init_ports,
     assert_pkg_version,
     configure_logger,
-    enable_show_time_cost,
-    is_hip,
+    is_port_available,
     kill_child_process,
     maybe_set_triton_cache_manager,
-    prepare_model,
-    prepare_tokenizer,
+    prepare_model_and_tokenizer,
     set_ulimit,
 )
 from sglang.utils import get_exception_traceback
@@ -127,6 +120,7 @@ async def health_generate(request: Request) -> Response:
 
 @app.get("/get_model_info")
 async def get_model_info():
+    """Get the model information."""
     result = {
         "model_path": tokenizer_manager.model_path,
         "is_generation": tokenizer_manager.is_generation,
@@ -136,11 +130,13 @@ async def get_model_info():
 
 @app.get("/get_server_args")
 async def get_server_args():
+    """Get the server arguments."""
     return dataclasses.asdict(tokenizer_manager.server_args)
 
 
 @app.get("/flush_cache")
 async def flush_cache():
+    """Flush the radix cache."""
     tokenizer_manager.flush_cache()
     return Response(
         content="Cache flushed.\nPlease check backend logs for more details. "
@@ -151,7 +147,7 @@ async def flush_cache():
 
 @app.post("/update_weights")
 async def update_weights(obj: UpdateWeightReqInput, request: Request):
-
+    """Update the weights inplace without re-launching the server."""
     success, message = await tokenizer_manager.update_weights(obj, request)
     content = {"success": success, "message": message}
     if success:
@@ -166,6 +162,7 @@ async def update_weights(obj: UpdateWeightReqInput, request: Request):
         )
 
 
+# fastapi implicitly converts json in the request to obj (dataclass)
 async def generate_request(obj: GenerateReqInput, request: Request):
     """Handle a generate request."""
     if obj.stream:
@@ -213,6 +210,21 @@ app.post("/encode")(encode_request)
 app.put("/encode")(encode_request)
 
 
+async def judge_request(obj: RewardReqInput, request: Request):
+    """Handle a reward model request."""
+    try:
+        ret = await tokenizer_manager.generate_request(obj, request).__anext__()
+        return ret
+    except ValueError as e:
+        return JSONResponse(
+            {"error": {"message": str(e)}}, status_code=HTTPStatus.BAD_REQUEST
+        )
+
+
+app.post("/judge")(judge_request)
+app.put("/judge")(judge_request)
+
+
 @app.post("/v1/completions")
 async def openai_v1_completions(raw_request: Request):
     return await v1_completions(tokenizer_manager, raw_request)
@@ -280,102 +292,95 @@ async def retrieve_file_content(file_id: str):
     return await v1_retrieve_file_content(file_id)
 
 
-def launch_server(
+def launch_engine(
     server_args: ServerArgs,
-    pipe_finish_writer: Optional[mp.connection.Connection] = None,
 ):
-    """Launch an HTTP server."""
+    """
+    Launch the Tokenizer Manager in the main process, the Scheduler in a subprocess, and the Detokenizer Manager in another subprocess.
+    """
+
     global tokenizer_manager
 
+    # Configure global environment
     configure_logger(server_args)
-
     server_args.check_server_args()
     _set_envs_and_config(server_args)
 
     # Allocate ports for inter-process communications
-    server_args.port, server_args.additional_ports = allocate_init_ports(
-        server_args.port,
-        server_args.additional_ports,
-        server_args.dp_size,
-    )
-    ports = server_args.additional_ports
-    port_args = PortArgs(
-        tokenizer_port=ports[0],
-        controller_port=ports[1],
-        detokenizer_port=ports[2],
-        nccl_ports=ports[3:],
-    )
+    port_args = PortArgs.init_new(server_args)
     logger.info(f"{server_args=}")
 
-    # Use model from www.modelscope.cn, first download the model.
-    server_args.model_path = prepare_model(server_args.model_path)
-    server_args.tokenizer_path = prepare_tokenizer(server_args.tokenizer_path)
-
-    # Launch processes for multi-node tensor parallelism
-    if server_args.nnodes > 1 and server_args.node_rank != 0:
-        tp_size_local = server_args.tp_size // server_args.nnodes
-        gpu_ids = [i for _ in range(server_args.nnodes) for i in range(tp_size_local)]
-        tp_rank_range = list(
-            range(
-                server_args.node_rank * tp_size_local,
-                (server_args.node_rank + 1) * tp_size_local,
-            )
-        )
-        procs = launch_tp_servers(
-            gpu_ids,
-            tp_rank_range,
-            server_args,
-            ports[3],
-        )
-
-        try:
-            for p in procs:
-                p.join()
-        finally:
-            kill_child_process(os.getpid(), including_parent=False)
-            return
-
-    # Launch processes
-    pipe_controller_reader, pipe_controller_writer = mp.Pipe(duplex=False)
+    # If using model from www.modelscope.cn, first download the model.
+    server_args.model_path, server_args.tokenizer_path = prepare_model_and_tokenizer(
+        server_args.model_path, server_args.tokenizer_path
+    )
 
-    if server_args.dp_size == 1:
-        start_controller_process = start_controller_process_single
-    else:
-        start_controller_process = start_controller_process_multi
-    proc_controller = mp.Process(
-        target=start_controller_process,
-        args=(server_args, port_args, pipe_controller_writer),
+    # 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(
+        tp_size_per_node * server_args.node_rank,
+        tp_size_per_node * (server_args.node_rank + 1),
     )
-    proc_controller.start()
+    for tp_rank in tp_rank_range:
+        reader, writer = mp.Pipe(duplex=False)
+        gpu_id = tp_rank % tp_size_per_node
+        proc = mp.Process(
+            target=run_scheduler_process,
+            args=(server_args, port_args, gpu_id, tp_rank, writer),
+        )
+        proc.start()
+        scheduler_procs.append(proc)
+        scheduler_pipe_readers.append(reader)
+
+    if server_args.node_rank >= 1:
+        # For other nodes, they do not need to run tokenizer or detokenizer,
+        # so they can just wait here.
+        while True:
+            pass
 
-    pipe_detoken_reader, pipe_detoken_writer = mp.Pipe(duplex=False)
-    proc_detoken = mp.Process(
-        target=start_detokenizer_process,
+    # Launch detokenizer process
+    detoken_proc = mp.Process(
+        target=run_detokenizer_process,
         args=(
             server_args,
             port_args,
-            pipe_detoken_writer,
         ),
     )
-    proc_detoken.start()
+    detoken_proc.start()
 
+    # Launch tokenizer process
     tokenizer_manager = TokenizerManager(server_args, port_args)
     if server_args.chat_template:
         load_chat_template_for_openai_api(tokenizer_manager, server_args.chat_template)
 
-    # Wait for the model to finish loading
-    controller_init_state = pipe_controller_reader.recv()
-    detoken_init_state = pipe_detoken_reader.recv()
-
-    if controller_init_state != "init ok" or detoken_init_state != "init ok":
-        proc_controller.kill()
-        proc_detoken.kill()
-        raise RuntimeError(
-            "Initialization failed. "
-            f"controller_init_state: {controller_init_state}, "
-            f"detoken_init_state: {detoken_init_state}"
-        )
-    assert proc_controller.is_alive() and proc_detoken.is_alive()
+    # Wait for model to finish loading
+    for i in range(len(scheduler_pipe_readers)):
+        scheduler_pipe_readers[i].recv()
+
+
+def launch_server(
+    server_args: ServerArgs,
+    pipe_finish_writer: Optional[mp.connection.Connection] = None,
+):
+    """
+    Launch SRT (SGLang Runtime) Server
+
+    The SRT server consists of an HTTP server and the SRT engine.
+
+    1. HTTP server: A FastAPI server that routes requests to the engine.
+    2. SRT engine:
+        1. Tokenizer Manager: Tokenizes the requests and sends them to the scheduler.
+        2. Scheduler (subprocess): Receives requests from the Tokenizer Manager, schedules batches, forwards them, and sends the output tokens to the Detokenizer Manager.
+        3. Detokenizer Manager (subprocess): Detokenizes the output tokens and sends the result back to the Tokenizer Manager.
+
+    Note:
+    1. The HTTP server and Tokenizer Manager both run in the main process.
+    2. Inter-process communication is done through ICP (each process uses a different port) via the ZMQ library.
+    """
+
+    launch_engine(server_args=server_args)
 
     # Add api key authorization
     if server_args.api_key:
@@ -388,7 +393,7 @@ def launch_server(
     t.start()
 
     try:
-        # Listen for requests
+        # Listen for HTTP requests
         uvicorn.run(
             app,
             host=server_args.host,
@@ -412,14 +417,6 @@ def _set_envs_and_config(server_args: ServerArgs):
     # Set ulimit
     set_ulimit()
 
-    # Enable show time cost for debugging
-    if server_args.show_time_cost:
-        enable_show_time_cost()
-
-    # Disable disk cache
-    if server_args.disable_disk_cache:
-        disable_cache()
-
     # Fix triton bugs
     if server_args.tp_size * server_args.dp_size > 1:
         # FIXME: remove this after https://github.com/triton-lang/triton/pull/4295 is used as a dependency.
@@ -435,9 +432,7 @@ def _set_envs_and_config(server_args: ServerArgs):
             "at https://docs.flashinfer.ai/installation.html.",
         )
 
-    if is_hip():
-        # to figure out a better method of not using fork later
-        mp.set_start_method("spawn", force=True)
+    mp.set_start_method("spawn", force=True)
 
 
 def _wait_and_warmup(server_args, pipe_finish_writer, pid):
@@ -467,7 +462,6 @@ def _wait_and_warmup(server_args, pipe_finish_writer, pid):
         return
 
     model_info = res.json()
-
     # Send a warmup request
     request_name = "/generate" if model_info["is_generation"] else "/encode"
     max_new_tokens = 8 if model_info["is_generation"] else 1
@@ -501,7 +495,7 @@ def _wait_and_warmup(server_args, pipe_finish_writer, pid):
 
     logger.info("The server is fired up and ready to roll!")
     if pipe_finish_writer is not None:
-        pipe_finish_writer.send("init ok")
+        pipe_finish_writer.send("ready")
 
 
 class Runtime:
@@ -520,18 +514,20 @@ class Runtime:
         """See the arguments in server_args.py::ServerArgs"""
         self.server_args = ServerArgs(*args, log_level=log_level, **kwargs)
 
+        # before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown()
+        atexit.register(self.shutdown)
+
         # Pre-allocate ports
-        self.server_args.port, self.server_args.additional_ports = allocate_init_ports(
-            self.server_args.port,
-            self.server_args.additional_ports,
-            self.server_args.dp_size,
-        )
+        for port in range(10000, 40000):
+            if is_port_available(port):
+                break
+            port += 1
+        self.server_args.port = port
 
         self.url = self.server_args.url()
-        self.generate_url = (
-            f"http://{self.server_args.host}:{self.server_args.port}/generate"
-        )
+        self.generate_url = self.url + "/generate"
 
+        # NOTE: We store pid instead of proc to fix some issues during __delete__
         self.pid = None
         pipe_reader, pipe_writer = mp.Pipe(duplex=False)
 
@@ -548,7 +544,7 @@ class Runtime:
         except EOFError:
             init_state = ""
 
-        if init_state != "init ok":
+        if init_state != "ready":
             self.shutdown()
             raise RuntimeError(
                 "Initialization failed. Please see the error messages above."
@@ -599,7 +595,7 @@ class Runtime:
                         if chunk == "data: [DONE]\n\n":
                             break
                         data = json.loads(chunk[5:].strip("\n"))
-                        if hasattr(data, "text"):
+                        if "text" in data:
                             cur = data["text"][pos:]
                             if cur:
                                 yield cur
@@ -635,16 +631,71 @@ class Runtime:
 
     def encode(
         self,
-        prompt: Union[str, List[str]],
+        prompt: Union[str, List[str], List[Dict], List[List[Dict]]],
     ):
-        json_data = {
-            "text": prompt,
-        }
-        response = requests.post(
-            self.url + "/encode",
-            json=json_data,
-        )
+        if isinstance(prompt, str) or isinstance(prompt[0], str):
+            # embedding
+            json_data = {
+                "text": prompt,
+            }
+            response = requests.post(
+                self.url + "/encode",
+                json=json_data,
+            )
+        else:
+            # reward
+            json_data = {
+                "conv": prompt,
+            }
+            response = requests.post(
+                self.url + "/judge",
+                json=json_data,
+            )
         return json.dumps(response.json())
 
     def __del__(self):
         self.shutdown()
+
+
+class Engine:
+    """
+    SRT Engine without an HTTP server layer.
+
+    This class provides a direct inference engine without the need for an HTTP server. It is designed for use cases where
+    launching the HTTP server adds unnecessary complexity or overhead,
+    """
+
+    def __init__(self, *args, **kwargs):
+
+        # before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown()
+        atexit.register(self.shutdown)
+
+        server_args = ServerArgs(*args, **kwargs)
+        launch_engine(server_args=server_args)
+
+    def generate(
+        self,
+        prompt: Union[str, List[str]],
+        sampling_params: Optional[Dict] = None,
+        return_logprob: Optional[Union[List[bool], bool]] = False,
+        logprob_start_len: Optional[Union[List[int], int]] = None,
+        top_logprobs_num: Optional[Union[List[int], int]] = None,
+        lora_path: Optional[List[Optional[str]]] = None,
+    ):
+        obj = GenerateReqInput(
+            text=prompt,
+            sampling_params=sampling_params,
+            return_logprob=return_logprob,
+            logprob_start_len=logprob_start_len,
+            top_logprobs_num=top_logprobs_num,
+            lora_path=lora_path,
+        )
+
+        # get the current event loop
+        loop = asyncio.get_event_loop()
+        return loop.run_until_complete(generate_request(obj, None))
+
+    def shutdown(self):
+        kill_child_process(os.getpid(), including_parent=False)
+
+    # TODO (ByronHsu): encode and async generate