release initial code

Co-authored-by: Ying Sheng <sqy1415@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>
Co-authored-by: parasol-aser <3848358+parasol-aser@users.noreply.github.com>
Co-authored-by: LiviaSun <33578456+ChuyueSun@users.noreply.github.com>
Co-authored-by: Cody Yu <hao.yu.cody@gmail.com>

python/sglang/backend/base_backend.py

+from typing import Callable, List, Optional, Union
+
+from sglang.lang.chat_template import get_chat_template
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import SamplingParams
+
+
+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: SamplingParams,
+    ):
+        raise NotImplementedError()
+
+    def generate_stream(
+        self,
+        s: StreamExecutor,
+        sampling_params: SamplingParams,
+    ):
+        raise NotImplementedError()
+
+    def select(
+        self,
+        s: StreamExecutor,
+        choices: List[str],
+        temperature: float,
+    ):
+        raise NotImplementedError()
+
+    def concatenate_and_append(self, src_rids: List[str], dst_rid: str):
+        raise NotImplementedError()
+
+    def shutdown(self):
+        pass

python/sglang/backend/huggingface.py

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

python/sglang/backend/openai.py

+from typing import Callable, List, Optional, Union
+
+import numpy as np
+from sglang.backend.base_backend import BaseBackend
+from sglang.lang.chat_template import get_chat_template
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import SamplingParams
+
+try:
+    import openai
+    import tiktoken
+except ImportError as e:
+    openai = tiktoken = e
+
+
+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
+
+
+CHAT_MODEL_NAMES = [
+    # GPT-4
+    "gpt-4",
+    "gpt-4-32k",
+    "gpt-4-1106-preview",
+    "gpt-4-vision-preview",
+    "gpt-4-0613",
+    "gpt-4-0314",
+    # GPT-3.5
+    "gpt-3.5-turbo",
+    "gpt-3.5-turbo-16k",
+    "gpt-3.5-turbo-1106",
+    "gpt-3.5-turbo-16k-0613",
+    "gpt-3.5-turbo-0613",
+    "gpt-3.5-turbo-0301",
+]
+
+
+class OpenAI(BaseBackend):
+    def __init__(self, model_name, *args, **kwargs):
+        super().__init__()
+        self.client = openai.OpenAI(*args, **kwargs)
+
+        if isinstance(openai, Exception):
+            raise e
+
+        self.model_name = model_name
+        self.tokenizer = tiktoken.encoding_for_model(model_name)
+        self.logit_bias_int = create_logit_bias_int(self.tokenizer)
+
+        if model_name in CHAT_MODEL_NAMES:
+            self.is_chat_model = True
+        else:
+            self.is_chat_model = False
+
+        self.chat_template = get_chat_template("default")
+
+    def get_chat_template(self):
+        return self.chat_template
+
+    def generate(
+        self,
+        s: StreamExecutor,
+        sampling_params: SamplingParams,
+    ):
+        if sampling_params.dtype is None:
+            if self.is_chat_model:
+                assert s.text_.endswith("ASSISTANT:")
+                prompt = s.messages_
+            else:
+                prompt = s.text_
+
+            kwargs = sampling_params.to_openai_kwargs()
+            comp = openai_completion(
+                client=self.client,
+                is_chat=self.is_chat_model,
+                model=self.model_name,
+                prompt=prompt,
+                **kwargs,
+            )
+        elif sampling_params.dtype in [str, "str", "string"]:
+            kwargs = sampling_params.to_openai_kwargs()
+            kwargs.pop("stop")
+            comp = openai_completion(
+                client=self.client,
+                is_chat=self.is_chat_model,
+                model=self.model_name,
+                prompt=s.text_ + '"',
+                stop='"',
+                **kwargs,
+            )
+            comp = '"' + comp + '"'
+        elif sampling_params.dtype in [int, "int"]:
+            kwargs = sampling_params.to_openai_kwargs()
+            kwargs.pop("stop")
+            comp = openai_completion(
+                client=self.client,
+                is_chat=self.is_chat_model,
+                model=self.model_name,
+                prompt=s.text_,
+                logit_bias=self.logit_bias_int,
+                stop=[" "],
+                **kwargs,
+            )
+        else:
+            raise ValueError(f"Unknown dtype: {dtype}")
+
+        return comp, {}
+
+    def generate_stream(
+        self,
+        s: StreamExecutor,
+        sampling_params: SamplingParams,
+    ):
+        if sampling_params.dtype is None:
+            if self.is_chat_model:
+                assert s.text_.endswith("ASSISTANT:")
+                prompt = s.messages_
+            else:
+                prompt = s.text_
+
+            kwargs = sampling_params.to_openai_kwargs()
+            generator = openai_completion_stream(
+                client=self.client,
+                is_chat=self.is_chat_model,
+                model=self.model_name,
+                prompt=prompt,
+                **kwargs,
+            )
+            return generator
+        else:
+            raise ValueError(f"Unknown dtype: {dtype}")
+
+    def select(
+        self,
+        s: StreamExecutor,
+        choices: List[str],
+        temperature: float,
+    ):
+        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]
+
+            # 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)
+
+        decision = choices[np.argmax(scores)]
+        return decision, scores
+
+
+def openai_completion(client, is_chat=None, prompt=None, **kwargs):
+    try:
+        if is_chat:
+            if kwargs["stop"] is None:
+                kwargs.pop("stop")
+            ret = client.chat.completions.create(messages=prompt, **kwargs)
+            comp = ret.choices[0].message.content
+        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
+    except openai.OpenAIError as e:
+        print(f"OpenAI Error: {e}")
+        raise e
+
+    return comp
+
+
+def openai_completion_stream(client, is_chat=None, prompt=None, **kwargs):
+    try:
+        if is_chat:
+            generator = client.chat.completions.create(
+                messages=prompt, stream=True, **kwargs
+            )
+            for ret in generator:
+                content = ret.choices[0].delta.content
+                yield content or "", {}
+        else:
+            generator = client.completions.create(prompt=prompt, stream=True, **kwargs)
+            for ret in generator:
+                content = ret.choices[0].text
+                yield content or "", {}
+    except openai.OpenAIError as e:
+        print(f"OpenAI Error: {e}")
+        raise e

python/sglang/backend/runtime_endpoint.py

+import json
+from typing import Callable, List, Optional, Union
+
+import numpy as np
+import requests
+from sglang.backend.base_backend import BaseBackend
+from sglang.global_config import global_config
+from sglang.lang.chat_template import get_chat_template_by_model_path
+from sglang.lang.interpreter import StreamExecutor
+from sglang.lang.ir import SamplingParams, SglArgument
+from sglang.utils import encode_image_base64, find_printable_text, http_request
+
+
+class RuntimeEndpoint(BaseBackend):
+    def __init__(self, base_url):
+        super().__init__()
+        self.support_concate_and_append = True
+
+        self.base_url = base_url
+
+        res = http_request(self.base_url + "/get_model_info")
+        assert res.status_code == 200
+        self.model_info = res.json()
+
+        self.chat_template = get_chat_template_by_model_path(
+            self.model_info["model_path"]
+        )
+
+    def get_model_name(self):
+        return self.model_info["model_path"]
+
+    def 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}},
+        )
+        assert res.status_code == 200
+
+    def commit_lazy_operations(self, s: StreamExecutor):
+        res = http_request(
+            self.base_url + "/generate",
+            json={"text": s.text_, "sampling_params": {"max_new_tokens": 0}},
+        )
+        assert res.status_code == 200
+
+    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)
+        assert res.status_code == 200
+
+    def generate(
+        self,
+        s: StreamExecutor,
+        sampling_params: SamplingParams,
+    ):
+        if sampling_params.dtype is None:
+            data = {
+                "text": s.text_,
+                "sampling_params": {
+                    "skip_special_tokens": global_config.skip_special_tokens_in_output,
+                    **sampling_params.to_srt_kwargs(),
+                },
+            }
+        elif sampling_params.dtype in [int, "int"]:
+            data = {
+                "text": s.text_,
+                "sampling_params": {
+                    "skip_special_tokens": global_config.skip_special_tokens_in_output,
+                    "dtype": "int",
+                    **sampling_params.to_srt_kwargs(),
+                },
+            }
+        else:
+            raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}")
+
+        self._add_images(s, data)
+
+        res = http_request(self.base_url + "/generate", json=data)
+        obj = res.json()
+        comp = obj["text"]
+        return comp, obj["meta_info"]
+
+    def generate_stream(
+        self,
+        s: StreamExecutor,
+        sampling_params: SamplingParams,
+    ):
+        if sampling_params.dtype is None:
+            data = {
+                "text": s.text_,
+                "sampling_params": {
+                    "skip_special_tokens": global_config.skip_special_tokens_in_output,
+                    **sampling_params.to_srt_kwargs(),
+                },
+            }
+        elif sampling_params.dtype in [int, "int"]:
+            data = {
+                "text": s.text_,
+                "sampling_params": {
+                    "skip_special_tokens": global_config.skip_special_tokens_in_output,
+                    "dtype": "int",
+                    **sampling_params.to_srt_kwargs(),
+                },
+            }
+        else:
+            raise RuntimeError(f"Invalid dtype: {sampling_params.dtype}")
+
+        data["stream"] = True
+        self._add_images(s, data)
+
+        response = http_request(self.base_url + "/generate", json=data, stream=True)
+        pos = 0
+
+        incomplete_text = ""
+        for chunk in response.iter_lines(decode_unicode=False, delimiter=b"\0"):
+            if chunk:
+                data = json.loads(chunk.decode())
+                text = find_printable_text(data["text"][pos:])
+                meta_info = data["meta_info"]
+                pos += len(text)
+                incomplete_text = data["text"][pos:]
+                yield text, meta_info
+
+        if len(incomplete_text) > 0:
+            yield incomplete_text, meta_info
+
+    def select(
+        self,
+        s: StreamExecutor,
+        choices: List[str],
+        temperature: float,
+    ):
+        assert temperature <= 1e-5
+
+        # Cache common prefix
+        data = {"text": s.text_, "sampling_params": {"max_new_tokens": 0}}
+        self._add_images(s, data)
+        res = http_request(self.base_url + "/generate", json=data)
+        assert res.status_code == 200
+        prompt_len = res.json()["meta_info"]["prompt_tokens"]
+
+        # Compute logprob
+        data = {
+            "text": [s.text_ + c for c in choices],
+            "sampling_params": {"max_new_tokens": 0},
+            "return_normalized_logprob": True,
+            "normalized_logprob_start_len": prompt_len,
+        }
+        self._add_images(s, data)
+        res = http_request(self.base_url + "/generate", json=data)
+        assert res.status_code == 200
+        logps = [r["meta_info"]["normalized_logprob"] for r in res.json()]
+
+        decision = choices[np.argmax(logps)]
+        return decision, logps
+
+    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},
+        )
+        assert res.status_code == 200
+
+    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]

python/sglang/backend/tgi.py

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

python/sglang/flush_cache.py

+"""Flush cache in the backend by sending random requests."""
+import argparse
+import random
+import string
+import time
+
+from sglang.test.test_utils import (
+    add_common_sglang_args_and_parse,
+    select_sglang_backend,
+)
+
+import sglang as sgl
+
+
+@sgl.function
+def flush_radix_cache(s, prompt):
+    s += prompt + sgl.gen("flush", max_tokens=1, stop="END")
+
+
+def main(args, max_total_tokens, context_length, print_flag):
+    backend = select_sglang_backend(args)
+    flush_length = int(context_length * 0.8)
+    batch_size = int(max_total_tokens / flush_length)
+    prompt_length = flush_length * 2
+    prompts = [
+        " ".join(random.choices(string.ascii_letters, k=int(prompt_length)))
+        for _ in range(batch_size)
+    ]
+    arguments = [{"prompt": prompts[i]} for i in range(batch_size)]
+
+    start_time = time.time()
+    flush_radix_cache.run_batch(
+        arguments, temperature=0, backend=backend, num_threads=1
+    )
+    end_time = time.time()
+
+    if print_flag:
+        print(
+            f"Flush length: {flush_length}\n",
+            f"Prompt length: {prompt_length}\n",
+            f"Total Prompt letters: {batch_size * prompt_length}\n",
+            f"Flush radix cache latency: {end_time - start_time:.3f}",
+            sep="",
+        )
+
+    # to prevent the backend still running
+    time.sleep(1)
+
+
+def run_flush(args, max_total_tokens=20000, context_length=1024, print_flag=False):
+    main(args, max_total_tokens, context_length, print_flag=print_flag)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--max-total-tokens", type=int, default=20000)
+    parser.add_argument("--context-length", type=int, default=1024)
+    args = add_common_sglang_args_and_parse(parser)
+    random.seed(0)
+    main(args, args.max_total_tokens, args.context_length, print_flag=True)

python/sglang/global_config.py

