chore: import upstream snapshot with attribution

python/sglang/lang/api.py

@@ -0,0 +1,292 @@
+"""Public APIs of the language."""
+
+import re
+from typing import Callable, List, Optional, Union
+
+from sglang.global_config import global_config
+from sglang.lang.backend.base_backend import BaseBackend
+from sglang.lang.choices import ChoicesSamplingMethod, token_length_normalized
+from sglang.lang.ir import (
+    SglExpr,
+    SglExprList,
+    SglFunction,
+    SglGen,
+    SglImage,
+    SglRoleBegin,
+    SglRoleEnd,
+    SglSelect,
+    SglSeparateReasoning,
+    SglVideo,
+)
+
+
+def function(
+    func: Optional[Callable] = None, num_api_spec_tokens: Optional[int] = None
+):
+    if func:
+        return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens)
+
+    def decorator(func):
+        return SglFunction(func, num_api_spec_tokens=num_api_spec_tokens)
+
+    return decorator
+
+
+def Runtime(*args, **kwargs):
+    # Avoid importing unnecessary dependency
+    from sglang.lang.backend.runtime_endpoint import Runtime
+
+    return Runtime(*args, **kwargs)
+
+
+def Engine(*args, **kwargs):
+    # Avoid importing unnecessary dependency
+    from sglang.srt.entrypoints.engine import Engine
+
+    return Engine(*args, **kwargs)
+
+
+def set_default_backend(backend: BaseBackend):
+    global_config.default_backend = backend
+
+
+def flush_cache(backend: Optional[BaseBackend] = None):
+    backend = backend or global_config.default_backend
+    if backend is None:
+        return False
+
+    # If backend is Runtime
+    if hasattr(backend, "endpoint"):
+        backend = backend.endpoint
+    return backend.flush_cache()
+
+
+def get_server_info(backend: Optional[BaseBackend] = None):
+    backend = backend or global_config.default_backend
+    if backend is None:
+        return None
+
+    # If backend is Runtime
+    if hasattr(backend, "endpoint"):
+        backend = backend.endpoint
+    return backend.get_server_info()
+
+
+def gen(
+    name: Optional[str] = None,
+    max_tokens: Optional[int] = None,
+    min_tokens: Optional[int] = None,
+    n: Optional[int] = None,
+    stop: Optional[Union[str, List[str]]] = None,
+    stop_token_ids: Optional[List[int]] = None,
+    stop_regex: Optional[Union[str, List[str]]] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    min_p: Optional[float] = None,
+    frequency_penalty: Optional[float] = None,
+    presence_penalty: Optional[float] = None,
+    ignore_eos: Optional[bool] = None,
+    return_logprob: Optional[bool] = None,
+    logprob_start_len: Optional[int] = None,
+    top_logprobs_num: Optional[int] = None,
+    return_text_in_logprobs: Optional[bool] = None,
+    dtype: Optional[Union[type, str]] = None,
+    choices: Optional[List[str]] = None,
+    choices_method: Optional[ChoicesSamplingMethod] = None,
+    regex: Optional[str] = None,
+    json_schema: Optional[str] = None,
+):
+    """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md"""
+
+    if choices:
+        return SglSelect(
+            name,
+            choices,
+            0.0 if temperature is None else temperature,
+            token_length_normalized if choices_method is None else choices_method,
+        )
+
+    # check regex is valid
+    if regex is not None:
+        try:
+            re.compile(regex)
+        except re.error as e:
+            raise e
+
+    return SglGen(
+        name,
+        max_tokens,
+        min_tokens,
+        n,
+        stop,
+        stop_token_ids,
+        stop_regex,
+        temperature,
+        top_p,
+        top_k,
+        min_p,
+        frequency_penalty,
+        presence_penalty,
+        ignore_eos,
+        return_logprob,
+        logprob_start_len,
+        top_logprobs_num,
+        return_text_in_logprobs,
+        dtype,
+        regex,
+        json_schema,
+    )
+
+
+def gen_int(
+    name: Optional[str] = None,
+    max_tokens: Optional[int] = None,
+    n: Optional[int] = None,
+    stop: Optional[Union[str, List[str]]] = None,
+    stop_token_ids: Optional[List[int]] = None,
+    stop_regex: Optional[Union[str, List[str]]] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    min_p: Optional[float] = None,
+    frequency_penalty: Optional[float] = None,
+    presence_penalty: Optional[float] = None,
+    ignore_eos: Optional[bool] = None,
+    return_logprob: Optional[bool] = None,
+    logprob_start_len: Optional[int] = None,
+    top_logprobs_num: Optional[int] = None,
+    return_text_in_logprobs: Optional[bool] = None,
+):
+    return SglGen(
+        name,
+        max_tokens,
+        None,
+        n,
+        stop,
+        stop_token_ids,
+        stop_regex,
+        temperature,
+        top_p,
+        top_k,
+        min_p,
+        frequency_penalty,
+        presence_penalty,
+        ignore_eos,
+        return_logprob,
+        logprob_start_len,
+        top_logprobs_num,
+        return_text_in_logprobs,
+        int,
+        None,
+    )
+
+
+def gen_string(
+    name: Optional[str] = None,
+    max_tokens: Optional[int] = None,
+    n: Optional[int] = None,
+    stop: Optional[Union[str, List[str]]] = None,
+    stop_token_ids: Optional[List[int]] = None,
+    stop_regex: Optional[Union[str, List[str]]] = None,
+    temperature: Optional[float] = None,
+    top_p: Optional[float] = None,
+    top_k: Optional[int] = None,
+    min_p: Optional[float] = None,
+    frequency_penalty: Optional[float] = None,
+    presence_penalty: Optional[float] = None,
+    ignore_eos: Optional[bool] = None,
+    return_logprob: Optional[bool] = None,
+    logprob_start_len: Optional[int] = None,
+    top_logprobs_num: Optional[int] = None,
+    return_text_in_logprobs: Optional[bool] = None,
+):
+    return SglGen(
+        name,
+        max_tokens,
+        None,
+        n,
+        stop,
+        stop_token_ids,
+        stop_regex,
+        temperature,
+        top_p,
+        top_k,
+        min_p,
+        frequency_penalty,
+        presence_penalty,
+        ignore_eos,
+        return_logprob,
+        logprob_start_len,
+        top_logprobs_num,
+        return_text_in_logprobs,
+        str,
+        None,
+    )
+
+
+def image(expr: SglExpr):
+    return SglImage(expr)
+
+
+def video(path: str, num_frames: int):
+    return SglVideo(path, num_frames)
+
+
+def select(
+    name: Optional[str] = None,
+    choices: Optional[List[str]] = None,
+    temperature: float = 0.0,
+    choices_method: ChoicesSamplingMethod = token_length_normalized,
+):
+    assert choices is not None
+    return SglSelect(name, choices, temperature, choices_method)
+
+
+def _role_common(name: str, expr: Optional[SglExpr] = None):
+    if expr is None:
+        return SglExprList([SglRoleBegin(name), SglRoleEnd(name)])
+    else:
+        return SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)])
+
+
+def system(expr: Optional[SglExpr] = None):
+    return _role_common("system", expr)
+
+
+def user(expr: Optional[SglExpr] = None):
+    return _role_common("user", expr)
+
+
+def assistant(expr: Optional[SglExpr] = None):
+    return _role_common("assistant", expr)
+
+
+def system_begin():
+    return SglRoleBegin("system")
+
+
+def system_end():
+    return SglRoleEnd("system")
+
+
+def user_begin():
+    return SglRoleBegin("user")
+
+
+def user_end():
+    return SglRoleEnd("user")
+
+
+def assistant_begin():
+    return SglRoleBegin("assistant")
+
+
+def assistant_end():
+    return SglRoleEnd("assistant")
+
+
+def separate_reasoning(
+    expr: Optional[SglExpr] = None, model_type: Optional[str] = None
+):
+    return SglExprList([expr, SglSeparateReasoning(model_type, expr=expr)])

python/sglang/lang/backend/anthropic.py

@@ -0,0 +1,73 @@
+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 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, {}

python/sglang/lang/backend/base_backend.py

@@ -0,0 +1,82 @@
+from typing import List, Optional, Union
+
+from sglang.lang.chat_template import get_chat_template
+from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod
+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,
+        choices_method: Optional[ChoicesSamplingMethod] = None,
+    ) -> ChoicesDecision:
+        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_info(self):
+        pass

python/sglang/lang/backend/crusoe.py

@@ -0,0 +1,43 @@
+import os
+from typing import Optional
+
+from sglang.lang.backend.openai import OpenAI
+from sglang.lang.chat_template import ChatTemplate
+
+CRUSOE_BASE_URL = "https://managed-inference-api-proxy.crusoecloud.com/v1/"
+
+
+class Crusoe(OpenAI):
+    """SGLang backend for Crusoe managed inference.
+
+    Crusoe exposes an OpenAI-compatible API, so this is a thin wrapper
+    around the OpenAI backend that handles Crusoe-specific defaults.
+
+    Args:
+        model_name: The model to use, e.g. "meta-llama/Llama-3.1-8B-Instruct".
+        api_key: Crusoe API key. Defaults to CRUSOE_API_KEY env var.
+        base_url: Override the Crusoe endpoint. Defaults to the Crusoe API.
+        chat_template: Optional custom chat template.
+    """
+
+    def __init__(
+        self,
+        model_name: str,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        chat_template: Optional[ChatTemplate] = None,
+        **kwargs,
+    ):
+        resolved_api_key = api_key or os.environ.get("CRUSOE_API_KEY")
+        if not resolved_api_key:
+            raise ValueError(
+                "Crusoe API key required. Pass api_key= or set CRUSOE_API_KEY."
+            )
+
+        super().__init__(
+            model_name=model_name,
+            chat_template=chat_template,
+            api_key=resolved_api_key,
+            base_url=base_url or CRUSOE_BASE_URL,
+            **kwargs,
+        )

python/sglang/lang/backend/litellm.py

@@ -0,0 +1,90 @@
+from typing import Mapping, Optional
+
+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
+
+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_litellm_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, {}

python/sglang/lang/backend/openai.py

