sglang 0.1.3__py3-none-any.whl → 0.1.5__py3-none-any.whl

sglang/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.3"
+__version__ = "0.1.5"
 
 from sglang.api import *
 from sglang.global_config import global_config

sglang/api.py

@@ -6,6 +6,7 @@ from sglang.backend.anthropic import Anthropic
 from sglang.backend.base_backend import BaseBackend
 from sglang.backend.openai import OpenAI
 from sglang.backend.runtime_endpoint import RuntimeEndpoint
+from sglang.backend.vertexai import VertexAI
 from sglang.global_config import global_config
 from sglang.lang.ir import (
     SglExpr,

sglang/backend/vertexai.py

@@ -0,0 +1,147 @@
+import os
+import warnings
+from typing import List, Optional, Union
+
+import numpy as np
+from sglang.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):
+        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")
+
+    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()),
+        )
+
+        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()),
+        )
+        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/interpreter.py

@@ -365,11 +365,10 @@ class StreamExecutor:
             for comp, meta_info in generator:
                 self.text_ += comp
                 self.variables[name] += comp
+                self.meta_info[name] = meta_info
                 self.stream_var_event[name].set()
                 self.stream_text_event.set()
 
-            self.meta_info[name] = meta_info
-
             self.variable_event[name].set()
             self.stream_var_event[name].set()
 
@@ -428,6 +427,7 @@ class StreamExecutor:
             self.messages_.append(last_msg)
             self.cur_images = []
         else:
+            # OpenAI chat API format
             self.messages_.append({"role": expr.role, "content": new_text})
 
         self.cur_role = None
@@ -582,7 +582,7 @@ class ProgramState:
             else:
                 yield self.get_var(name)
 
-    async def text_async_iter(self, var_name=None):
+    async def text_async_iter(self, var_name=None, return_meta_data=False):
         loop = asyncio.get_running_loop()
 
         if self.stream_executor.stream:
@@ -606,7 +606,10 @@ class ProgramState:
                     out = str(self.stream_executor.variables[var_name][prev:])
                     prev += len(out)
                     if out:
-                        yield out
+                        if return_meta_data:
+                            yield out, self.stream_executor.meta_info[var_name]
+                        else:
+                            yield out
                     if self.stream_executor.variable_event[var_name].is_set():
                         break
         else:
@@ -632,11 +635,7 @@ class ProgramState:
         self.stream_executor.end()
 
     def __repr__(self) -> str:
-        msgs = self.messages()
-        ret = ""
-        for msg in msgs:
-            ret += msg["role"] + ":\n" + msg["content"] + "\n"
-        return ret
+        return f"ProgramState({self.text()})"
 
 
 class ProgramStateGroup:

sglang/lang/ir.py

@@ -2,6 +2,7 @@
 
 import dataclasses
 import inspect
+import warnings
 from typing import List, Optional, Union
 
 from sglang.global_config import global_config
@@ -40,6 +41,8 @@ class SglSamplingParams:
 
     def to_openai_kwargs(self):
         # OpenAI does not support top_k, so we drop it here
+        if self.regex is not None:
+            warnings.warn("Regular expression is not supported in the OpenAI backend.")
         return {
             "max_tokens": self.max_new_tokens,
             "stop": self.stop or None,
@@ -49,8 +52,26 @@ class SglSamplingParams:
             "presence_penalty": self.presence_penalty,
         }
 
+    def to_vertexai_kwargs(self):
+        if self.regex is not None:
+            warnings.warn(
+                "Regular expression is not supported in the VertexAI backend."
+            )
+        return {
+            "candidate_count": 1,
+            "max_output_tokens": self.max_new_tokens,
+            "stop_sequences": self.stop,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k if self.top_k > 0 else None,
+        }
+
     def to_anthropic_kwargs(self):
         # Anthropic does not support frequency_penalty or presence_penalty, so we drop it here