+"""Global configurations"""
+
+
+class GlobalConfig:
+    def __init__(self):
+        # Verbosity level
+        # 0: do not output anything
+        # 2: output final text after every run
+        self.verbosity = 0
+
+        self.default_backend = None
+
+        # Output configs
+        self.skip_special_tokens_in_output = True
+
+        # Optimization configs
+        self.eager_fill_image = False
+        self.enable_prefix_sharing = True
+        self.enable_parallel_encoding = True
+        self.enable_parallel_decoding = True
+
+        # Choices: ["no_adjust", "adjust_cache"]
+        # no_adjust: Do not adjust the position embedding of KV cache.
+        # adjust_cache: Adjust the position embedding of KV cache.
+        self.concate_and_append_mode = "no_adjust"
+
+
+global_config = GlobalConfig()

python/sglang/lang/chat_template.py

+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]]
+    image_token: str = "<image>"
+    style: ChatTemplateStyle = ChatTemplateStyle.PLAIN
+
+    def get_prefix_and_suffix(self, role, hist_messages):
+        if self.style == ChatTemplateStyle.PLAIN:
+            return self.role_prefix_and_suffix[role]
+        elif self.style == ChatTemplateStyle.LLAMA2:
+            if len(hist_messages) == 0 and role == "system":
+                return (
+                    self.role_prefix_and_suffix["user"][0]
+                    + self.role_prefix_and_suffix["system"][0],
+                    self.role_prefix_and_suffix["system"][1],
+                )
+            elif (
+                len(hist_messages) == 1
+                and role == "user"
+                and hist_messages[0]["content"] is not None
+            ):
+                return ("", self.role_prefix_and_suffix["user"][1])
+            return self.role_prefix_and_suffix[role]
+        else:
+            raise ValueError(f"Invalid style: {self.style}")
+
+    def get_prompt(self, messages):
+        prompt = ""
+        for i in range(len(messages)):
+            role, content = messages[i]["role"], messages[i]["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 += 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 = matching_func(model_path)
+        if template is not None:
+            return template
+    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", "\n<|im_end|>\n"),
+            "user": ("<|im_start|>user\n", "\n<|im_end|>\n"),
+            "assistant": ("<|im_start|>assistant\n", "\n<|im_end|>\n"),
+        },
+        style=ChatTemplateStyle.PLAIN,
+    )
+)
+
+
+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,
+    )
+)
+
+
+@register_chat_template_matching_function
+def match_vicuna(model_path: str):
+    if "vicuna" in model_path.lower():
+        return get_chat_template("vicuna_v1.1")
+    if "llava" in model_path.lower():
+        return get_chat_template("vicuna_v1.1")
+
+
+@register_chat_template_matching_function
+def match_llama2_chat(model_path: str):
+    model_path = model_path.lower()
+    if "llama-2" in model_path and "chat" in model_path:
+        return get_chat_template("llama-2-chat")
+    if (
+        "mistral" in model_path or "mixtral" in model_path
+    ) and "instruct" in model_path:
+        return get_chat_template("llama-2-chat")
+    if "codellama" in model_path and "instruct" in model_path:
+        return get_chat_template("llama-2-chat")
+
+
+@register_chat_template_matching_function
+def match_chat_ml(model_path: str):
+    if "tinyllama" in model_path.lower():
+        return get_chat_template("chatml")
+
+
+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/compiler.py

+import multiprocessing
+from concurrent.futures import ThreadPoolExecutor
+from queue import Queue
+from typing import List, Union
+
+from sglang.global_config import global_config
+from sglang.lang.interpreter import ProgramState, StreamExecutor, pin_program
+from sglang.lang.ir import (
+    SamplingParams,
+    SglArgument,
+    SglConstantText,
+    SglExpr,
+    SglVariable,
+)
+
+
+def compile_func(function, backend):
+    tracer = function.trace(backend=backend)
+    compiler = CompiledFunction(tracer, function)
+    return compiler
+
+
+class CompiledFunction:
+    def __init__(self, tracer, function):
+        self.function = function
+
+        self.last_node = CompGraphNode(tracer.last_node)
+        self.expr_to_node = {}
+        self.build_graph(tracer)
+        self.topological_sort()
+
+    def build_graph(self, tracer):
+        self.nodes = [self.last_node]
+        self.expr_to_node[tracer.last_node] = self.nodes[-1]
+
+        rename_pid = {}
+
+        visited = set([tracer.last_node])
+        head = 0
+        while head < len(self.nodes):
+            cur_node = self.nodes[head]
+
+            # add prev node
+            prev_node = cur_node.expr.prev_node
+            if prev_node is not None:
+                if prev_node not in visited:
+                    visited.add(prev_node)
+                    self.nodes.append(CompGraphNode(prev_node))
+                    self.expr_to_node[prev_node] = self.nodes[-1]
+                cur_node.prev_node = self.expr_to_node[prev_node]
+                self.expr_to_node[prev_node].add_next_node(cur_node)
+
+            # add source node
+            if isinstance(cur_node.expr, SglVariable):
+                if cur_node.expr.name in tracer.variables:
+                    source = tracer.variables[cur_node.expr.name].source
+                else:
+                    source = cur_node.expr.source
+                if source not in visited:
+                    visited.add(source)
+                    self.nodes.append(CompGraphNode(source))
+                    self.expr_to_node[source] = self.nodes[-1]
+                cur_node.source_node = self.expr_to_node[source]
+                self.expr_to_node[source].add_next_node(cur_node)
+            head += 1
+
+            # rename pid
+            if cur_node.expr.pid not in rename_pid:
+                rename_pid[cur_node.expr.pid] = len(rename_pid)
+            cur_node.expr.pid = rename_pid[cur_node.expr.pid]
+
+    def topological_sort(self):
+        prevd = {}
+        cand = Queue()
+        for x in self.nodes:
+            prevd[x] = (x.prev_node is not None) + (x.source_node is not None)
+            if prevd[x] == 0:
+                cand.put(x)
+        new_list = []
+        while cand.qsize() > 0:
+            head = cand.get()
+            new_list.append(head)
+            for x in head.next_nodes:
+                prevd[x] -= 1
+                if prevd[x] == 0:
+                    cand.put(x)
+        self.nodes = new_list
+
+    def print_graph(
+        self,
+    ):
+        for node in self.nodes:
+            print(node)
+
+    def run_internal(
+        self,
+        backend,
+        kwargs,
+        default_sampling_para,
+    ):
+        stream_executor_ids = set([x.expr.pid for x in self.nodes])
+        stream_executors = {}
+        for x in stream_executor_ids:
+            arguments = kwargs if x == self.last_node.expr.pid else {}
+            stream_executors[x] = StreamExecutor(
+                backend, arguments, default_sampling_para, None, False
+            )
+        for node in self.nodes:
+            se_id = node.expr.pid
+            expr = node.expr
+            if isinstance(expr, SglVariable):
+                # Make a copy for SglVariable
+                expr = SglVariable(expr.name, expr.source)
+                expr.source_stream_executor = stream_executors[
+                    node.source_node.expr.pid
+                ]
+            elif isinstance(expr, SglArgument):
+                # Substitute SglArgument
+                expr = kwargs[expr.name]
+            stream_executors[se_id].submit(expr)
+        for stream_executor in stream_executors.values():
+            stream_executor.end()
+        return ProgramState(stream_executors[self.last_node.expr.pid])
+
+    def run(
+        self,
+        *,
+        max_new_tokens: int = 16,
+        stop: Union[str, List[str]] = (),
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        frequency_penalty: float = 0.0,
+        presence_penalty: float = 0.0,
+        backend=None,
+        **kwargs,
+    ):
+        backend = backend or global_config.default_backend
+
+        kwargs = {k: SglArgument(k, v) for k, v in kwargs.items()}
+        kwargs.update(self.function.bind_arguments)
+
+        default_sampling_para = SamplingParams(
+            max_new_tokens=max_new_tokens,
+            stop=stop,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+        )
+
+        return self.run_internal(backend, kwargs, default_sampling_para)
+
+    def run_batch(
+        self,
+        batch_kwargs,
+        *,
+        max_new_tokens: int = 16,
+        stop: Union[str, List[str]] = (),
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        frequency_penalty: float = 0.0,
+        presence_penalty: float = 0.0,
+        backend=None,
+        num_threads: Union[str, int] = "auto",
+    ):
+        assert isinstance(batch_kwargs, (list, tuple))
+        if len(batch_kwargs) == 0:
+            return []
+        assert isinstance(batch_kwargs[0], dict)
+
+        backend = backend or global_config.default_backend
+
+        default_sampling_para = SamplingParams(
+            max_new_tokens=max_new_tokens,
+            stop=stop,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+        )
+        batch_kwargs = [
+            {k: SglArgument(k, v) for k, v in kwargs.items()} for kwargs in batch_kwargs
+        ]
+
+        # Extract prefix by tracing and cache it
+        if len(batch_kwargs) > 1:
+            pin_program(self.function, backend)
+
+        # Run all programs
+        if num_threads == "auto":
+            num_threads = multiprocessing.cpu_count()
+        num_threads = min(num_threads, len(batch_kwargs))
+
+        if num_threads == 1:
+            rets = []
+            for arguments in batch_kwargs:
+                rets.append(
+                    self.run_internal(backend, arguments, default_sampling_para)
+                )
+        else:
+            with ThreadPoolExecutor(num_threads) as executor:
+                futures = []
+                for arguments in batch_kwargs:
+                    futures.append(
+                        executor.submit(
+                            self.run_internal, backend, arguments, default_sampling_para
+                        )
+                    )
+                rets = [f.result() for f in futures]
+            rets[-1].sync()
+
+        return rets
+
+
+class CompGraphNode:
+    def __init__(
+        self, expr: SglExpr, prev_node=None, next_nodes=None, source_node=None
+    ):
+        self.expr = expr
+        self.next_nodes = next_nodes or []
+        self.prev_node = prev_node
+        self.source_node = source_node
+
+    def add_next_node(self, other):
+        self.next_nodes.append(other)
+
+    def __repr__(self):
+        re = f"stream {self.expr.pid:2d}: "
+        re += f"%{self.expr.node_id} = "
+        if self.prev_node is not None:
+            re += f"%{self.prev_node.expr.node_id} + "
+        re += repr(self.expr)
+        return re

python/sglang/lang/interpreter.py

