sglang 0.2.9.post1__py3-none-any.whl → 0.2.10__py3-none-any.whl

sglang/srt/openai_api/adapter.py

@@ -251,7 +251,9 @@ async def process_batch(tokenizer_manager, batch_id: str, batch_request: BatchRe
             if end_point == "/v1/chat/completions":
                 responses = v1_chat_generate_response(request, ret, to_file=True)
             else:
-                responses = v1_generate_response(request, ret, to_file=True)
+                responses = v1_generate_response(
+                    request, ret, tokenizer_manager, to_file=True
+                )
 
         except Exception as e:
             error_json = {
@@ -339,6 +341,7 @@ def v1_generate_request(all_requests):
     return_logprobs = []
     top_logprobs_nums = []
     first_prompt_type = type(all_requests[0].prompt)
+
     for request in all_requests:
         prompt = request.prompt
         assert (
@@ -364,7 +367,7 @@ def v1_generate_request(all_requests):
         )
         if len(all_requests) > 1 and request.n > 1:
             raise ValueError(
-                "Batch operation is not supported for completions from files"
+                "Parallel sampling is not supported for completions from files"
             )
 
     if len(all_requests) == 1:
@@ -377,10 +380,11 @@ def v1_generate_request(all_requests):
         else:
             prompt_kwargs = {"input_ids": prompt}
     else:
-        if isinstance(prompts[0], str):
+        if isinstance(prompts[0], str) or isinstance(propmt[0][0], str):
             prompt_kwargs = {"text": prompts}
         else:
             prompt_kwargs = {"input_ids": prompts}
+
     adapted_request = GenerateReqInput(
         **prompt_kwargs,
         sampling_params=sampling_params_list,
@@ -389,35 +393,52 @@ def v1_generate_request(all_requests):
         return_text_in_logprobs=True,
         stream=all_requests[0].stream,
     )
+
     if len(all_requests) == 1:
         return adapted_request, all_requests[0]
     return adapted_request, all_requests
 
 
-def v1_generate_response(request, ret, to_file=False):
+def v1_generate_response(request, ret, tokenizer_manager, to_file=False):
     choices = []
     echo = False
 
-    if (not isinstance(request, List)) and request.echo:
+    if (not isinstance(request, list)) and request.echo:
         # TODO: handle the case propmt is token ids
-        if isinstance(request.prompt, list):
+        if isinstance(request.prompt, list) and isinstance(request.prompt[0], str):
+            # for the case of multiple str prompts
             prompts = request.prompt
+        elif isinstance(request.prompt, list) and isinstance(request.prompt[0], list):
+            # for the case of multiple token ids prompts
+            prompts = [
+                tokenizer_manager.tokenizer.decode(prompt, skip_special_tokens=True)
+                for prompt in request.prompt
+            ]
+        elif isinstance(request.prompt, list) and isinstance(request.prompt[0], int):
+            # for the case of single token ids prompt
+            prompts = [
+                tokenizer_manager.tokenizer.decode(
+                    request.prompt, skip_special_tokens=True
+                )
+            ]
         else:
+            # for the case of single str prompt
             prompts = [request.prompt]
         echo = True
 
     for idx, ret_item in enumerate(ret):
         text = ret_item["text"]
-        if isinstance(request, List) and request[idx].echo:
+        if isinstance(request, list) and request[idx].echo:
             echo = True
             text = request[idx].prompt + text
-        if (not isinstance(request, List)) and echo:
-            text = prompts[idx] + text
+        if (not isinstance(request, list)) and echo:
+            prompt_index = idx // request.n
+            text = prompts[prompt_index] + text
 
         logprobs = False
-        if isinstance(request, List) and request[idx].logprobs:
+        if isinstance(request, list) and request[idx].logprobs:
             logprobs = True
-        elif (not isinstance(request, List)) and request.logprobs:
+        elif (not isinstance(request, list)) and request.logprobs:
             logprobs = True
         if logprobs:
             if echo:
@@ -479,15 +500,16 @@ def v1_generate_response(request, ret, to_file=False):
             responses.append(response)
         return responses
     else:
+        prompt_tokens = sum(item["meta_info"]["prompt_tokens"] for item in ret)
         completion_tokens = sum(item["meta_info"]["completion_tokens"] for item in ret)
         response = CompletionResponse(
             id=ret[0]["meta_info"]["id"],
             model=request.model,
             choices=choices,
             usage=UsageInfo(
-                prompt_tokens=ret[0]["meta_info"]["prompt_tokens"],
+                prompt_tokens=prompt_tokens,
                 completion_tokens=completion_tokens,
-                total_tokens=ret[0]["meta_info"]["prompt_tokens"] + completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens,
             ),
         )
     return response
@@ -513,8 +535,18 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
 
                     if not stream_buffer:  # The first chunk
                         if request.echo:
+                            if isinstance(request.prompt, str):
+                                # for the case of single str prompts
+                                prompts = request.prompt
+                            elif isinstance(request.prompt, list) and isinstance(
+                                request.prompt[0], int
+                            ):
+                                prompts = tokenizer_manager.tokenizer.decode(
+                                    request.prompt, skip_special_tokens=True
+                                )
+
                             # Prepend prompt in response text.
-                            text = request.prompt + text
+                            text = prompts + text
 
                     if request.logprobs:
                         # The first chunk and echo is enabled.
@@ -539,7 +571,6 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
                                 "output_top_logprobs"
                             ][n_prev_token:],
                         )
