sglang 0.1.24__py3-none-any.whl → 0.1.26__py3-none-any.whl

sglang/__init__.py

@@ -1,5 +1,3 @@
-__version__ = "0.1.24"
-
 # SGL API Components
 from sglang.api import (
     Runtime,
@@ -32,6 +30,8 @@ from sglang.lang.backend.openai import OpenAI
 from sglang.lang.backend.runtime_endpoint import RuntimeEndpoint
 from sglang.lang.backend.vertexai import VertexAI
 
+from .version import __version__
+
 # public APIs management
 __all__ = [
     "global_config",

sglang/srt/managers/controller/model_runner.py

@@ -15,6 +15,7 @@ from flashinfer import (
     BatchPrefillWithRaggedKVCacheWrapper,
 )
 from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
+from torch.nn.parameter import Parameter
 from vllm.config import DeviceConfig, LoadConfig
 from vllm.config import ModelConfig as VllmModelConfig
 from vllm.distributed import (
@@ -22,6 +23,7 @@ from vllm.distributed import (
     init_distributed_environment,
     initialize_model_parallel,
 )
+from vllm.model_executor.layers.linear import QKVParallelLinear
 from vllm.model_executor.models import ModelRegistry
 
 from sglang.global_config import global_config
@@ -38,6 +40,18 @@ from sglang.srt.utils import (
 logger = logging.getLogger("srt.model_runner")
 
 
+def is_llama3_405b_fp8(model_config):
+    if (
+        model_config.hf_config.architectures[0] == "LlamaForCausalLM"
+        and model_config.hf_config.hidden_size == 16384
+        and model_config.hf_config.intermediate_size == 53248
+        and model_config.hf_config.num_hidden_layers == 126
+        and model_config.hf_config.quantization_config["quant_method"] == "fbgemm_fp8"
+    ):
+        return True
+    return False
+
+
 class ModelRunner:
     def __init__(
         self,
@@ -118,6 +132,9 @@ class ModelRunner:
             seed=42,
             skip_tokenizer_init=True,
         )
+        if is_llama3_405b_fp8(self.model_config):
+            self.model_config.hf_config.num_key_value_heads = 8
+            vllm_model_config.hf_config.num_key_value_heads = 8
         self.dtype = vllm_model_config.dtype
         if self.model_config.model_overide_args is not None:
             vllm_model_config.hf_config.update(self.model_config.model_overide_args)
@@ -370,5 +387,39 @@ def load_model_cls_srt(model_arch: str) -> Optional[Type[nn.Module]]:
     return model_arch_name_to_cls[model_arch]
 
 
+def get_original_weight(loaded_weight, head_dim):
+    n_kv_head = loaded_weight.shape[0] // (2 * head_dim)
+    dim = loaded_weight.shape[1]
+    for i in range(n_kv_head):
+        loaded_weight[i * head_dim : (i + 1) * head_dim, :] = loaded_weight[
+            2 * i * head_dim : (2 * i + 1) * head_dim, :
+        ]
+    original_kv_weight = loaded_weight[: n_kv_head * head_dim, :]
+    assert original_kv_weight.shape == (n_kv_head * head_dim, dim)
+    return original_kv_weight
+
+
+def get_weight_loader_srt(weight_loader):
+    def weight_loader_srt(
+        self,
+        param: Parameter,
+        loaded_weight: torch.Tensor,
+        loaded_shard_id: Optional[str] = None,
+    ):
+        if (
+            loaded_shard_id in ["k", "v"]
+            and loaded_weight.shape[0] == self.head_size * self.total_num_kv_heads * 2
+        ):
+            loaded_weight = get_original_weight(loaded_weight, self.head_size)
+
+        weight_loader(self, param, loaded_weight, loaded_shard_id)
+
+    return weight_loader_srt
+
+
 # Monkey patch model loader
 setattr(ModelRegistry, "load_model_cls", load_model_cls_srt)
+original_weight_loader = QKVParallelLinear.weight_loader
+setattr(
+    QKVParallelLinear, "weight_loader", get_weight_loader_srt(original_weight_loader)
+)

sglang/srt/server.py

@@ -52,6 +52,7 @@ from sglang.srt.utils import (
     allocate_init_ports,
     assert_pkg_version,
     enable_show_time_cost,
+    maybe_set_triton_cache_manager,
     set_ulimit,
 )
 from sglang.utils import get_exception_traceback
@@ -201,6 +202,11 @@ def launch_server(
             "reinstall the latest version by following the instructions "
             "at https://docs.flashinfer.ai/installation.html.",
         )
+
+    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()
+
     if server_args.chat_template:
         # TODO: replace this with huggingface transformers template
         load_chat_template_for_openai_api(server_args.chat_template)

sglang/srt/utils.py

@@ -18,10 +18,15 @@ import psutil
 import requests
 import torch
 import torch.distributed as dist
-import triton
 from fastapi.responses import JSONResponse
 from packaging import version as pkg_version
 from starlette.middleware.base import BaseHTTPMiddleware
+from triton.runtime.cache import (
+    FileCacheManager,
+    default_cache_dir,
+    default_dump_dir,
+    default_override_dir,
+)
 
 logger = logging.getLogger(__name__)
 
@@ -460,6 +465,44 @@ def monkey_patch_vllm_all_gather(reverse: bool = False):
         setattr(GroupCoordinator, "all_gather", all_gather)
 
 
+def maybe_set_triton_cache_manager() -> None:
+    """Set environment variable to tell Triton to use a
+    custom cache manager"""
+    cache_manger = os.environ.get("TRITON_CACHE_MANAGER", None)
+    if cache_manger is None:
+        manager = "sglang.srt.utils:CustomCacheManager"
+        logger.info("Setting Triton cache manager to: %s", manager)
+        os.environ["TRITON_CACHE_MANAGER"] = manager
+
+
+class CustomCacheManager(FileCacheManager):
+    # Adapted from: https://github.com/tdoublep/vllm/blob/3307522289fdfefe323b6c00d0db696651989a2f/vllm/triton_utils/custom_cache_manager.py
+    def __init__(self, key, override=False, dump=False):
+
+        self.key = key
+        self.lock_path = None
+        if dump:
+            self.cache_dir = default_dump_dir()
+            self.cache_dir = os.path.join(self.cache_dir, self.key)
+            self.lock_path = os.path.join(self.cache_dir, "lock")
+            os.makedirs(self.cache_dir, exist_ok=True)
+        elif override:
+            self.cache_dir = default_override_dir()
+            self.cache_dir = os.path.join(self.cache_dir, self.key)
+        else:
+            # create cache directory if it doesn't exist
+            self.cache_dir = (
+                os.getenv("TRITON_CACHE_DIR", "").strip() or default_cache_dir()
+            )
+            if self.cache_dir:
+                self.cache_dir = f"{self.cache_dir}_{os.getpid()}"
+                self.cache_dir = os.path.join(self.cache_dir, self.key)
+                self.lock_path = os.path.join(self.cache_dir, "lock")
+                os.makedirs(self.cache_dir, exist_ok=True)
+            else:
+                raise RuntimeError("Could not create or locate cache dir")
+
+
 API_KEY_HEADER_NAME = "X-API-Key"
 
 

sglang/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.26"

{sglang-0.1.24.dist-info → sglang-0.1.26.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sglang
-Version: 0.1.24
+Version: 0.1.26
 Summary: SGLang is yet another fast serving framework for large language models and vision language models.
 License: Apache License
                                    Version 2.0, January 2004
@@ -244,7 +244,7 @@ Requires-Dist: vllm ==0.5.3.post1 ; extra == 'srt'
 Requires-Dist: outlines >=0.0.44 ; extra == 'srt'
 
 <div align="center">
-<img src="assets/logo.png" alt="logo" width="400"></img>
+<img src="https://raw.githubusercontent.com/sgl-project/sglang/main/assets/logo.png" alt="logo" width="400"></img>
 </div>
 
 --------------------------------------------------------------------------------
@@ -282,7 +282,7 @@ The core features include:
 
 ### Method 1: With pip
 ```
-pip install --upgrade pip setuptools wheel
+pip install --upgrade pip
 pip install "sglang[all]"
 
 # Install FlashInfer CUDA kernels
@@ -405,7 +405,7 @@ python -m sglang.launch_server --model-path meta-llama/Meta-Llama-3-8B-Instruct
 
 ### Supported Models
 
-- Llama / Llama 2 / Llama 3
+- Llama / Llama 2 / Llama 3 / Llama 3.1
 - Mistral / Mixtral
 - Gemma / Gemma 2
 - Qwen / Qwen 2 / Qwen 2 MoE

{sglang-0.1.24.dist-info → sglang-0.1.26.dist-info}/RECORD

@@ -1,6 +1,5 @@
-sglang/__init__.py,sha256=nMs6lYeKcQpYArIaZLQ2VGNleY1dVvdBFaHyG7fpOsA,1141
+sglang/__init__.py,sha256=UV7VlXhXrwi00Zg45iNB9KcnmrwLjdMtjMz06AiafY0,1151
 sglang/api.py,sha256=1JARbc1wNYF6tODdUpgmNgTyLOvMnxdTBctLvEwzGTY,5565
-sglang/bench.py,sha256=p34wnfMRdiedOUf9GKGZkkNxehmyTzK6Q1O20q_SGjY,21841
 sglang/bench_latency.py,sha256=UPy6WhrddMTDX7HqIeHNhCn5vF0YMOKxJlQRvhMC8zU,10552
 sglang/bench_serving.py,sha256=zKGgVX3S-ggUvOxvEM4AszzXRPRVU6NGNnBG5vAAvRY,34577
 sglang/check_env.py,sha256=CscuPMlf68dkgZf0m-FiLpUisNNDoihMck4qhLOeV1Q,4124
@@ -8,13 +7,7 @@ sglang/global_config.py,sha256=QG-ABVJksKK_llvUx7fSZcmK4GGCs-hBUVcM4LCr7Nw,1749
 sglang/launch_server.py,sha256=Gg8CwNlTCCfg1dF65ZT9ePLxOT9LKtY79GhIPG6PCrU,358
 sglang/launch_server_llavavid.py,sha256=40uaazMsavKuk6YXFa5v37kdUpFGuealgJJeph1g8gU,1025
 sglang/utils.py,sha256=arJuwOAEX445M2NL9SAOi6jBNu0-cfU04PLAr-hIH3U,8168
-sglang/backend/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sglang/backend/anthropic.py,sha256=iJjXiDMZbtvX2XNG78MG9kM7SpZq9hmXVuzT_T18elw,2076
-sglang/backend/base_backend.py,sha256=APiMht4WYECLCOGRPCEUF6lX-an1vjVe2dWoMSgymWY,1831
-sglang/backend/litellm.py,sha256=ZqsEZXgxLge-Fh3SMr1XkVPU7z3FKntpRppNwd1a12s,2447
-sglang/backend/openai.py,sha256=Id4vDzfefG9R7AqJBMXqYmKHv2FMu0PBSYEGbK7Q510,14803
-sglang/backend/runtime_endpoint.py,sha256=PAdnQBj3yQNtgw8GH9F1ecGE7HhxGa2T7Tz_c--H2aE,9203
-sglang/backend/vertexai.py,sha256=98toR-L0OTi4dYHaSmmzJdlQ2qN_0lImoKZFlVgYLRE,4850
+sglang/version.py,sha256=3_QdGLpuk_SDY7k9PpNcHpSTjlPdhadPiEgF82wzkqk,23
 sglang/lang/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sglang/lang/chat_template.py,sha256=psIlhaDo70twgLrx5Lgln03metLEA3-FZuixeI0Y7Ao,13309
 sglang/lang/compiler.py,sha256=UiXUmPR9wBAPtnORrLcyQX8Uh0ZL0nKeV8ZgBozAJPw,7531
@@ -34,12 +27,10 @@ sglang/srt/hf_transformers_utils.py,sha256=94mOI93B2xOmXKqfJfEoGxqHgwwlWNbPHgsA4
 sglang/srt/memory_pool.py,sha256=FhJk5GtYortO3MJIsMMQ-o49agwDHVX1aEQH2LITq6c,3949
 sglang/srt/mm_utils.py,sha256=OptgAHDX-73Bk4jAdr2BOAJtiEXJNzPrMhaM-dy275c,8889
 sglang/srt/model_config.py,sha256=lZu1D-XLVMETHS6FBMoPn8Uowa9QFGe95d3SuWrr2q8,5282
-sglang/srt/openai_api_adapter.py,sha256=iw-FquXQeM2Z4nxOoYGFPjTkIdgA8rQkh_IcmJRy-R0,15143
-sglang/srt/openai_protocol.py,sha256=lGBhfxG6jmgUkMOh2NpBK9w9TUTRZKrsfHdW7XYhKKI,5700
 sglang/srt/sampling_params.py,sha256=OI11asr1Bd_E5soDjih614v4flgWxdMZU9HAF0aBafQ,3062
-sglang/srt/server.py,sha256=JC6rs8mkWg2mWwriwZvYEZyO514_HJFOUNda-pu8U_4,14369
+sglang/srt/server.py,sha256=DXhcJt0V24a7yhydP1abPrK1qqV3qt7r8cyOMVOAI4M,14611
 sglang/srt/server_args.py,sha256=aF6L35mEB-FU3BL_ooKuCIcOXLhYLxA9-MjpaOTQRCo,13189
-sglang/srt/utils.py,sha256=ZB9WLlZ_GpKVpPJiETrYkqH10J8iWrN_4buxDnQoA88,18568
+sglang/srt/utils.py,sha256=bUp3SLzbDms0dvuETaccDPAGRHOIGW5A61pqH62XiT0,20370
 sglang/srt/constrained/__init__.py,sha256=5LB3_mDTMW6wcRkFA5J2Rd5HPHHEKRyiELhe4gtlBYM,1472
 sglang/srt/constrained/base_cache.py,sha256=QQjmFEiT8jlOskJoZobhrDl2TKB-B4b1LPQo9JQCP_w,1405
 sglang/srt/constrained/fsm_cache.py,sha256=P4qNDHHxpKpTnYL_8V1R6OFXlUwbM6ZcBdzddpcBgb4,1135
@@ -57,11 +48,10 @@ sglang/srt/managers/detokenizer_manager.py,sha256=8rN2cdMr61LWy07lingEqLnNy0W5Re
 sglang/srt/managers/io_struct.py,sha256=Y6jW3p0cNg0jcrEQNki1H8MMEWxwWA4p6Y-xVgUVWaI,5404
 sglang/srt/managers/tokenizer_manager.py,sha256=SbivhFhZUR9HU9pLTe93MlYprAFAHzOU3KMBA2piQUk,19308
 sglang/srt/managers/controller/cuda_graph_runner.py,sha256=0aRqA1_34oJ557Zn8PjpJecex5bBWJdnCmBlcDVvYO0,8509
-sglang/srt/managers/controller/dp_worker.py,sha256=ES3-jyxGfHzpgVoXub_3qjVygwfWYWpfN4vuVWU23Gs,3675
 sglang/srt/managers/controller/infer_batch.py,sha256=SKwCwhnZ_CNlG0mVCEc4X0e4HNjJFke-c8zdWP3TzjQ,34186
 sglang/srt/managers/controller/manager_multi.py,sha256=DT8Y9RF5OyTxlrLEZYz4claNWir3UrVztdOZaVPiA6g,6077
 sglang/srt/managers/controller/manager_single.py,sha256=2xO_iWK6tWvc0B31nKbe2N3klxwQBJmPTnFhNjzhVSI,4566
-sglang/srt/managers/controller/model_runner.py,sha256=927tf6nJjLjEDgz2wCDj2kvpZ-E_rAVm8PVKFVfP4p8,13951
+sglang/srt/managers/controller/model_runner.py,sha256=FwZ7FU7nhJsYhtoTNxYFc4e6oMEwSqOh8ohXOKtFPKc,15828
 sglang/srt/managers/controller/radix_cache.py,sha256=tx8LEQpqLxipw9UUVj4D1YQLMMDmWnjDYv8oDlOl-co,8210
 sglang/srt/managers/controller/schedule_heuristic.py,sha256=SQAGzPS3aB_TPj7rnPBhewwyR6W1sVwW4D3zG3JUY00,2714
 sglang/srt/managers/controller/tp_worker.py,sha256=yjz-Xzl0zEy4QSU-EYneZH5vi3oHtBuXTtYe4VuDp2g,30517
@@ -90,16 +80,13 @@ sglang/srt/models/qwen2_moe.py,sha256=oHNoo45myV5kitkls2GWVzuGt1Q4pRHN2nLlXEltFI
 sglang/srt/models/stablelm.py,sha256=Z_XCDSHY_QMz3lZwwkZdIZjEOizZjLYJU9GDi8o08qQ,10802
 sglang/srt/models/yivl.py,sha256=55KPrQ-dVplI0hh2WCSugjc1luE0J2UAafjZxu_7Xuc,4367
 sglang/srt/openai_api/adapter.py,sha256=eirFYVGIp5D-UrQLqW5dRJOQYKmzF9nmgCzFeUOb2z8,15737
-sglang/srt/openai_api/api_adapter.py,sha256=eirFYVGIp5D-UrQLqW5dRJOQYKmzF9nmgCzFeUOb2z8,15737
-sglang/srt/openai_api/openai_api_adapter.py,sha256=5pDaktIEteHxp3qN89U_U3ndd7N0FIfUZAM06YeziUY,15687
-sglang/srt/openai_api/openai_protocol.py,sha256=lGBhfxG6jmgUkMOh2NpBK9w9TUTRZKrsfHdW7XYhKKI,5700
 sglang/srt/openai_api/protocol.py,sha256=j7ifIR2SFQxTwaHAd9ksM096vfffcNltzTH4sg7H0RA,5739
 sglang/test/test_conversation.py,sha256=gF_AyOxQgpPQBPnA57-kq-M0p_zFu-rBDMFgAq655Rw,1596
 sglang/test/test_openai_protocol.py,sha256=DVx3r6hrb8oRqbo5AYIleldxbqMBTtb-gtORM6t_Y1c,1661
 sglang/test/test_programs.py,sha256=uefeHUFKT2NJESOujj-CsnPXdw1aQQN2TzUbPCHJjGs,13654
 sglang/test/test_utils.py,sha256=kD_fQe3WroZ9Kc3NBRKPiZOFJ_JD2uEE9XIvPp6AD9Y,11048
-sglang-0.1.24.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-sglang-0.1.24.dist-info/METADATA,sha256=_HKFljParVedu-eht7OKKb_RpEkVcB-Wh_P_jRW3TJk,30933
-sglang-0.1.24.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
-sglang-0.1.24.dist-info/top_level.txt,sha256=yxhh3pYQkcnA7v3Bg889C2jZhvtJdEincysO7PEB09M,7
-sglang-0.1.24.dist-info/RECORD,,
+sglang-0.1.26.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+sglang-0.1.26.dist-info/METADATA,sha256=QnzTK6blFTHKTDw9ULRpaJVvXyg0MuzkdqwYkk0zPb0,30986
+sglang-0.1.26.dist-info/WHEEL,sha256=Wyh-_nZ0DJYolHNn1_hMa4lM7uDedD_RGVwbmTjyItk,91
+sglang-0.1.26.dist-info/top_level.txt,sha256=yxhh3pYQkcnA7v3Bg889C2jZhvtJdEincysO7PEB09M,7
+sglang-0.1.26.dist-info/RECORD,,

sglang/backend/anthropic.py

@@ -1,77 +0,0 @@
-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 anthropic
-except ImportError as e:
-    anthropic = e
-
-
-class Anthropic(BaseBackend):
-    def __init__(self, model_name, *args, **kwargs):
-        super().__init__()
-
-        if isinstance(anthropic, Exception):
-            raise anthropic
-
-        self.model_name = model_name
-        self.chat_template = get_chat_template("claude")
-        self.client = anthropic.Anthropic(*args, **kwargs)
-
-    def get_chat_template(self):
-        return self.chat_template
-
-    def generate(
-        self,
-        s: StreamExecutor,
-        sampling_params: SglSamplingParams,
-    ):
-        if s.messages_:
-            messages = s.messages_
-        else:
-            messages = [{"role": "user", "content": s.text_}]
-
-        if messages and messages[0]["role"] == "system":
-            system = messages.pop(0)["content"]
-        else:
-            system = ""
-
-        ret = self.client.messages.create(
-            model=self.model_name,
-            system=system,
-            messages=messages,
-            **sampling_params.to_anthropic_kwargs(),
-        )
-        comp = ret.content[0].text
-
-        return comp, {}
-
-    def generate_stream(
-        self,
-        s: StreamExecutor,
-        sampling_params: SglSamplingParams,
-    ):
-        if s.messages_:
-            messages = s.messages_
-        else:
-            messages = [{"role": "user", "content": s.text_}]
-
-        if messages and messages[0]["role"] == "system":
-            system = messages.pop(0)["content"]
-        else:
-            system = ""
-
-        with self.client.messages.stream(
-            model=self.model_name,
-            system=system,
-            messages=messages,
-            **sampling_params.to_anthropic_kwargs(),
-        ) as stream:
-            for text in stream.text_stream:
-                yield text, {}

sglang/backend/base_backend.py

@@ -1,80 +0,0 @@
-from typing import Callable, List, Optional, Union
-
-from sglang.lang.chat_template import get_chat_template
-from sglang.lang.interpreter import StreamExecutor
-from sglang.lang.ir import SglSamplingParams
-
-
-class BaseBackend:
-    def __init__(self) -> None:
-        self.support_concate_and_append = False
-        self.chat_template = get_chat_template("default")
-
-    def get_model_name(self):
-        raise NotImplementedError()
-
-    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: Union[str, List[str]]):
-        pass
-
-    def begin_program(self, s: StreamExecutor):
-        pass
-
-    def end_program(self, s: Union[StreamExecutor, List[StreamExecutor]]):
-        pass
-
-    def commit_lazy_operations(self, s: StreamExecutor):
-        pass
-
-    def fork_program(
-        self,
-        src: StreamExecutor,
-        dst: List[StreamExecutor],
-        position_ids_offset: Optional[List[int]] = None,
-    ):
-        pass
-
-    def fill_image(self, s: StreamExecutor):
-        pass
-
-    def generate(
-        self,
-        s: StreamExecutor,
-        sampling_params: SglSamplingParams,
-    ):
-        raise NotImplementedError()
-
-    def generate_stream(
-        self,
-        s: StreamExecutor,
-        sampling_params: SglSamplingParams,
-    ):
-        raise NotImplementedError()
-
-    def select(
-        self,
-        s: StreamExecutor,
-        choices: List[str],
-        temperature: float,
-    ):
-        raise NotImplementedError()
-
-    def concatenate_and_append(self, src_rids: List[str], dst_rid: str):
-        raise NotImplementedError()
-
-    def shutdown(self):
-        pass
-
-    def flush_cache(self):
-        pass
-
-    def get_server_args(self):
-        pass

sglang/backend/litellm.py

@@ -1,90 +0,0 @@
-from typing import Mapping, Optional
-
-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
-
-try:
-    import litellm
-except ImportError as e:
-    litellm = e
-    litellm.num_retries = 1
-
-
-class LiteLLM(BaseBackend):
-    def __init__(
-        self,
-        model_name,
-        chat_template=None,
-        api_key=None,
-        organization: Optional[str] = None,
-        base_url: Optional[str] = None,
-        timeout: Optional[float] = 600,
-        max_retries: Optional[int] = litellm.num_retries,
-        default_headers: Optional[Mapping[str, str]] = None,
-    ):
-        super().__init__()
-
-        if isinstance(litellm, Exception):
-            raise litellm
-
-        self.model_name = model_name
-
-        self.chat_template = chat_template or get_chat_template_by_model_path(
-            model_name
-        )
-
-        self.client_params = {
-            "api_key": api_key,
-            "organization": organization,
-            "base_url": base_url,
-            "timeout": timeout,
-            "max_retries": max_retries,
-            "default_headers": default_headers,
-        }
-
-    def get_chat_template(self):
-        return self.chat_template
-
-    def generate(
-        self,
-        s: StreamExecutor,
-        sampling_params: SglSamplingParams,
-    ):
-        if s.messages_:
-            messages = s.messages_
-        else:
-            messages = [{"role": "user", "content": s.text_}]
-
-        ret = litellm.completion(
-            model=self.model_name,
-            messages=messages,
-            **self.client_params,
-            **sampling_params.to_anthropic_kwargs(),
-        )
-        comp = ret.choices[0].message.content
-
-        return comp, {}
-
-    def generate_stream(
-        self,
-        s: StreamExecutor,
-        sampling_params: SglSamplingParams,
-    ):
-        if s.messages_:
-            messages = s.messages_
-        else:
-            messages = [{"role": "user", "content": s.text_}]
-
-        ret = litellm.completion(
-            model=self.model_name,
-            messages=messages,
-            stream=True,
-            **self.client_params,
-            **sampling_params.to_litellm_kwargs(),
-        )
-        for chunk in ret:
-            text = chunk.choices[0].delta.content
-            if text is not None:
-                yield text, {}