@@ -0,0 +1,475 @@
+import dataclasses
+import logging
+import time
+import warnings
+from typing import List, Optional, Union
+
+import numpy as np
+
+from sglang.lang.backend.base_backend import BaseBackend
+from sglang.lang.chat_template import ChatTemplate, get_chat_template_by_model_path
+from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import SglSamplingParams
+
+try:
+    import openai
+    import tiktoken
+except ImportError as e:
+    openai = tiktoken = e
+
+
+logger = logging.getLogger(__name__)
+
+
+def create_logit_bias_int(tokenizer):
+    """Get logit bias for integer numbers."""
+    int_token_ids = []
+
+    tokens = tokenizer._mergeable_ranks
+    for token, token_id in tokens.items():
+        s = tokenizer.decode([token_id])
+        if all([c.isdigit() for c in s]) or s in [" "]:
+            int_token_ids.append(token_id)
+            if len(int_token_ids) >= 300:  # OpenAI API limit
+                break
+    special_tokens = tokenizer._special_tokens
+    mask = {t: 100 for t in int_token_ids[:299]}
+    mask[special_tokens["<|endoftext|>"]] = 100
+    return mask
+
+
+INSTRUCT_MODEL_NAMES = [
+    "gpt-3.5-turbo-instruct",
+]
+
+
+@dataclasses.dataclass
+class TokenUsage:
+    prompt_tokens: int
+    completion_tokens: int
+
+    def reset(self):
+        self.prompt_tokens = self.completion_tokens = 0
+
+
+class OpenAI(BaseBackend):
+    def __init__(
+        self,
+        model_name: str,
+        is_chat_model: Optional[bool] = None,
+        chat_template: Optional[ChatTemplate] = None,
+        is_azure: bool = False,
+        *args,
+        **kwargs,
+    ):
+        super().__init__()
+
+        if isinstance(openai, Exception):
+            raise openai
+
+        if is_azure:
+            self.client = openai.AzureOpenAI(*args, **kwargs)
+        else:
+            self.client = openai.OpenAI(*args, **kwargs)
+
+        self.model_name = model_name
+        try:
+            self.tokenizer = tiktoken.encoding_for_model(model_name)
+        except KeyError:
+            self.tokenizer = tiktoken.get_encoding("cl100k_base")
+        self.logit_bias_int = create_logit_bias_int(self.tokenizer)
+
+        self.chat_template = chat_template or get_chat_template_by_model_path(
+            model_name
+        )
+
+        if is_chat_model is not None:
+            self.is_chat_model = is_chat_model
+        else:
+            if model_name in INSTRUCT_MODEL_NAMES:
+                self.is_chat_model = False
+            else:
+                self.is_chat_model = True
+
+        self.chat_prefix = self.chat_template.role_prefix_and_suffix["assistant"][0]
+
+        # Usage
+        self.token_usage = TokenUsage(0, 0)
+
+        # API speculative execution
+        # TODO(ying): This does not support multi-threading (run_batch)
+        self.spec_kwargs = {}
+        self.spec_format = []
+        self.spec_max_num_tries = 3
+
+    def get_chat_template(self):
+        return self.chat_template
+
+    def _prepare_spec_execution(
+        self,
+        sampling_params: SglSamplingParams,
+        num_api_spec_tokens: int,
+        spec_var_name: str,
+    ):
+        if "max_tokens" not in self.spec_kwargs:
+            self.spec_kwargs["max_tokens"] = num_api_spec_tokens
+        else:
+            assert self.spec_kwargs["max_tokens"] == num_api_spec_tokens
+
+        params = sampling_params.to_openai_kwargs()
+        for key, value in params.items():
+            if key in ["stop"]:
+                continue
+            if key in ["max_tokens"]:
+                warnings.warn(
+                    "The parameter max_tokens will be overwritten by speculated number of tokens."
+                )
+                continue
+            if key not in self.spec_kwargs:
+                self.spec_kwargs[key] = value
+            else:
+                assert (
+                    value == self.spec_kwargs[key]
+                ), "sampling parameters should be consistent if turn on api speculative execution."
+        self.spec_format.append(
+            {"text": "", "stop": params["stop"], "name": spec_var_name}
+        )
+        return "", {}
+
+    def generate(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+        spec_var_name: str = None,
+    ):
+        if sampling_params.dtype is None:
+            if self.is_chat_model:
+                if s.num_api_spec_tokens is None:
+                    if not s.text_.endswith(self.chat_prefix):
+                        raise RuntimeError(
+                            "This use case is not supported if api speculative execution is off. "
+                            "For OpenAI chat models, sgl.gen must be right after sgl.assistant. "
+                            "Example of adding api speculative execution: @function(num_api_spec_tokens=128)."
+                        )
+                    prompt = s.messages_
+                else:
+                    return self._prepare_spec_execution(
+                        sampling_params, s.num_api_spec_tokens, spec_var_name
+                    )
+            else:
+                prompt = s.text_
+
+            kwargs = sampling_params.to_openai_kwargs()
+            if (
+                self.model_name.startswith("o1")
+                or self.model_name.startswith("o3")
+                or "o1" in self.model_name
+            ):
+                kwargs.pop("max_tokens", None)
+            else:
+                kwargs.pop("max_completion_tokens", None)
+
+            comp = openai_completion(
+                client=self.client,
+                token_usage=self.token_usage,
+                is_chat=self.is_chat_model,
+                model=self.model_name,
+                prompt=prompt,
+                **kwargs,
+            )
+            # Keep the returned list (or string) as is.
+        elif sampling_params.dtype in [str, "str", "string"]:
+            assert (
+                not self.is_chat_model
+            ), "constrained type not supported on chat model"
+            kwargs = sampling_params.to_openai_kwargs()
+            kwargs.pop("stop")
+            comp = openai_completion(
+                client=self.client,
+                token_usage=self.token_usage,
+                is_chat=self.is_chat_model,
+                model=self.model_name,
+                prompt=s.text_ + '"',
+                stop='"',
+                **kwargs,
+            )
+            # Wrap each element in quotes if we have a list.
+            if isinstance(comp, list):
+                comp = ['"' + x + '"' for x in comp]
+            else:
+                comp = '"' + comp + '"'
+        elif sampling_params.dtype in [int, "int"]:
+            assert (
+                not self.is_chat_model
+            ), "constrained type not supported on chat model"
+            kwargs = sampling_params.to_openai_kwargs()
+            kwargs.pop("stop")
+            comp = openai_completion(
+                client=self.client,
+                token_usage=self.token_usage,
+                is_chat=self.is_chat_model,
+                model=self.model_name,
+                prompt=s.text_,
+                logit_bias=self.logit_bias_int,
+                stop=[" "],
+                **kwargs,
+            )
+            # Leave as a list if that's what is returned.
+        else:
+            raise ValueError(f"Unknown dtype: {sampling_params.dtype}")
+
+        return comp, {}
+
+    def spec_fill(self, value: str):
+        assert self.is_chat_model
+        self.spec_format.append({"text": value, "stop": None, "name": None})
+
+    def spec_pattern_match(self, comp):
+        for i, term in enumerate(self.spec_format):
+            text = term["text"]
+            if text != "":
+                if comp.startswith(text):
+                    comp = comp[len(text) :]
+                else:
+                    return False
+            else:
+                pos = comp.find(term["stop"])
+                if pos != -1:
+                    term["text"] = comp[:pos]
+                    comp = comp[pos:]
+                else:
+                    if i == len(self.spec_format) - 1:
+                        term["text"] = comp
+                    else:
+                        return False
+        return True
+
+    def role_end_generate(
+        self,
+        s: StreamExecutor,
+    ):
+        if s.num_api_spec_tokens is None or not s.text_.endswith(self.chat_prefix):
+            return
+
+        comp = ""
+        if not all(x["name"] is None for x in self.spec_format):
+            # TODO(ying): throw errors or warnings
+            for i in range(self.spec_max_num_tries):
+                comp = openai_completion(
+                    client=self.client,
+                    token_usage=self.token_usage,
+                    is_chat=self.is_chat_model,
+                    model=self.model_name,
+                    prompt=s.messages_,
+                    **self.spec_kwargs,
+                )
+                # Use a string for pattern matching.
+                comp_for_match = comp[0] if isinstance(comp, list) else comp
+                if self.spec_pattern_match(comp_for_match):
+                    break
+
+        for term in self.spec_format:
+            s.text_ += term["text"]
+            name = term["name"]
+            if name is not None:
+                s.variables[name] = term["text"]
+                s.meta_info[name] = {}
+                s.variable_event[name].set()
+
+        self.spec_kwargs = {}
+        self.spec_format = []
+
+    def generate_stream(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+    ):
+        if sampling_params.dtype is None:
+            if self.is_chat_model:
+                if not s.text_.endswith(self.chat_prefix):
+                    raise RuntimeError(
+                        "This use case is not supported. "
+                        "For OpenAI chat models, sgl.gen must be right after sgl.assistant"
+                    )
+                prompt = s.messages_
+            else:
+                prompt = s.text_
+
+            kwargs = sampling_params.to_openai_kwargs()
+            generator = openai_completion_stream(
+                client=self.client,
+                token_usage=self.token_usage,
+                is_chat=self.is_chat_model,
+                model=self.model_name,
+                prompt=prompt,
+                **kwargs,
+            )
+            return generator
+        else:
+            raise ValueError(f"Unknown dtype: {sampling_params.dtype}")
+
+    def select(
+        self,
+        s: StreamExecutor,
+        choices: List[str],
+        temperature: float,
+        choices_method: ChoicesSamplingMethod,
+    ) -> ChoicesDecision:
+        """Note: `choices_method` is not used by the OpenAI backend."""
+        if self.is_chat_model:
+            raise NotImplementedError(
+                "select/choices is not supported for chat models. "
+                "Please try to use a non-chat model such as gpt-3.5-turbo-instruct"
+            )
+
+        n_choices = len(choices)
+        token_ids = [self.tokenizer.encode(x) for x in choices]
+        scores = [0] * n_choices
+        valid = [len(x) > 0 for x in token_ids]
+        prompt_tokens = self.tokenizer.encode(s.text_)
+
+        max_len = max([len(x) for x in token_ids])
+        for step in range(max_len):
+            # Build logit bias
+            logit_bias = {}
+            for i in range(n_choices):
+                if valid[i]:
+                    logit_bias[token_ids[i][step]] = 100
+
+            # Call API
+            ret = self.client.completions.create(
+                model=self.model_name,
+                prompt=prompt_tokens,
+                logit_bias=logit_bias,
+                max_tokens=1,
+                temperature=temperature,
+            )
+            ret_str = ret.choices[0].text
+            ret_token = self.tokenizer.encode(ret_str)[0]
+            self.token_usage.prompt_tokens += ret.usage.prompt_tokens
+            self.token_usage.completion_tokens = ret.usage.completion_tokens
+
+            # TODO:
+            # 1. return logits as the scores
+            # 2. compute logits of the full choice
+            # 3. consider chunk-based decoding
+
+            # Update valid
+            hit = False
+            for i in range(n_choices):
+                if valid[i]:
+                    if step == len(token_ids[i]) - 1:
+                        valid[i] = False
+
+                    if ret_token == token_ids[i][step]:
+                        scores[i] += 1
+                        hit = True
+                    else:
+                        valid[i] = False
+            assert hit
+
+            if np.sum(valid) <= 1:
+                break
+
+            prompt_tokens.append(ret_token)
+
+        return ChoicesDecision(
+            decision=choices[np.argmax(scores)],
+            meta_info={"scores": scores},
+        )
+
+
+def openai_completion(
+    client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs
+) -> Union[str, List[str]]:
+    # if "ebnf" is in kwargs, warn and remove
+    if "ebnf" in kwargs:
+        warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.")
+        del kwargs["ebnf"]
+
+    for attempt in range(retries):
+        try:
+            if is_chat:
+                if "stop" in kwargs and kwargs["stop"] is None:
+                    kwargs.pop("stop")
+                ret = client.chat.completions.create(messages=prompt, **kwargs)
+                if len(ret.choices) == 1:
+                    comp = ret.choices[0].message.content
+                else:
+                    comp = [c.message.content for c in ret.choices]
+            else:
+                ret = client.completions.create(prompt=prompt, **kwargs)
+                if isinstance(prompt, (list, tuple)):
+                    comp = [c.text for c in ret.choices]
+                else:
+                    comp = ret.choices[0].text
+                    if len(ret.choices) > 1:
+                        comp = [c.text for c in ret.choices]
+
+            token_usage.prompt_tokens += ret.usage.prompt_tokens
+            token_usage.completion_tokens += ret.usage.completion_tokens
+            break
+        except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e:
+            logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...")
+            time.sleep(5)
+            if attempt == retries - 1:
+                raise e
+        except Exception as e:
+            logger.error(f"RuntimeError {e}.")
+            raise e
+
+    return comp
+
+
+def openai_completion_stream(
+    client, token_usage, is_chat=None, retries=3, prompt=None, **kwargs
+):
+    # if "ebnf" is in kwargs, warn and remove
+    if "ebnf" in kwargs:
+        warnings.warn("EBNF is not officially supported by OpenAI endpoints. Ignoring.")
+        del kwargs["ebnf"]
+
+    for attempt in range(retries):
+        try:
+            if is_chat:
+                if "stop" in kwargs and kwargs["stop"] is None:
+                    kwargs.pop("stop")
+                generator = client.chat.completions.create(
+                    messages=prompt,
+                    stream=True,
+                    stream_options={"include_usage": True},
+                    **kwargs,
+                )
+                for ret in generator:
+                    if len(ret.choices) == 0:
+                        continue
+                    try:
+                        content = ret.choices[0].delta.content
+                    except IndexError:
+                        content = None
+                    yield content or "", {}
+            else:
+                generator = client.completions.create(
+                    prompt=prompt,
+                    stream=True,
+                    stream_options={"include_usage": True},
+                    **kwargs,
+                )
+                for ret in generator:
+                    if len(ret.choices) == 0:
+                        continue
+                    content = ret.choices[0].text
+                    yield content or "", {}
+
+            token_usage.prompt_tokens += ret.usage.prompt_tokens
+            token_usage.completion_tokens += ret.usage.completion_tokens
+            break
+        except (openai.APIError, openai.APIConnectionError, openai.RateLimitError) as e:
+            logger.error(f"OpenAI Error: {e}. Waiting 5 seconds...")
+            time.sleep(5)
+            if attempt == retries - 1:
+                raise e
+        except Exception as e:
+            logger.error(f"RuntimeError {e}.")
+            raise e