+"""The interpreter that executes SGL programs"""
+
+import asyncio
+import multiprocessing
+import queue
+import threading
+import uuid
+from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
+from typing import Any, Callable, Dict, List, Optional, Union
+
+import tqdm
+from sglang.global_config import global_config
+from sglang.lang.ir import (
+    SglArgument,
+    SglCommitLazy,
+    SglConcateAndAppend,
+    SglConstantText,
+    SglExpr,
+    SglExprList,
+    SglFunction,
+    SglGen,
+    SglImage,
+    SglRoleBegin,
+    SglRoleEnd,
+    SglSelect,
+    SglVariable,
+    SglVarScopeBegin,
+    SglVarScopeEnd,
+)
+from sglang.utils import encode_image_base64
+
+
+def run_internal(state, program, func_args, func_kwargs, sync):
+    try:
+        state.ret_value = program.func(state, *func_args, **func_kwargs)
+    except Exception as e:
+        raise e
+    finally:
+        state.stream_executor.end()
+
+    if sync:
+        state.stream_executor.sync()
+
+    if global_config.verbosity >= 2:
+        print(state.text())
+
+
+def run_program(
+    program, backend, func_args, func_kwargs, default_sampling_para, stream, sync=False
+):
+    assert backend is not None, "Please specify a backend"
+    func_kwargs.update(program.bind_arguments)
+    stream_executor = StreamExecutor(
+        backend, func_kwargs, default_sampling_para, chat_template=None, stream=stream
+    )
+    state = ProgramState(stream_executor)
+
+    if stream:
+        t = threading.Thread(
+            target=run_internal, args=(state, program, func_args, func_kwargs, sync)
+        )
+        t.start()
+        return state
+    else:
+        run_internal(state, program, func_args, func_kwargs, sync)
+        return state
+
+
+def run_program_batch(
+    program,
+    backend,
+    batch_arguments,
+    default_sampling_para,
+    num_threads,
+    progress_bar,
+):
+    # Extract prefix by tracing and cache it
+    if len(batch_arguments) > 1:
+        pin_program(program, backend)
+
+    # Run all programs
+    if num_threads == "auto":
+        num_threads = multiprocessing.cpu_count()
+    num_threads = min(num_threads, len(batch_arguments))
+
+    if num_threads == 1:
+        rets = []
+        for arguments in batch_arguments:
+            rets.append(
+                run_program(
+                    program, backend, (), arguments, default_sampling_para, False, False
+                )
+            )
+    else:
+        if progress_bar:
+            pbar = tqdm.tqdm(total=len(batch_arguments))
+
+        with ThreadPoolExecutor(num_threads) as executor:
+            futures = []
+            for arguments in batch_arguments:
+                futures.append(
+                    executor.submit(
+                        run_program,
+                        program,
+                        backend,
+                        (),
+                        arguments,
+                        default_sampling_para,
+                        False,
+                        False,
+                    )
+                )
+                if progress_bar:
+                    futures[-1].add_done_callback(lambda _: pbar.update())
+
+            rets = [f.result() for f in futures]
+        rets[-1].sync()
+
+        if progress_bar:
+            pbar.close()
+
+    return rets
+
+
+def pin_program(program, backend):
+    if global_config.enable_prefix_sharing and program.pin_prefix_rid is None:
+        # TODO: handle multiple backends
+        from sglang.lang.tracer import extract_prefix_by_tracing
+
+        prefix = extract_prefix_by_tracing(program, backend)
+        if prefix and len(prefix) > 64:
+            prefix_rid = backend.cache_prefix(prefix)
+            program.pin_prefix_rid = prefix_rid
+            return prefix_rid
+    return None
+
+
+def unpin_program(program, backend):
+    pass
+
+
+class StreamExecutor:
+    """A stream executor that executes SGL expressions in a background thread."""
+
+    def __init__(
+        self,
+        backend,
+        arguments,
+        default_sampling_para,
+        chat_template,
+        stream,
+        use_thread=True,
+    ):
+        self.sid = uuid.uuid4().hex
+        self.backend = backend
+        self.arguments: Dict[str, Any] = arguments
+        self.default_sampling_para = default_sampling_para
+        self.stream = stream
+
+        if hasattr(backend, "endpoint"):
+            self.backend = backend.endpoint
+
+        self.variables = {}  # Dict[name: str -> value: str]
+        self.variable_event = {}  # Dict[name: str -> event: threading.Event]
+        self.meta_info = {}  # Dict[name: str -> info: str]
+        self.is_finished = False
+
+        # For completion
+        self.text_ = ""  # The full text
+
+        # For chat
+        self.messages_ = []  # The messages in the OpenAI API format
+        self.chat_template = chat_template or self.backend.get_chat_template()
+        self.cur_role = None
+        self.cur_role_begin_pos = None
+
+        # For vision
+        self.images_ = []
+        self.cur_images = []
+
+        # For fork/join
+        self.fork_start_text_pos = None
+
+        # Worker thread
+        self.use_thread = use_thread
+        if self.use_thread:
+            self.queue = queue.Queue()
+            self.worker = threading.Thread(target=self._thread_worker_func)
+            self.worker.start()
+
+        # For streaming
+        if stream:
+            self.stream_text_event = threading.Event()
+            self.stream_var_event = {}
+        else:
+            self.stream_text_event = None
+            self.stream_var_event = None
+
+    def submit(self, expr: SglExpr):
+        if isinstance(expr, (SglGen, SglSelect, SglVarScopeBegin)):
+            self.variable_event[expr.name] = threading.Event()
+            if self.stream:
+                self.stream_var_event[expr.name] = threading.Event()
+        elif isinstance(expr, SglExprList):
+            for e in expr.expr_list:
+                if isinstance(e, (SglGen, SglSelect, SglVarScopeBegin)):
+                    self.variable_event[e.name] = threading.Event()
+                    if self.stream:
+                        self.stream_var_event[e.name] = threading.Event()
+
+        if self.use_thread:
+            self.queue.put(expr)
+        else:
+            self._execute(expr)
+
+    def sync(self):
+        if self.use_thread:
+            self.queue.join()
+
+    def get_var(self, name):
+        if name in self.variable_event:
+            self.variable_event[name].wait()
+        return self.variables[name]
+
+    def get_meta_info(self, name):
+        if name in self.variable_event:
+            self.variable_event[name].wait()
+        ret = self.meta_info.get(name, None)
+        return ret
+
+    def fork(self, number: int, position_ids_offset: Optional[List[int]] = None):
+        if number > 1:
+            self.submit(SglCommitLazy())
+            self.sync()
+
+        number = int(number)
+
+        exes = [
+            StreamExecutor(
+                self.backend,
+                self.arguments,
+                self.default_sampling_para,
+                self.chat_template,
+                self.stream,
+            )
+            for _ in range(number)
+        ]
+        for i in range(number):
+            exes[i].variables = dict(self.variables)
+            exes[i].text_ = str(self.text_)
+            exes[i].messages_ = list(self.messages_)
+            exes[i].cur_role = self.cur_role
+            exes[i].fork_start_text_pos = len(self.text_)
+
+        return exes
+
+    def text(self):
+        self.sync()
+        return self.text_
+
+    def messages(self):
+        self.sync()
+        return self.messages_
+
+    def end(self):
+        if self.use_thread:
+            if self.worker.is_alive():
+                self.queue.put(None)
+        self.backend.end_program(self)
+
+    def _thread_worker_func(self):
+        while True:
+            expr = self.queue.get()
+            if expr is None:
+                self.queue.task_done()
+                break
+
+            self._execute(expr)
+            self.queue.task_done()
+            if self.stream_text_event:
+                self.stream_text_event.set()
+
+        if self.stream_text_event:
+            self.stream_text_event.set()
+
+        self.is_finished = True
+
+    def _execute(self, other):
+        if isinstance(other, str):
+            other = SglConstantText(other)
+
+        assert isinstance(other, SglExpr), f"{other}"
+
+        if isinstance(other, (SglConstantText, SglArgument)):
+            self._execute_fill(other.value)
+        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, SglImage):
+            self._execute_image(other)
+        elif isinstance(other, SglVariable):
+            self._execute_variable(other)
+        elif isinstance(other, SglVarScopeBegin):
+            self._execute_var_scope_begin(other)
+        elif isinstance(other, SglVarScopeEnd):
+            self._execute_var_scope_end(other)
+        elif isinstance(other, SglCommitLazy):
+            self._execute_commit_lazy_operations(other)
+        elif isinstance(other, SglConcateAndAppend):
+            if (
+                global_config.enable_parallel_encoding
+                and self.backend.support_concate_and_append
+            ):
+                self._execute_concatenate_and_append_kv_cache(other)
+            else:
+                self._execute_concatenate_and_append_text(other)
+        else:
+            raise ValueError(f"Unknown type: {type(other)}")
+
+    def _execute_fill(self, value: str):
+        value = str(value)
+        self.text_ += value
+
+    def _execute_image(self, expr: SglImage):
+        path = expr.path
+        if isinstance(path, SglArgument):
+            path = path.value
+
+        base64_data = encode_image_base64(path)
+
+        self.images_.append((path, base64_data))
+        self.cur_images.append((path, base64_data))
+        self.text_ += self.chat_template.image_token
+
+        # if global_config.eager_fill_image:
+        #     self.backend.fill_image(self)
+
+    def _execute_gen(self, expr: SglGen):
+        sampling_params = self._resolve_sampling_params(expr.sampling_params)
+        name = expr.name
+
+        if not self.stream:
+            comp, meta_info = self.backend.generate(
+                self, sampling_params=sampling_params
+            )
+            self.text_ += comp
+
+            self.variables[name] = comp
+            self.meta_info[name] = meta_info
+            self.variable_event[name].set()
+        else:
+            generator = self.backend.generate_stream(
+                self, sampling_params=sampling_params
+            )
+
+            self.stream_var_event[name].set()
+
+            self.variables[name] = ""
+            for comp, meta_info in generator:
+                self.text_ += comp
+                self.variables[name] += comp
+                self.stream_var_event[name].set()
+                self.stream_text_event.set()
+
+            self.meta_info[name] = meta_info
+
+            self.variable_event[name].set()
+            self.stream_var_event[name].set()
+
+    def _execute_select(self, expr: SglSelect):
+        decision, scores = self.backend.select(self, expr.choices, expr.temperature)
+        if expr.name is not None:
+            name = expr.name
+            self.variables[name] = decision
+            self.variable_event[name].set()
+        self.text_ += decision
+
+    def _execute_variable(self, expr: SglVariable):
+        src_executor = expr.source_stream_executor
+        value = src_executor.get_var(expr.name)
+        self._execute_fill(value)
+
+    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 the 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, _ = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_)
+
+        self._execute_fill(prefix)
+        self.cur_role_begin_pos = len(self.text_)
+
+    def _execute_role_end(self, expr: SglRoleEnd):
+        new_text = self.text_[self.cur_role_begin_pos :].lstrip()
+
+        _, suffix = self.chat_template.get_prefix_and_suffix(expr.role, self.messages_)
+        self._execute_fill(suffix)
+
+        if self.cur_images:
+            # OpenAI vision API format
+            last_msg = {
+                "role": expr.role,
+                "content": [{"type": "text", "text": new_text}],
+            }
+            for (image_path, image_base64_data) in self.cur_images:
+                last_msg["content"].append(
+                    {
+                        "type": "image_url",
+                        "image_url": {
+                            "url": f"data:image/jpeg;base64,{image_base64_data}"
+                        },
+                    }
+                )
+            self.messages_.append(last_msg)
+            self.cur_images = []
+        else:
+            self.messages_.append({"role": expr.role, "content": new_text})
+
+        self.cur_role = None
+
+    def _execute_var_scope_begin(self, expr: SglVarScopeBegin):
+        self.variables[expr.name] = int(len(self.text_))
+
+    def _execute_var_scope_end(self, expr: SglVarScopeEnd):
+        self.variables[expr.name] = self.text_[self.variables[expr.name] :]
+        self.variable_event[expr.name].set()
+
+    def _execute_commit_lazy_operations(self, expr: SglCommitLazy):
+        self.backend.commit_lazy_operations(self)
+
+    def _execute_concatenate_and_append_text(self, expr: SglConcateAndAppend):
+        new_text = ""
+        for s in expr.states:
+            exe = s.stream_executor
+            exe.sync()
+            new_text += exe.text_[exe.fork_start_text_pos :]
+
+        self._execute_fill(new_text)
+
+    def _execute_concatenate_and_append_kv_cache(self, expr: SglConcateAndAppend):
+        self_len = len(self.text_)
+
+        for i, s in enumerate(expr.states):
+            exe = s.stream_executor
+            exe.submit(SglCommitLazy())
+
+        for i, s in enumerate(expr.states):
+            exe = s.stream_executor
+            exe.sync()
+            assert exe.fork_start_text_pos == self_len
+            self.text_ += exe.text_[exe.fork_start_text_pos :]
+
+        src_rids = [state.stream_executor.sid for state in expr.states]
+        self.backend.concatenate_and_append(src_rids, self.sid)
+
+    def _resolve_sampling_params(self, sampling_params):
+        clone = None
+        for item in [
+            "max_new_tokens",
+            "stop",
+            "temperature",
+            "top_p",
+            "top_k",
+            "frequency_penalty",
+            "presence_penalty",
+            "dtype",
+            "regex",
+        ]:
+            value = getattr(sampling_params, item, None)
+            if value is not None:
+                if clone is None:
+                    clone = self.default_sampling_para.clone()
+                setattr(clone, item, value)
+        return clone or self.default_sampling_para
+
+    def __del__(self):
+        self.end()
+
+
+class ProgramState:
+    """The state of an SGL program."""
+
+    def __init__(self, stream_executor: StreamExecutor):
+        self.stream_executor = stream_executor
+
+    def _role_common(self, name: str, expr: Optional[SglExpr] = None):
+        if expr is not None:
+            self.stream_executor.submit(
+                SglExprList([SglRoleBegin(name), expr, SglRoleEnd(name)])
+            )
+        else:
+
+            @contextmanager
+            def role_scope():
+                self.stream_executor.submit(SglRoleBegin(name))
+                yield
+                self.stream_executor.submit(SglRoleEnd(name))
+
+            return role_scope()
+
+    def system(self, expr: Optional[SglExpr] = None):
+        return self._role_common("system", expr)
+
+    def user(self, expr: Optional[SglExpr] = None):
+        return self._role_common("user", expr)
+
+    def assistant(self, expr: Optional[SglExpr] = None):
+        return self._role_common("assistant", expr)
+
+    @contextmanager
+    def var_scope(self, name: str):
+        self.stream_executor.submit(SglVarScopeBegin(name))
+        yield
+        self.stream_executor.submit(SglVarScopeEnd(name))
+
+    def fork(self, number: int = 1, position_ids_offset: Optional[List[int]] = None):
+        stream_executors = self.stream_executor.fork(number, position_ids_offset)
+        states = [ProgramState(x) for x in stream_executors]
+        state_group = ProgramStateGroup(states, self)
+        return state_group
+
+    @contextmanager
+    def copy(self, position_ids_offset: Optional[List[int]] = None):
+        state_group = self.fork(1, position_ids_offset)
+        try:
+            yield state_group[0]
+        finally:
+            state_group.join()
+
+    def text(self):
+        return self.stream_executor.text()
+
+    def messages(self):
+        return self.stream_executor.messages()
+
+    def sync(self):
+        return self.stream_executor.sync()
+
+    def text_iter(self, var_name=None):
+        if self.stream_executor.stream:
+            prev = 0
+            if var_name is None:
+                event = self.stream_executor.stream_text_event
+                while True:
+                    event.wait()
+                    event.clear()
+                    out = str(self.stream_executor.text_[prev:])
+                    prev += len(out)
+                    if out:
+                        yield out
+                    if self.stream_executor.is_finished:
+                        break
+            else:
+                event = self.stream_executor.stream_var_event[var_name]
+                while True:
+                    event.wait()
+                    event.clear()
+                    out = str(self.stream_executor.variables[var_name][prev:])
+                    prev += len(out)
+                    if out:
+                        yield out
+                    if self.stream_executor.variable_event[var_name].is_set():
+                        break
+        else:
+            if var_name is None:
+                yield self.text()
+            else:
+                yield self.get_var(name)
+
+    async def text_async_iter(self, var_name=None):
+        loop = asyncio.get_running_loop()
+
+        if self.stream_executor.stream:
+            prev = 0
+            if var_name is None:
+                event = self.stream_executor.stream_text_event
+                while True:
+                    await loop.run_in_executor(None, event.wait)
+                    event.clear()
+                    out = str(self.stream_executor.text_[prev:])
+                    prev += len(out)
+                    if out:
+                        yield out
+                    if self.stream_executor.is_finished:
+                        break
+            else:
+                event = self.stream_executor.stream_var_event[var_name]
+                while True:
+                    await loop.run_in_executor(None, event.wait)
+                    event.clear()
+                    out = str(self.stream_executor.variables[var_name][prev:])
+                    prev += len(out)
+                    if out:
+                        yield out
+                    if self.stream_executor.variable_event[var_name].is_set():
+                        break
+        else:
+            if var_name is None:
+                yield self.text()
+            else:
+                yield self.get_var(name)
+
+    def get_var(self, name):
+        return self.stream_executor.get_var(name)
+
+    def get_meta_info(self, name):
+        return self.stream_executor.get_meta_info(name)
+
+    def __iadd__(self, other):
+        self.stream_executor.submit(other)
+        return self
+
+    def __getitem__(self, name):
+        return self.get_var(name)
+
+    def __del__(self):
+        self.stream_executor.end()
+
+    def __repr__(self) -> str:
+        msgs = self.messages()
+        ret = ""
+        for msg in msgs:
+            ret += msg["role"] + ":\n" + msg["content"] + "\n"
+        return ret
+
+
+class ProgramStateGroup:
+    def __init__(
+        self, states: List[ProgramState], src_state: Optional[ProgramState] = None
+    ):
+        self.states = states
+        self.src_state = src_state
+
+    def join(self, mode: str = "gather_variable"):
+        if mode == "gather_variable":
+            # Copy variables back
+            src_vars = self.src_state.stream_executor.variables
+            src_var_set = set(src_vars.keys())
+            for child_state in self.states:
+                child_state.stream_executor.sync()
+                child_vars = child_state.stream_executor.variables
+                new_vars = set(child_vars.keys()) - src_var_set
+
+                for k in new_vars:
+                    if k in src_vars:
+                        src_vars[k].append(child_vars[k])
+                    else:
+                        src_vars[k] = [child_vars[k]]
+        elif mode == "concate_and_append":
+            # Concatenate and append KV cache
+            self.src_state += SglConcateAndAppend(self.states)
+            # Need a sync here. Otherwise, `states` can be deleted.
+            self.src_state.stream_executor.sync()
+        else:
+            raise ValueError(f"Invalid join mode: {mode}")
+
+        for s in self.states:
+            s.stream_executor.end()
+
+    def __getitem__(self, i: int):
+        return self.states[i]
+
+    def __setitem__(self, i: int, value):
+        assert self.states[i] == value
+
+    def __iadd__(self, other):
+        if isinstance(other, Callable):
+            # lambda function
+            for i in range(len(self.states)):
+                self.states[i] += other(i)
+        elif isinstance(other, SglExpr):
+            for i in range(len(self.states)):
+                self.states[i] += other
+        elif isinstance(other, (list, tuple)):
+            for i in range(len(self.states)):
+                self.states[i] += other[i]
+        else:
+            raise ValueError(f"Invalid value: {other}")
+
+        return self