-
                         n_prev_token = len(
                             content["meta_info"]["output_token_logprobs"]
                         )
@@ -588,7 +619,7 @@ async def v1_completions(tokenizer_manager, raw_request: Request):
     if not isinstance(ret, list):
         ret = [ret]
 
-    response = v1_generate_response(request, ret)
+    response = v1_generate_response(request, ret, tokenizer_manager)
     return response
 
 
@@ -626,7 +657,7 @@ def v1_chat_generate_request(all_requests, tokenizer_manager):
                 prompt_ids = tokenizer_manager.tokenizer.encode(prompt)
         else:
             # Use the raw prompt and stop strings if the messages is already a string.
-            prompt = request.messages
+            prompt_ids = request.messages
             stop = request.stop
             image_data = None
         input_ids.append(prompt_ids)
@@ -647,12 +678,21 @@ def v1_chat_generate_request(all_requests, tokenizer_manager):
         image_data_list.append(image_data)
     if len(all_requests) == 1:
         input_ids = input_ids[0]
+        if isinstance(input_ids, str):
+            prompt_kwargs = {"text": input_ids}
+        else:
+            prompt_kwargs = {"input_ids": input_ids}
         sampling_params_list = sampling_params_list[0]
         image_data = image_data_list[0]
         return_logprobs = return_logprobs[0]
         top_logprobs_nums = top_logprobs_nums[0]
+    else:
+        if isinstance(input_ids[0], str):
+            prompt_kwargs = {"text": input_ids}
+        else:
+            prompt_kwargs = {"input_ids": input_ids}
     adapted_request = GenerateReqInput(
-        input_ids=input_ids,
+        **prompt_kwargs,
         image_data=image_data,
         sampling_params=sampling_params_list,
         return_logprob=return_logprobs,
@@ -672,9 +712,9 @@ def v1_chat_generate_response(request, ret, to_file=False):
 
     for idx, ret_item in enumerate(ret):
         logprobs = False
-        if isinstance(request, List) and request[idx].logprobs:
+        if isinstance(request, list) and request[idx].logprobs:
             logprobs = True
-        elif (not isinstance(request, List)) and request.logprobs:
+        elif (not isinstance(request, list)) and request.logprobs:
             logprobs = True
         if logprobs:
             logprobs = to_openai_style_logprobs(
@@ -779,10 +819,58 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
             is_first = True
 
             stream_buffer = ""
+            n_prev_token = 0
             try:
                 async for content in tokenizer_manager.generate_request(
                     adapted_request, raw_request
                 ):
+                    prompt_tokens = content["meta_info"]["prompt_tokens"]
+                    completion_tokens = content["meta_info"]["completion_tokens"]
+                    if request.logprobs:
+                        logprobs = to_openai_style_logprobs(
+                            output_token_logprobs=content["meta_info"][
+                                "output_token_logprobs"
+                            ][n_prev_token:],
+                            output_top_logprobs=content["meta_info"][
+                                "output_top_logprobs"
+                            ][n_prev_token:],
+                        )
+
+                        n_prev_token = len(
+                            content["meta_info"]["output_token_logprobs"]
+                        )
+                        token_logprobs = []
+                        for token, logprob in zip(
+                            logprobs.tokens, logprobs.token_logprobs
+                        ):
+                            token_bytes = list(token.encode("utf-8"))
+                            top_logprobs = []
+                            if logprobs.top_logprobs:
+                                for top_token, top_logprob in logprobs.top_logprobs[
+                                    0
+                                ].items():
+                                    top_token_bytes = list(top_token.encode("utf-8"))
+                                    top_logprobs.append(
+                                        TopLogprob(
+                                            token=top_token,
+                                            bytes=top_token_bytes,
+                                            logprob=top_logprob,
+                                        )
+                                    )
+                            token_logprobs.append(
+                                ChatCompletionTokenLogprob(
+                                    token=token,
+                                    bytes=token_bytes,
+                                    logprob=logprob,
+                                    top_logprobs=top_logprobs,
+                                )
+                            )
+
+                        choice_logprobs = ChoiceLogprobs(content=token_logprobs)
+
+                    else:
+                        choice_logprobs = None
+
                     if is_first:
                         # First chunk with role
                         is_first = False
@@ -790,11 +878,17 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
                             index=0,
                             delta=DeltaMessage(role="assistant"),
                             finish_reason=content["meta_info"]["finish_reason"],
+                            logprobs=choice_logprobs,
                         )
                         chunk = ChatCompletionStreamResponse(
                             id=content["meta_info"]["id"],
                             choices=[choice_data],
                             model=request.model,
+                            usage=UsageInfo(
+                                prompt_tokens=prompt_tokens,
+                                completion_tokens=completion_tokens,
+                                total_tokens=prompt_tokens + completion_tokens,
+                            ),
                         )
                         yield f"data: {chunk.model_dump_json()}\n\n"
 
@@ -805,11 +899,17 @@ async def v1_chat_completions(tokenizer_manager, raw_request: Request):
                         index=0,
                         delta=DeltaMessage(content=delta),
                         finish_reason=content["meta_info"]["finish_reason"],
+                        logprobs=choice_logprobs,
                     )
                     chunk = ChatCompletionStreamResponse(
                         id=content["meta_info"]["id"],
                         choices=[choice_data],
                         model=request.model,
+                        usage=UsageInfo(
+                            prompt_tokens=prompt_tokens,
+                            completion_tokens=completion_tokens,
+                            total_tokens=prompt_tokens + completion_tokens,
+                        ),
                     )
                     yield f"data: {chunk.model_dump_json()}\n\n"
             except ValueError as e:

sglang/srt/openai_api/protocol.py

@@ -278,7 +278,7 @@ class DeltaMessage(BaseModel):
 class ChatCompletionResponseStreamChoice(BaseModel):
     index: int
     delta: DeltaMessage
-    logprobs: Optional[LogProbs] = None
+    logprobs: Optional[Union[LogProbs, ChoiceLogprobs]] = None
     finish_reason: Optional[str] = None
 
 

sglang/srt/server.py

@@ -28,7 +28,7 @@ import sys
 import threading
 import time
 from http import HTTPStatus
-from typing import Dict, Optional
+from typing import Dict, List, Optional, Union
 
 # Fix a bug of Python threading
 setattr(threading, "_register_atexit", lambda *args, **kwargs: None)
@@ -67,13 +67,13 @@ from sglang.srt.openai_api.adapter import (
 from sglang.srt.openai_api.protocol import ModelCard, ModelList
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.utils import (
-    API_KEY_HEADER_NAME,
-    APIKeyValidatorMiddleware,
+    add_api_key_middleware,
     allocate_init_ports,
     assert_pkg_version,
     enable_show_time_cost,
     kill_child_process,
     maybe_set_triton_cache_manager,
+    set_torch_compile_config,
     set_ulimit,
 )
 from sglang.utils import get_exception_traceback
@@ -158,6 +158,16 @@ async def openai_v1_chat_completions(raw_request: Request):
     return await v1_chat_completions(tokenizer_manager, raw_request)
 
 
+@app.get("/v1/models")
+def available_models():
+    """Show available models."""
+    served_model_names = [tokenizer_manager.served_model_name]
+    model_cards = []
+    for served_model_name in served_model_names:
+        model_cards.append(ModelCard(id=served_model_name, root=served_model_name))
+    return ModelList(data=model_cards)
+
+
 @app.post("/v1/files")
 async def openai_v1_files(file: UploadFile = File(...), purpose: str = Form("batch")):
     return await v1_files_create(
@@ -187,69 +197,11 @@ async def retrieve_file_content(file_id: str):
     return await v1_retrieve_file_content(file_id)
 
 
-@app.get("/v1/models")
-def available_models():
-    """Show available models."""
-    served_model_names = [tokenizer_manager.served_model_name]
-    model_cards = []
-    for served_model_name in served_model_names:
-        model_cards.append(ModelCard(id=served_model_name, root=served_model_name))
-    return ModelList(data=model_cards)
-
-
-def _set_torch_compile_config():
-    # The following configurations are for torch compile optimizations
-    import torch._dynamo.config
-    import torch._inductor.config
-
-    torch._inductor.config.coordinate_descent_tuning = True
-    torch._inductor.config.triton.unique_kernel_names = True
-    torch._inductor.config.fx_graph_cache = True  # Experimental feature to reduce compilation times, will be on by default in future
-
-    # FIXME: tmp workaround
-    torch._dynamo.config.accumulated_cache_size_limit = 256
-
-
-def set_envs_and_config(server_args: ServerArgs):
-    # Set global environments
-    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
-    os.environ["NCCL_CUMEM_ENABLE"] = "0"
-    os.environ["NCCL_NVLS_ENABLE"] = "0"
-    os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
-
-    # 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.
-        maybe_set_triton_cache_manager()
-
-    # Set torch compile config
-    if server_args.enable_torch_compile:
-        _set_torch_compile_config()
-
-    # Set global chat template
-    if server_args.chat_template:
-        # TODO: replace this with huggingface transformers template
-        load_chat_template_for_openai_api(server_args.chat_template)
-
-
 def launch_server(
     server_args: ServerArgs,
     model_overide_args: Optional[dict] = None,
     pipe_finish_writer: Optional[mp.connection.Connection] = None,
 ):
-    server_args.check_server_args()
-
     """Launch an HTTP server."""
     global tokenizer_manager
 
@@ -258,16 +210,8 @@ def launch_server(
         format="%(message)s",
     )
 
-    if not server_args.disable_flashinfer:
-        assert_pkg_version(
-            "flashinfer",
-            "0.1.3",
-            "Please uninstall the old version and "
-            "reinstall the latest version by following the instructions "
-            "at https://docs.flashinfer.ai/installation.html.",
-        )
-
-    set_envs_and_config(server_args)
+    server_args.check_server_args()
+    _set_envs_and_config(server_args)
 
     # Allocate ports
     server_args.port, server_args.additional_ports = allocate_init_ports(
@@ -284,7 +228,7 @@ def launch_server(
     )
     logger.info(f"{server_args=}")
 
-    # Handle multi-node tensor parallelism
+    # Launch processes for multi-node tensor parallelism
     if server_args.nnodes > 1:
         if server_args.node_rank != 0:
             tp_size_local = server_args.tp_size // server_args.nnodes
@@ -349,8 +293,9 @@ def launch_server(
         sys.exit(1)
     assert proc_controller.is_alive() and proc_detoken.is_alive()
 
-    if server_args.api_key and server_args.api_key != "":
-        app.add_middleware(APIKeyValidatorMiddleware, api_key=server_args.api_key)
+    # Add api key authorization
+    if server_args.api_key:
+        add_api_key_middleware(app, server_args.api_key)
 
     # Send a warmup request
     t = threading.Thread(
@@ -372,15 +317,58 @@ def launch_server(
         t.join()
 
 
+def _set_envs_and_config(server_args: ServerArgs):
+    # Set global environments
+    os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
+    os.environ["NCCL_CUMEM_ENABLE"] = "0"
+    os.environ["NCCL_NVLS_ENABLE"] = "0"
+    os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1"
+
+    # 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.
+        maybe_set_triton_cache_manager()
+
+    # Set torch compile config
+    if server_args.enable_torch_compile:
+        set_torch_compile_config()
+
+    # Set global chat template
+    if server_args.chat_template:
+        # TODO: replace this with huggingface transformers template
+        load_chat_template_for_openai_api(server_args.chat_template)
+
+    # Check flashinfer version
+    if not server_args.disable_flashinfer:
+        assert_pkg_version(
+            "flashinfer",
+            "0.1.3",
+            "Please uninstall the old version and "
+            "reinstall the latest version by following the instructions "
+            "at https://docs.flashinfer.ai/installation.html.",
+        )
+
+
 def _wait_and_warmup(server_args, pipe_finish_writer):
     headers = {}
     url = server_args.url()
     if server_args.api_key:
-        headers[API_KEY_HEADER_NAME] = server_args.api_key
+        headers["Authorization"] = f"Bearer {server_args.api_key}"
 
     # Wait until the server is launched
     for _ in range(120):
-        time.sleep(0.5)
+        time.sleep(1)
         try:
             requests.get(url + "/get_model_info", timeout=5, headers=headers)
             break
@@ -481,10 +469,10 @@ class Runtime:
             trust_remote_code=self.server_args.trust_remote_code,
         )
 
-    async def add_request(
+    async def async_generate(
         self,
         prompt: str,
-        sampling_params: Dict,
+        sampling_params: Optional[Dict] = None,
     ):
         json_data = {
             "text": prompt,
@@ -507,5 +495,26 @@ class Runtime:
                             yield cur
                         pos += len(cur)
 
+    add_request = async_generate
+
+    def generate(
+        self,
+        prompt: str,
+        sampling_params: Optional[Dict] = None,
+        return_logprob: Optional[Union[List[bool], bool]] = False,
+        top_logprobs_num: Optional[Union[List[int], int]] = None,
+    ):
+        json_data = {
+            "text": prompt,
+            "sampling_params": sampling_params,
+            "return_logprob": return_logprob,
+            "top_logprobs_num": top_logprobs_num,
+        }
+        response = requests.post(
+            self.url + "/generate",
+            json=json_data,
+        )
+        return json.dumps(response.json())
+
     def __del__(self):
         self.shutdown()

sglang/srt/server_args.py

@@ -61,7 +61,7 @@ class ServerArgs:
     show_time_cost: bool = False
 
     # Other
-    api_key: str = ""
+    api_key: Optional[str] = None
     file_storage_pth: str = "SGlang_storage"
 
     # Data parallelism
@@ -80,6 +80,7 @@ class ServerArgs:
     disable_disk_cache: bool = False
     enable_torch_compile: bool = False
     enable_p2p_check: bool = False
+    enable_mla: bool = False
     attention_reduce_in_fp32: bool = False
     efficient_weight_load: bool = False
 
@@ -306,7 +307,7 @@ class ServerArgs:
             "--api-key",
             type=str,
             default=ServerArgs.api_key,
-            help="Set API key of the server.",
+            help="Set API key of the server. It is also used in the OpenAI API compatible server.",
         )
         parser.add_argument(
             "--file-storage-pth",
@@ -393,6 +394,11 @@ class ServerArgs:
             action="store_true",
             help="Enable P2P check for GPU access, otherwise the p2p access is allowed by default.",
         )
+        parser.add_argument(
+            "--enable-mla",
+            action="store_true",
+            help="Enable Multi-head Latent Attention (MLA) for DeepSeek-V2",
+        )
         parser.add_argument(
             "--attention-reduce-in-fp32",
             action="store_true",

sglang/srt/utils.py

@@ -539,26 +539,6 @@ class CustomCacheManager(FileCacheManager):
                 raise RuntimeError("Could not create or locate cache dir")
 
 
-API_KEY_HEADER_NAME = "X-API-Key"
-
-
-class APIKeyValidatorMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app, api_key: str):
-        super().__init__(app)
-        self.api_key = api_key
-
-    async def dispatch(self, request, call_next):
-        # extract API key from the request headers
-        api_key_header = request.headers.get(API_KEY_HEADER_NAME)
-        if not api_key_header or api_key_header != self.api_key:
-            return JSONResponse(
-                status_code=403,
-                content={"detail": "Invalid API Key"},
-            )
-        response = await call_next(request)
-        return response
-
-
 def get_ip_address(ifname):
     """
     Get the IP address of a network interface.
@@ -642,6 +622,19 @@ def receive_addrs(model_port_args, server_args):
     dist.destroy_process_group()
 
 
+def set_torch_compile_config():
+    # The following configurations are for torch compile optimizations
+    import torch._dynamo.config
+    import torch._inductor.config
+
+    torch._inductor.config.coordinate_descent_tuning = True
+    torch._inductor.config.triton.unique_kernel_names = True
+    torch._inductor.config.fx_graph_cache = True  # Experimental feature to reduce compilation times, will be on by default in future
+
+    # FIXME: tmp workaround
+    torch._dynamo.config.accumulated_cache_size_limit = 256
+
+
 def set_ulimit(target_soft_limit=65535):
     resource_type = resource.RLIMIT_NOFILE
     current_soft, current_hard = resource.getrlimit(resource_type)
@@ -700,3 +693,15 @@ def monkey_patch_vllm_qvk_linear_loader():
         origin_weight_loader(self, param, loaded_weight, loaded_shard_id)
 
     setattr(QKVParallelLinear, "weight_loader", weight_loader_srt)
+
+
+def add_api_key_middleware(app, api_key):
+    @app.middleware("http")
+    async def authentication(request, call_next):
+        if request.method == "OPTIONS":
+            return await call_next(request)
+        if request.url.path.startswith("/health"):
+            return await call_next(request)
+        if request.headers.get("Authorization") != "Bearer " + api_key:
+            return JSONResponse(content={"error": "Unauthorized"}, status_code=401)
+        return await call_next(request)