python/sglang/lang/backend/runtime_endpoint.py

@@ -0,0 +1,549 @@
+import atexit
+import json
+import multiprocessing
+import time
+import warnings
+from typing import Dict, List, Optional, Union
+
+import aiohttp
+import requests
+
+from sglang.global_config import global_config
+from sglang.lang.backend.base_backend import BaseBackend
+from sglang.lang.chat_template import get_chat_template, get_chat_template_by_model_path
+from sglang.lang.choices import ChoicesDecision, ChoicesSamplingMethod
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import (
+    REGEX_BOOL,
+    REGEX_FLOAT,
+    REGEX_INT,
+    REGEX_STR,
+    SglSamplingParams,
+)
+from sglang.utils import http_request
+
+
+class RuntimeEndpoint(BaseBackend):
+    def __init__(
+        self,
+        base_url: str,
+        api_key: Optional[str] = None,
+        verify: Optional[str] = None,
+        chat_template_name: Optional[str] = None,
+    ):
+        super().__init__()
+        self.support_concate_and_append = True
+
+        self.base_url = base_url
+        self.api_key = api_key
+        self.verify = verify
+
+        res = http_request(
+            self.base_url + "/get_model_info",
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+        self.model_info = res.json()
+
+        if chat_template_name:
+            self.chat_template = get_chat_template(chat_template_name)
+        else:
+            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",
+            api_key=self.api_key,
+            verify=self.verify,
+            method="POST",
+        )
+        self._assert_success(res)
+
+    def get_server_info(self):
+        res = http_request(
+            self.base_url + "/server_info",
+            api_key=self.api_key,
+            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}},
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def start_profile(self):
+        res = http_request(
+            self.base_url + "/start_profile",
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def stop_profile(self):
+        res = http_request(
+            self.base_url + "/stop_profile",
+            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,
+            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,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def _handle_dtype_to_regex(self, sampling_params: SglSamplingParams):
+        if sampling_params.dtype is None:
+            return
+
+        if sampling_params.stop == ():
+            sampling_params.stop = []
+
+        dtype_regex = None
+        if sampling_params.dtype in ["int", int]:
+
+            dtype_regex = REGEX_INT
+            sampling_params.stop.extend([" ", "\n"])
+        elif sampling_params.dtype in ["float", float]:
+
+            dtype_regex = REGEX_FLOAT
+            sampling_params.stop.extend([" ", "\n"])
+        elif sampling_params.dtype in ["str", str]:
+
+            dtype_regex = REGEX_STR
+        elif sampling_params.dtype in ["bool", bool]:
+
+            dtype_regex = REGEX_BOOL
+        else:
+            raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}")
+
+        if dtype_regex is not None and sampling_params.regex is not None:
+            warnings.warn(
+                f"Both dtype and regex are set. Only dtype will be used. dtype: {sampling_params.dtype}, regex: {sampling_params.regex}"
+            )
+
+        sampling_params.regex = dtype_regex
+
+    def generate(
+        self,
+        s: StreamExecutor,
+        sampling_params: SglSamplingParams,
+    ):
+        self._handle_dtype_to_regex(sampling_params)
+        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(),
+            },
+        }
+
+        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,
+            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,
+    ):
+        self._handle_dtype_to_regex(sampling_params)
+
+        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(),
+            },
+        }
+
+        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,
+            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,
+        choices_method: ChoicesSamplingMethod,
+    ) -> ChoicesDecision:
+        assert temperature <= 1e-5
+
+        # Cache common prefix
+        data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
+        obj = self._generate_http_request(s, data)
+        prompt_len = obj["meta_info"]["prompt_tokens"]
+        logprob_start_len = max(prompt_len - 2, 0)  # For token healing
+
+        # Compute logprob
+        data = {
+            "text": [s.text_ + c for c in choices],
+            "sampling_params": {
+                "max_new_tokens": 0,
+                "temperature": 0,
+            },
+            "return_logprob": True,
+            "return_text_in_logprobs": True,
+            "logprob_start_len": logprob_start_len,
+        }
+        obj = self._generate_http_request(s, data)
+
+        input_token_logprobs = [r["meta_info"]["input_token_logprobs"] for r in obj]
+        output_token_logprobs = [r["meta_info"]["output_token_logprobs"] for r in obj]
+        normalized_prompt_logprobs = [
+            compute_normalized_prompt_logprobs(r["meta_info"]["input_token_logprobs"])
+            for r in obj
+        ]
+
+        # Remove extra token if no token healing occurred
+        for i in range(len(input_token_logprobs)):
+            healed_token_str = input_token_logprobs[i][0][-1]
+            if s.text_.endswith(healed_token_str):
+                healed_token_logprob = input_token_logprobs[i][0][0]
+                normalized_prompt_logprobs[i] = (
+                    normalized_prompt_logprobs[i] * len(input_token_logprobs[i])
+                    - healed_token_logprob
+                ) / (len(input_token_logprobs[i]) - 1)
+                input_token_logprobs[i] = input_token_logprobs[i][1:]
+
+        # Compute unconditional logprobs if required
+        if choices_method.requires_unconditional_logprobs:
+            input_ids = [[el[1] for el in subl] for subl in input_token_logprobs]
+            data = {
+                "input_ids": input_ids,
+                "sampling_params": {"max_new_tokens": 0},
+                "return_logprob": True,
+            }
+            obj = self._generate_http_request(s, data)
+            unconditional_token_logprobs = [
+                r["meta_info"]["input_token_logprobs"] for r in obj
+            ]
+        else:
+            unconditional_token_logprobs = None
+
+        return choices_method(
+            choices=choices,
+            normalized_prompt_logprobs=normalized_prompt_logprobs,
+            input_token_logprobs=input_token_logprobs,
+            output_token_logprobs=output_token_logprobs,
+            unconditional_token_logprobs=unconditional_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},
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+
+    def _generate_http_request(self, s: StreamExecutor, data):
+        self._add_images(s, data)
+        res = http_request(
+            self.base_url + "/generate",
+            json=data,
+            api_key=self.api_key,
+            verify=self.verify,
+        )
+        self._assert_success(res)
+        return res.json()
+
+    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:
+            try:
+                content = res.json()
+            except json.JSONDecodeError:
+                content = res.text
+            raise RuntimeError(content)
+
+
+def compute_normalized_prompt_logprobs(input_logprobs):
+    values = [x[0] for x in input_logprobs if x[0]]
+    return sum(values) / len(values)
+
+
+class Runtime:
+    """
+    A wrapper for the HTTP server.
+    This is used for launching the server in a python program without
+    using the command line interface.
+
+    It is mainly used for the frontend language.
+    You should use the Engine class if you want to do normal offline processing without the frontend language.
+    """
+
+    def __init__(
+        self,
+        log_level: str = "error",
+        launch_timeout: float = 300.0,
+        *args,
+        **kwargs,
+    ):
+        """See the arguments in server_args.py::ServerArgs
+
+        Args:
+            log_level: Log level for the server.
+            timeout: Timeout in seconds for waiting for the server to start.
+            *args: Additional arguments passed to ServerArgs.
+            **kwargs: Additional keyword arguments passed to ServerArgs.
+        """
+        # We delay the import of any `sglang.srt` components in `sglang.lang`, so users can run
+        # client code without installing SRT server and its dependency if they want.
+        from sglang.srt.entrypoints.http_server import launch_server
+        from sglang.srt.server_args import ServerArgs
+        from sglang.srt.utils.network import is_port_available
+
+        self.server_args = ServerArgs(*args, log_level=log_level, **kwargs)
+
+        # Pre-allocate ports
+        for port in range(self.server_args.port, 40000):
+            if is_port_available(port):
+                break
+        self.server_args.override("runtime_endpoint.port_alloc", port=port)
+
+        self.url = self.server_args.url()
+        self.generate_url = self.url + "/generate"
+
+        # NOTE: We store pid instead of proc to fix some issues during __delete__
+        self.pid = None
+
+        ctx = multiprocessing.get_context("spawn")
+        proc = ctx.Process(
+            target=launch_server,
+            args=(self.server_args,),
+        )
+        proc.start()
+        self.pid = proc.pid
+
+        # Before python program terminates, call shutdown implicitly. Therefore, users don't have to explicitly call .shutdown()
+        atexit.register(self.shutdown)
+
+        # Wait for server to be ready by polling /health_generate
+        start_time = time.time()
+        with requests.Session() as session:
+            while time.time() - start_time < launch_timeout:
+                try:
+                    response = session.get(f"{self.url}/health_generate")
+                    if response.status_code == 200:
+                        break
+                except requests.RequestException:
+                    pass
+
+                if not proc.is_alive():
+                    self.shutdown()
+                    raise RuntimeError(
+                        "Initialization failed. Please see the error messages above."
+                    )
+
+                time.sleep(2)
+            else:
+                self.shutdown()
+                raise TimeoutError("Server failed to start within the timeout period.")
+
+        self.endpoint = RuntimeEndpoint(self.url)
+
+    def shutdown(self):
+        from sglang.srt.utils import kill_process_tree
+
+        if self.pid is not None:
+            kill_process_tree(self.pid)
+            self.pid = None
+
+    def start_profile(self):
+        self.endpoint.start_profile()
+
+    def stop_profile(self):
+        self.endpoint.stop_profile()
+
+    def cache_prefix(self, prefix: str):
+        self.endpoint.cache_prefix(prefix)
+
+    def get_tokenizer(self):
+        from sglang.srt.utils.hf_transformers_utils import get_tokenizer
+
+        return get_tokenizer(
+            self.server_args.tokenizer_path,
+            tokenizer_mode=self.server_args.tokenizer_mode,
+            trust_remote_code=self.server_args.trust_remote_code,
+            revision=self.server_args.revision,
+        )
+
+    async def async_generate(
+        self,
+        prompt: str,
+        sampling_params: Optional[Dict] = None,
+        session_id: Optional[str] = None,
+    ):
+        if self.server_args.skip_tokenizer_init:
+            json_data = {
+                "input_ids": prompt,
+                "sampling_params": sampling_params,
+                "stream": True,
+                "session_id": session_id,
+            }
+        else:
+            json_data = {
+                "text": prompt,
+                "sampling_params": sampling_params,
+                "stream": True,
+                "session_id": session_id,
+            }
+        pos = 0
+
+        timeout = aiohttp.ClientTimeout(total=3 * 3600)
+        async with aiohttp.ClientSession(timeout=timeout, trust_env=True) as session:
+            async with session.post(self.generate_url, json=json_data) as response:
+                async for chunk, _ in response.content.iter_chunks():
+                    chunk = chunk.decode("utf-8")
+                    if chunk and chunk.startswith("data:"):
+                        if chunk == "data: [DONE]\n\n":
+                            break
+                        data = json.loads(chunk[5:].strip("\n"))
+                        if "text" in data:
+                            cur = data["text"][pos:]
+                            if cur:
+                                yield cur
+                            pos += len(cur)
+                        else:
+                            yield data
+
+    add_request = async_generate
+
+    def generate(
+        self,
+        prompt: Union[str, List[str]],
+        sampling_params: Optional[Dict] = None,
+        return_logprob: Optional[Union[List[bool], bool]] = False,
+        logprob_start_len: Optional[Union[List[int], int]] = None,
+        top_logprobs_num: Optional[Union[List[int], int]] = None,
+        lora_path: Optional[List[Optional[str]]] = None,
+        session_id: Optional[str] = None,
+    ):
+        json_data = {
+            "text": prompt,
+            "sampling_params": sampling_params,
+            "return_logprob": return_logprob,
+            "logprob_start_len": logprob_start_len,
+            "top_logprobs_num": top_logprobs_num,
+            "lora_path": lora_path,
+            "session_id": session_id,
+        }
+        assert not isinstance(lora_path, list) or len(lora_path) == len(prompt)
+        response = requests.post(
+            self.url + "/generate",
+            json=json_data,
+        )
+        return json.dumps(response.json())
+
+    def encode(
+        self,
+        prompt: Union[str, List[str], List[Dict], List[List[Dict]]],
+    ):
+        json_data = {"text": prompt}
+        response = requests.post(self.url + "/encode", json=json_data)
+        return json.dumps(response.json())
+
+    async def get_server_info(self):
+        async with aiohttp.ClientSession() as session:
+            async with session.get(f"{self.url}/server_info") as response:
+                if response.status == 200:
+                    return await response.json()
+                else:
+                    error_data = await response.json()
+                    raise RuntimeError(
+                        f"Failed to get server info. {error_data['error']['message']}"
+                    )
+
+    def __del__(self):
+        self.shutdown()