python/sglang/lang/ir.py

+"""The intermediate representation."""
+
+import dataclasses
+import inspect
+from typing import List, Optional, Union
+
+from sglang.global_config import global_config
+
+
+@dataclasses.dataclass
+class SamplingParams:
+    max_new_tokens: int = 16
+    stop: Union[str, List[str]] = ()
+    temperature: float = 1.0
+    top_p: float = 1.0
+    top_k: int = -1  # -1 means disable
+    frequency_penalty: float = 0.0
+    presence_penalty: float = 0.0
+
+    # for constrained generation, not included in to_xxx_kwargs
+    dtype: Optional[str] = None
+    regex: Optional[str] = None
+
+    def clone(self):
+        return SamplingParams(
+            self.max_new_tokens,
+            self.stop,
+            self.temperature,
+            self.top_p,
+            self.top_k,
+            self.frequency_penalty,
+            self.presence_penalty,
+        )
+
+    def to_openai_kwargs(self):
+        # OpenAI does not support top_k, so we drop it here
+        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_anthropic_kwargs(self):
+        # Anthropic does not support frequency_penalty or presence_penalty, so we drop it here
+        return {
+            "max_tokens_to_sample": self.max_new_tokens,
+            "stop_sequences": self.stop,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+        }
+
+    def to_srt_kwargs(self):
+        return {
+            "max_new_tokens": self.max_new_tokens,
+            "stop": self.stop,
+            "temperature": self.temperature,
+            "top_p": self.top_p,
+            "top_k": self.top_k,
+            "frequency_penalty": self.frequency_penalty,
+            "presence_penalty": self.presence_penalty,
+            "regex": self.regex,
+        }
+
+
+class SglFunction:
+    def __init__(self, func, bind_arguments=None):
+        self.func = func
+        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:]
+
+    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 = 16,
+        stop: Union[str, List[str]] = (),
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        frequency_penalty: float = 0.0,
+        presence_penalty: float = 0.0,
+        stream: bool = False,
+        backend=None,
+        **kwargs,
+    ):
+        from sglang.lang.interpreter import run_program
+
+        default_sampling_para = SamplingParams(
+            max_new_tokens=max_new_tokens,
+            stop=stop,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+        )
+        backend = backend or global_config.default_backend
+        kwargs = {k: SglArgument(k, v) for k, v in kwargs.items()}
+        return run_program(self, backend, args, kwargs, default_sampling_para, stream)
+
+    def run_batch(
+        self,
+        batch_kwargs,
+        *,
+        max_new_tokens: int = 16,
+        stop: Union[str, List[str]] = (),
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        frequency_penalty: float = 0.0,
+        presence_penalty: float = 0.0,
+        backend=None,
+        num_threads: Union[str, int] = "auto",
+        progress_bar: bool = False,
+    ):
+        from sglang.lang.interpreter import run_program_batch
+
+        assert isinstance(batch_kwargs, (list, tuple))
+        if len(batch_kwargs) == 0:
+            return []
+        assert isinstance(batch_kwargs[0], dict)
+
+        default_sampling_para = SamplingParams(
+            max_new_tokens=max_new_tokens,
+            stop=stop,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+        )
+        backend = backend or global_config.default_backend
+        batch_kwargs = [
+            {k: SglArgument(k, v) for k, v in kwargs.items()} for kwargs in batch_kwargs
+        ]
+        return run_program_batch(
+            self,
+            backend,
+            batch_kwargs,
+            default_sampling_para,
+            num_threads,
+            progress_bar,
+        )
+
+    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 pin(self, backend=None):
+        from sglang.lang.interpreter import pin_program
+
+        backend = backend or global_config.default_backend
+        return pin_program(self, backend)
+
+    def unpin(self, backend=None):
+        from sglang.lang.interpreter import unpin_program
+
+        backend = backend or global_config.default_backend
+        return unpin_program(self, backend)
+
+    def compile(self, *, backend=None):
+        from sglang.lang.compiler import compile_func
+
+        return compile_func(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):
+        self.path = path
+
+    def __repr__(self) -> str:
+        return f"SglImage({self.path})"
+
+
+class SglGen(SglExpr):
+    def __init__(
+        self,
+        name,
+        max_new_tokens,
+        stop,
+        temperature,
+        top_p,
+        top_k,
+        frequency_penalty,
+        presence_penalty,
+        dtype,
+        regex,
+    ):
+        super().__init__()
+        self.name = name
+        self.sampling_params = SamplingParams(
+            max_new_tokens=max_new_tokens,
+            stop=stop,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+            dtype=dtype,
+            regex=regex,
+        )
+
+    def __repr__(self):
+        return f"Gen('{self.name}')"
+
+
+class SglConstantText(SglExpr):
+    def __init__(self, value):
+        super().__init__()
+        self.value = value
+
+    def __repr__(self):
+        return f"Constant({repr(self.value)})"
+
+
+class SglRoleBegin(SglExpr):
+    def __init__(self, role):
+        super().__init__()
+        self.role = role
+
+    def __repr__(self):
+        return f"RoleBegin({self.role})"
+
+
+class SglRoleEnd(SglExpr):
+    def __init__(self, role):
+        super().__init__()
+        self.role = role
+
+    def __repr__(self):
+        return f"RoleEnd({self.role})"
+
+
+class SglSelect(SglExpr):
+    def __init__(self, name, choices, temperature):
+        super().__init__()
+        self.name = name
+        self.choices = choices
+        self.temperature = temperature
+
+    def __repr__(self):
+        return f"Select({self.name}, choices={self.choices})"
+
+
+class SglFork(SglExpr):
+    def __init__(self, number, 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):
+        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, 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):
+        super().__init__()
+        self.name = name
+
+    def __repr__(self):
+        return f"VarScopeBegin('{self.name}')"
+
+
+class SglVarScopeEnd(SglExpr):
+    def __init__(self, name):
+        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 f"CommitLazy()"

python/sglang/lang/tracer.py