+        if self.regex is not None:
+            warnings.warn(
+                "Regular expression is not supported in the Anthropic backend."
+            )
         return {
             "max_tokens_to_sample": self.max_new_tokens,
             "stop_sequences": self.stop,

sglang/srt/layers/context_flashattention_nopad.py

@@ -5,6 +5,8 @@ import triton
 import triton.language as tl
 from sglang.srt.utils import wrap_kernel_launcher
 
+CUDA_CAPABILITY = torch.cuda.get_device_capability()
+
 
 @triton.jit
 def _fwd_kernel(
@@ -120,7 +122,11 @@ cached_kernel = None
 
 
 def context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len):
-    BLOCK = 128
+    if CUDA_CAPABILITY[0] >= 8:
+        BLOCK = 128
+    else:
+        BLOCK = 64
+
     Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
     assert Lq == Lk and Lk == Lv
     assert Lk in {16, 32, 64, 128}

sglang/srt/layers/extend_attention.py

@@ -2,6 +2,9 @@ import torch
 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()
 
 
 @triton.jit
@@ -153,6 +156,9 @@ 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,
@@ -175,7 +181,11 @@ def extend_attention_fwd(
 
     k_buffer, v_buffer: (prefix + extend) tensors in mem_manager
     """
-    BLOCK_M, BLOCK_N = 128, 128
+    if CUDA_CAPABILITY[0] >= 8:
+        BLOCK_M, BLOCK_N = 128, 128
+    else:
+        BLOCK_M, BLOCK_N = 64, 64
+
     Lq, Lk, Lv, Lo = (
         q_extend.shape[-1],
         k_extend.shape[-1],
@@ -193,6 +203,40 @@ 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,
@@ -226,6 +270,7 @@ def extend_attention_fwd(
         num_warps=num_warps,
         num_stages=num_stages,
     )
+    cached_kernel = wrap_kernel_launcher(_fwd_kernel)
 
 
 def redundant_attention(

sglang/srt/managers/router/manager.py

@@ -28,7 +28,7 @@ class RouterManager:
         self.model_client = model_client
         self.recv_reqs = []
 
-        # Init Some Configs
+        # Init some configs
         self.extend_dependency_time = GLOBAL_BACKEND_CONFIG.extend_dependency_time
 
     async def loop_for_forward(self):
@@ -46,7 +46,7 @@ class RouterManager:
                 if has_finished:
                     await asyncio.sleep(self.extend_dependency_time)
 
-            await asyncio.sleep(0.001)
+            await asyncio.sleep(0.0006)
 
     async def loop_for_recv_requests(self):
         while True:

sglang/srt/managers/router/model_rpc.py

@@ -2,6 +2,7 @@ import asyncio
 import logging
 import multiprocessing
 import time
+import warnings
 from concurrent.futures import ThreadPoolExecutor
 from enum import Enum, auto
 from typing import Dict, List, Optional, Tuple, Union
@@ -44,6 +45,7 @@ class ModelRpcServer(rpyc.Service):
         self.tp_rank = tp_rank
         self.tp_size = server_args.tp_size
         self.schedule_heuristic = server_args.schedule_heuristic
+        self.schedule_conservativeness = server_args.schedule_conservativeness
 
         # Init model and tokenizer
         self.model_config = ModelConfig(
@@ -107,7 +109,7 @@ class ModelRpcServer(rpyc.Service):
         self.running_batch: Batch = None
         self.out_pyobjs = []
         self.decode_forward_ct = 0
-        self.stream_interval = 2
+        self.stream_interval = server_args.stream_interval
 
         # Init the FSM cache for constrained generation
         self.regex_fsm_cache = FSMCache(self.tokenizer)
@@ -164,7 +166,7 @@ class ModelRpcServer(rpyc.Service):
                     + self.tree_cache.evictable_size()
                 )
                 if available_size != self.max_total_num_token:
-                    logger.warning(
+                    warnings.warn(
                         "Warning: "
                         f"available_size={available_size}, max_total_num_token={self.max_total_num_token}\n"
                         "KV cache pool leak detected!"
@@ -247,7 +249,9 @@ class ModelRpcServer(rpyc.Service):
         available_size = (
             self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size()
         )
-        new_ratio = self.scheduler.new_token_estimation_ratio()
+        new_ratio = (
+            self.scheduler.new_token_estimation_ratio() * self.schedule_conservativeness
+        )
         if self.running_batch:
             available_size -= sum(
                 [

sglang/srt/managers/router/model_runner.py

@@ -278,7 +278,7 @@ class ModelRunner:
                 load_format=self.load_format,
                 revision=None,
             )
-        self.model = model
+        self.model = model.eval()
 
     def profile_max_num_token(self, total_gpu_memory):
         available_gpu_memory = get_available_gpu_memory(

sglang/srt/models/mixtral.py

@@ -355,7 +355,7 @@ class MixtralForCausalLM(nn.Module):
         ):
             if "rotary_emb.inv_freq" in name:
                 continue
-            for (param_name, weight_name, shard_id) in stacked_params_mapping:
+            for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
                 name = name.replace(weight_name, param_name)

sglang/srt/server_args.py

@@ -16,7 +16,9 @@ class ServerArgs:
     tp_size: int = 1
     model_mode: List[str] = ()
     schedule_heuristic: str = "lpm"
+    schedule_conservativeness: float = 1.0
     random_seed: int = 42
+    stream_interval: int = 2
     disable_log_stats: bool = False
     log_stats_interval: int = 10
     log_level: str = "info"
@@ -25,10 +27,14 @@ class ServerArgs:
         if self.tokenizer_path is None:
             self.tokenizer_path = self.model_path
         if self.mem_fraction_static is None:
-            if self.tp_size > 1:
-                self.mem_fraction_static = 0.8
+            if self.tp_size >= 8:
+                self.mem_fraction_static = 0.80
+            elif self.tp_size >= 4:
+                self.mem_fraction_static = 0.82
+            elif self.tp_size >= 2:
+                self.mem_fraction_static = 0.85
             else:
-                self.mem_fraction_static = 0.9
+                self.mem_fraction_static = 0.90
 
     @staticmethod
     def add_cli_args(parser: argparse.ArgumentParser):
@@ -80,7 +86,7 @@ class ServerArgs:
             "--mem-fraction-static",
             type=float,
             default=ServerArgs.mem_fraction_static,
-            help="The fraction of the memory used for static allocation (model weights and KV cache memory pool)",
+            help="The fraction of the memory used for static allocation (model weights and KV cache memory pool). Use a smaller value if you see out-of-memory errors.",
         )
         parser.add_argument(
             "--tp-size",
@@ -102,12 +108,24 @@ class ServerArgs:
             default=ServerArgs.schedule_heuristic,
             help="Schudule mode: [lpm, weight, random, fcfs]",
         )
+        parser.add_argument(
+            "--schedule-conservativeness",
+            type=float,
+            default=ServerArgs.schedule_conservativeness,
+            help="How conservative the schedule policy is. A larger value means more conservative scheduling. Use a larger value if you see out-of-memory errors.",
+        )
         parser.add_argument(
             "--random-seed",
             type=int,
             default=ServerArgs.random_seed,
             help="Random seed.",
         )
+        parser.add_argument(
+            "--stream-interval",
+            type=int,
+            default=ServerArgs.stream_interval,
+            help="The interval in terms of token length for streaming",
+        )
         parser.add_argument(
             "--log-level",
             type=str,

sglang/srt/utils.py

@@ -209,7 +209,7 @@ def load_image(image_file):
     elif image_file.lower().endswith(("png", "jpg", "jpeg", "webp", "gif")):
         image = Image.open(image_file)
     elif image_file.startswith("data:"):
-        image_file = image_url.split(",")[1]
+        image_file = image_file.split(",")[1]
         image = Image.open(BytesIO(base64.b64decode(image_file)))
     else:
         image = Image.open(BytesIO(base64.b64decode(image_file)))

sglang/test/test_programs.py

@@ -304,7 +304,10 @@ def test_image_qa():
         temperature=0,
         max_new_tokens=64,
     )
-    assert "taxi" in state.messages()[-1]["content"]
+    assert (
+        "taxi" in state.messages()[-1]["content"]
+        or "car" in state.messages()[-1]["content"]
+    )
 
 
 def test_stream():

{sglang-0.1.3.dist-info → sglang-0.1.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sglang
-Version: 0.1.3
+Version: 0.1.5
 Summary: A structured generation langauge for LLMs.
 License: Apache License
                                    Version 2.0, January 2004
@@ -234,6 +234,7 @@ Requires-Dist: lark ; extra == 'srt'
 Requires-Dist: numba ; extra == 'srt'
 
 # SGLang
+| [**Blog**](https://lmsys.org/blog/2024-01-17-sglang/) | [**Paper**](https://arxiv.org/abs/2312.07104) |
 
 SGLang is a structured generation language designed for large language models (LLMs).
 It makes your interaction with LLMs faster and more controllable by co-designing the frontend language and the runtime system.
@@ -267,10 +268,20 @@ pip install --upgrade pip
 pip install -e "python[all]"
 ```
 
+### Notes
+- If you are using older GPUs (NVIDIA T4, V100), please use `pip install "triton>=2.2.0"` to avoid some bugs in the triton compiler
+- If you only need to use the OpenAI backend, you can avoid installing other dependencies by using `pip install sglang[openai]`
+
 ## Quick Start
 The example below shows how to use sglang to answer a mulit-turn question.
 
 ### Using OpenAI Models
+Set the OpenAI API Key
+```
+export OPENAI_API_KEY=sk-******
+```
+
+Then, answer a multi-turn question.
 ```python
 from sglang import function, system, user, assistant, gen, set_default_backend, OpenAI
 
@@ -325,6 +336,7 @@ for m in state.messages():
 
 ### More Examples
 
+Anthropic and VertexAI (Gemini) models are also supported.
 You can find more examples at [examples/quick_start](examples/quick_start).
 
 ## Frontend: Structured Generation Langauge (SGLang)
@@ -334,19 +346,20 @@ To begin with, import sglang.
 import sglang as sgl
 ```
 
-`sglang` provides some simple primitives such as `gen`, `select`, `fork`.
+`sglang` provides some simple primitives such as `gen`, `select`, `fork`, `image`.
 You can implement your prompt flow in a function decorated by `sgl.function`.
 You can then invoke the function with `run` or `run_batch`.
 The system will manage the state, chat template, and parallelism for you.
 
 ### Control Flow
+You can use any Python code within the function body, including control flow, nested function calls, and external libraries.
+
 ```python
 @sgl.function
 def control_flow(s, question):
     s += "To answer this question: " + question + ", "
     s += "I need to use a " + sgl.gen("tool", choices=["calculator", "web browser"]) + ". "
 
-    # You can use if or nested function calls
     if s["tool"] == "calculator":
         s += "The math expression is" + sgl.gen("expression")
     elif s["tool"] == "web browser":
@@ -354,6 +367,9 @@ def control_flow(s, question):
 ```
 
 ### Parallelism
+Use `fork` to launch parallel prompts.
+Because `sgl.gen` is non-blocking, the for loop below issues two generation calls in parallel.
+
 ```python
 @sgl.function
 def tip_suggestion(s):
@@ -362,7 +378,7 @@ def tip_suggestion(s):
         "1. Balanced Diet. 2. Regular Exercise.\n\n"
     )
 
-    forks = s.fork(2)  # Launch parallel prompts
+    forks = s.fork(2)
     for i, f in enumerate(forks):
         f += f"Now, expand tip {i+1} into a paragraph:\n"
         f += sgl.gen(f"detailed_tip", max_tokens=256, stop="\n\n")
@@ -373,6 +389,8 @@ def tip_suggestion(s):
 ```
 
 ### Multi Modality
+Use `sgl.image` to pass an image as input.
+
 ```python
 @sgl.function
 def image_qa(s, image_file, question):
@@ -381,11 +399,13 @@ def image_qa(s, image_file, question):
 ```
 
 ### Constrained Decoding
+Use `regex=` to specify a regular expression as a decoding constraint.
+
 ```python
-@function
+@sgl.function
 def regular_expression_gen(s):
     s += "Q: What is the IP address of the Google DNS servers?\n"
-    s += "A: " + gen(
+    s += "A: " + sgl.gen(
         "answer",
         temperature=0,
         regex=r"((25[0-5]|2[0-4]\d|[01]?\d\d?).){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)",
@@ -393,6 +413,8 @@ def regular_expression_gen(s):
 ```
 
 ### Batching
+Use `run_batch` to run a batch of requests with continuous batching.
+
 ```python
 @sgl.function
 def text_qa(s, question):
@@ -405,10 +427,13 @@ states = text_qa.run_batch(
         {"question": "What is the capital of France?"},
         {"question": "What is the capital of Japan?"},
     ],
+    progress_bar=True
 )
 ```
 
 ### Streaming
+Add `stream=True` to enable streaming.
+
 ```python
 @sgl.function
 def text_qa(s, question):
@@ -417,7 +442,9 @@ def text_qa(s, question):
 
 states = text_qa.run(
     question="What is the capital of France?",
-    temperature=0.1)
+    temperature=0.1,
+    stream=True
+)
 
 for out in state.text_iter():
     print(out, end="", flush=True)
@@ -426,7 +453,7 @@ for out in state.text_iter():
 ## Backend: SGLang Runtime (SRT)
 The SGLang Runtime (SRT) is designed to work best with the SGLang frontend.
 However, it can also be used as a standalone API server.
-In this case, the RadixAttention can still greatly accelerate many use cases.
+In this case, the [RadixAttention](https://arxiv.org/abs/2312.07104) can still greatly accelerate many use cases with automatic KV cache reuse.
 
 ### Usage
 Launch a server
@@ -450,6 +477,10 @@ curl http://localhost:30000/v1/completions \
 ```
 python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --tp 2
 ```
+- If you see out-of-memory errors during serving, please try to reduce the memory usage of the KV cache pool by setting a smaller value of `--mem-fraction-static`. The default value is `0.9`
+```
+python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port 30000 --mem-fraction-static 0.7
+```
 
 ### Supported Models
 - Llama
@@ -457,6 +488,7 @@ python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port
 - Mixtral
 - LLaVA
   - `python3 -m sglang.launch_server --model-path liuhaotian/llava-v1.5-7b --tokenizer-path llava-hf/llava-1.5-7b-hf --port 30000`
+- AWQ quantization
 
 ## Benchmark And Performance
 
@@ -466,13 +498,13 @@ python -m sglang.launch_server --model-path meta-llama/Llama-2-7b-chat-hf --port
 - Mixtral-8x7B on NVIDIA A10G, FP16, Tensor Parallelism=8
 ![mixtral_8x7b](assets/mixtral_8x7b.jpg)
 
-Learn more [here]().
+Learn more [here](docs/benchmark_results.md).
 
 ## Roadmap
-- [ ] Function call
-- [ ] Quantization
+- [ ] Function call APIs
 - [ ] S-LoRA
-- [ ] More models
+- [ ] Support more models
+- [ ] Support more hardware backends
 
 ## Citation And Acknowledgment
 ```

{sglang-0.1.3.dist-info → sglang-0.1.5.dist-info}/RECORD

@@ -1,5 +1,5 @@
-sglang/__init__.py,sha256=U_vIUJQoAKKm3mK9wNlAiUFO4rk5G0epSNmOO43IQrI,95
-sglang/api.py,sha256=tJuEyB28BUQfl0-dQr4vi6UMHBhUbmyu9Z3iAE5xFcU,3883
+sglang/__init__.py,sha256=G73L_PWJ_6mF3NIE4ZAOWcb1CUbETSeRWr3wDTePrZ4,95
+sglang/api.py,sha256=SxmPP_PMYi4DfUcwz_V9UvYOwGmQdHPgpMV6jDDJq68,3928
 sglang/flush_cache.py,sha256=cCD_MTlQ5qEv__w0nOthDnVitdAfyscYjksBljwC5Mw,1835
 sglang/global_config.py,sha256=PAX7TWeFcq0HBzNUWyCONAOjqIokWqw8vT7I6sBSKTc,797
 sglang/launch_server.py,sha256=jKPZRDN5bUe8Wgz5eoDkqeePhmKa8DLD4DpXQLT5auo,294
@@ -7,15 +7,14 @@ sglang/utils.py,sha256=tvJs95QGZ_PcnTjvm-CDGQ8dJe84qUUOfG7BeF79nsA,5670
 sglang/backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sglang/backend/anthropic.py,sha256=y5TN9EDrJtOH4JEUxpXu-endloeYBy7xMUr3r7Ah3MA,1462
 sglang/backend/base_backend.py,sha256=pPalZfoezxnUBs752j7lm0uMwa8tZuCWd-ijSdStMO8,1745
-sglang/backend/huggingface.py,sha256=roQlt8y41PQbmnAY47CXiR0KJaxhtljH6j8RhbsR4f0,10533
 sglang/backend/openai.py,sha256=umTWzC2p4PypDaXHe6Kc8By5IM_Doi0Ob97vK_fFWDc,7367
 sglang/backend/runtime_endpoint.py,sha256=rIhwtKJaLLCJAc6q6kqxEVC8xO_NNjmJs7BnxlOydLM,5860
-sglang/backend/tgi.py,sha256=2wlfparGJNLN806bvPi_8jsk6ezJG1QviSZu2IBf1No,5935
+sglang/backend/vertexai.py,sha256=BLfWf_tEgoHY9srCufJM5PLe3tql2j0G6ia7cPykxCM,4713
 sglang/lang/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sglang/lang/chat_template.py,sha256=1x4724K2oxu7VID40-5Megk7SbZb97PQCbRjLpoescU,5599
 sglang/lang/compiler.py,sha256=wNn_UqV6Sxl22mv-PpzFUtRgiFFV-Y4OYpO4LshEoRM,7527
-sglang/lang/interpreter.py,sha256=YqCqsVonZt_xwL1ZMBNXHRSyxnGUVQr736ESn1Q7NWE,22339
-sglang/lang/ir.py,sha256=9vUL68VkT3gmcDaLjTiwerM21UwlQg-95FRaIt32jSM,12380
+sglang/lang/interpreter.py,sha256=0WTJxCB57WDBr_E6kW39wByhcG2nRFjEMTzOjAaNhrY,22453
+sglang/lang/ir.py,sha256=uUnBRyaM-8suVOEb2qf4EAt_VN2pWbXV6V88jLk6wsI,13160
 sglang/lang/tracer.py,sha256=zH9DENdJBPEvWkThgwqvHOW7aC1EPC8xha_WpEj-SRs,8243
 sglang/srt/backend_config.py,sha256=7MdHjNsZeAKB9IWWxyrvyOjJJAdI5tl9hWl-MV7yHrI,226
 sglang/srt/hf_transformers_utils.py,sha256=soRyYLoCn7GxgxvonufGFkdFBA3eH5i3Izk_wi7p1l0,5285
@@ -23,14 +22,14 @@ sglang/srt/memory_pool.py,sha256=cN3Lrs9fn0DFmt67_IN4g06mPzKUxpbAJGUw4O33xbo,360
 sglang/srt/model_config.py,sha256=R7YaR8H8AmCJl_1XcSP0zII_5ebZNl0wMXNVANGWd2c,997
 sglang/srt/sampling_params.py,sha256=Sd9l_uIIuS_mhbzljKwTGDO9ESMviNOYGxOifc71RrY,2895
 sglang/srt/server.py,sha256=XxTS1K4N5y-ZknLBQefxk1UxC50l6DABVqJOrJ-NG74,6388
-sglang/srt/server_args.py,sha256=Fpj3To5hEgmWn9qCS-pfypOEh34x9xVmiHBoEx5Smbo,4932
-sglang/srt/utils.py,sha256=YtTLEtVnOTrjub0Ct_xjrKGtHIajiQ57FB38l6Dw3a4,5691
+sglang/srt/server_args.py,sha256=ojox8nu2tgPEy_JlKKEvRenby4HKkmWk-1MpHy3PmnI,5771
+sglang/srt/utils.py,sha256=-2F99bqYT99x1jScMjciJxgQec6CaH6PcCHSmrKHhhY,5692
 sglang/srt/constrained/fsm.py,sha256=H4kXSsV4IX2ow5TMmnmd-8ho4qqJ5mpVZ4MOH5FUtnY,12900
 sglang/srt/constrained/fsm_cache.py,sha256=KX4bFX5hj0W66SC9pSvst1ew7etaOMTtTC75z0enRME,1087
 sglang/srt/constrained/regex.py,sha256=CcV7KBOKS2ZxGoEr6BHG5okagNIGEXYvGvhKXu5gtDA,18689
 sglang/srt/constrained/tokenizer.py,sha256=rei9yKHFETcbDPOpI7bpIYdrBFgIBhGr_U-zb3r5Beo,7951
-sglang/srt/layers/context_flashattention_nopad.py,sha256=qQc35BVOYPoZlLbbTUWB3a43Zwd3v5ZKR_uFRoypUIU,5084
-sglang/srt/layers/extend_attention.py,sha256=X-3nrQBeUyA3_cp2vZH1dC85x-EF9rppiK95FocMnKA,11423
+sglang/srt/layers/context_flashattention_nopad.py,sha256=GkjLiTkS4px_uLcW0aDocE3_OBXtujZ-SlsN2b2U7ng,5204
+sglang/srt/layers/extend_attention.py,sha256=pWVE6ySnPiVLFON__bie73eDhmXHk4tECMK8zTiJNbI,12558
 sglang/srt/layers/get_selected_logprob.py,sha256=CpMXM9WXMSB-AVaxBB_aVl1Qx_ZtAFFnjDTm4CgNDpU,2199
 sglang/srt/layers/logits_processor.py,sha256=rwcXwdZ7-dW9zvJX3MF_EHSxMLbU7TIQ9xUIYRu-WAs,3013
 sglang/srt/layers/radix_attention.py,sha256=hmPNFg2TkN4EAVUj376N_89RRtUYRwFgUpjj5SydnRk,6170
@@ -40,18 +39,18 @@ sglang/srt/managers/io_struct.py,sha256=5jMWj6_U8yTQd5V3tpDtThnoFyF0A3ln-4Z5bSL3
 sglang/srt/managers/openai_protocol.py,sha256=Eid_734Wup4jsL1ZS2Op0vwRuzvNbF4mV2UcwFxqEvI,327
 sglang/srt/managers/tokenizer_manager.py,sha256=jVwr0lM18RFJLhDb5TWlUpQ4Q8tALT4L6GY0jmaZkLw,7861
 sglang/srt/managers/router/infer_batch.py,sha256=UfS1uVhGnM-62Xv1cfu_IoTeIUxkjkKc4W3trtGbadc,11541
-sglang/srt/managers/router/manager.py,sha256=H-T-LlnIssHw-FXMHbs3yDQewkTMBCqG6jTYjugopCA,2527
-sglang/srt/managers/router/model_rpc.py,sha256=ZLK5izxMGpfCs4uT7DJ8u-aww5UG_jwjr7eJdbWGZ3Y,19271
-sglang/srt/managers/router/model_runner.py,sha256=U-SBnEeLvwolLcaxyxrPgVG7PnR2rRvuXWV50t9y0Fo,16480
+sglang/srt/managers/router/manager.py,sha256=AVCdYKKYcIQsIwpudkfFY4jh6M--ubLjXeYGzfi2ebw,2528
+sglang/srt/managers/router/model_rpc.py,sha256=CR3qbHvShttlC19qAZ8B8nhT6UPobeu2Dy3Z0n6WdC8,19448
+sglang/srt/managers/router/model_runner.py,sha256=IhSdpBcd54HN01HDi_PAkJztFxEGDcnktdoPZDWEx3s,16487
 sglang/srt/managers/router/radix_cache.py,sha256=ZQPm9HhQ7vD3Gl5nhuvw3ZW4ZRARcplqWed1GYUvHCg,6441
 sglang/srt/managers/router/scheduler.py,sha256=ejuIRwqqMZVXFKUionRJxy5AtNvK25YoGRO9rFY-rc8,2926
 sglang/srt/models/llama2.py,sha256=D3j-NtyM8PA74UhXM7wSPogI2HKX-JcQAWcOusrZZo0,11320
 sglang/srt/models/llava.py,sha256=COS0IC6Yo-QiwKe5emgCbtEe9HgaSu5tt6CQA7UtV38,8533
-sglang/srt/models/mixtral.py,sha256=j91xOt6NZ5tJiyTPqmUSzgJqFAw7vTDnfBtEs5x0jDM,13714
-sglang/test/test_programs.py,sha256=ua3wufnS3x6d_U3aboY4ivqoglrRPZj18j96vuiUtiE,11348
+sglang/srt/models/mixtral.py,sha256=frd2XsNZwP0XsQtRiYhgy4PErLNLgtIsLakmNrOKBAU,13712
+sglang/test/test_programs.py,sha256=EovA2xL7fODcTbFj2wAAmYKlg1mLZ1x1BRU6nrXFRdE,11416
 sglang/test/test_utils.py,sha256=Knxg3BTA6d_7XSlprbBCdvfDr2SN5x7LhkT-tZFk5EQ,4828
-sglang-0.1.3.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sglang-0.1.3.dist-info/METADATA,sha256=SSRJ09MVErF7DrD5lJLm2oBDkk7sySET3AVaxJMciKs,21885
-sglang-0.1.3.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-sglang-0.1.3.dist-info/top_level.txt,sha256=yxhh3pYQkcnA7v3Bg889C2jZhvtJdEincysO7PEB09M,7
-sglang-0.1.3.dist-info/RECORD,,
+sglang-0.1.5.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sglang-0.1.5.dist-info/METADATA,sha256=aepmAL6VoXRcxZBIDKvxwikCYSbvWFm_JFGTxb3Mgfw,23345
+sglang-0.1.5.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+sglang-0.1.5.dist-info/top_level.txt,sha256=yxhh3pYQkcnA7v3Bg889C2jZhvtJdEincysO7PEB09M,7
+sglang-0.1.5.dist-info/RECORD,,

sglang/backend/huggingface.py

@@ -1,349 +0,0 @@
-import functools
-from enum import Enum, auto
-from typing import Callable, List, Optional, Union
-
-import numpy as np
-import torch
-import transformers
-from sglang.backend.base_backend import BaseBackend
-from sglang.lang.chat_template import get_chat_template_by_model_path
-from sglang.lang.interpreter import ProgramState
-from sglang.utils import get_available_gpu_memory
-from transformers import (
-    AutoModelForCausalLM,
-    AutoTokenizer,
-    StoppingCriteria,
-    StoppingCriteriaList,
-)
-from transformersgl.generation.logits_process import (
-    LogitsProcessorList,
-    RepetitionPenaltyLogitsProcessor,
-    TemperatureLogitsWarper,
-    TopKLogitsWarper,
-    TopPLogitsWarper,
-)
-
-
-class StopReason(Enum):
-    EOS_TOKEN = auto()
-    STOP_STR = auto()
-    LENGTH = auto()
-
-
-def load_model(
-    model_name: str,
-    device,
-    num_gpus,
-    max_gpu_memory,
-    model_kwargs=None,
-    tokenizer_kwargs=None,
-):
-    model_kwargs = model_kwargs or {}
-    tokenizer_kwargs = tokenizer_kwargs or {}
-
-    if device == "cuda":
-        model_kwargs["torch_dtype"] = torch.float16
-        if num_gpus != 1:
-            model_kwargs["device_map"] = "auto"
-            if max_gpu_memory is None:
-                model_kwargs[
-                    "device_map"
-                ] = "sequential"  # This is important for not the same VRAM sizes
-                available_gpu_memory = [
-                    get_available_gpu_memory(i, False) for i in range(num_gpus)
-                ]
-                model_kwargs["max_memory"] = {
-                    i: str(int(available_gpu_memory[i] * 0.85)) + "GiB"
-                    for i in range(num_gpus)
-                }
-            else:
-                model_kwargs["max_memory"] = {
-                    i: max_gpu_memory for i in range(num_gpus)
-                }
-    elif device == "cpu":
-        model_kwargs["torch_dtype"] = torch.float32
-    else:
-        raise ValueError(f"Invalid device: {device}")
-
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name, low_cpu_mem_usage=True, **model_kwargs
-    )
-    tokenizer = AutoTokenizer.from_pretrained(model_name, **tokenizer_kwargs)
-
-    if num_gpus == 1:
-        model.to(device).eval()
-
-    return model, tokenizer
-
-
-def prepare_logits_processor(
-    temperature: float, repetition_penalty: float, top_p: float, top_k: int
-) -> LogitsProcessorList:
-    processor_list = LogitsProcessorList()
-    # TemperatureLogitsWarper doesn't accept 0.0, 1.0 makes it a no-op so we skip two cases.
-    if temperature >= 1e-5 and temperature != 1.0:
-        processor_list.append(TemperatureLogitsWarper(temperature))
-    if repetition_penalty > 1.0:
-        processor_list.append(RepetitionPenaltyLogitsProcessor(repetition_penalty))
-    if 1e-8 <= top_p < 1.0:
-        processor_list.append(TopPLogitsWarper(top_p))
-    if top_k > 0:
-        processor_list.append(TopKLogitsWarper(top_k))
-    return processor_list
-
-
-@functools.lru_cache
-def get_token_healing_mask(tokenizer, prompt_last_token):
-    last_str = tokenizer.convert_ids_to_tokens(prompt_last_token)
-    disallowed = torch.zeros(len(tokenizer), dtype=bool)
-    for s, t_id in tokenizer.get_vocab().items():
-        if not s.startswith(last_str):
-            disallowed[t_id] = 1
-    return disallowed
-
-
-@functools.lru_cache
-def get_int_token_mask(tokenizer):
-    disallowed = torch.zeros(len(tokenizer), dtype=bool)
-    for s, t_id in tokenizer.get_vocab().items():
-        s = s.replace("▁", "").strip()
-        if not (s.isdigit() or len(s) == 0 or s == ","):
-            disallowed[t_id] = 1
-    disallowed[tokenizer.eos_token_id] = 0
-    return disallowed
-
-
-@torch.inference_mode()
-def generate_stream(
-    model,
-    tokenizer,
-    prompt,
-    max_new_tokens,
-    stop: List[str],
-    temperature,
-    top_p,
-    token_healing,
-    logit_mask=None,
-):
-    logits_processor = prepare_logits_processor(
-        temperature=temperature, repetition_penalty=1.0, top_p=top_p, top_k=0
-    )
-    device = model.device
-    input_ids = tokenizer.encode(prompt)
-    output_ids = list(input_ids)
-    prompt_len = len(prompt)
-
-    # Resolve stop
-    stop_token_ids = [tokenizer.eos_token_id]
-
-    # Token healing
-    token_healing = token_healing and len(input_ids) > 0
-    if token_healing:
-        token_healing_mask = get_token_healing_mask(tokenizer, input_ids[-1])
-        del output_ids[-1]
-
-    # Generate
-    past_key_values = None
-    stop_reason = None
-    for i in range(max_new_tokens):
-        # Forward
-        if i == 0:  # prefill
-            out = model(torch.as_tensor([output_ids], device=device), use_cache=True)
-        else:  # decoding
-            out = model(
-                input_ids=torch.as_tensor([[token]], device=device),
-                use_cache=True,
-                past_key_values=past_key_values,
-            )
-        logits = out.logits
-        past_key_values = out.past_key_values
-
-        # Logit mask
-        if token_healing and i == 0:
-            logits[0, -1, token_healing_mask] = -1e4
-        if logit_mask is not None:
-            logits[0, -1, logit_mask] = -1e4
-
-        # Sample next token
-        last_token_logits = logits_processor(None, logits[:, -1, :])[0]
-        if temperature < 1e-5 or top_p < 1e-8:  # greedy
-            token = int(torch.argmax(last_token_logits))
-        else:
-            probs = torch.softmax(last_token_logits, dim=-1)
-            token = int(torch.multinomial(probs, num_samples=1))
-        output_ids.append(token)
-
-        # Stop condition
-        if token in stop_token_ids:
-            stop_reason = StopReason.EOS_TOKEN
-            break
-
-        output_str = tokenizer.decode(output_ids, skip_special_tokens=True)
-        for stop_str in stop:
-            pos = output_str[prompt_len:].find(stop_str)
-            if pos != -1:
-                stop_reason = StopReason.STOP_STR
-                output_str = output_str[: prompt_len + pos]
-                break
-
-        if stop_reason:
-            break
-
-    return output_str[prompt_len:]
-
-
-class HuggingFaceTransformers(BaseBackend):
-    def __init__(
-        self,
-        model_name,
-        device="cuda",
-        num_gpus=1,
-        max_gpu_memory=None,
-        model_kwargs=None,
-        tokenizer_kwargs=None,
-    ):
-        self.model_name = model_name
-        self.device = device
-
-        self.model, self.tokenizer = load_model(
-            model_name, device, num_gpus, max_gpu_memory, model_kwargs, tokenizer_kwargs
-        )
-
-        self.chat_template = get_chat_template_by_model_path(model_name)
-
-    def get_chat_template(self):
-        return self.chat_template
-
-    def cache_prefix(self, prefix_str: str):
-        pass
-
-    def uncache_prefix(self, rid: str):
-        pass
-
-    def end_request(self, rid: str):
-        pass
-
-    def begin_program(self, s: ProgramState):
-        pass
-
-    def end_program(self, s: ProgramState):
-        pass
-
-    def fill(self, s: ProgramState, text: str):
-        return False
-
-    def generate_internal(
-        self,
-        prompt: str,
-        max_tokens: int,
-        stop: Union[str, List[str]],
-        temperature: float,
-        top_p: float,
-        dtype: Optional[str] = None,
-    ):
-        if dtype is None:
-            comp = generate_stream(
-                self.model,
-                self.tokenizer,
-                prompt,
-                max_new_tokens=max_tokens,
-                stop=stop,
-                temperature=temperature,
-                top_p=top_p,
-                token_healing=True,
-            )
-        elif dtype in [str, "str", "string"]:
-            comp = generate_stream(
-                self.model,
-                self.tokenizer,
-                prompt + '"',
-                max_new_tokens=max_tokens,
-                stop=['"'],
-                temperature=temperature,
-                top_p=top_p,
-                token_healing=False,
-            )
-            comp = '"' + comp + '"'
-        elif dtype in [int, "int"]:
-            logit_mask = get_int_token_mask(self.tokenizer)
-            comp = generate_stream(
-                self.model,
-                self.tokenizer,
-                prompt,
-                max_new_tokens=max_tokens,
-                stop=stop + [" ", ","],
-                temperature=temperature,
-                top_p=top_p,
-                token_healing=False,
-                logit_mask=logit_mask,
-            )
-        return comp
-
-    def generate(
-        self,
-        s: ProgramState,
-        max_tokens: int,
-        stop: Union[str, List[str]],
-        temperature: float,
-        top_p: float,
-        dtype: Optional[str] = None,
-    ):
-        prompt = s.text
-        comp = self.generate_internal(
-            prompt, max_tokens, stop, temperature, top_p, dtype
-        )
-        return comp
-
-    def parallel_generate(
-        self,
-        s: ProgramState,
-        prefixes: List[str],
-        join_func: Callable,
-        max_tokens: int,
-        stop: Union[str, List[str]],
-        temperature: float,
-        top_p: float,
-        dtype: Optional[str] = None,
-    ):
-        prompt = s.text
-        parallel_prompts = [prompt + prefix for prefix in prefixes]
-
-        comps = []
-        for i in range(len(parallel_prompts)):
-            comps.append(
-                self.generate_internal(
-                    parallel_prompts[i], max_tokens, stop, temperature, top_p, dtype
-                )
-            )
-
-        joined = join_func([p + c for p, c in zip(prefixes, comps)])
-        return joined, comps
-
-    @torch.inference_mode()
-    def select(
-        self, s: ProgramState, choices: List[str], temperature: float, top_p: float
-    ):
-        loss_fct = torch.nn.CrossEntropyLoss()
-        prompt = s.text
-
-        prompt_len = self.tokenizer.encode(prompt, return_tensors="pt").shape[1]
-        prompt_choices = [prompt + choice for choice in choices]
-
-        scores = []
-        for i in range(len(choices)):
-            choice_ids = self.tokenizer.encode(
-                prompt_choices[i], return_tensors="pt"
-            ).to(self.model.device)
-            logits = self.model(choice_ids).logits
-
-            # score = -loss_fct(logits[0, :-1, :], choice_ids[0, 1:]).item()
-
-            logprobs = torch.log(torch.softmax(logits, dim=-1))
-            idx1 = torch.arange(0, logits.shape[1] - 1, device=logits.device)
-            idx2 = choice_ids[0, 1:]
-            selected_logprobs = logprobs[0, idx1, idx2]
-            score = selected_logprobs.mean().item()
-            scores.append(score)
-
-        decision = choices[np.argmax(scores)]
-        return decision, scores

sglang/backend/tgi.py

@@ -1,190 +0,0 @@
-import re
-from concurrent.futures import ThreadPoolExecutor
-from functools import partial
-from itertools import repeat
-from typing import List, Optional, Union
-
-from sglang.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 TGI(BaseBackend):
-    def __init__(self, base_url):
-        super().__init__()
-
-        self.base_url = base_url
-
-        res = http_request(self.base_url + "/info")
-        assert res.status_code == 200
-        self.model_info = res.json()
-        self.chat_template = get_chat_template_by_model_path(
-            self.model_info["model_id"]
-        )
-
-    def get_model_name(self):
-        return self.model_info["model_id"]
-
-    def get_chat_template(self):
-        return self.chat_template
-
-    @staticmethod
-    def adapt_params(max_tokens, stop, sampling_params, **override_params):
-        temperature = sampling_params.temperature
-        do_sample = True
-        if temperature == 0:
-            do_sample = False
-            temperature = None
-
-        if stop is None:
-            stop = []
-        elif isinstance(stop, str):
-            stop = [stop]
-
-        top_p = sampling_params.top_p
-        if top_p == 0:
-            top_p = 0.001
-        if top_p == 1:
-            top_p = 0.999
-
-        top_k = sampling_params.top_k
-        if top_k == -1:
-            top_k = None
-
-        params = {
-            "decoder_input_details": False,
-            "details": False,
-            "do_sample": do_sample,
-            "max_new_tokens": max_tokens,
-            "stop": stop,
-            "temperature": temperature,
-            "top_p": top_p,
-            "top_k": top_k,
-            "return_full_text": False,
-        }
-        params.update(override_params)
-        return params
-
-    @staticmethod
-    def _extract_int(text):
-        words = re.split("\ |'|\/|\(|\)|\n|\.|,", text)
-        for word in words:
-            try:
-                int(word)
-                return word
-            except ValueError:
-                continue
-        raise ValueError
-
-    @staticmethod
-    def _extract_choice(choices, text):
-        # FIXME: Current only support the case where the choices are single words.
-        words = re.split("\ |'|\/|\(|\)|\n|\.|,", text)
-        for word in words:
-            if word in choices:
-                return word
-        raise ValueError
-
-    @staticmethod
-    def _truncate_to_stop(text, stop):
-        # The stop sequence may not be a single token. In this case TGI will generate
-        # too many tokens so we need to truncate the output.
-        if stop:
-            stop = [stop] if isinstance(stop, str) else stop
-            for stop_seq in stop:
-                pos = text.find(stop_seq)
-                if pos != -1:
-                    return text[:pos]
-        return text
-
-    def _make_request(self, params):
-        res = http_request(self.base_url + "/generate", json=params)
-        if res.status_code != 200:
-            raise ValueError(f"Error from TGI backend: {res.text}")
-        return res.json()
-
-    def retry_for_expected(self, prompt, params, extract_fn, retry=5):
-        # TGI does not support logis_bias (yet), so we have to use an inefficient hack.
-        failed = []
-        while retry > 0:
-            res_json = self._make_request(
-                {
-                    "inputs": prompt,
-                    "parameters": params,
-                }
-            )
-            text = res_json["generated_text"]
-            try:
-                return extract_fn(text)
-            except ValueError:
-                retry -= 1
-            failed.append(text)
-
-        msg = "=" * 20 + "\n"
-        msg += f"Prompt:\n{prompt}\n"
-        msg += "=" * 20 + "\n"
-        for i, text in enumerate(failed):
-            msg += f"====== Try {i+1}:\n{text}\n"
-
-        raise ValueError(
-            f"Model {self.model_info['model_id']} served by TGI backend does not generate"
-            "expected output. Please improve the prompt, increase the temperature, or "
-            f"use different models.\n{msg}"
-        )
-
-    def select(
-        self,
-        s: StreamExecutor,
-        choices: List[str],
-        sampling_params: SglSamplingParams,
-    ):
-        decision = self.retry_for_expected(
-            s.text_,
-            self.adapt_params(16, [], sampling_params),
-            partial(self._extract_choice, choices),
-        )
-        return decision, [1 if choice == decision else 0 for choice in choices]
-
-    def generate(
-        self,
-        s: StreamExecutor,
-        max_tokens: int,
-        stop: Union[str, List[str]],
-        sampling_params: SglSamplingParams,
-        dtype: Optional[str] = None,
-    ):
-        if dtype is None:
-            res_json = self._make_request(
-                {
-                    "inputs": s.text_,
-                    "parameters": self.adapt_params(max_tokens, stop, sampling_params),
-                }
-            )
-            return self._truncate_to_stop(res_json["generated_text"], stop), {}
-
-        if dtype in [str, "str", "string"]:
-            stop = ['"']
-            res_json = self._make_request(
-                {
-                    "inputs": f'{s.text_}"',
-                    "parameters": self.adapt_params(max_tokens, stop, sampling_params),
-                }
-            )
-            return (
-                '"' + self._truncate_to_stop(res_json["generated_text"], stop) + '"',
-                {},
-            )
-
-        if dtype in [int, "int"]:
-            return (
-                self.retry_for_expected(
-                    s.text_,
-                    self.adapt_params(max_tokens, stop, sampling_params),
-                    self._extract_int,
-                ),
-                {},
-            )
-
-        raise ValueError(f"Unknown dtype: {dtype}")