sglang 0.1.21__py3-none-any.whl → 0.1.22__py3-none-any.whl

sglang/lang/backend/runtime_endpoint.py

@@ -0,0 +1,283 @@
+import json
+from typing import List, Optional
+
+import numpy as np
+
+from sglang.global_config import global_config
+from sglang.lang.backend.base_backend import BaseBackend
+from sglang.lang.chat_template import get_chat_template_by_model_path
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import SglSamplingParams
+from sglang.utils import http_request
+
+
+class RuntimeEndpoint(BaseBackend):
+    def __init__(
+        self,
+        base_url: str,
+        auth_token: Optional[str] = None,
+        api_key: Optional[str] = None,
+        verify: Optional[str] = None,
+    ):
+        super().__init__()
+        self.support_concate_and_append = True
+
+        self.base_url = base_url
+        self.auth_token = auth_token
+        self.api_key = api_key
+        self.verify = verify
+
+        res = http_request(
+            self.base_url + "/get_model_info",
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+        self.model_info = res.json()
+
+        self.chat_template = get_chat_template_by_model_path(
+            self.model_info["model_path"]
+        )
+
+    def get_model_name(self):
+        return self.model_info["model_path"]
+
+    def flush_cache(self):
+        res = http_request(
+            self.base_url + "/flush_cache",
+            auth_token=self.auth_token,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def get_server_args(self):
+        res = http_request(
+            self.base_url + "/get_server_args",
+            auth_token=self.auth_token,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+        return res.json()
+
+    def get_chat_template(self):
+        return self.chat_template
+
+    def cache_prefix(self, prefix_str: str):
+        res = http_request(
+            self.base_url + "/generate",
+            json={"text": prefix_str, "sampling_params": {"max_new_tokens": 0}},
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def commit_lazy_operations(self, s: StreamExecutor):
+        data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
+        self._add_images(s, data)
+        res = http_request(
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def fill_image(self, s: StreamExecutor):
+        data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
+        self._add_images(s, data)
+        res = http_request(
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def generate(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+    ):
+        if sampling_params.dtype is None:
+            data = {
+                "text": s.text_,
+                "sampling_params": {
+                    "skip_special_tokens": global_config.skip_special_tokens_in_output,
+                    "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out,
+                    **sampling_params.to_srt_kwargs(),
+                },
+            }
+        elif sampling_params.dtype in [int, "int"]:
+            data = {
+                "text": s.text_,
+                "sampling_params": {
+                    "skip_special_tokens": global_config.skip_special_tokens_in_output,
+                    "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out,
+                    "dtype": "int",
+                    **sampling_params.to_srt_kwargs(),
+                },
+            }
+        else:
+            raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}")
+
+        for item in [
+            "return_logprob",
+            "logprob_start_len",
+            "top_logprobs_num",
+            "return_text_in_logprobs",
+        ]:
+            value = getattr(sampling_params, item, None)
+            if value is not None:
+                data[item] = value
+
+        self._add_images(s, data)
+
+        res = http_request(
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+        obj = res.json()
+        comp = obj["text"]
+        return comp, obj["meta_info"]
+
+    def generate_stream(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+    ):
+        if sampling_params.dtype is None:
+            data = {
+                "text": s.text_,
+                "sampling_params": {
+                    "skip_special_tokens": global_config.skip_special_tokens_in_output,
+                    "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out,
+                    **sampling_params.to_srt_kwargs(),
+                },
+            }
+        elif sampling_params.dtype in [int, "int"]:
+            data = {
+                "text": s.text_,
+                "sampling_params": {
+                    "skip_special_tokens": global_config.skip_special_tokens_in_output,
+                    "spaces_between_special_tokens": global_config.spaces_between_special_tokens_in_out,
+                    "dtype": "int",
+                    **sampling_params.to_srt_kwargs(),
+                },
+            }
+        else:
+            raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}")
+
+        for item in [
+            "return_logprob",
+            "logprob_start_len",
+            "top_logprobs_num",
+            "return_text_in_logprobs",
+        ]:
+            value = getattr(sampling_params, item, None)
+            if value is not None:
+                data[item] = value
+
+        data["stream"] = True
+        self._add_images(s, data)
+
+        res = http_request(
+            self.base_url + "/generate",
+            json=data,
+            stream=True,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+        pos = 0
+
+        for chunk in res.iter_lines(decode_unicode=False):
+            chunk = chunk.decode("utf-8")
+            if chunk and chunk.startswith("data:"):
+                if chunk == "data: [DONE]":
+                    break
+                data = json.loads(chunk[5:].strip("\n"))
+                chunk_text = data["text"][pos:]
+                meta_info = data["meta_info"]
+                pos += len(chunk_text)
+                yield chunk_text, meta_info
+
+    def select(
+        self,
+        s: StreamExecutor,
+        choices: List[str],
+        temperature: float,
+    ):
+        assert temperature <= 1e-5
+
+        # Cache common prefix
+        data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
+        self._add_images(s, data)
+        res = http_request(
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+        prompt_len = res.json()["meta_info"]["prompt_tokens"]
+
+        # Compute logprob
+        data = {
+            "text": [s.text_ + c for c in choices],
+            "sampling_params": {"max_new_tokens": 0},
+            "return_logprob": True,
+            "logprob_start_len": max(prompt_len - 2, 0),
+        }
+        self._add_images(s, data)
+        res = http_request(
+            self.base_url + "/generate",
+            json=data,
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+        obj = res.json()
+        normalized_prompt_logprobs = [
+            r["meta_info"]["normalized_prompt_logprob"] for r in obj
+        ]
+        decision = choices[np.argmax(normalized_prompt_logprobs)]
+        prefill_token_logprobs = [r["meta_info"]["prefill_token_logprobs"] for r in obj]
+        decode_token_logprobs = [r["meta_info"]["decode_token_logprobs"] for r in obj]
+
+        return (
+            decision,
+            normalized_prompt_logprobs,
+            prefill_token_logprobs,
+            decode_token_logprobs,
+        )
+
+    def concatenate_and_append(self, src_rids: List[str], dst_rid: str):
+        res = http_request(
+            self.base_url + "/concate_and_append_request",
+            json={"src_rids": src_rids, "dst_rid": dst_rid},
+            auth_token=self.auth_token,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def _add_images(self, s: StreamExecutor, data):
+        if s.images_:
+            assert len(s.images_) == 1, "Only support one image."
+            data["image_data"] = s.images_[0][1]
+
+    def _assert_success(self, res):
+        if res.status_code != 200:
+            raise RuntimeError(res.json())

sglang/lang/backend/vertexai.py

@@ -0,0 +1,149 @@
+import os
+import warnings
+from typing import Optional
+
+from sglang.lang.backend.base_backend import BaseBackend
+from sglang.lang.chat_template import get_chat_template
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import SglSamplingParams
+
+try:
+    import vertexai
+    from vertexai.preview.generative_models import (
+        GenerationConfig,
+        GenerativeModel,
+        Image,
+    )
+except ImportError as e:
+    GenerativeModel = e
+
+
+class VertexAI(BaseBackend):
+    def __init__(self, model_name, safety_settings=None):
+        super().__init__()
+
+        if isinstance(GenerativeModel, Exception):
+            raise GenerativeModel
+
+        project_id = os.environ["GCP_PROJECT_ID"]
+        location = os.environ.get("GCP_LOCATION")
+        vertexai.init(project=project_id, location=location)
+
+        self.model_name = model_name
+        self.chat_template = get_chat_template("default")
+        self.safety_settings = safety_settings
+
+    def get_chat_template(self):
+        return self.chat_template
+
+    def generate(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+    ):
+        if s.messages_:
+            prompt = self.messages_to_vertexai_input(s.messages_)
+        else:
+            # single-turn
+            prompt = (
+                self.text_to_vertexai_input(s.text_, s.cur_images)
+                if s.cur_images
+                else s.text_
+            )
+        ret = GenerativeModel(self.model_name).generate_content(
+            prompt,
+            generation_config=GenerationConfig(**sampling_params.to_vertexai_kwargs()),
+            safety_settings=self.safety_settings,
+        )
+
+        comp = ret.text
+
+        return comp, {}
+
+    def generate_stream(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+    ):
+        if s.messages_:
+            prompt = self.messages_to_vertexai_input(s.messages_)
+        else:
+            # single-turn
+            prompt = (
+                self.text_to_vertexai_input(s.text_, s.cur_images)
+                if s.cur_images
+                else s.text_
+            )
+        generator = GenerativeModel(self.model_name).generate_content(
+            prompt,
+            stream=True,
+            generation_config=GenerationConfig(**sampling_params.to_vertexai_kwargs()),
+            safety_settings=self.safety_settings,
+        )
+        for ret in generator:
+            yield ret.text, {}
+
+    def text_to_vertexai_input(self, text, images):
+        input = []
+        # split with image token
+        text_segs = text.split(self.chat_template.image_token)
+        for image_path, image_base64_data in images:
+            text_seg = text_segs.pop(0)
+            if text_seg != "":
+                input.append(text_seg)
+            input.append(Image.from_bytes(image_base64_data))
+        text_seg = text_segs.pop(0)
+        if text_seg != "":
+            input.append(text_seg)
+        return input
+
+    def messages_to_vertexai_input(self, messages):
+        vertexai_message = []
+        # from openai message format to vertexai message format
+        for msg in messages:
+            if isinstance(msg["content"], str):
+                text = msg["content"]
+            else:
+                text = msg["content"][0]["text"]
+
+            if msg["role"] == "system":
+                warnings.warn("Warning: system prompt is not supported in VertexAI.")
+                vertexai_message.append(
+                    {
+                        "role": "user",
+                        "parts": [{"text": "System prompt: " + text}],
+                    }
+                )
+                vertexai_message.append(
+                    {
+                        "role": "model",
+                        "parts": [{"text": "Understood."}],
+                    }
+                )
+                continue
+            if msg["role"] == "user":
+                vertexai_msg = {
+                    "role": "user",
+                    "parts": [{"text": text}],
+                }
+            elif msg["role"] == "assistant":
+                vertexai_msg = {
+                    "role": "model",
+                    "parts": [{"text": text}],
+                }
+
+            # images
+            if isinstance(msg["content"], list) and len(msg["content"]) > 1:
+                for image in msg["content"][1:]:
+                    assert image["type"] == "image_url"
+                    vertexai_msg["parts"].append(
+                        {
+                            "inline_data": {
+                                "data": image["image_url"]["url"].split(",")[1],
+                                "mime_type": "image/jpeg",
+                            }
+                        }
+                    )
+
+            vertexai_message.append(vertexai_msg)
+        return vertexai_message

sglang/lang/tracer.py

@@ -3,8 +3,8 @@
 import uuid
 from typing import Any, Callable, Dict, List, Optional, Union
 
-from sglang.backend.base_backend import BaseBackend
 from sglang.global_config import global_config
+from sglang.lang.backend.base_backend import BaseBackend
 from sglang.lang.interpreter import ProgramState, ProgramStateGroup
 from sglang.lang.ir import (
     SglArgument,

sglang/launch_server.py

@@ -11,4 +11,4 @@ if __name__ == "__main__":
     args = parser.parse_args()
     server_args = ServerArgs.from_cli_args(args)
 
-    launch_server(server_args, None)
+    launch_server(server_args)

sglang/launch_server_llavavid.py

@@ -1,7 +1,6 @@
 """Launch the inference server for Llava-video model."""
 
 import argparse
-import multiprocessing as mp
 
 from sglang.srt.server import ServerArgs, launch_server
 
@@ -27,6 +26,4 @@ if __name__ == "__main__":
 
     server_args = ServerArgs.from_cli_args(args)
 
-    pipe_reader, pipe_writer = mp.Pipe(duplex=False)
-
-    launch_server(server_args, pipe_writer, model_overide_args)
+    launch_server(server_args, model_overide_args, None)

sglang/srt/conversation.py

@@ -6,7 +6,7 @@ import dataclasses
 from enum import IntEnum, auto
 from typing import Dict, List, Optional, Tuple, Union
 
-from sglang.srt.openai_protocol import ChatCompletionRequest
+from sglang.srt.openai_api.protocol import ChatCompletionRequest
 
 
 class SeparatorStyle(IntEnum):

sglang/srt/layers/context_flashattention_nopad.py

@@ -4,8 +4,6 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.utils import wrap_kernel_launcher
-
 CUDA_CAPABILITY = torch.cuda.get_device_capability()
 
 
@@ -119,9 +117,6 @@ def _fwd_kernel(
     tl.store(out_ptrs, acc, mask=offs_m[:, None] < cur_batch_seq_len)
 
 
-cached_kernel = None
-
-
 def context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len):
     if CUDA_CAPABILITY[0] >= 8:
         BLOCK = 128
@@ -139,29 +134,6 @@ def context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len):
     grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
     num_warps = 4 if Lk <= 64 else 8
 
-    global cached_kernel
-    if cached_kernel:
-        cached_kernel(
-            grid,
-            num_warps,
-            q,
-            k,
-            v,
-            sm_scale,
-            b_start_loc,
-            b_seq_len,
-            o,
-            q.stride(0),
-            q.stride(1),
-            k.stride(0),
-            k.stride(1),
-            v.stride(0),
-            v.stride(1),
-            o.stride(0),
-            o.stride(1),
-        )
-        return
-
     _fwd_kernel[grid](
         q,
         k,
@@ -185,4 +157,3 @@ def context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len):
         num_warps=num_warps,
         num_stages=1,
     )
-    cached_kernel = wrap_kernel_launcher(_fwd_kernel)

sglang/srt/layers/extend_attention.py

@@ -3,7 +3,6 @@ import triton
 import triton.language as tl
 
 from sglang.srt.layers.context_flashattention_nopad import context_attention_fwd
-from sglang.srt.utils import wrap_kernel_launcher
 
 CUDA_CAPABILITY = torch.cuda.get_device_capability()
 
@@ -172,9 +171,6 @@ def _fwd_kernel(
     tl.store(O_Extend + offs_o, acc / deno[:, None], mask=mask_m[:, None])
 
 
-cached_kernel = None
-
-
 def extend_attention_fwd(
     q_extend,
     k_extend,
@@ -222,40 +218,6 @@ def extend_attention_fwd(
     num_warps = 4 if Lk <= 64 else 8
     num_stages = 1
 
-    global cached_kernel
-    if cached_kernel:
-        cached_kernel(
-            grid,
-            num_warps,
-            q_extend,
-            k_extend,
-            v_extend,
-            o_extend,
-            k_buffer,
-            v_buffer,
-            req_to_tokens,
-            b_req_idx,
-            b_seq_len,
-            b_start_loc_extend,
-            b_seq_len_extend,
-            sm_scale,
-            kv_group_num,
-            q_extend.stride(0),
-            q_extend.stride(1),
-            k_extend.stride(0),
-            k_extend.stride(1),
-            v_extend.stride(0),
-            v_extend.stride(1),
-            o_extend.stride(0),
-            o_extend.stride(1),
-            k_buffer.stride(0),
-            k_buffer.stride(1),
-            v_buffer.stride(0),
-            v_buffer.stride(1),
-            req_to_tokens.stride(0),
-        )
-        return
-
     _fwd_kernel[grid](
         q_extend,
         k_extend,
@@ -290,7 +252,6 @@ def extend_attention_fwd(
         num_stages=num_stages,
         logit_cap=logit_cap,
     )
-    cached_kernel = wrap_kernel_launcher(_fwd_kernel)
 
 
 def redundant_attention(