+"""Tracing a program."""
+import uuid
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from sglang.backend.base_backend import BaseBackend
+from sglang.global_config import global_config
+from sglang.lang.interpreter import ProgramState, ProgramStateGroup
+from sglang.lang.ir import (
+    SglArgument,
+    SglCommitLazy,
+    SglConcateAndAppend,
+    SglConstantText,
+    SglExpr,
+    SglExprList,
+    SglFork,
+    SglFunction,
+    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:
+        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, number: int, position_ids_offset: Optional[List[int]] = None):
+        if self.only_trace_prefix:
+            raise StopTracing()
+
+        fork_node = SglFork(number)
+        fork_node.prev_node = self.last_node
+
+        states = [
+            TracerProgramState(self.backend, self.arguments, self.only_trace_prefix)
+            for _ in range(number)
+        ]
+
+        for i in range(number):
+            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(name, source=self.last_node)
+        self.variables[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 != None:
+            cur_scope.tracer_state.child_states.append(state)
+            cur_scope = cur_scope.last_scope

python/sglang/launch_server.py

+import argparse
+
+from sglang.srt.server import ServerArgs, launch_server
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    ServerArgs.add_cli_args(parser)
+    args = parser.parse_args()
+    server_args = ServerArgs.from_cli_args(args)
+
+    launch_server(server_args, None)

python/sglang/srt/constrained/fsm.py

+# Adapted from:
+# https://github.com/outlines-dev/outlines/blob/0355ab4272a5d7e4d94c4a53a52593f885b81a61/outlines/fsm/fsm.py
+from typing import List, NewType, Protocol
+
+import interegular
+from lark import Lark
+
+# from outlines.fsm.parsing import PartialLark
+from sglang.srt.constrained.regex import (
+    create_fsm_index_tokenizer,
+    make_deterministic_fsm,
+)
+from sglang.srt.constrained.tokenizer import Tokenizer
+
+FSMState = NewType("FSMState", int)
+
+
+class FSM(Protocol):
+    def allowed_token_ids(self, state: FSMState, idx: int = 0) -> List[int]:
+        ...
+
+    def next_state(self, state: FSMState, token_id: int, idx: int = 0) -> FSMState:
+        ...
+
+    def is_final_state(self, state: FSMState, idx: int = 0) -> bool:
+        ...
+
+    def reset(self) -> None:
+        ...
+
+
+class StopAtTokenFSM(FSM):
+    """FSM to generate text until a specified token id is generated or
+    a specified number of tokens has been generated.
+
+    Text is usually produced until the EOS token is generated by the
+    model.
+
+    """
+
+    def __init__(
+        self,
+        tokenizer: "Tokenizer",
+        stop_token_id: int,
+    ):
+        self.stop_token_id = stop_token_id
+        self.num_tokens_generated = 0
+        self.vocabulary = tokenizer.vocabulary.values()
+        self.final_states = {1}
+
+    def allowed_token_ids(self, state: FSMState, idx: int = 0) -> List[int]:
+        """Generate a list of allowed tokens for the next step.
+
+        When in the initial state we allow every token to be generated.
+        In the final state the only allowed token is `stop_token_id`.
+
+        Parameters
+        ----------
+        state
+            The current state of the FSM.
+        idx
+            The index of the current input in the batch.
+
+        Returns
+        -------
+        A list that contains the tokens to mask.
+
+        """
+        if state == 0:
+            return list(self.vocabulary)
+        else:
+            return [self.stop_token_id]
+
+    def next_state(self, state: FSMState, token_id: int, idx: int = 0) -> FSMState:
+        """Update the state of the FSM.
+
+        The FSM stays in the initial state `0` unless the specified stop token
+        has been generated or the maximum number of tokens has been reached. In
+        which case the FSM moves to the final state `1`.
+
+        Parameters
+        ----------
+        state
+            The current state of the FSM.
+        token_id
+            The id of the token that was just generated.
+        idx
+            The index of the current input in the batch.
+
+        Returns
+        -------
+        The new state of the FSM.
+
+        """
+        if idx == 0:
+            self.num_tokens_generated += 1
+
+        if token_id == self.stop_token_id:
+            return FSMState(1)
+
+        return FSMState(0)
+
+    def is_final_state(self, state: FSMState, idx: int = 0) -> bool:
+        """Determine whether the current state of the FSM is a final state."""
+        return state in self.final_states
+
+    def reset(self) -> None:
+        """Reset the FSM to its initial state. Here this only resets the token counter."""
+        self.num_tokens_generated = 0
+
+
+class RegexFSM(FSM):
+    """FSM to generate text that is in the language of a regular expression."""
+
+    def __init__(
+        self,
+        regex_string: str,
+        tokenizer: "Tokenizer",
+    ):
+        regex_pattern = interegular.parse_pattern(regex_string)
+        regex_fsm, _ = make_deterministic_fsm(regex_pattern.to_fsm().reduce())
+        (
+            self.states_to_token_maps,
+            self.empty_token_ids,
+        ) = create_fsm_index_tokenizer(regex_fsm, tokenizer)
+
+        # We make sure that it is possible to generate strings in the language
+        # of the regular expression with the tokens present in the model's
+        # vocabulary.
+        if not any(
+            regex_fsm.finals.intersection(v.values())
+            for v in self.states_to_token_maps.values()
+        ):
+            raise ValueError(
+                "The vocabulary does not allow us to build a sequence that matches the input regex"
+            )
+
+        self.final_states = regex_fsm.finals | {
+            -1
+        }  # Include the EOS token in final states
+        self.num_tokens_generated = 0
+        self.vocabulary = tokenizer.vocabulary.values()
+        self.end_token_id = tokenizer.eos_token_id
+
+    def allowed_token_ids(self, state: FSMState, idx: int = 0) -> List[int]:
+        """Generate a list of allowed tokens for the next step.
+
+        The initialization of the FSM builds an index which maps FSM states to a
+        map from authorized tokens to the state in which the FSM needs to move
+        if said token is generated. Therefore the authorized tokens at the
+        current state are the keys of the map returned by the value of the index
+        for current state.
+
+        If the current state is not contained in the end this means that we are
+        in a final state of the FSM. We only authorize EOS tokens in the final
+        state.
+
+        Parameters
+        ----------
+        state
+            The current state of the FSM.
+        idx
+            The index of the current input in the batch.
+
+        Returns
+        -------
+        A list that contains the tokens to mask.
+
+        """
+        next_tokens_to_end_states = self.states_to_token_maps.get(state)
+
+        if next_tokens_to_end_states is None:
+            return [self.end_token_id]
+        else:
+            return list(next_tokens_to_end_states.keys())
+
+    def next_state(self, state: FSMState, token_id: int, idx: int = 0) -> FSMState:
+        """Update the state of the FSM.
+
+        We use the index to determine to which state the FSM should transition
+        given the token that was just generated.
+
+        Parameters
+        ----------
+        state
+            The current state of the FSM.
+        token_id
+            The id of the token that was just generated.
+        idx
+            The index of the current input in the batch.
+
+        Returns
+        -------
+        The new state of the FSM.
+
+        """
+        if idx == 0:
+            self.num_tokens_generated += 1
+
+        if token_id == self.end_token_id:
+            return FSMState(-1)
+
+        last_token_to_end_state = self.states_to_token_maps[state]
+        next_state = last_token_to_end_state.get(token_id)
+        if next_state is None:
+            next_state = -1
+
+        return FSMState(next_state)
+
+    def is_final_state(self, state: FSMState, idx: int = 0) -> bool:
+        """Determine whether the current state of the FSM is a final state."""
+        return state in self.final_states
+
+    def reset(self) -> None:
+        """Reset the FSM to its initial state. Here this only resets the token counter."""
+        self.num_tokens_generated = 0
+
+
+class CFGFSM(FSM):
+    """FSM to generate text that is in the language of a context-free grammar."""
+
+    def __init__(
+        self,
+        cfg_string: str,
+        tokenizer: "Tokenizer",
+    ):
+        # self.parser = PartialLark(cfg_string, parser="lalr")
+        self.parser = Lark(
+            cfg_string,
+            parser="lalr",
+            lexer="contextual",
+            propagate_positions=False,
+            maybe_placeholders=False,
+            regex=True,
+        )
+        self.terminal_regexps = dict()
+        for terminal in self.parser.terminals:
+            if terminal.pattern is not None:
+                self.terminal_regexps[terminal.name] = terminal.pattern.to_regexp()
+        self.terminal_regexps["$END"] = tokenizer.eos_token
+
+        self.tokenizer = tokenizer
+        self.num_tokens_generated = 0
+        self.generations: List[str] = []
+        self.regex_fsms: List[RegexFSM] = []
+        self.reset_state: List[bool] = []
+        self.allow_eos: List[bool] = []
+        self.done: List[bool] = []
+
+    def _set_next_regex_fsm(self, idx: int = 0) -> None:
+        """Use the CFG incremental parser to set the next regex FSM.
+
+        Check what the CFG incremental parser proposes next.
+        If the only proposal is the EOS token,
+            we set the state to done and return.
+        If there are other proposals,
+            we set a new regex FSM and return.
+
+        """
+        interactive = self.parser.parse_interactive(self.generations[idx])
+        interactive.exhaust_lexer()
+        options = {self.terminal_regexps[x] for x in interactive.accepts()}
+
+        if self.terminal_regexps["$END"] in options:
+            options.remove(self.terminal_regexps["$END"])
+            if len(options) == 0:
+                self.done[idx] = True
+                return
+            self.allow_eos[idx] = True
+            options.add("")
+            assert len(options) > 1
+
+        regex_string = r"(" + r"|".join([r"(" + x + r")" for x in options]) + r")"
+        args = (
+            regex_string,
+            self.tokenizer,
+        )
+        if len(self.regex_fsms) <= idx:
+            self.regex_fsms.append(RegexFSM(*args))
+        else:
+            self.regex_fsms[idx] = RegexFSM(*args)
+        self.reset_state[idx] = True
+
+    def allowed_token_ids(self, state: FSMState, idx: int = 0) -> List[int]:
+        """Generate a list of allowed tokens for the next step.
+
+        Upon initialization, the CFG incremental parser is used to determine the first regex.
+
+        This regex is used for proposals until either:
+        - the regex is exhausted, and its only remaining option is the EOS token,
+            in which case we always transition to the next regex
+        - the regex can be exhausted, but the EOS token is not the only remaining option,
+            in which case we transition to the next regex with probability P (TODO)
+            or remove the possibility of generating the EOS token and continue with the current regex
+
+        The CFG incremental parser is allowed to propose the EOS token from any final state,
+        and once it is generated, the FSM will continue to always generate the EOS token.
+
+        Parameters
+        ----------
+        state
+            The current state of the FSM.
+        idx
+            The index of the current input in the batch.
+
+        Returns
+        -------
+        A list that contains the tokens to mask.
+
+        """
+        if len(self.generations) <= idx:
+            self.generations.append("")
+            self.reset_state.append(False)
+            self.allow_eos.append(False)
+            self.done.append(False)
+
+        if len(self.regex_fsms) > idx:
+            proposal = self.regex_fsms[idx].allowed_token_ids(state)
+            if self.tokenizer.eos_token_id not in proposal:
+                return proposal
+            if set(proposal) != {self.tokenizer.eos_token_id}:
+                if False:  # TODO: THIS NEEDS TO BE SAMPLED
+                    proposal = [x for x in proposal if x != self.tokenizer.eos_token_id]
+                    return proposal
+
+        self._set_next_regex_fsm(idx)
+
+        if self.done[idx]:
+            return [self.tokenizer.eos_token_id]
+
+        if self.reset_state[idx]:
+            state = FSMState(0)
+
+        proposal = self.regex_fsms[idx].allowed_token_ids(state)
+        if self.allow_eos[idx]:
+            self.allow_eos[idx] = False
+        else:
+            proposal = [x for x in proposal if x != self.tokenizer.eos_token_id]
+            assert len(proposal) > 0
+        return proposal
+
+    def next_state(self, state: FSMState, token_id: int, idx: int = 0) -> FSMState:
+        """Update the state of the FSM.
+
+        Transitions the underlying regex FSM to its next state.
+        If at max tokens or EOS token, transition permanently to the final state.
+        Update stored partial generations for subsequent incremental parsing.
+
+        Parameters
+        ----------
+        state
+            The current state of the FSM.
+        token_id
+            The id of the token that was just generated.
+        idx
+            The index of the current input in the batch.
+
+        Returns
+        -------
+        The new state of the FSM.
+        """
+        if idx == 0:
+            self.num_tokens_generated += 1
+        if token_id == self.tokenizer.eos_token_id:
+            self.done[idx] = True
+            return FSMState(-1)
+        if self.reset_state[idx]:
+            self.reset_state[idx] = False
+            state = FSMState(0)
+
+        self.generations[idx] += self.tokenizer.decode([token_id])[0]
+
+        return self.regex_fsms[idx].next_state(state, token_id, idx)
+
+    def is_final_state(self, state: FSMState, idx: int = 0) -> bool:
+        """Return whether the current state of the FSM is a final state."""
+        return self.done[idx]
+
+    def reset(self) -> None:
+        """Reset the FSM to its initial state, so it can be called on a fresh batch on inputs."""
+        self.num_tokens_generated = 0
+        self.generations = []
+        self.regex_fsms = []
+        self.reset_state = []
+        self.done = []

python/sglang/srt/constrained/fsm_cache.py

+import threading
+
+from sglang.srt.constrained.fsm import RegexFSM
+from sglang.srt.constrained.tokenizer import TransformerTokenizer
+
+
+def get_fsm(regex, tokenizer, fsm_cache_entry):
+    outlines_tokenizer = TransformerTokenizer(tokenizer)
+    fsm = RegexFSM(regex, outlines_tokenizer)
+    fsm_cache_entry.fsm = fsm
+    fsm_cache_entry.event.set()
+
+
+class FSMCacheEntry:
+    def __init__(self):
+        self.fsm = None
+        self.event = threading.Event()
+
+
+class FSMCache:
+    def __init__(self, tokenizer):
+        self.cache = {}
+        self.tokenizer = tokenizer
+
+    def init_fsm_in_background(self, regex):
+        if regex not in self.cache:
+            self.cache[regex] = FSMCacheEntry()
+            threading.Thread(
+                target=get_fsm,
+                args=(
+                    regex,
+                    self.tokenizer,
+                    self.cache[regex],
+                ),
+            ).start()
+
+    def get_fsm(self, regex):
+        self.init_fsm_in_background(regex)
+        entry = self.cache[regex]
+        entry.event.wait()
+        return entry.fsm

python/sglang/srt/constrained/regex.py

+# Adapted from:
+# https://github.com/outlines-dev/outlines/blob/0355ab4272a5d7e4d94c4a53a52593f885b81a61/outlines/fsm/regex.py
+from collections import namedtuple
+from functools import lru_cache
+from typing import Dict, Generator, List, Sequence, Set, Tuple
+
+import numba
+import numpy as np
+from interegular.fsm import FSM, Alphabet, OblivionError, anything_else
+from numba.typed.typedobjectutils import _nonoptional
+from sglang.srt.constrained.tokenizer import Tokenizer
+
+
+class BetterAlphabet(Alphabet):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        assert anything_else in self._symbol_mapping
+        self.anything_value = self._symbol_mapping[anything_else]
+
+    def __getitem__(self, item):
+        return self._symbol_mapping.get(item, self.anything_value)
+
+    def copy(self):
+        return BetterAlphabet(self._symbol_mapping.copy())
+
+
+class BetterFSM(FSM):
+    flat_transition_map: Dict[Tuple[int, int], int]
+    trans_key_to_states: Dict[int, List[int]]
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        if not isinstance(self.alphabet, BetterAlphabet):
+            self.__dict__["alphabet"] = BetterAlphabet(self.alphabet._symbol_mapping)
+
+        flat_transition_map = {}
+        trans_key_to_states = {}
+        for from_state, trans_map in self.map.items():
+            for trans_key, to_state in trans_map.items():
+                flat_transition_map[(from_state, trans_key)] = to_state
+                trans_key_to_states.setdefault(trans_key, set()).add(from_state)
+
+        self.__dict__["trans_key_to_states"] = trans_key_to_states
+        self.__dict__["flat_transition_map"] = flat_transition_map
+        self.__dict__["_fsm_info"] = None
+
+    def copy(self):
+        return BetterFSM(
+            alphabet=self.alphabet.copy(),
+            states=self.states.copy(),
+            initial=self.initial,
+            finals=self.finals.copy(),
+            map=self.map.copy(),
+            __no_validation__=True,
+        )
+
+    @property
+    def fsm_info(self):
+        if self._fsm_info is None:
+            flat_transition_map_items = np.fromiter(
+                ((a[0], a[1], b) for a, b in self.flat_transition_map.items()),
+                dtype=np.dtype("i8, i8, i8"),
+            )
+            trans_key_to_states_items = np.fromiter(
+                ((k, z) for k, v in self.trans_key_to_states.items() for z in v),
+                dtype=np.dtype("i8, i8"),
+            )
+            alphabet_symbol_mapping_items = np.fromiter(
+                (
+                    it
+                    for it in self.alphabet._symbol_mapping.items()
+                    if it[0] != anything_else
+                ),
+                dtype=np.dtype("U1, i8"),
+            )
+            nb_finals = np.fromiter(self.finals, dtype=np.dtype("i8"))
+            self.__dict__["_fsm_info"] = create_fsm_info(
+                self.initial,
+                nb_finals,
+                flat_transition_map_items,
+                trans_key_to_states_items,
+                self.alphabet.anything_value,
+                alphabet_symbol_mapping_items,
+            )
+
+        return self._fsm_info
+
+
+nb_int_list_type = numba.types.ListType(numba.int64)
+nb_int_pair_type = numba.types.UniTuple(numba.int64, 2)
+nb_unichar_1_type = numba.types.UnicodeCharSeq(1)
+
+
+@numba.njit(cache=True)
+def create_fsm_info(
+    py_initial,
+    py_finals,
+    flat_transition_map_items,
+    trans_key_to_states_items,
+    py_anything_value,
+    alphabet_symbol_mapping_items,
+):
+    trans_key_to_states = numba.typed.Dict.empty(numba.int64, nb_int_list_type)
+    for trans_key_and_state in trans_key_to_states_items:
+        trans_key_to_states.setdefault(
+            trans_key_and_state[0], numba.typed.List.empty_list(numba.int64)
+        ).append(trans_key_and_state[1])
+
+    flat_transition_map = numba.typed.Dict.empty(nb_int_pair_type, numba.int64)
+    for trans_key_and_state in flat_transition_map_items:
+        flat_transition_map[
+            (trans_key_and_state[0], trans_key_and_state[1])
+        ] = trans_key_and_state[2]
+
+    alphabet_symbol_map = numba.typed.Dict.empty(nb_unichar_1_type, numba.int64)
+    for symbol_and_trans_key in alphabet_symbol_mapping_items:
+        alphabet_symbol_map[symbol_and_trans_key[0]] = symbol_and_trans_key[1]
+
+    initial = numba.int64(py_initial)
+
+    finals = set()
+    for final in py_finals:
+        finals.add(final)
+
+    anything_value = numba.int64(py_anything_value)
+
+    return FSMInfo(
+        initial,
+        finals,
+        flat_transition_map,
+        trans_key_to_states,
+        anything_value,
+        alphabet_symbol_map,
+    )
+
+
+FSMInfo = namedtuple(
+    "FSMInfo",
+    [
+        "initial",
+        "finals",
+        "transitions",
+        "trans_key_to_states",
+        "alphabet_anything_value",
+        "alphabet_symbol_mapping",
+    ],
+)
+
+
+def make_deterministic_fsm(fsm: FSM) -> Tuple[BetterFSM, Dict[int, int]]:
+    """Construct an equivalent FSM with deterministic state labels."""
+    old_to_new_trans_keys = {
+        trans_key: i
+        for i, (trans_key, _) in enumerate(
+            sorted(fsm.alphabet.by_transition.items(), key=lambda x: sorted(x[1]))
+        )
+    }
+
+    new_symbol_mapping = {
+        symbol: old_to_new_trans_keys[trans_key]
+        for symbol, trans_key in fsm.alphabet._symbol_mapping.items()
+    }
+
+    new_alphabet = BetterAlphabet(new_symbol_mapping)
+
+    new_map = {
+        from_state: {
+            old_to_new_trans_keys[trans_key]: to_state
+            for trans_key, to_state in trans_map.items()
+        }
+        for from_state, trans_map in fsm.map.items()
+    }
+
+    old_to_new_states = {}
+    old_to_new_states[fsm.initial] = 0
+
+    i = 0
+    seen = {fsm.initial}
+    old_state_queue = [fsm.initial]
+    while old_state_queue:
+        old_state = old_state_queue.pop(-1)
+        transitions = new_map[old_state]
+        sorted_transitions = sorted(transitions.items(), key=lambda v: v[0])
+        for _, old_state in sorted_transitions:
+            if old_state not in seen:
+                old_state_queue.append(old_state)
+                seen.add(old_state)
+            if old_state not in old_to_new_states:
+                i += 1
+                old_to_new_states[old_state] = i
+
+    new_map = dict(
+        sorted(
+            (
+                (
+                    old_to_new_states[from_state],
+                    dict(
+                        sorted(
+                            (
+                                (trans_key, old_to_new_states[to_state])
+                                for trans_key, to_state in trans_map.items()
+                            ),
+                            key=lambda v: v[0],
+                        )
+                    ),
+                )
+                for from_state, trans_map in new_map.items()
+            ),
+            key=lambda v: v[0],
+        )
+    )
+
+    new_initial = 0
+    new_finals = frozenset(
+        sorted(old_to_new_states[old_state] for old_state in fsm.finals)
+    )
+    new_states = frozenset(sorted(new_map.keys()))
+
+    new_fsm = BetterFSM(new_alphabet, new_states, new_initial, new_finals, new_map)
+
+    return new_fsm, old_to_new_states
+
+
+@numba.njit(nogil=True, cache=True)
+def _walk_fsm(
+    fsm_transitions: Dict[Tuple[int, int], int],
+    alphabet_symbol_mapping: Dict[str, int],
+    alphabet_anything_value: int,
+    fsm_initial: int,
+    fsm_finals: Set[int],
+    input_string: str,
+    start_state: int,
+    full_match: bool = True,
+) -> List[int]:
+    state = start_state
+    accepted_states: List[int] = numba.typed.List.empty_list(numba.int64)
+    last_final_idx: int = numba.uint64(0)
+
+    for i, symbol in enumerate(input_string):
+        trans_key = alphabet_symbol_mapping.get(symbol, alphabet_anything_value)
+
+        new_state = fsm_transitions.get((state, trans_key))
+
+        if new_state is None:
+            if not full_match and last_final_idx > 0:
+                return accepted_states[:last_final_idx]
+
+            return numba.typed.List.empty_list(numba.int64)
+
+        state = new_state
+
+        if state in fsm_finals:
+            last_final_idx = numba.uint64(i + 1)
+
+        accepted_states.append(_nonoptional(state))
+
+    if full_match and last_final_idx - 1 != i:
+        return numba.typed.List.empty_list(numba.int64)
+
+    return accepted_states
+
+
+def walk_fsm(
+    fsm: BetterFSM,
+    input_string: str,
+    start_state: int,
+    full_match: bool = True,
+) -> List[int]:
+    fsm_finals = fsm.finals
+
+    state = start_state
+    accepted_states: List[int] = []
+    last_final_idx: int = 0
+
+    alphabet_symbol_mapping = fsm.alphabet._symbol_mapping
+    alphabet_anything_value = fsm.alphabet.anything_value
+    fsm_transitions = fsm.flat_transition_map
+
+    for i, symbol in enumerate(input_string):
+        trans_key = alphabet_symbol_mapping.get(symbol, alphabet_anything_value)
+
+        new_state = fsm_transitions.get((state, trans_key))
+
+        if new_state is None:
+            if not full_match and last_final_idx > 0:
+                return accepted_states[:last_final_idx]
+
+            return []
+
+        state = new_state
+
+        if state in fsm_finals:
+            last_final_idx = i + 1
+
+        accepted_states.append(state)
+
+    if full_match and last_final_idx - 1 != i:
+        return []
+
+    return accepted_states
+
+
+def fsm_union(
+    fsms: Sequence[FSM],
+) -> Tuple[FSM, Dict[int, Tuple[Set[Tuple[int, int]], Set[int], Dict[int, Set[int]]]]]:
+    """Construct an FSM representing the union of the FSMs in `fsms`.
+
+    This is an updated version of `interegular.fsm.FSM.union` made to return an
+    extra map of component FSMs to the sets of state transitions that
+    correspond to them in the new FSM.
+
+    """
+
+    alphabet, new_to_old = Alphabet.union(*[fsm.alphabet for fsm in fsms])
+
+    indexed_fsms = tuple(enumerate(fsms))
+
+    initial = {i: fsm.initial for (i, fsm) in indexed_fsms}
+
+    # Dedicated function accepting a "superset" and returning the next
+    # "superset" obtained by following this transition in the new FSM
+    def follow(current_state, new_transition: int):
+        next = {}
+        for i, f in indexed_fsms:
+            old_transition = new_to_old[i][new_transition]
+            if (
+                i in current_state
+                and current_state[i] in f.map
+                and old_transition in f.map[current_state[i]]
+            ):
+                next[i] = f.map[current_state[i]][old_transition]
+        if not next:
+            raise OblivionError
+        return next
+
+    states = [initial]
+    finals: Set[int] = set()
+    map: Dict[int, Dict[int, int]] = {}
+
+    # Map component FSMs to their new state-to-state transitions, finals, and a
+    # map translating component FSM states to aggregate FSM states
+    fsms_to_trans_finals: Dict[
+        int, Tuple[Set[Tuple[int, int]], Set[int], Dict[int, Set[int]]]
+    ] = {}
+
+    i = 0
+    while i < len(states):
+        state = states[i]
+
+        # Add to the finals of the aggregate FSM whenever we hit a final in a
+        # component FSM
+        if any(state.get(j, -1) in fsm.finals for (j, fsm) in indexed_fsms):
+            finals.add(i)
+
+        # Compute the map for this state
+        map[i] = {}
+        for transition in alphabet.by_transition:
+            try:
+                next = follow(state, transition)
+            except OblivionError:
+                # Reached an oblivion state; don't list it
+                continue
+            else:
+                try:
+                    # TODO: Seems like this could--and should--be avoided
+                    j = states.index(next)
+                except ValueError:
+                    j = len(states)
+                    states.append(next)
+
+                map[i][transition] = j
+
+                for fsm_id, fsm_state in next.items():
+                    (
+                        fsm_transitions,
+                        fsm_finals,
+                        fsm_old_to_new,
+                    ) = fsms_to_trans_finals.setdefault(fsm_id, (set(), set(), {}))
+                    old_from = state[fsm_id]
+                    old_to = fsm_state
+                    fsm_old_to_new.setdefault(old_from, set()).add(i)
+                    fsm_old_to_new.setdefault(old_to, set()).add(j)
+                    fsm_transitions.add((i, j))
+                    if fsm_state in fsms[fsm_id].finals:
+                        fsm_finals.add(j)
+
+        i += 1
+
+    fsm = FSM(
+        alphabet=alphabet,
+        states=range(len(states)),
+        initial=0,
+        finals=finals,
+        map=map,
+        __no_validation__=True,
+    )
+
+    fsm, old_to_new_states = make_deterministic_fsm(fsm)
+    _fsms_to_trans_finals = {
+        fsm_id: (
+            {(old_to_new_states[s1], old_to_new_states[s2]) for s1, s2 in transitions},
+            {old_to_new_states[s] for s in finals},
+            {
+                old_state: {old_to_new_states[new_state] for new_state in new_states}
+                for old_state, new_states in old_to_new.items()
+            },
+        )
+        for fsm_id, (transitions, finals, old_to_new) in sorted(
+            fsms_to_trans_finals.items(), key=lambda x: x[0]
+        )
+    }
+
+    return (
+        fsm,
+        _fsms_to_trans_finals,
+    )
+
+
+def get_sub_fsms_from_seq(
+    state_seq: Sequence[int],
+    fsms_to_trans_finals: Dict[
+        int, Tuple[Set[Tuple[int, int]], Set[int], Dict[int, Set[int]]]
+    ],
+) -> Generator[Tuple[int, bool, bool], None, None]:
+    """Get the indices of the sub-FSMs in `fsm` that could have matched the state sequence `state_seq`.
+
+    Parameters
+    ----------
+    state_seq
+        A state sequence.
+    fsms_to_trans_finals
+        A map from FSM indices to tuples containing sets of their state transitions
+        and sets of the final/accept states.
+
+    Returns
+    -------
+    A generator returning tuples containing each sub-FSM index (in the order
+    they were union-ed to construct `fsm`) and booleans indicating whether or
+    not there is another valid transition from the last state in the sequence
+    for the associated sub-FSM (i.e. if the FSM can continue
+    accepting/matching) and whether or not the sequence ends in a final state
+    of the sub-FSM.
+    """
+    state_seq_transitions = set(zip(state_seq[:-1], state_seq[1:]))
+    last_fsm_state = state_seq[-1]
+    yield from (
+        (
+            # The sub-FMS index
+            fsm_idx,
+            # Is there another possible transition in this sub-FSM?
+            any(last_fsm_state == from_s for (from_s, to_s) in transitions),
+            # Is this sub-FSM in a final state?
+            state_seq[-1] in finals,
+        )
+        for fsm_idx, (transitions, finals, _) in fsms_to_trans_finals.items()
+        if state_seq_transitions.issubset(transitions)
+    )
+
+
+@numba.njit(cache=True, nogil=True)
+def state_scan_tokens(
+    fsm_transitions: Dict[Tuple[int, int], int],
+    alphabet_symbol_mapping: Dict[str, int],
+    alphabet_anything_value: int,
+    fsm_initial: int,
+    fsm_finals: Set[int],
+    vocabulary: Dict[str, List[int]],
+    start_state: int,
+) -> Set[Tuple[int, int]]:
+    res = set()
+
+    for token, token_ids in vocabulary.items():
+        state_seq = _walk_fsm(
+            fsm_transitions,
+            alphabet_symbol_mapping,
+            alphabet_anything_value,
+            fsm_initial,
+            fsm_finals,
+            token,
+            start_state,
+            False,
+        )
+
+        if state_seq is not None and len(state_seq) < len(token):
+            continue
+
+        for token_id in token_ids:
+            res.add((token_id, state_seq[-1]))
+
+    return res
+
+
+def create_fsm_index_end_to_end(
+    fsm_info: FSMInfo,
+    vocabulary: Dict[str, List[int]],
+) -> Dict[int, Set[Tuple[int, int]]]:
+    """Create an FSM state-to-vocabulary map/index through end-to-end token parsing."""
+
+    # TODO: Consider using a `List` of `Set`s instead; that way we can JIT this
+    # code, too.
+    states_to_token_subsets: Dict[int, Set[Tuple[int, int]]] = {}
+    seen: Set[int] = set()
+    next_states = {fsm_info.initial}
+
+    while next_states:
+        start_state = next_states.pop()
+
+        token_ids_end_states = state_scan_tokens(
+            fsm_info.transitions,
+            fsm_info.alphabet_symbol_mapping,
+            fsm_info.alphabet_anything_value,
+            fsm_info.initial,
+            fsm_info.finals,
+            vocabulary,
+            start_state,
+        )
+
+        for token_id_and_end_state in token_ids_end_states:
+            states_to_token_subsets.setdefault(start_state, set()).add(
+                token_id_and_end_state
+            )
+            end_state = token_id_and_end_state[1]
+            if end_state not in seen:
+                next_states.add(end_state)
+
+        seen.add(start_state)
+
+    return states_to_token_subsets
+
+
+# TODO: Cannot cache typed collections to disk, yet.  See
+# https://github.com/numba/numba/issues/4698
+@lru_cache
+def reduced_vocabulary(tokenizer: "Tokenizer"):
+    """Create a map from decoded vocabulary tokens to lists of equivalent token ids."""
+    vocabulary = numba.typed.Dict.empty(
+        numba.types.string, numba.types.ListType(numba.int64)
+    )
+    empty_token_ids = set()
+    for token, token_idx in tokenizer.vocabulary.items():
+        if token in tokenizer.special_tokens:
+            continue
+
+        token_str = tokenizer.convert_token_to_string(token)
+
+        if token_str:
+            vocabulary.setdefault(
+                token_str,
+                numba.typed.List.empty_list(numba.int64),
+            ).append(numba.int64(token_idx))
+        else:
+            empty_token_ids.add(numba.int64(token_idx))
+
+    return vocabulary, empty_token_ids
+
+
+def create_fsm_index_tokenizer(
+    fsm: BetterFSM,
+    tokenizer: "Tokenizer",
+) -> Tuple[Dict[int, Dict[int, int]], Set[int]]:
+    """Construct an FMS index from a tokenizer.
+
+    This uses the end-to-end approach of `create_fsm_index_end_to_end`.
+
+    .. warning::
+
+        `fsm` needs to be deterministically ordered so that future caching makes sense.
+
+    """
+    vocabulary, empty_token_ids = reduced_vocabulary(tokenizer)
+
+    states_to_token_subsets = create_fsm_index_end_to_end(fsm.fsm_info, vocabulary)
+
+    # Allow transitions to EOS from all terminals FSM states that are
+    # reachable
+    # TODO: Do we really need this anymore?
+    for state in fsm.fsm_info.finals:
+        subset = states_to_token_subsets.get(state)
+        if subset is not None:
+            subset.add((tokenizer.eos_token_id, state))
+
+    # Convert to token-to-end-state maps
+    states_to_token_subsets = {k: dict(v) for k, v in states_to_token_subsets.items()}
+
+    return states_to_token_subsets, empty_token_ids

python/sglang/srt/constrained/tokenizer.py

+# Adapted from:
+# https://github.com/outlines-dev/outlines/blob/0355ab4272a5d7e4d94c4a53a52593f885b81a61/outlines/models/tokenizer.py
+# https://github.com/outlines-dev/outlines/blob/0355ab4272a5d7e4d94c4a53a52593f885b81a61/outlines/models/transformers.py
+from abc import abstractmethod
+from typing import (
+    TYPE_CHECKING,
+    Dict,
+    Hashable,
+    List,
+    Optional,
+    Protocol,
+    Set,
+    Tuple,
+    Union,
+)
+
+import numpy as np
+import torch
+from numpy.typing import NDArray
+
+
+class Tokenizer(Protocol, Hashable):
+    eos_token: str
+    eos_token_id: int
+    pad_token_id: int
+    vocabulary: Dict[str, int]
+    special_tokens: Set[int]
+
+    @abstractmethod
+    def encode(
+        self, prompt: Union[str, List[str]]
+    ) -> Tuple[NDArray[np.int64], NDArray[np.int64]]:
+        """Translate the input prompts into NumPy arrays of token ids and attention mask."""
+        ...
+
+    @abstractmethod
+    def decode(self, token_ids: NDArray[np.int64]) -> List[str]:
+        """Translate an array of token ids to a string or list of strings."""
+        ...
+
+    @abstractmethod
+    def convert_token_to_string(self, token: str) -> str:
+        """Convert a token to its equivalent string.
+
+        This is for instance useful for BPE tokenizers where whitespaces are
+        represented by the special characted `Ġ`. This prevents matching a raw
+        token that includes `Ġ` with a string.
+
+        """
+        ...
+
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedModel, PreTrainedTokenizer
+
+__all__ = ["transformers"]
+
+
+KVCacheType = Tuple[Tuple[torch.DoubleTensor, torch.DoubleTensor], ...]
+
+
+def get_llama_tokenizer_types():
+    """Get all the Llama tokenizer types/classes that need work-arounds.
+
+    When they can't be imported, a dummy class is created.
+
+    """
+    try:
+        from transformers.models.llama import LlamaTokenizer
+    except ImportError:
+
+        class LlamaTokenizer:  # type: ignore
+            pass
+
+    try:
+        from transformers.models.llama import LlamaTokenizerFast
+    except ImportError:
+
+        class LlamaTokenizerFast:  # type: ignore
+            pass
+
+    try:
+        from transformers.models.code_llama import CodeLlamaTokenizer
+    except ImportError:
+
+        class CodeLlamaTokenizer:  # type: ignore
+            pass
+
+    try:
+        from transformers.models.code_llama import CodeLlamaTokenizerFast
+    except ImportError:
+
+        class CodeLlamaTokenizerFast:  # type: ignore
+            pass
+
+    return (
+        LlamaTokenizer,
+        LlamaTokenizerFast,
+        CodeLlamaTokenizer,
+        CodeLlamaTokenizerFast,
+    )
+
+
+class Transformer:
+    """Represents a `transformers` model."""
+
+    def __init__(
+        self,
+        model: "PreTrainedModel",
+        tokenizer: "PreTrainedTokenizer",
+    ):
+        self.device = model.device
+        self.model = model
+        self.tokenizer = tokenizer
+
+    @torch.inference_mode
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.LongTensor,
+        past_key_values: Optional[Tuple] = None,
+    ) -> Tuple[torch.FloatTensor, Optional[KVCacheType]]:
+        """Compute a forward pass through the transformer model.
+
+        Parameters
+        ----------
+        input_ids
+            The input token ids.  Must be one or two dimensional.
+        attention_mask
+            The attention mask.  Must be one or two dimensional.
+        past_key_values
+            A tuple of tuples containing the cached key and value tensors for each
+            attention head.
+
+        Returns
+        -------
+        The computed logits and the new cached key and value tensors.
+
+        """
+        assert 0 < input_ids.ndim < 3
+
+        if past_key_values:
+            input_ids = input_ids[..., -1].unsqueeze(-1)
+
+        output = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            return_dict=True,
+            output_attentions=False,
+            output_hidden_states=False,
+            past_key_values=past_key_values,
+        )
+
+        return output.logits, output.past_key_values
+
+    def __call__(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: torch.LongTensor,
+        past_key_values: Optional[Tuple] = None,
+    ) -> torch.FloatTensor:
+        logits, kv_cache = self.forward(input_ids, attention_mask, past_key_values)
+        next_token_logits = logits[..., -1, :]
+
+        return next_token_logits, kv_cache
+
+
+class TransformerTokenizer(Tokenizer):
+    """Represents a tokenizer for models in the `transformers` library."""
+
+    def __init__(self, tokenizer):
+        # TODO: Do something to make this hashable?
+        self.tokenizer = tokenizer
+        self.eos_token_id = self.tokenizer.eos_token_id
+        self.eos_token = self.tokenizer.eos_token
+
+        if not self.tokenizer.pad_token_id:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+            self.pad_token_id = self.eos_token_id
+        else:
+            self.pad_token_id = self.tokenizer.pad_token_id
+            self.pad_token = self.tokenizer.pad_token
+
+        self.special_tokens = set(self.tokenizer.all_special_tokens)
+
+        self.vocabulary = self.tokenizer.get_vocab()
+        self.is_llama = isinstance(self.tokenizer, get_llama_tokenizer_types())
+
+    def encode(
+        self, prompt: Union[str, List[str]], **kwargs
+    ) -> Tuple[torch.LongTensor, torch.LongTensor]:
+        kwargs["padding"] = True
+        kwargs["return_tensors"] = "pt"
+        output = self.tokenizer(prompt, **kwargs)
+        return output["input_ids"], output["attention_mask"]
+
+    def decode(self, token_ids: torch.LongTensor) -> List[str]:
+        text = self.tokenizer.batch_decode(token_ids, skip_special_tokens=True)
+        return text
+
+    def convert_token_to_string(self, token: str) -> str:
+        from transformers.file_utils import SPIECE_UNDERLINE
+
+        string = self.tokenizer.convert_tokens_to_string([token])
+
+        if self.is_llama:
+            # A hack to handle missing spaces to HF's Llama tokenizers
+            if token.startswith(SPIECE_UNDERLINE) or token == "<0x20>":
+                return " " + string
+
+        return string
+
+    def __eq__(self, other):
+        if isinstance(other, type(self)):
+            return False
+            # TODO(lsyin): the lru_cache for the TransoformerTokenizer is useless ?
+            # return other.model_name == self.model_name and other.kwargs == self.kwargs
+        return NotImplemented
+
+    def __hash__(self):
+        from datasets.fingerprint import Hasher
+
+        return hash(Hasher.hash(self.tokenizer))
+
+
+def transformers(
+    model_name: str,
+    device: Optional[str] = None,
+    model_kwargs: dict = {},
+    tokenizer_kwargs: dict = {},
+):
+    """Instantiate a model from the `transformers` library and its tokenizer.
+
+    Parameters
+    ----------
+    model_name
+        The name of the model as listed on Hugging Face's model page.
+    device
+        The device(s) on which the model should be loaded. This overrides
+        the `device_map` entry in `model_kwargs` when provided.
+    model_kwargs
+        A dictionary that contains the keyword arguments to pass to the
+        `from_pretrained` method when loading the model.
+    tokenizer_kwargs
+        A dictionary that contains the keyword arguments to pass to the
+        `from_pretrained` method when loading the tokenizer.
+
+    Returns
+    -------
+    A `TransformersModel` model instance.
+
+    """
+    try:
+        from transformers import AutoModelForCausalLM
+    except ImportError:
+        raise ImportError(
+            "The `transformers` library needs to be installed in order to use `transformers` models."
+        )
+
+    if device is not None:
+        model_kwargs["device_map"] = device
+
+    model = AutoModelForCausalLM.from_pretrained(model_name, **model_kwargs)
+    tokenizer = TransformerTokenizer(model_name, **tokenizer_kwargs)
+
+    return Transformer(model, tokenizer)