python/sglang/lang/backend/vertexai.py

@@ -0,0 +1,148 @@
+import os
+import warnings
+
+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

python/sglang/lang/chat_template.py

@@ -0,0 +1,679 @@
+import re
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import Callable, Dict, List, Tuple
+
+
+class ChatTemplateStyle(Enum):
+    PLAIN = auto()
+    LLAMA2 = auto()
+
+
+@dataclass
+class ChatTemplate:
+    name: str
+    default_system_prompt: str
+    role_prefix_and_suffix: Dict[str, Tuple[str, str]]
+    stop_str: List[str] = ()
+    image_token: str = "<image>"
+    audio_token: str = "<audio>"
+    style: ChatTemplateStyle = ChatTemplateStyle.PLAIN
+
+    def get_prefix_and_suffix(
+        self, role: str, hist_messages: List[Dict]
+    ) -> Tuple[str, str]:
+        prefix, suffix = self.role_prefix_and_suffix.get(role, ("", ""))
+
+        if self.style == ChatTemplateStyle.LLAMA2:
+            if role == "system" and not hist_messages:
+                user_prefix, _ = self.role_prefix_and_suffix.get("user", ("", ""))
+                system_prefix, system_suffix = self.role_prefix_and_suffix.get(
+                    "system", ("", "")
+                )
+                return (user_prefix + system_prefix, system_suffix)
+            elif (
+                role == "user"
+                and len(hist_messages) == 1
+                and hist_messages[0]["content"] is not None
+            ):
+                return ("", suffix)
+
+        return prefix, suffix
+
+    def get_prompt(self, messages: List[Dict]) -> str:
+        prompt = ""
+        for i, message in enumerate(messages):
+            role, content = message["role"], message["content"]
+            if role == "system" and content is None:
+                content = self.default_system_prompt
+                if content is None:
+                    continue
+
+            prefix, suffix = self.get_prefix_and_suffix(role, messages[:i])
+            prompt += f"{prefix}{content}{suffix}"
+        return prompt
+
+
+chat_template_registry: Dict[str, ChatTemplate] = {}
+matching_function_registry: List[Callable] = []
+
+
+def register_chat_template(template):
+    chat_template_registry[template.name] = template
+
+
+def register_chat_template_matching_function(func):
+    matching_function_registry.append(func)
+
+
+def get_chat_template(name):
+    return chat_template_registry[name]
+
+
+def get_chat_template_by_model_path(model_path):
+    for matching_func in matching_function_registry:
+        template_name = matching_func(model_path)
+        if template_name is not None:
+            return get_chat_template(template_name)
+    return get_chat_template("default")
+
+
+register_chat_template(
+    ChatTemplate(
+        name="default",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("SYSTEM:", "\n"),
+            "user": ("USER:", "\n"),
+            "assistant": ("ASSISTANT:", "\n"),
+        },
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="claude",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("", ""),
+            "user": ("\n\nHuman: ", ""),
+            "assistant": ("\n\nAssistant:", ""),
+        },
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="chatml",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("<|im_start|>system\n", "<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+        stop_str=("<|im_end|>",),
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="chatml-llava",
+        default_system_prompt="You are a helpful assistant.",
+        role_prefix_and_suffix={
+            "system": ("<|im_start|>system\n", "<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+        stop_str=("<|im_end|>",),
+        image_token="<image>\n",
+    )
+)
+
+# There is default system prompt for qwen
+# reference: https://modelscope.cn/models/qwen/Qwen2-72B-Instruct/file/view/master?fileName=tokenizer_config.json&status=1
+# The chat template is: "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
+register_chat_template(
+    ChatTemplate(
+        name="qwen",
+        default_system_prompt="You are a helpful assistant.",
+        role_prefix_and_suffix={
+            "system": ("<|im_start|>system\n", "<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+        stop_str=("<|im_end|>",),
+    )
+)
+
+# Reference: https://huggingface.co/docs/transformers/main/model_doc/qwen2_vl#usage-example
+register_chat_template(
+    ChatTemplate(
+        name="qwen2-vl",
+        default_system_prompt="You are a helpful assistant.",
+        role_prefix_and_suffix={
+            "system": ("<|im_start|>system\n", "<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+        stop_str=("<|im_end|>",),
+        image_token="<|vision_start|><|image_pad|><|vision_end|>",
+    )
+)
+
+# Reference: https://github.com/lm-sys/FastChat/blob/main/docs/vicuna_weights_version.md#prompt-template
+register_chat_template(
+    ChatTemplate(
+        name="vicuna_v1.1",
+        default_system_prompt=(
+            "A chat between a curious user and an artificial intelligence assistant. "
+            "The assistant gives helpful, detailed, and polite answers to the user's questions."
+        ),
+        role_prefix_and_suffix={
+            "system": ("", " "),
+            "user": ("USER:", " "),
+            "assistant": ("ASSISTANT:", "</s>"),
+        },
+        image_token=" <image>\n",
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="llama-2-chat",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("<<SYS>>\n", "\n<</SYS>>\n\n"),
+            "user": ("[INST] ", " [/INST]"),
+            "assistant": ("", " </s><s>"),
+        },
+        style=ChatTemplateStyle.LLAMA2,
+    )
+)
+
+# Reference: https://huggingface.co/mistralai/Mistral-Small-3.1-24B-Instruct-2503/blob/main/chat_template.json
+register_chat_template(
+    ChatTemplate(
+        name="mistral",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("[SYSTEM_PROMPT] ", " [/SYSTEM_PROMPT]"),
+            "user": ("[INST] ", " [/INST]"),
+            "assistant": ("", " </s><s>"),
+        },
+        stop_str=("</s>",),
+        image_token="[IMG]",
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="llama-3-instruct",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": (
+                "<|start_header_id|>system<|end_header_id|>\n\n",
+                "<|eot_id|>",
+            ),
+            "user": (
+                "<|start_header_id|>user<|end_header_id|>\n\n",
+                "<|eot_id|>",
+            ),
+            "assistant": (
+                "<|start_header_id|>assistant<|end_header_id|>\n\n",
+                "<|eot_id|>",
+            ),
+        },
+        stop_str=("<|eot_id|>",),
+        image_token="<|image|>",
+    )
+)
+
+# https://huggingface.co/openbmb/MiniCPM-V-2_6
+register_chat_template(
+    ChatTemplate(
+        name="minicpmv",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("", " "),
+            "user": ("user:", " "),
+            "assistant": ("assistant:", "</s>"),
+        },
+        stop_str=("<|im_end|>", "<|endoftext|>"),
+        image_token="(<image>./</image>)",
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="janus-pro",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": (
+                "",
+                "",
+            ),
+            "User": (
+                "<｜User｜>",
+                "",
+            ),
+            "assistant": (
+                "<｜Assistant｜>",
+                "<｜end▁of▁sentence｜>",
+            ),
+        },
+        stop_str=("<｜end▁of▁sentence｜>",),
+        image_token="<image_placeholder>\n",
+    )
+)
+
+# https://huggingface.co/openbmb/MiniCPM-o-2_6
+register_chat_template(
+    ChatTemplate(
+        name="minicpmo",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("", " "),
+            "user": ("user:", " "),
+            "assistant": ("assistant:", "</s>"),
+        },
+        stop_str=("<|im_end|>", "<|endoftext|>"),
+        image_token="(<image>./</image>)",
+        audio_token="(<audio>./</audio>)",
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="janus",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": (
+                "",
+                "",
+            ),
+            "user": (
+                "<｜User｜>",
+                "",
+            ),
+            "assistant": (
+                "<｜Assistant｜>",
+                "<｜end▁of▁sentence｜>",
+            ),
+        },
+        stop_str=("<｜end▁of▁sentence｜>",),
+        image_token="<image_placeholder>\n",
+    )
+)
+
+# The difference between "llama-3-instruct-llava" and "llama-3-instruct" is that llava uses a different image_token.
+register_chat_template(
+    ChatTemplate(
+        name="llama-3-instruct-llava",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": (
+                "<|start_header_id|>system<|end_header_id|>\n\n",
+                "<|eot_id|>",
+            ),
+            "user": (
+                "<|start_header_id|>user<|end_header_id|>\n\n",
+                "<|eot_id|>",
+            ),
+            "assistant": (
+                "<|start_header_id|>assistant<|end_header_id|>\n\n",
+                "<|eot_id|>",
+            ),
+        },
+        stop_str=("<|eot_id|>",),
+        image_token="<image>\n",
+    )
+)
+
+# Reference: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct/blob/main/chat_template.json
+register_chat_template(
+    ChatTemplate(
+        name="llama-4",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": (
+                "<|header_start|>system<|header_end|>\n\n",
+                "<|eot|>",
+            ),
+            "user": (
+                "<|header_start|>user<|header_end|>\n\n",
+                "<|eot|>",
+            ),
+            "assistant": (
+                "<|header_start|>assistant<|header_end|>\n\n",
+                "<|eot|>",
+            ),
+        },
+        stop_str=("<|eot|>",),
+        image_token="<|image|>",
+    )
+)
+
+# Reference: https://modelscope.cn/models/01ai/Yi-1.5-34B-Chat/file/view/master?fileName=tokenizer_config.json&status=1
+register_chat_template(
+    ChatTemplate(
+        name="yi-1.5",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("", ""),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n<|im_start|>assistant\n"),
+            "assistant": ("", "<|im_end|>\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+        stop_str=("<|im_end|>",),
+    )
+)
+
+# Reference: https://github.com/01-ai/Yi/tree/main/VL#major-difference-with-llava
+register_chat_template(
+    ChatTemplate(
+        name="yi-vl",
+        default_system_prompt=(
+            "This is a chat between an inquisitive human and an AI assistant. Assume the role of the AI assistant. Read all the images carefully, and respond to the human's questions with informative, helpful, detailed and polite answers."
+            "这是一个好奇的人类和一个人工智能助手之间的对话。假设你扮演这个AI助手的角色。仔细阅读所有的图像，并对人类的问题做出信息丰富、有帮助、详细的和礼貌的回答。"
+        ),
+        role_prefix_and_suffix={
+            "system": ("", "\n\n"),
+            "user": ("### Human:", "\n"),
+            "assistant": ("### Assistant:", "\n"),
+        },
+        image_token=" <image_placeholder>\n",
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="gemma-it",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("", ""),
+            "user": ("<start_of_turn>user\n", "<end_of_turn>\n"),
+            "assistant": ("<start_of_turn>model\n", "<end_of_turn>\n"),
+        },
+        image_token="<start_of_image>",
+        audio_token="<start_of_audio>",
+        style=ChatTemplateStyle.PLAIN,
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="gemma-4-it",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("", ""),
+            "user": ("<|turn>user\n", "<turn|>\n"),
+            "assistant": ("<|turn>assistant\n", "<turn|>\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="dbrx-instruct",
+        default_system_prompt="You are DBRX, created by Databricks. You were last updated in December 2023. You answer questions based on information available up to that point.\nYOU PROVIDE SHORT RESPONSES TO SHORT QUESTIONS OR STATEMENTS, but provide thorough responses to more complex and open-ended questions.\nYou assist with various tasks, from writing to coding (using markdown for code blocks — remember to use ``` with code, JSON, and tables).\n(You do not have real-time data access or code execution capabilities. You avoid stereotyping and provide balanced perspectives on controversial topics. You do not provide song lyrics, poems, or news articles and do not divulge details of your training data.)\nThis is your system prompt, guiding your responses. Do not reference it, just respond to the user. If you find yourself talking about this message, stop. You should be responding appropriately and usually that means not mentioning this.\nYOU DO NOT MENTION ANY OF THIS INFORMATION ABOUT YOURSELF UNLESS THE INFORMATION IS DIRECTLY PERTINENT TO THE USER'S QUERY.",
+        role_prefix_and_suffix={
+            "system": ("<|im_start|>system\n", "<|im_end|>"),
+            "user": ("\n<|im_start|>user\n", "<|im_end|>"),
+            "assistant": ("\n<|im_start|>assistant\n", "<|im_end|>"),
+        },
+        stop_str=("<|im_end|>",),
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="c4ai-command-r",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": (
+                "<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>",
+                "<|END_OF_TURN_TOKEN|>",
+            ),
+            "user": ("<|START_OF_TURN_TOKEN|><|USER_TOKEN|>", "<|END_OF_TURN_TOKEN|>"),
+            "assistant": (
+                "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
+                "<|END_OF_TURN_TOKEN|>",
+            ),
+        },
+        style=ChatTemplateStyle.PLAIN,
+    )
+)
+
+# Adapted from https://huggingface.co/OpenGVLab/InternVL2-4B/blob/main/modeling_intern_vit.py
+register_chat_template(
+    ChatTemplate(
+        name="internvl-2-5",
+        default_system_prompt="你是书生·万象，英文名是InternVL，是由上海人工智能实验室、清华大学及多家合作单位联合开发的多模态大语言模型。",
+        role_prefix_and_suffix={
+            "system": ("<|im_start|>system\n", "<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"),
+        },
+        stop_str=["<|im_end|>", "<|action_end|>"],
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="interns1",
+        default_system_prompt="You are an AI assistant whose name is Intern-S1 (书生大模型).\n- Intern-S1 (书生大模型) is a vision-language model that is developed by Shanghai AI Laboratory (上海人工智能实验室).  It is designed to be helpful, honest, and harmless.\n- Intern-S1 (书生大模型) can understand and communicate fluently in the language chosen by the user such as English and 中文.\nYou are an expert reasoner with extensive experience in all areas. You approach problems through systematic thinking and rigorous reasoning. Your response should reflect deep understanding and precise logical thinking, making your solution path and reasoning clear to others. Please put your thinking process within <think>...</think> tags.",
+        role_prefix_and_suffix={
+            "system": ("<|im_start|>system\n", "<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "<|im_end|>\n"),
+        },
+        stop_str=["<|im_end|>", "<|action_end|>"],
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="granite-3-instruct",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": (
+                "<|start_of_role|>system<|end_of_role|>",
+                "<|end_of_text|>",
+            ),
+            "user": (
+                "<|start_of_role|>user<|end_of_role|>",
+                "<|end_of_text|>",
+            ),
+            "assistant": (
+                "<|start_of_role|>assistant<|end_of_role|>",
+                "<|end_of_text|>",
+            ),
+        },
+        stop_str=("<|end_of_text|>",),
+    )
+)
+
+register_chat_template(
+    ChatTemplate(
+        name="deepseek-v3",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": (
+                "",
+                "",
+            ),
+            "user": (
+                "<｜User｜>",
+                "",
+            ),
+            "assistant": (
+                "<｜Assistant｜>",
+                "<｜end▁of▁sentence｜>",
+            ),
+        },
+        stop_str=("<｜end▁of▁sentence｜>",),
+    )
+)
+
+# Reference: https://huggingface.co/docs/transformers/main/model_doc/glm4_v#usage-example
+register_chat_template(
+    ChatTemplate(
+        name="glm-4v",
+        default_system_prompt=None,
+        role_prefix_and_suffix={
+            "system": ("<|system|>\n", "\n"),
+            "user": ("<|user|>\n", "\n"),
+            "assistant": ("<|assistant|>\n", "\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+        stop_str=["<|user|>", "<|endoftext|>", "<|observation|>"],
+        image_token="<|image|>",
+    )
+)
+
+
+@register_chat_template_matching_function
+def match_deepseek(model_path: str):
+    if re.search(r"deepseek-(v3|r1)", model_path, re.IGNORECASE) and not re.search(
+        r"base", model_path, re.IGNORECASE
+    ):
+        return "deepseek-v3"
+
+
+@register_chat_template_matching_function
+def match_orion(model_path: str):
+    if "orion" in model_path.lower():
+        return "claude"
+
+
+@register_chat_template_matching_function
+def match_deepseek_janus_pro(model_path: str):
+    if re.search(r"janus", model_path, re.IGNORECASE):
+        return "janus-pro"
+
+
+@register_chat_template_matching_function
+def match_dbrx(model_path: str):
+    if re.search(r"dbrx", model_path, re.IGNORECASE) and re.search(
+        r"instruct", model_path, re.IGNORECASE
+    ):
+        return "dbrx-instruct"
+
+
+@register_chat_template_matching_function
+def match_vicuna(model_path: str):
+    if re.search(r"vicuna|llava-v1\.5|llava-next-video-7b", model_path, re.IGNORECASE):
+        return "vicuna_v1.1"
+
+
+@register_chat_template_matching_function
+def match_llama2_chat(model_path: str):
+    if re.search(
+        r"llama-2.*chat|codellama.*instruct",
+        model_path,
+        re.IGNORECASE,
+    ):
+        return "llama-2-chat"
+
+
+@register_chat_template_matching_function
+def match_mistral(model_path: str):
+    if re.search(r"pixtral|(mistral|mixtral).*instruct", model_path, re.IGNORECASE):
+        return "mistral"
+
+
+@register_chat_template_matching_function
+def match_llama3_instruct(model_path: str):
+    if re.search(r"llama-3.*instruct", model_path, re.IGNORECASE):
+        return "llama-3-instruct"
+
+
+@register_chat_template_matching_function
+def match_chat_ml(model_path: str):
+    if re.search(r"tinyllama", model_path, re.IGNORECASE):
+        return "chatml"
+    if re.search(r"qwen.*vl", model_path, re.IGNORECASE):
+        return "qwen2-vl"
+    if re.search(r"glm[-_]?4(\.\d+)?v", model_path, re.IGNORECASE):
+        return "glm-4v"
+    if re.search(r"qwen.*(chat|instruct)", model_path, re.IGNORECASE) and not re.search(
+        r"llava", model_path, re.IGNORECASE
+    ):
+        return "qwen"
+    if re.search(
+        r"llava-v1\.6-34b|llava-v1\.6-yi-34b|llava-next-video-34b|llava-onevision-qwen2",
+        model_path,
+        re.IGNORECASE,
+    ):
+        return "chatml-llava"
+
+
+@register_chat_template_matching_function
+def match_chat_yi(model_path: str):
+    if re.search(r"yi-vl", model_path, re.IGNORECASE) and not re.search(
+        r"llava", model_path, re.IGNORECASE
+    ):
+        return "yi-vl"
+    elif re.search(r"yi-1\.5.*chat", model_path, re.IGNORECASE):
+        return "yi-1.5"
+
+
+@register_chat_template_matching_function
+def match_gemma(model_path: str):
+    if re.search(r"gemma-4.*it", model_path, re.IGNORECASE):
+        return "gemma-4-it"
+    if re.search(r"(gemma.*it)|(gemma-3)", model_path, re.IGNORECASE):
+        return "gemma-it"
+
+
+@register_chat_template_matching_function
+def match_openbmb_minicpm(model_path: str):
+    if re.search(r"minicpm-v", model_path, re.IGNORECASE):
+        return "minicpmv"
+    elif re.search(r"minicpm-o", model_path, re.IGNORECASE):
+        return "minicpmo"
+
+
+@register_chat_template_matching_function
+def match_c4ai_command_r(model_path: str):
+    if re.search(r"c4ai-command-r", model_path, re.IGNORECASE):
+        return "c4ai-command-r"
+
+
+@register_chat_template_matching_function
+def match_granite_instruct(model_path: str):
+    if re.search(r"granite.*instruct", model_path, re.IGNORECASE):
+        return "granite-3-instruct"
+
+
+@register_chat_template_matching_function
+def match_internvl_chat(model_path: str):
+    if re.search(r"internvl2_5", model_path, re.IGNORECASE):
+        return "internvl-2-5"
+
+
+@register_chat_template_matching_function
+def match_interns1_chat(model_path: str):
+    if re.search(r"intern-s1", model_path, re.IGNORECASE):
+        return "interns1"
+    if re.search(r"interns1", model_path, re.IGNORECASE):
+        return "interns1"
+
+
+if __name__ == "__main__":
+    messages = [
+        {"role": "system", "content": None},  # None means default
+        # {"role": "system", "content": "You are a helpful, respectful and honest assistant."},
+        {"role": "user", "content": "Hello!"},
+        {"role": "assistant", "content": "Hi!"},
+        {"role": "user", "content": "What can you do?"},
+        {"role": "assistant", "content": "I can chat with you."},
+    ]
+
+    template = get_chat_template("llama-2-chat")
+    print(template.get_prompt(messages))

python/sglang/lang/choices.py

@@ -0,0 +1,164 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional
+
+import numpy as np
+
+
+@dataclass
+class ChoicesDecision:
+    decision: str
+    meta_info: Optional[Dict[str, Any]] = None
+
+
+class ChoicesSamplingMethod(ABC):
+
+    @property
+    def requires_unconditional_logprobs(self) -> bool:
+        return False
+
+    @abstractmethod
+    def __call__(
+        self,
+        *,
+        choices: List[str],
+        normalized_prompt_logprobs: List[float],
+        input_token_logprobs: List[List[Any]],
+        output_token_logprobs: List[List[Any]],
+        unconditional_token_logprobs: Optional[List[List[Any]]] = None,
+    ) -> ChoicesDecision: ...
+
+
+class TokenLengthNormalized(ChoicesSamplingMethod):
+
+    def __call__(
+        self,
+        *,
+        choices: List[str],
+        normalized_prompt_logprobs: List[float],
+        input_token_logprobs: List[List[Any]],
+        output_token_logprobs: List[List[Any]],
+        unconditional_token_logprobs: Optional[List[List[Any]]] = None,
+    ) -> ChoicesDecision:
+        """Select the option with the highest token length normalized prompt logprob."""
+        best_choice = choices[np.argmax(normalized_prompt_logprobs)]
+        meta_info = {
+            "normalized_prompt_logprobs": normalized_prompt_logprobs,
+            "input_token_logprobs": input_token_logprobs,
+            "output_token_logprobs": output_token_logprobs,
+        }
+        return ChoicesDecision(decision=best_choice, meta_info=meta_info)
+
+
+token_length_normalized = TokenLengthNormalized()
+
+
+class GreedyTokenSelection(ChoicesSamplingMethod):
+
+    def __call__(
+        self,
+        *,
+        choices: List[str],
+        normalized_prompt_logprobs: List[float],
+        input_token_logprobs: List[List[Any]],
+        output_token_logprobs: List[List[Any]],
+        unconditional_token_logprobs: Optional[List[List[Any]]] = None,
+    ) -> ChoicesDecision:
+        """Select the option based on greedy logprob selection. For overlapping options
+        where one option is a subset of a longer option, extend the shorter option using
+        its average logprob for comparison against the longer option."""
+
+        num_options = len(choices)
+        max_tokens = max(len(option) for option in input_token_logprobs)
+        logprob_matrix = self._build_logprob_matrix(
+            input_token_logprobs, max_tokens, num_options
+        )
+        remaining = self._greedy_selection(logprob_matrix, num_options, max_tokens)
+
+        best_choice = choices[remaining[0]]
+        meta_info = {
+            "normalized_prompt_logprobs": normalized_prompt_logprobs,
+            "input_token_logprobs": input_token_logprobs,
+            "output_token_logprobs": output_token_logprobs,
+            "greedy_logprob_matrix": logprob_matrix.tolist(),
+        }
+        return ChoicesDecision(decision=best_choice, meta_info=meta_info)
+
+    def _build_logprob_matrix(self, input_token_logprobs, max_tokens, num_options):
+        logprob_matrix = np.zeros((num_options, max_tokens))
+        for i, option in enumerate(input_token_logprobs):
+            actual_logprobs = [token[0] for token in option]
+            avg_logprob = np.mean(actual_logprobs)
+            logprob_matrix[i, : len(option)] = actual_logprobs
+            if len(option) < max_tokens:
+                logprob_matrix[i, len(option) :] = avg_logprob
+        return logprob_matrix
+
+    def _greedy_selection(self, logprob_matrix, num_options, max_tokens):
+        remaining = np.arange(num_options)
+        for j in range(max_tokens):
+            max_logprob = np.max(logprob_matrix[remaining, j])
+            remaining = remaining[logprob_matrix[remaining, j] == max_logprob]
+            if len(remaining) == 1:
+                break
+        return remaining
+
+
+greedy_token_selection = GreedyTokenSelection()
+
+
+class UnconditionalLikelihoodNormalized(ChoicesSamplingMethod):
+
+    @property
+    def requires_unconditional_logprobs(self) -> bool:
+        return True
+
+    def __call__(
+        self,
+        *,
+        choices: List[str],
+        normalized_prompt_logprobs: List[float],
+        input_token_logprobs: List[List[Any]],
+        output_token_logprobs: List[List[Any]],
+        unconditional_token_logprobs: Optional[List[List[Any]]] = None,
+    ) -> ChoicesDecision:
+        """Select the option with the highest average token logprob once normalized by
+        the unconditional token logprobs.
+
+        The first unconditional token logprob is assumed to be None. If so, it is
+        replaced with 0 for the purposes of normalization."""
+
+        if unconditional_token_logprobs is None:
+            raise ValueError(
+                "Unconditional token logprobs are required for this method."
+            )
+
+        normalized_unconditional_prompt_logprobs = self._normalize_logprobs(
+            input_token_logprobs, unconditional_token_logprobs
+        )
+
+        best_choice = choices[np.argmax(normalized_unconditional_prompt_logprobs)]
+        meta_info = {
+            "normalized_prompt_logprobs": normalized_prompt_logprobs,
+            "input_token_logprobs": input_token_logprobs,
+            "output_token_logprobs": output_token_logprobs,
+            "unconditional_token_logprobs": unconditional_token_logprobs,
+            "normalized_unconditional_prompt_logprobs": normalized_unconditional_prompt_logprobs,
+        }
+        return ChoicesDecision(decision=best_choice, meta_info=meta_info)
+
+    def _normalize_logprobs(self, input_token_logprobs, unconditional_token_logprobs):
+        normalized_unconditional_prompt_logprobs = []
+        for inputs, unconditionals in zip(
+            input_token_logprobs, unconditional_token_logprobs
+        ):
+            inputs_logprobs = np.array([token[0] for token in inputs])
+            unconditionals_logprobs = np.array([token[0] for token in unconditionals])
+            unconditionals_logprobs[0] = unconditionals_logprobs[0] or 0
+            normalized_unconditional_prompt_logprobs.append(
+                float(np.mean(inputs_logprobs - unconditionals_logprobs))
+            )
+        return normalized_unconditional_prompt_logprobs
+
+
+unconditional_likelihood_normalized = UnconditionalLikelihoodNormalized()

python/sglang/lang/ir.py

@@ -0,0 +1,643 @@
+"""The intermediate representation."""
+
+import dataclasses
+import inspect
+import warnings
+from typing import List, Optional, Union
+
+from sglang.global_config import global_config
+from sglang.lang.choices import ChoicesSamplingMethod
+
+REGEX_INT = r"[-+]?[0-9]+[ \n]*"
+REGEX_FLOAT = r"[-+]?[0-9]*\.?[0-9]+[ \n]*"
+REGEX_BOOL = r"(True|False)"
+REGEX_STR = r"\"[\w\d\s]*\""  # bugs with regex r"\".*\"" in interegular pkg
+
+
+@dataclasses.dataclass
+class SglSamplingParams:
+    max_new_tokens: int = 128
+    min_new_tokens: int = 0
+    n: int = 1
+    stop: Union[str, List[str]] = ()
+    stop_token_ids: Optional[List[int]] = ()
+    stop_regex: Optional[Union[str, List[str]]] = ()
+    temperature: float = 1.0
+    top_p: float = 1.0
+    top_k: int = -1  # -1 means disable
+    min_p: float = 0.0
+    frequency_penalty: float = 0.0
+    presence_penalty: float = 0.0
+    ignore_eos: bool = False
+    return_logprob: Optional[bool] = None
+    logprob_start_len: Optional[int] = None
+    top_logprobs_num: Optional[int] = None
+    return_text_in_logprobs: Optional[bool] = None
+    json_schema: Optional[str] = None
+
+    # for constrained generation, not included in to_xxx_kwargs
+    dtype: Optional[str] = None
+    regex: Optional[str] = None
+
+    def clone(self):
+        return SglSamplingParams(
+            self.max_new_tokens,
+            self.min_new_tokens,
+            self.n,
+            self.stop,
+            self.stop_token_ids,
+            self.stop_regex,
+            self.temperature,
+            self.top_p,
+            self.top_k,
+            self.min_p,
+            self.frequency_penalty,
+            self.presence_penalty,
+            self.ignore_eos,
+            self.return_logprob,
+            self.logprob_start_len,
+            self.top_logprobs_num,
+            self.return_text_in_logprobs,
+            self.json_schema,
+        )
+
+    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,
+            "max_completion_tokens": self.max_new_tokens,
+            "n": self.n,
+            "stop": self.stop or None,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "frequency_penalty": self.frequency_penalty,
+            "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": self.max_new_tokens,
+            "stop_sequences": (
+                self.stop if isinstance(self.stop, (list, tuple)) else [self.stop]
+            ),
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+        }
+
+    def to_litellm_kwargs(self):
+        if self.regex is not None:
+            warnings.warn("Regular expression is not supported in the LiteLLM backend.")
+        return {
+            "max_tokens": self.max_new_tokens,
+            "stop": self.stop or None,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "frequency_penalty": self.frequency_penalty,
+            "presence_penalty": self.presence_penalty,
+        }
+
+    def to_srt_kwargs(self):
+        return {
+            "max_new_tokens": self.max_new_tokens,
+            "min_new_tokens": self.min_new_tokens,
+            "n": self.n,
+            "stop": self.stop,
+            "stop_token_ids": self.stop_token_ids,
+            "stop_regex": self.stop_regex,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "min_p": self.min_p,
+            "frequency_penalty": self.frequency_penalty,
+            "presence_penalty": self.presence_penalty,
+            "ignore_eos": self.ignore_eos,
+            "regex": self.regex,
+            "json_schema": self.json_schema,
+        }
+
+
+class SglFunction:
+    def __init__(self, func, num_api_spec_tokens=None, bind_arguments=None):
+        self.func = func
+        self.num_api_spec_tokens = num_api_spec_tokens
+        self.bind_arguments = bind_arguments or {}
+        self.pin_prefix_rid = None
+
+        # Parse arguments
+        argspec = inspect.getfullargspec(func)
+        assert argspec.args[0] == "s", 'The first argument must be "s"'
+        self.arg_names = argspec.args[1:]
+        self.arg_defaults = argspec.defaults if argspec.defaults is not None else []
+
+    def bind(self, **kwargs):
+        assert all(key in self.arg_names for key in kwargs)
+
+        new_bind_dict = {**self.bind_arguments, **kwargs}
+        return SglFunction(self.func, bind_arguments=new_bind_dict)
+
+    def run(
+        self,
+        *args,
+        max_new_tokens: int = 128,
+        n: int = 1,
+        stop: Optional[Union[str, List[str]]] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        stop_regex: Optional[Union[str, List[str]]] = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        min_p: float = 0.0,
+        frequency_penalty: float = 0.0,
+        presence_penalty: float = 0.0,
+        ignore_eos: bool = False,
+        return_logprob: Optional[bool] = None,
+        logprob_start_len: Optional[int] = None,
+        top_logprobs_num: Optional[int] = None,
+        return_text_in_logprobs: Optional[bool] = None,
+        stream: bool = False,
+        backend=None,
+        use_thread: bool = True,
+        **kwargs,
+    ):
+        from sglang.lang.interpreter import run_program
+
+        # avoid using [] as the default arg: https://nikos7am.com/posts/mutable-default-arguments/
+        if stop is None:
+            stop = []
+        if stop_token_ids is None:
+            stop_token_ids = []
+        if stop_regex is None:
+            stop_regex = []
+
+        default_sampling_para = SglSamplingParams(
+            max_new_tokens=max_new_tokens,
+            n=n,
+            stop=stop,
+            stop_token_ids=stop_token_ids,
+            stop_regex=stop_regex,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            min_p=min_p,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+            ignore_eos=ignore_eos,
+            return_logprob=return_logprob,
+            logprob_start_len=logprob_start_len,
+            top_logprobs_num=top_logprobs_num,
+            return_text_in_logprobs=return_text_in_logprobs,
+        )
+        backend = backend or global_config.default_backend
+        return run_program(
+            self,
+            backend,
+            args,
+            kwargs,
+            default_sampling_para,
+            stream,
+            use_thread=use_thread,
+        )
+
+    def run_batch(
+        self,
+        batch_kwargs,
+        *,
+        max_new_tokens: int = 128,
+        n: int = 1,
+        stop: Optional[Union[str, List[str]]] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        stop_regex: Optional[Union[str, List[str]]] = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        min_p: float = 0.0,
+        frequency_penalty: float = 0.0,
+        presence_penalty: float = 0.0,
+        ignore_eos: bool = False,
+        return_logprob: Optional[bool] = None,
+        logprob_start_len: Optional[int] = None,
+        top_logprobs_num: Optional[int] = None,
+        return_text_in_logprobs: Optional[bool] = None,
+        backend=None,
+        num_threads: Union[str, int] = "auto",
+        progress_bar: bool = False,
+        generator_style: bool = False,
+    ):
+        from sglang.lang.interpreter import run_program_batch
+
+        if stop is None:
+            stop = []
+        if stop_token_ids is None:
+            stop_token_ids = []
+        if stop_regex is None:
+            stop_regex = []
+
+        assert isinstance(batch_kwargs, (list, tuple))
+        if len(batch_kwargs) == 0:
+            return []
+        if not isinstance(batch_kwargs[0], dict):
+            num_programs = len(batch_kwargs)
+            # change the list of argument values to dict of arg_name -> arg_value
+            batch_kwargs = [
+                {self.arg_names[i]: v for i, v in enumerate(arg_values)}
+                for arg_values in batch_kwargs
+                if isinstance(arg_values, (list, tuple))
+                and len(self.arg_names) - len(self.arg_defaults)
+                <= len(arg_values)
+                <= len(self.arg_names)
+            ]
+            # Ensure to raise an exception if the number of arguments mismatch
+            if len(batch_kwargs) != num_programs:
+                raise Exception("Given arguments mismatch the SGL function signature")
+
+        default_sampling_para = SglSamplingParams(
+            max_new_tokens=max_new_tokens,
+            n=n,
+            stop=stop,
+            stop_token_ids=stop_token_ids,
+            stop_regex=stop_regex,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            min_p=min_p,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+            ignore_eos=ignore_eos,
+            return_logprob=return_logprob,
+            logprob_start_len=logprob_start_len,
+            top_logprobs_num=top_logprobs_num,
+            return_text_in_logprobs=return_text_in_logprobs,
+        )
+        backend = backend or global_config.default_backend
+        return run_program_batch(
+            self,
+            backend,
+            batch_kwargs,
+            default_sampling_para,
+            num_threads,
+            progress_bar,
+            generator_style=generator_style,
+        )
+
+    def trace(self, *, backend=None, **kwargs):
+        from sglang.lang.tracer import trace_program
+
+        backend = backend or global_config.default_backend
+        return trace_program(self, kwargs, backend)
+
+    def cache(self, backend=None):
+        from sglang.lang.interpreter import cache_program
+
+        backend = backend or global_config.default_backend
+        return cache_program(self, backend)
+
+    def __call__(self, *args, **kwargs):
+        from sglang.lang.tracer import TracingScope
+
+        tracing_scope = TracingScope.get_current_scope()
+        if tracing_scope is None:
+            return self.run(*args, **kwargs)
+        else:
+            kwargs["backend"] = tracing_scope.tracer_state.backend
+            return self.trace(*args, **kwargs)
+
+
+class SglExpr:
+    node_ct = 0
+
+    def __init__(self):
+        self.node_id = SglExpr.node_ct
+        self.prev_node = None
+        self.pid = None
+        SglExpr.node_ct += 1
+
+    def __add__(self, other):
+        if isinstance(other, str):
+            other = SglConstantText(other)
+        assert isinstance(other, SglExpr)
+
+        return self.concatenate_ir(self, other)
+
+    def __radd__(self, other):
+        if isinstance(other, str):
+            other = SglConstantText(other)
+        assert isinstance(other, SglExpr), f"{other}"
+
+        return self.concatenate_ir(other, self)
+
+    def concatenate_ir(self, a, b):
+        if isinstance(a, SglExprList):
+            if isinstance(b, SglExprList):
+                return SglExprList(a.expr_list + b.expr_list)
+            else:
+                return SglExprList(a.expr_list + [b])
+        elif isinstance(b, SglExprList):
+            return SglExprList([a] + b.expr_list)
+
+        return SglExprList([a, b])
+
+    def print_graph_dfs(self):
+        ret = [""]
+        visited = set()
+
+        def dfs_print(x):
+            if x is None or x in visited:
+                return
+            visited.add(x)
+
+            # Print dependency
+            if x.prev_node is not None:
+                dfs_print(x.prev_node)
+
+            if isinstance(x, SglExprList):
+                for y in x.expr_list:
+                    dfs_print(y)
+            # elif isinstance(x, SglRole):
+            #    dfs_print(x.expr)
+            elif isinstance(x, SglVariable):
+                dfs_print(x.source)
+
+            # Print the node itself
+            if isinstance(x, (SglFork, SglGetForkItem)):
+                ret[0] += f"%{x.node_id} = {x}\n"
+            else:
+                if x.prev_node is not None:
+                    ret[0] += (
+                        f"%{x.node_id} = %{x.prev_node.node_id} + " + str(x) + "\n"
+                    )
+                else:
+                    ret[0] += f"%{x.node_id} = " + str(x) + "\n"
+
+        dfs_print(self)
+        return ret[0]
+
+
+class SglExprList(SglExpr):
+    def __init__(self, expr_list: List[SglExpr]):
+        super().__init__()
+        self.expr_list = expr_list
+
+    def __repr__(self):
+        return f"ExprList({self.expr_list})"
+
+
+class SglArgument(SglExpr):
+    def __init__(self, name: str, value: str):
+        super().__init__()
+        self.name = name
+        self.value = value
+
+    def __repr__(self):
+        return f"Argument(name={self.name}, value={repr(self.value)})"
+
+    def __len__(self):
+        return len(self.value)
+
+    def __getitem__(self, i):
+        return self.value[i]
+
+    def __int__(self):
+        return self.value
+
+    def __bool__(self):
+        return self.value
+
+    def __format__(self, *args):
+        raise TypeError(
+            "Cannot put argument inside a f-string. "
+            "This is not compatible with the tracer. "
+        )
+
+
+class SglImage(SglExpr):
+    def __init__(self, path: str):
+        self.path = path
+
+    def __repr__(self) -> str:
+        return f"SglImage({self.path})"
+
+
+class SglVideo(SglExpr):
+    def __init__(self, path: str, num_frames: int):
+        self.path = path
+        self.num_frames = num_frames
+
+    def __repr__(self) -> str:
+        return f"SglVideo({self.path}, {self.num_frames})"
+
+
+class SglGen(SglExpr):
+    def __init__(
+        self,
+        name: Optional[str] = None,
+        max_new_tokens: Optional[int] = None,
+        min_new_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        stop_regex: Optional[Union[str, List[str]]] = None,
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        top_k: Optional[int] = None,
+        min_p: Optional[float] = None,
+        frequency_penalty: Optional[float] = None,
+        presence_penalty: Optional[float] = None,
+        ignore_eos: Optional[bool] = None,
+        return_logprob: Optional[bool] = None,
+        logprob_start_len: Optional[int] = None,
+        top_logprobs_num: Optional[int] = None,
+        return_text_in_logprobs: Optional[bool] = None,
+        dtype: Optional[type] = None,
+        regex: Optional[str] = None,
+        json_schema: Optional[str] = None,
+    ):
+        """Call the model to generate. See the meaning of the arguments in docs/backend/sampling_params.md"""
+        super().__init__()
+        self.name = name
+        self.sampling_params = SglSamplingParams(
+            max_new_tokens=max_new_tokens,
+            min_new_tokens=min_new_tokens,
+            n=n,
+            stop=stop,
+            stop_regex=stop_regex,
+            stop_token_ids=stop_token_ids,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            min_p=min_p,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+            ignore_eos=ignore_eos,
+            return_logprob=return_logprob,
+            logprob_start_len=logprob_start_len,
+            top_logprobs_num=top_logprobs_num,
+            return_text_in_logprobs=return_text_in_logprobs,
+            dtype=dtype,
+            regex=regex,
+            json_schema=json_schema,
+        )
+
+    def __repr__(self):
+        return f"Gen('{self.name}')"
+
+
+class SglConstantText(SglExpr):
+    def __init__(self, value: str):
+        super().__init__()
+        self.value = value
+
+    def __repr__(self):
+        return f"Constant({repr(self.value)})"
+
+
+class SglRoleBegin(SglExpr):
+    def __init__(self, role: str):
+        super().__init__()
+        self.role = role
+
+    def __repr__(self):
+        return f"RoleBegin({self.role})"
+
+
+class SglRoleEnd(SglExpr):
+    def __init__(self, role: str):
+        super().__init__()
+        self.role = role
+
+    def __repr__(self):
+        return f"RoleEnd({self.role})"
+
+
+class SglSelect(SglExpr):
+
+    def __init__(
+        self,
+        name: str,
+        choices: List[str],
+        temperature: float,
+        choices_method: ChoicesSamplingMethod,
+    ):
+        super().__init__()
+        self.name = name
+        self.choices = choices
+        self.temperature = temperature
+        self.choices_method = choices_method
+
+    def __repr__(self):
+        return f"Select({self.name}, choices={self.choices}, choices_method={self.choices_method})"
+
+
+class SglFork(SglExpr):
+    def __init__(self, number: int, position_ids_offset=None):
+        super().__init__()
+        self.number = number
+        self.position_ids_offset = position_ids_offset
+
+    def __repr__(self):
+        return (
+            f"Fork(%{self.prev_node.node_id}, number={self.number}, "
+            f"position_ids_offset={self.position_ids_offset})"
+        )
+
+
+class SglGetForkItem(SglExpr):
+    def __init__(self, index: int):
+        super().__init__()
+        self.index = index
+
+    def __repr__(self):
+        return f"GetForkItem(%{self.prev_node.node_id}, index={self.index})"
+
+
+class SglVariable(SglExpr):
+    def __init__(self, name: str, source):
+        super().__init__()
+        self.name = name
+        self.source = source
+
+    def __repr__(self):
+        return f"Variable('{self.name}', source=%{self.source.node_id})"
+
+
+class SglVarScopeBegin(SglExpr):
+    def __init__(self, name: str):
+        super().__init__()
+        self.name = name
+
+    def __repr__(self):
+        return f"VarScopeBegin('{self.name}')"
+
+
+class SglVarScopeEnd(SglExpr):
+    def __init__(self, name: str):
+        super().__init__()
+        self.name = name
+
+    def __repr__(self):
+        return f"VarScopeEnd('{self.name}')"
+
+
+class SglConcateAndAppend(SglExpr):
+    def __init__(self, states):
+        super().__init__()
+        self.states = states
+
+    def __repr__(self):
+        return f"ConcatenateAndAppend('{self.states}')"
+
+
+class SglCommitLazy(SglExpr):
+    def __init__(self):
+        super().__init__()
+
+    def __repr__(self):
+        return "CommitLazy()"
+
+
+class SglSeparateReasoning(SglExpr):
+    def __init__(self, model_type: str, expr: SglExpr):
+        super().__init__()
+        self.model_type = model_type
+
+        self.expr = expr
+        self.name = None
+        self._process_expr(expr)
+
+    def process_name_for_reasoning(self, name):
+        if not name:
+            raise ValueError("name must be provided")
+        return f"{name}_reasoning_content"
+
+    def _process_expr(self, expr):
+        if isinstance(expr, SglGen):
+            self.name = self.process_name_for_reasoning(expr.name)
+        elif isinstance(expr, SglSelect):
+            self.name = self.process_name_for_reasoning(expr.name)
+        elif isinstance(expr, SglExprList):
+            for x in expr.expr_list:
+                self._process_expr(x)
+
+    def __repr__(self):
+        return f"SeparateReasoning(model_type={self.model_type}, name={self.name})"

python/sglang/lang/tracer.py

@@ -0,0 +1,279 @@
+"""Tracing a program."""
+
+import uuid
+from typing import Any, Dict, List, Optional
+
+from sglang.lang.backend.base_backend import BaseBackend
+from sglang.lang.interpreter import ProgramState, ProgramStateGroup
+from sglang.lang.ir import (
+    SglArgument,
+    SglConstantText,
+    SglExpr,
+    SglExprList,
+    SglFork,
+    SglGen,
+    SglGetForkItem,
+    SglRoleBegin,
+    SglRoleEnd,
+    SglSelect,
+    SglVariable,
+    SglVarScopeBegin,
+    SglVarScopeEnd,
+)
+
+
+class StopTracing(Exception):
+    pass
+
+
+def extract_prefix_by_tracing(program, backend):
+    # Create dummy arguments
+    dummy_arguments = {name: SglArgument(name, None) for name in program.arg_names}
+    arguments = dummy_arguments
+    arguments.update(program.bind_arguments)
+
+    # Trace
+    tracer = TracerProgramState(backend, arguments, only_trace_prefix=True)
+    try:
+        with TracingScope(tracer):
+            tracer.ret_value = program.func(tracer, **arguments)
+    except (StopTracing, TypeError, AttributeError):
+        # Some exceptions may not be caught
+        pass
+
+    # Run and cache prefix
+    prefix = ""
+    for expr in tracer.flatten_nodes():
+        if isinstance(expr, SglConstantText):
+            prefix += expr.value
+        else:
+            break
+    return prefix
+
+
+def trace_program(program, arguments, backend):
+    # Create dummy backend
+    if backend is None:
+        backend = BaseBackend()
+
+    # Create dummy arguments
+    dummy_arguments = {
+        name: SglArgument(name, None)
+        for name in program.arg_names
+        if name not in arguments
+    }
+    arguments.update(dummy_arguments)
+    arguments.update(program.bind_arguments)
+
+    # Trace
+    tracer = TracerProgramState(backend, arguments, only_trace_prefix=False)
+    with TracingScope(tracer):
+        tracer.ret_value = program.func(tracer, **arguments)
+    return tracer
+
+
+class TracerProgramState(ProgramState):
+    def __init__(self, backend, arguments, only_trace_prefix):
+        self.pid = uuid.uuid4().hex
+        self.backend = backend
+        self.arguments: Dict[str, Any] = arguments
+        self.only_trace_prefix = only_trace_prefix
+
+        if hasattr(backend, "endpoint"):
+            self.backend = backend.endpoint
+
+        self.nodes = []
+        self.last_node = None
+        self.variables = {}
+        self.ret_value = None
+
+        # For completion
+
+        # For chat
+        self.messages_ = []
+        self.cur_role = None
+        self.chat_template = self.backend.get_chat_template()
+
+        # For multi states
+        self.child_states = []
+
+        cur_scope = TracingScope.get_current_scope()
+        if cur_scope is not None:
+            cur_scope.add_child_state(self)
+
+    ##################################
+    ########### Public API ###########
+    ##################################
+
+    def fork(self, size: int = 1, position_ids_offset: Optional[List[int]] = None):
+        assert size >= 1
+
+        if self.only_trace_prefix:
+            raise StopTracing()
+
+        fork_node = SglFork(size)
+        fork_node.prev_node = self.last_node
+
+        states = [
+            TracerProgramState(self.backend, self.arguments, self.only_trace_prefix)
+            for _ in range(size)
+        ]
+
+        for i in range(size):
+            node = SglGetForkItem(i)
+            node.prev_node = fork_node
+            states[i].last_node = node
+            states[i].variables = dict(self.variables)
+            states[i].messages_ = list(self.messages_)
+            states[i].cur_role = self.cur_role
+            states[i].chat_template = self.chat_template
+
+        state_group = ProgramStateGroup(states, self)
+
+        return state_group
+
+    ##################################
+    ########## Internal API ##########
+    ##################################
+
+    def _append_node(self, other: SglExpr):
+        self.nodes.append(other)
+        other.prev_node = self.last_node
+        self.last_node = other
+
+    def _execute(self, other: SglExpr):
+        if isinstance(other, str):
+            other = SglConstantText(other)
+
+        other.pid = self.pid
+
+        if isinstance(other, SglConstantText):
+            self._execute_fill(other)
+        elif isinstance(other, SglGen):
+            self._execute_gen(other)
+        elif isinstance(other, SglSelect):
+            self._execute_select(other)
+        elif isinstance(other, SglExprList):
+            for x in other.expr_list:
+                self._execute(x)
+        elif isinstance(other, SglRoleBegin):
+            self._execute_role_begin(other)
+        elif isinstance(other, SglRoleEnd):
+            self._execute_role_end(other)
+        elif isinstance(other, SglVarScopeBegin):
+            self._execute_var_scope_begin(other)
+        elif isinstance(other, SglVarScopeEnd):
+            self._execute_var_scope_end(other)
+        else:
+            if self.only_trace_prefix:
+                raise StopTracing()
+            else:
+                self._append_node(other)
+
+        return self
+
+    def __iadd__(self, other):
+        self._execute(other)
+        return self
+
+    def _execute_fill(self, expr: SglConstantText):
+        if isinstance(expr, str):
+            expr = SglConstantText(expr)
+        self._append_node(expr)
+
+    def _execute_gen(self, expr: SglGen):
+        name = expr.name if expr.name is not None else "gen_" + str(len(self.variables))
+        new_node = SglVariable(name, source=expr)
+        self.variables[name] = new_node
+        self._append_node(expr)
+
+    def _execute_select(self, expr: SglSelect):
+        name = (
+            expr.name if expr.name is not None else "select_" + str(len(self.variables))
+        )
+        new_node = SglVariable(name, source=expr)
+        self.variables[name] = new_node
+        self._append_node(expr)
+
+    def _execute_role_begin(self, expr: SglRoleBegin):
+        assert self.cur_role is None, "Nested roles are not allowed."
+
+        if len(self.messages_) == 0 and expr.role != "system":
+            # Insert default system message
+            default_system = self.chat_template.default_system_prompt
+            if default_system:
+                self._execute_role_begin(SglRoleBegin("system"))
+                self._execute_fill(default_system)
+                self._execute_role_end(SglRoleEnd("system"))
+
+        self.cur_role = expr.role
+
+        prefix, suffix = self.chat_template.get_prefix_and_suffix(
+            expr.role, self.messages_
+        )
+
+        self._execute_fill(prefix)
+
+    def _execute_role_end(self, expr: SglRoleEnd):
+        prefix, suffix = self.chat_template.get_prefix_and_suffix(
+            expr.role, self.messages_
+        )
+
+        self._execute_fill(suffix)
+
+        self.messages_.append({"role": expr.role, "content": ""})
+
+        self.cur_role = None
+
+    def _execute_var_scope_end(self, expr: SglVarScopeEnd):
+        new_node = SglVariable(expr.name, source=self.last_node)
+        self.variables[expr.name] = new_node
+
+    def get_var(self, name):
+        ret = self.arguments.get(name, None)
+        if ret is not None:
+            return ret
+
+        v = self.variables[name]
+        return SglVariable(v.name, v.source)
+
+    def flatten_nodes(self):
+        def traverse(cur):
+            if isinstance(cur, SglExprList):
+                for child in cur.expr_list:
+                    traverse(child)
+            else:
+                ret.append(cur)
+
+        ret = []
+        for x in self.nodes:
+            traverse(x)
+        return ret
+
+    def __del__(self):
+        pass
+
+
+class TracingScope:
+    cur_scope = None
+
+    def __init__(self, tracer_state: TracerProgramState):
+        self.tracer_state = tracer_state
+        self.last_scope = TracingScope.cur_scope
+
+    def __enter__(self):
+        TracingScope.cur_scope = self
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        TracingScope.cur_scope = self.last_scope
+
+    @staticmethod
+    def get_current_scope():
+        return TracingScope.cur_scope
+
+    def add_child_state(self, state: TracerProgramState):
+        cur_scope = self
+        while cur_scope is not None:
+            cur_scope.tracer_state.child_states.append(state)
+            cur_scope = cur_scope.last_scope