python/sglang/srt/hf_transformers_utils.py

+"""Utilities for Huggingface Transformers."""
+
+import json
+import os
+import warnings
+from typing import List, Optional, Tuple, Union
+
+from huggingface_hub import snapshot_download
+from sglang.srt.utils import is_multimodal_model
+from transformers import (
+    AutoConfig,
+    AutoProcessor,
+    AutoTokenizer,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+)
+
+
+def download_from_hf(model_path: str):
+    if os.path.exists(model_path):
+        return model_path
+
+    return snapshot_download(model_path, allow_patterns=["*.json", "*.bin", "*.model"])
+
+
+def get_config_json(model_path: str):
+    with open(os.path.join(model_path, "config.json")) as f:
+        config = json.load(f)
+    return config
+
+
+def get_config(model: str, trust_remote_code: bool, revision: Optional[str] = None):
+    config = AutoConfig.from_pretrained(
+        model, trust_remote_code=trust_remote_code, revision=revision
+    )
+    return config
+
+
+# Models don't use the same configuration key for determining the maximum
+# context length.  Store them here so we can sanely check them.
+# NOTE: The ordering here is important. Some models have two of these and we
+# have a preference for which value gets used.
+CONTEXT_LENGTH_KEYS = [
+    "max_sequence_length",
+    "seq_length",
+    "max_position_embeddings",
+    "max_seq_len",
+    "model_max_length",
+]
+
+
+def get_context_length(config):
+    """Get the context length of a model from a huggingface model config."""
+    rope_scaling = getattr(config, "rope_scaling", None)
+    if rope_scaling:
+        rope_scaling_factor = config.rope_scaling["factor"]
+    else:
+        rope_scaling_factor = 1
+
+    for key in CONTEXT_LENGTH_KEYS:
+        val = getattr(config, key, None)
+        if val is not None:
+            return int(rope_scaling_factor * val)
+    return 2048
+
+
+# A fast LLaMA tokenizer with the pre-processed `tokenizer.json` file.
+_FAST_LLAMA_TOKENIZER = "hf-internal-testing/llama-tokenizer"
+
+
+def get_tokenizer(
+    tokenizer_name: str,
+    *args,
+    tokenizer_mode: str = "auto",
+    trust_remote_code: bool = False,
+    tokenizer_revision: Optional[str] = None,
+    **kwargs,
+) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]:
+    """Gets a tokenizer for the given model name via Huggingface."""
+    if is_multimodal_model(tokenizer_name):
+        processor = get_processor(
+            tokenizer_name,
+            *args,
+            trust_remote_code=trust_remote_code,
+            tokenizer_revision=tokenizer_revision,
+            **kwargs,
+        )
+        tokenizer = processor.tokenizer
+        return tokenizer
+
+    if tokenizer_mode == "slow":
+        if kwargs.get("use_fast", False):
+            raise ValueError("Cannot use the fast tokenizer in slow tokenizer mode.")
+        kwargs["use_fast"] = False
+
+    if (
+        "llama" in tokenizer_name.lower()
+        and kwargs.get("use_fast", True)
+        and tokenizer_name != _FAST_LLAMA_TOKENIZER
+    ):
+        pass
+        # warnings.warn(
+        #    "For some LLaMA V1 models, initializing the fast tokenizer may "
+        #    "take a long time. To reduce the initialization time, consider "
+        #    f"using '{_FAST_LLAMA_TOKENIZER}' instead of the original "
+        #    "tokenizer."
+        # )
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_name,
+            *args,
+            trust_remote_code=trust_remote_code,
+            tokenizer_revision=tokenizer_revision,
+            **kwargs,
+        )
+    except TypeError as e:
+        # The LLaMA tokenizer causes a protobuf error in some environments.
+        err_msg = (
+            "Failed to load the tokenizer. If you are using a LLaMA V1 model "
+            f"consider using '{_FAST_LLAMA_TOKENIZER}' instead of the "
+            "original tokenizer."
+        )
+        raise RuntimeError(err_msg) from e
+    except ValueError as e:
+        # If the error pertains to the tokenizer class not existing or not
+        # currently being imported, suggest using the --trust-remote-code flag.
+        if not trust_remote_code and (
+            "does not exist or is not currently imported." in str(e)
+            or "requires you to execute the tokenizer file" in str(e)
+        ):
+            err_msg = (
+                "Failed to load the tokenizer. If the tokenizer is a custom "
+                "tokenizer not yet available in the HuggingFace transformers "
+                "library, consider setting `trust_remote_code=True` in LLM "
+                "or using the `--trust-remote-code` flag in the CLI."
+            )
+            raise RuntimeError(err_msg) from e
+        else:
+            raise e
+
+    if not isinstance(tokenizer, PreTrainedTokenizerFast):
+        warnings.warn(
+            "Using a slow tokenizer. This might cause a significant "
+            "slowdown. Consider using a fast tokenizer instead."
+        )
+    return tokenizer
+
+
+def get_processor(
+    tokenizer_name: str,
+    *args,
+    tokenizer_mode: str = "auto",
+    trust_remote_code: bool = False,
+    tokenizer_revision: Optional[str] = None,
+    **kwargs,
+):
+    processor = AutoProcessor.from_pretrained(
+        tokenizer_name,
+        *args,
+        trust_remote_code=trust_remote_code,
+        tokenizer_revision=tokenizer_revision,
+        **kwargs,
+    )
+    return processor

python/sglang/srt/layers/context_flashattention_nopad.py

+# Adapted from
+# https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/context_flashattention_nopad.py#L1
+import torch
+import triton
+import triton.language as tl
+from sglang.srt.utils import wrap_kernel_launcher
+
+
+@triton.jit
+def _fwd_kernel(
+    Q,
+    K,
+    V,
+    sm_scale,
+    B_Start_Loc,
+    B_Seqlen,
+    Out,
+    stride_qbs,
+    stride_qh,
+    stride_kbs,
+    stride_kh,
+    stride_vbs,
+    stride_vh,
+    stride_obs,
+    stride_oh,
+    kv_group_num: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_DMODEL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    cur_batch = tl.program_id(0)
+    cur_head = tl.program_id(1)
+    start_m = tl.program_id(2)
+
+    cur_kv_head = cur_head // kv_group_num
+
+    cur_batch_seq_len = tl.load(B_Seqlen + cur_batch)
+    cur_batch_in_all_start_index = tl.load(B_Start_Loc + cur_batch)
+
+    block_start_loc = BLOCK_M * start_m
+
+    # initialize offsets
+    offs_n = tl.arange(0, BLOCK_N)
+    offs_d = tl.arange(0, BLOCK_DMODEL)
+    offs_m = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    off_q = (
+        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_qbs
+        + cur_head * stride_qh
+        + offs_d[None, :]
+    )
+    off_k = offs_n[None, :] * stride_kbs + cur_kv_head * stride_kh + offs_d[:, None]
+    off_v = offs_n[:, None] * stride_vbs + cur_kv_head * stride_vh + offs_d[None, :]
+
+    q = tl.load(Q + off_q, mask=offs_m[:, None] < cur_batch_seq_len, other=0.0)
+
+    k_ptrs = K + off_k
+    v_ptrs = V + off_v
+
+    # initialize pointer to m and l
+    m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf")
+    l_i = tl.zeros([BLOCK_M], dtype=tl.float32)
+    acc = tl.zeros([BLOCK_M, BLOCK_DMODEL], dtype=tl.float32)
+
+    block_mask = tl.where(block_start_loc < cur_batch_seq_len, 1, 0)
+
+    for start_n in range(0, block_mask * (start_m + 1) * BLOCK_M, BLOCK_N):
+        start_n = tl.multiple_of(start_n, BLOCK_N)
+        # -- compute qk ----
+        k = tl.load(
+            k_ptrs + (cur_batch_in_all_start_index + start_n) * stride_kbs,
+            mask=(start_n + offs_n[None, :]) < cur_batch_seq_len,
+            other=0.0,
+        )
+        # mask = tl.load(mask_ptrs + start_n, mask=start_n + offs_n < cur_batch_end_loc, other=0.0)
+
+        qk = tl.zeros([BLOCK_M, BLOCK_N], dtype=tl.float32)
+        qk += tl.dot(q, k)
+        qk *= sm_scale
+        qk = tl.where(offs_m[:, None] >= (start_n + offs_n[None, :]), qk, float("-inf"))
+
+        # -- compute m_ij, p, l_ij
+        m_ij = tl.max(qk, 1)
+        p = tl.exp(qk - m_ij[:, None])
+        l_ij = tl.sum(p, 1)
+        # -- update m_i and l_i
+        m_i_new = tl.maximum(m_i, m_ij)
+        alpha = tl.exp(m_i - m_i_new)
+        beta = tl.exp(m_ij - m_i_new)
+        l_i_new = alpha * l_i + beta * l_ij
+        # -- update output accumulator --
+        # scale p
+        p_scale = beta / l_i_new
+        p = p * p_scale[:, None]
+        # scale acc
+        acc_scale = l_i / l_i_new * alpha
+        acc = acc * acc_scale[:, None]
+        # update acc
+        v = tl.load(
+            v_ptrs + (cur_batch_in_all_start_index + start_n) * stride_vbs,
+            mask=(start_n + offs_n[:, None]) < cur_batch_seq_len,
+            other=0.0,
+        )
+
+        p = p.to(v.dtype)
+        acc += tl.dot(p, v)
+        # update m_i and l_i
+        l_i = l_i_new
+        m_i = m_i_new
+    # initialize pointers to output
+    off_o = (
+        (cur_batch_in_all_start_index + offs_m[:, None]) * stride_obs
+        + cur_head * stride_oh
+        + offs_d[None, :]
+    )
+    out_ptrs = Out + off_o
+    tl.store(out_ptrs, acc, mask=offs_m[:, None] < cur_batch_seq_len)
+
+
+cached_kernel = None
+
+
+def context_attention_fwd(q, k, v, o, b_start_loc, b_seq_len, max_input_len):
+    BLOCK = 128
+    Lq, Lk, Lv = q.shape[-1], k.shape[-1], v.shape[-1]
+    assert Lq == Lk and Lk == Lv
+    assert Lk in {16, 32, 64, 128}
+
+    sm_scale = 1.0 / (Lq**0.5)
+    batch, head = b_seq_len.shape[0], q.shape[1]
+    kv_group_num = q.shape[1] // k.shape[1]
+
+    grid = (batch, head, triton.cdiv(max_input_len, BLOCK))
+    num_warps = 4 if Lk <= 64 else 8
+
+    global cached_kernel
+    if cached_kernel:
+        cached_kernel(
+            grid,
+            num_warps,
+            q,
+            k,
+            v,
+            sm_scale,
+            b_start_loc,
+            b_seq_len,
+            o,
+            q.stride(0),
+            q.stride(1),
+            k.stride(0),
+            k.stride(1),
+            v.stride(0),
+            v.stride(1),
+            o.stride(0),
+            o.stride(1),
+        )
+        return
+
+    _fwd_kernel[grid](
+        q,
+        k,
+        v,
+        sm_scale,
+        b_start_loc,
+        b_seq_len,
+        o,
+        q.stride(0),
+        q.stride(1),
+        k.stride(0),
+        k.stride(1),
+        v.stride(0),
+        v.stride(1),
+        o.stride(0),
+        o.stride(1),
+        kv_group_num=kv_group_num,
+        BLOCK_M=BLOCK,
+        BLOCK_DMODEL=Lk,
+        BLOCK_N=BLOCK,
+        num_warps=num_warps,
+        num_stages=1,
+    )
+    cached_kernel = wrap_kernel_launcher(_fwd_kernel)