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/srt/managers/detokenizer_manager.py

@@ -0,0 +1,85 @@
+import asyncio
+
+import uvloop
+import zmq
+import zmq.asyncio
+from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.managers.io_struct import BatchStrOut, BatchTokenIDOut
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import get_exception_traceback
+
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
+
+class DetokenizerManager:
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+    ):
+        context = zmq.asyncio.Context(2)
+        self.recv_from_router = context.socket(zmq.PULL)
+        self.recv_from_router.bind(f"tcp://127.0.0.1:{port_args.detokenizer_port}")
+
+        self.send_to_tokenizer = context.socket(zmq.PUSH)
+        self.send_to_tokenizer.connect(f"tcp://127.0.0.1:{port_args.tokenizer_port}")
+
+        self.tokenizer = get_tokenizer(
+            server_args.tokenizer_path,
+            tokenizer_mode=server_args.tokenizer_mode,
+            trust_remote_code=server_args.trust_remote_code,
+        )
+
+    async def handle_loop(self):
+        while True:
+            recv_obj = await self.recv_from_router.recv_pyobj()
+
+            if isinstance(recv_obj, BatchTokenIDOut):
+                output_tokens = recv_obj.output_tokens
+
+                # TODO(lmzheng): handle skip_special_tokens per request
+                output_strs = self.tokenizer.batch_decode(
+                    output_tokens,
+                    skip_special_tokens=recv_obj.skip_special_tokens[0],
+                )
+
+                # Trim stop str
+                # TODO(lmzheng): handle the case where multiple stop strs are hit
+                for i in range(len(output_strs)):
+                    if recv_obj.hit_stop_str[i] is not None:
+                        pos = output_strs[i].find(recv_obj.hit_stop_str[i])
+                        if pos != -1:
+                            output_strs[i] = output_strs[i][:pos]
+
+                    if len(output_tokens[i]) > 0:
+                        first_token = self.tokenizer.convert_ids_to_tokens(
+                            int(output_tokens[i][0])
+                        )
+                        if first_token.startswith("▁"):
+                            output_strs[i] = " " + output_strs[i]
+
+                self.send_to_tokenizer.send_pyobj(
+                    BatchStrOut(
+                        recv_obj.rids,
+                        output_strs,
+                        recv_obj.meta_info,
+                        recv_obj.finished,
+                    )
+                )
+            else:
+                raise ValueError(f"Invalid object: {recv_obj}")
+
+
+def start_detokenizer_process(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    pipe_writer,
+):
+    try:
+        manager = DetokenizerManager(server_args, port_args)
+    except Exception as e:
+        pipe_writer.send(get_exception_traceback())
+        raise
+    pipe_writer.send("init ok")
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(manager.handle_loop())

python/sglang/srt/managers/io_struct.py

@@ -0,0 +1,88 @@
+import uuid
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Union
+
+from sglang.srt.sampling_params import SamplingParams
+
+
+@dataclass
+class GenerateReqInput:
+    text: Union[List[str], str]
+    image_data: Optional[Union[List[str], str]] = None
+    sampling_params: Union[List[Dict], Dict] = None
+    rid: Optional[Union[List[str], str]] = None
+    return_normalized_logprob: Optional[Union[List[bool], bool]] = None
+    normalized_logprob_start_len: Optional[Union[List[int], int]] = None
+    stream: bool = False
+
+    def post_init(self):
+        is_single = isinstance(self.text, str)
+
+        if is_single:
+            if self.sampling_params is None:
+                self.sampling_params = {}
+            if self.rid is None:
+                self.rid = uuid.uuid4().hex
+            if self.return_normalized_logprob is None:
+                self.return_normalized_logprob = False
+            if self.normalized_logprob_start_len is None:
+                self.normalized_logprob_start_len = 0
+        else:
+            num = len(self.text)
+
+            if self.image_data is None:
+                self.image_data = [None] * num
+            elif not isinstance(self.image_data, list):
+                self.image_data = [self.image_data] * num
+
+            if self.sampling_params is None:
+                self.sampling_params = [{}] * num
+            elif not isinstance(self.sampling_params, list):
+                self.sampling_params = [self.sampling_params] * num
+
+            if self.rid is None:
+                self.rid = [uuid.uuid4().hex for _ in range(num)]
+            else:
+                assert isinstance(self.rid, list)
+
+            if self.return_normalized_logprob is None:
+                self.return_normalized_logprob = [False] * num
+            elif not isinstance(self.return_normalized_logprob, list):
+                self.return_normalized_logprob = [self.return_normalized_logprob] * num
+
+            if self.normalized_logprob_start_len is None:
+                self.normalized_logprob_start_len = [0] * num
+            elif not isinstance(self.normalized_logprob_start_len, list):
+                self.normalized_logprob_start_len = [
+                    self.normalized_logprob_start_len
+                ] * num
+
+
+@dataclass
+class TokenizedGenerateReqInput:
+    rid: str
+    input_ids: List[int]
+    pixel_values: List[float]
+    image_hash: int
+    sampling_params: SamplingParams
+    return_normalized_logprob: bool
+    normalized_logprob_start_len: int
+    stream: bool
+
+
+@dataclass
+class BatchTokenIDOut:
+    rids: List[str]
+    output_tokens: List[List[int]]
+    hit_stop_str: List[Optional[str]]
+    skip_special_tokens: List[bool]
+    meta_info: List[Dict]
+    finished: List[bool]
+
+
+@dataclass
+class BatchStrOut:
+    rids: List[str]
+    output_str: List[str]
+    meta_info: List[Dict]
+    finished: List[bool]

python/sglang/srt/managers/openai_protocol.py

@@ -0,0 +1,12 @@
+from dataclasses import dataclass
+from typing import Any, List, Optional, Union
+
+
+@dataclass
+class CompletionRequest:
+    prompt: Union[str, List[Any]]
+    model: str = "default"
+    temperature: Optional[float] = 0.7
+    max_tokens: Optional[int] = 16
+    n: Optional[int] = 1
+    stop: Optional[Union[str, List[str]]] = None

python/sglang/srt/managers/router/infer_batch.py

@@ -0,0 +1,326 @@
+from enum import Enum, auto
+from typing import List
+
+import numpy as np
+import torch
+from sglang.srt.managers.router.radix_cache import RadixCache
+from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
+
+
+class ForwardMode(Enum):
+    PREFILL = auto()
+    EXTEND = auto()
+    DECODE = auto()
+
+
+class FinishReason(Enum):
+    LENGTH = auto()
+    EOS_TOKEN = auto()
+    STOP_STR = auto()
+
+
+class Req:
+    def __init__(self, rid):
+        self.rid = rid
+        self.input_ids = []
+        self.output_ids = []
+        self.pixel_values = None
+        self.image_offset = 0
+        self.sampling_params = None
+        self.return_normalized_logprob = False
+        self.normalized_logprob_start_len = 0
+        self.stream = False
+
+        self.tokenizer = None
+        self.finished = False
+        self.finish_reason = None
+        self.hit_stop_str = None
+
+        self.adjust_input_len = 0
+        self.prefix_indices = []
+
+        self.normalized_logprob = None
+
+        # for constrained decoding
+        self.regex_fsm = None
+        self.regex_fsm_state = None
+
+    def max_new_tokens(self):
+        return self.sampling_params.max_new_tokens
+
+    def check_finished(self):
+        if self.finished:
+            return
+
+        if len(self.output_ids) >= self.sampling_params.max_new_tokens:
+            self.finished = True
+            self.finish_reason = FinishReason.LENGTH
+            return
+
+        if (
+            self.output_ids[-1] == self.tokenizer.eos_token_id
+            and self.sampling_params.ignore_eos == False
+        ):
+            self.finished = True
+            self.finish_reason = FinishReason.EOS_TOKEN
+            return
+
+        if len(self.sampling_params.stop_strs) > 0:
+            tail_str = self.tokenizer.decode(
+                self.output_ids[-(self.sampling_params.stop_str_max_len + 1) :]
+            )
+
+            for stop_str in self.sampling_params.stop_strs:
+                if stop_str in tail_str:
+                    self.finished = True
+                    self.finish_reason = FinishReason.STOP_STR
+                    self.hit_stop_str = stop_str
+                    return
+
+    def __repr__(self):
+        return f"rid(n={self.rid}, " f"input_ids={self.input_ids}, "
+
+
+class Batch:
+    def __init__(
+        self,
+        reqs: List[Req],
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool: TokenToKVPool,
+        tree_cache: RadixCache,
+    ):
+        self.reqs = reqs
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool = token_to_kv_pool
+        self.tree_cache = tree_cache
+
+        self.return_normalized_logprob = any(
+            req.return_normalized_logprob for req in reqs
+        )
+
+    def is_empty(self):
+        return len(self.reqs) == 0
+
+    def init_extend_batch(self, vocab_size: int, int_token_logit_bias: torch.Tensor):
+        device = "cuda"
+        bs = len(self.reqs)
+        reqs = self.reqs
+        input_ids = [r.input_ids[len(r.prefix_indices) :] for r in reqs]
+        prefix_indices = [r.prefix_indices for r in reqs]
+
+        # Handle prefix
+        flatten_input_ids = []
+        extend_lens = []
+        prefix_lens = []
+        seq_lens = []
+
+        req_pool_indices = self.req_to_token_pool.alloc(bs)
+        req_pool_indices_cpu = req_pool_indices.cpu().numpy()
+        for i in range(bs):
+            flatten_input_ids.extend(input_ids[i])
+            extend_lens.append(len(input_ids[i]))
+
+            if len(prefix_indices[i]) == 0:
+                prefix_lens.append(0)
+            else:
+                prefix_lens.append(len(prefix_indices[i]))
+                self.req_to_token_pool.req_to_token[req_pool_indices_cpu[i]][
+                    : len(prefix_indices[i])
+                ] = prefix_indices[i]
+
+            seq_lens.append(prefix_lens[-1] + extend_lens[-1])
+
+        position_ids_offsets = torch.zeros((bs,), dtype=torch.int32, device=device)
+
+        # Alloc mem
+        seq_lens, prefix_lens = np.array(seq_lens), np.array(prefix_lens)
+        extend_num_tokens = seq_lens.sum() - prefix_lens.sum()
+        out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
+        if out_cache_loc is None:
+            self.tree_cache.evict(extend_num_tokens, self.token_to_kv_pool.free)
+            out_cache_loc = self.token_to_kv_pool.alloc(extend_num_tokens)
+
+            if out_cache_loc is None:
+                print("Prefill out of memory.")
+                self.tree_cache.pretty_print()
+                exit()
+
+        pt = 0
+        for i in range(bs):
+            self.req_to_token_pool.req_to_token[req_pool_indices_cpu[i]][
+                prefix_lens[i] : prefix_lens[i] + extend_lens[i]
+            ] = out_cache_loc[pt : pt + extend_lens[i]]
+            pt += extend_lens[i]
+
+        # Handle logit bias
+        logit_bias = torch.zeros((bs, vocab_size), dtype=torch.float32, device=device)
+        for i in range(bs):
+            if reqs[i].sampling_params.dtype == "int":
+                logit_bias[i] = int_token_logit_bias
+
+        # Set fields
+        self.input_ids = torch.tensor(
+            flatten_input_ids, dtype=torch.int32, device=device
+        )
+        self.pixel_values = [r.pixel_values for r in reqs]
+        self.image_offsets = [
+            r.image_offset - p_len for r, p_len in zip(reqs, prefix_lens)
+        ]
+        self.req_pool_indices = req_pool_indices
+        self.seq_lens = torch.tensor(seq_lens, dtype=torch.int32, device=device)
+        self.prefix_lens = torch.tensor(prefix_lens, dtype=torch.int32, device=device)
+        self.position_ids_offsets = position_ids_offsets
+        self.extend_num_tokens = extend_num_tokens
+        self.out_cache_loc = out_cache_loc
+
+        self.temperatures = torch.tensor(
+            [r.sampling_params.temperature for r in reqs],
+            dtype=torch.float,
+            device=device,
+        ).view(-1, 1)
+        self.top_ps = torch.tensor(
+            [r.sampling_params.top_p for r in reqs], dtype=torch.float, device=device
+        ).view(-1, 1)
+        self.top_ks = torch.tensor(
+            [r.sampling_params.top_k for r in reqs], dtype=torch.int, device=device
+        ).view(-1, 1)
+        self.frequency_penalties = torch.tensor(
+            [r.sampling_params.frequency_penalty for r in reqs],
+            dtype=torch.float,
+            device=device,
+        )
+        self.presence_penalties = torch.tensor(
+            [r.sampling_params.presence_penalty for r in reqs],
+            dtype=torch.float,
+            device=device,
+        )
+        self.logit_bias = logit_bias
+
+    def update_for_decode(self, input_ids=None):
+        if input_ids is None:
+            input_ids = [
+                r.output_ids[-1] if r.output_ids else r.input_ids[-1] for r in self.reqs
+            ]
+        self.input_ids = torch.tensor(input_ids, dtype=torch.int32, device="cuda")
+        self.seq_lens.add_(1)
+        self.prefix_lens = None
+
+        # Alloc mem
+        bs = len(self.reqs)
+        alloc_res = self.token_to_kv_pool.alloc_contiguous(bs)
+        if alloc_res is None:
+            self.out_cache_loc = self.token_to_kv_pool.alloc(bs)
+
+            if self.out_cache_loc is None:
+                self.tree_cache.evict(bs, self.token_to_kv_pool.free)
+                self.out_cache_loc = self.token_to_kv_pool.alloc(bs)
+
+                if self.out_cache_loc is None:
+                    print("Decode out of memory.")
+                    self.tree_cache.pretty_print()
+                    exit()
+
+            self.out_cache_cont_start = None
+            self.out_cache_cont_end = None
+        else:
+            self.out_cache_loc = alloc_res[0]
+            self.out_cache_cont_start = alloc_res[1]
+            self.out_cache_cont_end = alloc_res[2]
+
+        self.req_to_token_pool.req_to_token[
+            self.req_pool_indices, self.seq_lens - 1
+        ] = self.out_cache_loc
+
+    def filter_batch(self, unfinished_indices: List[int]):
+        self.reqs = [self.reqs[i] for i in unfinished_indices]
+        new_indices = torch.tensor(unfinished_indices, dtype=torch.int32, device="cuda")
+        self.seq_lens = self.seq_lens[new_indices]
+        self.input_ids = None
+        self.req_pool_indices = self.req_pool_indices[new_indices]
+        self.prefix_lens = None
+        self.position_ids_offsets = self.position_ids_offsets[new_indices]
+        self.out_cache_loc = self.out_cache_cont_start = self.out_cache_cont_end = None
+
+        for item in [
+            "temperatures",
+            "top_ps",
+            "top_ks",
+            "frequency_penalties",
+            "presence_penalties",
+            "logit_bias",
+        ]:
+            setattr(self, item, getattr(self, item)[new_indices])
+
+    def merge(self, other):
+        self.reqs.extend(other.reqs)
+
+        self.req_pool_indices = torch.concat(
+            [self.req_pool_indices, other.req_pool_indices]
+        )
+        self.seq_lens = torch.concat([self.seq_lens, other.seq_lens])
+        self.prefix_lens = None
+        self.position_ids_offsets = torch.concat(
+            [self.position_ids_offsets, other.position_ids_offsets]
+        )
+        self.out_cache_loc = self.out_cache_cont_start = self.out_cache_cont_end = None
+
+        for item in [
+            "temperatures",
+            "top_ps",
+            "top_ks",
+            "frequency_penalties",
+            "presence_penalties",
+            "logit_bias",
+        ]:
+            setattr(
+                self, item, torch.concat([getattr(self, item), getattr(other, item)])
+            )
+
+    def sample(self, logits: torch.Tensor):
+        # Post process logits
+        logits = logits.contiguous()
+        logits.div_(self.temperatures)
+        logits.add_(self.logit_bias)
+
+        has_regex = any(req.regex_fsm is not None for req in self.reqs)
+        if has_regex:
+            allowed_mask = torch.empty_like(logits[0], dtype=torch.bool)
+            for i, req in enumerate(self.reqs):
+                if req.regex_fsm is not None:
+                    allowed_mask.zero_()
+                    allowed_mask[
+                        req.regex_fsm.allowed_token_ids(req.regex_fsm_state)
+                    ] = 1
+                    logits[i].masked_fill_(~allowed_mask, float("-inf"))
+
+        # TODO(lmzheng): apply penalty
+        probs = torch.softmax(logits, dim=-1)
+        probs_sort, probs_idx = _top_p_top_k(probs, self.top_ps, self.top_ks)
+        sampled_index = torch.multinomial(probs_sort, num_samples=1)
+        batch_next_token_ids = torch.gather(probs_idx, dim=1, index=sampled_index).view(
+            -1
+        )
+        batch_next_token_probs = torch.gather(
+            probs_sort, dim=1, index=sampled_index
+        ).view(-1)
+
+        if has_regex:
+            batch_next_token_ids_cpu = batch_next_token_ids.cpu().numpy()
+            for i, req in enumerate(self.reqs):
+                if req.regex_fsm is not None:
+                    req.regex_fsm_state = req.regex_fsm.next_state(
+                        req.regex_fsm_state, batch_next_token_ids_cpu[i]
+                    )
+
+        return batch_next_token_ids, batch_next_token_probs
+
+
+def _top_p_top_k(probs: torch.Tensor, top_ps: torch.Tensor, top_ks: torch.Tensor):
+    probs_sort, probs_idx = probs.sort(dim=-1, descending=True)
+    probs_sum = torch.cumsum(probs_sort, dim=-1)
+    probs_sort[(probs_sum - probs_sort) > top_ps] = 0.0
+    probs_sort[
+        torch.arange(0, probs.shape[-1], device=probs.device).view(1, -1) >= top_ks
+    ] = 0.0
+    probs_sort.div_(probs_sort.max(dim=-1, keepdim=True)[0])
+    return probs_sort, probs_idx

python/sglang/srt/managers/router/manager.py

@@ -0,0 +1,71 @@
+import asyncio
+import logging
+from typing import List, Tuple
+
+import uvloop
+import zmq
+import zmq.asyncio
+from sglang.srt.managers.router.model_rpc import ModelRpcClient
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import get_exception_traceback
+
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
+
+class RouterManager:
+    def __init__(self, model_client: ModelRpcClient, port_args: PortArgs):
+        # Init communication
+        context = zmq.asyncio.Context(2)
+        self.recv_from_tokenizer = context.socket(zmq.PULL)
+        self.recv_from_tokenizer.bind(f"tcp://127.0.0.1:{port_args.router_port}")
+
+        self.send_to_detokenizer = context.socket(zmq.PUSH)
+        self.send_to_detokenizer.connect(
+            f"tcp://127.0.0.1:{port_args.detokenizer_port}"
+        )
+
+        # Init status
+        self.model_client = model_client
+        self.recv_reqs = []
+
+    async def loop_for_forward(self):
+        while True:
+            next_step_input = list(self.recv_reqs)
+            self.recv_reqs = []
+            out_pyobjs = await self.model_client.step(next_step_input)
+
+            for obj in out_pyobjs:
+                self.send_to_detokenizer.send_pyobj(obj)
+
+            # await for a while to accept input requests
+            await asyncio.sleep(0.001)
+
+    async def loop_for_recv_requests(self):
+        while True:
+            recv_req = await self.recv_from_tokenizer.recv_pyobj()
+            self.recv_reqs.append(recv_req)
+
+
+def start_router_process(
+    server_args: ServerArgs,
+    port_args: PortArgs,
+    pipe_writer,
+):
+    logging.basicConfig(
+        level=getattr(logging, server_args.log_level.upper()),
+        format="%(message)s",
+    )
+
+    try:
+        model_client = ModelRpcClient(server_args, port_args)
+        router = RouterManager(model_client, port_args)
+    except Exception:
+        pipe_writer.send(get_exception_traceback())
+        raise
+
+    pipe_writer.send("init ok")
+
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+    loop.create_task(router.loop_for_recv_requests())
+    loop.run_until_complete(router.loop_for_forward())

python/sglang/srt/managers/router/model_rpc.py

@@ -0,0 +1,497 @@
+import asyncio
+import logging
+import multiprocessing
+import time
+from concurrent.futures import ThreadPoolExecutor
+from enum import Enum, auto
+from typing import Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import rpyc
+import torch
+from rpyc.utils.classic import obtain
+from rpyc.utils.server import ThreadedServer
+from sglang.srt.constrained.fsm_cache import FSMCache
+from sglang.srt.hf_transformers_utils import get_processor, get_tokenizer
+from sglang.srt.managers.io_struct import BatchTokenIDOut, TokenizedGenerateReqInput
+from sglang.srt.managers.router.infer_batch import Batch, ForwardMode, Req
+from sglang.srt.managers.router.model_runner import ModelRunner
+from sglang.srt.managers.router.radix_cache import RadixCache
+from sglang.srt.managers.router.scheduler import Scheduler
+from sglang.srt.model_config import ModelConfig
+from sglang.srt.sampling_params import SamplingParams
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import (
+    get_exception_traceback,
+    get_int_token_logit_bias,
+    is_multimodal_model,
+    set_random_seed,
+)
+
+logger = logging.getLogger("model_rpc")
+
+
+class ModelRpcServer(rpyc.Service):
+    def exposed_init_model(
+        self,
+        tp_rank: int,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+    ):
+        server_args, port_args = [obtain(x) for x in [server_args, port_args]]
+
+        # Copy arguments
+        self.model_mode = server_args.model_mode
+        self.tp_rank = tp_rank
+        self.tp_size = server_args.tp_size
+        self.schedule_heuristic = server_args.schedule_heuristic
+
+        # Init model and tokenizer
+        self.model_config = ModelConfig(
+            server_args.model_path, server_args.trust_remote_code
+        )
+        self.model_runner = ModelRunner(
+            self.model_config,
+            server_args.mem_fraction_static,
+            tp_rank,
+            server_args.tp_size,
+            port_args.nccl_port,
+            server_args.load_format,
+            server_args.trust_remote_code,
+            server_args.model_mode,
+        )
+        if is_multimodal_model(server_args.model_path):
+            self.processor = get_processor(
+                server_args.tokenizer_path,
+                tokenizer_mode=server_args.tokenizer_mode,
+                trust_remote_code=server_args.trust_remote_code,
+            )
+            self.tokenizer = self.processor.tokenizer
+        else:
+            self.tokenizer = get_tokenizer(
+                server_args.tokenizer_path,
+                tokenizer_mode=server_args.tokenizer_mode,
+                trust_remote_code=server_args.trust_remote_code,
+            )
+        self.eos_token_id = self.tokenizer.eos_token_id
+        self.max_total_num_token = self.model_runner.max_total_num_token
+        self.max_num_running_seq = self.max_total_num_token // 2
+        self.max_prefill_num_token = max(
+            self.model_config.context_len, self.max_total_num_token // 6
+        )
+        self.int_token_logit_bias = torch.tensor(
+            get_int_token_logit_bias(self.tokenizer, self.model_config.vocab_size)
+        )
+        set_random_seed(server_args.random_seed)
+        logger.info(
+            f"Rank {self.tp_rank}: "
+            f"max_total_num_token={self.max_total_num_token}, "
+            f"max_prefill_num_token={self.max_prefill_num_token}, "
+            f"context_len={self.model_config.context_len}, "
+            f"model_mode={self.model_mode}"
+        )
+
+        # Init cache
+        self.tree_cache = RadixCache(disable="no-cache" in self.model_mode)
+        self.scheduler = Scheduler(
+            self.schedule_heuristic,
+            self.max_num_running_seq,
+            self.max_prefill_num_token,
+            self.max_total_num_token,
+            self.tree_cache,
+        )
+        self.req_to_token_pool = self.model_runner.req_to_token_pool
+        self.token_to_kv_pool = self.model_runner.token_to_kv_pool
+
+        # Init running status
+        self.forward_queue: List[Req] = []
+        self.running_batch: Batch = None
+        self.out_pyobjs = []
+        self.decode_forward_ct = 0
+        self.stream_interval = 2
+
+        # Init the FSM cache for constrained generation
+        self.regex_fsm_cache = FSMCache(self.tokenizer)
+
+    def exposed_step(self, recv_reqs):
+        if self.tp_size != 1:
+            recv_reqs = obtain(recv_reqs)
+
+        try:
+            # Recv requests
+            for recv_req in recv_reqs:
+                if isinstance(recv_req, TokenizedGenerateReqInput):
+                    self.handle_generate_request(recv_req)
+                else:
+                    raise ValueError(f"Invalid request: {recv_req}")
+
+            # Forward
+            self.forward_step()
+        except Exception:
+            logger.error("Exception in ModelRpcClient:\n" + get_exception_traceback())
+
+        # Return results
+        ret = self.out_pyobjs
+        self.out_pyobjs = []
+        return ret
+
+    @torch.inference_mode()
+    def forward_step(self):
+        new_batch = self.get_new_fill_batch()
+
+        if new_batch is not None:
+            # Run new fill batch
+            self.forward_fill_batch(new_batch)
+
+            if not new_batch.is_empty():
+                if self.running_batch is None:
+                    self.running_batch = new_batch
+                else:
+                    self.running_batch.merge(new_batch)
+        else:
+            # Run decode batch
+            if self.running_batch is not None:
+                # Run a few decode batches continuously for reducing overhead
+                for _ in range(10):
+                    self.forward_decode_batch(self.running_batch)
+
+                    if self.running_batch.is_empty():
+                        self.running_batch = None
+                        break
+
+        if self.running_batch is not None and self.tp_rank == 0:
+            if self.decode_forward_ct >= 20:
+                self.decode_forward_ct = 0
+                num_used = self.max_total_num_token - (
+                    self.token_to_kv_pool.available_size()
+                    + self.tree_cache.evictable_size()
+                )
+                logger.info(
+                    f"#running-req: {len(self.running_batch.reqs)}, "
+                    f"#token: {num_used}, "
+                    f"token usage: {num_used / self.max_total_num_token:.2f}, "
+                    f"#queue-req: {len(self.forward_queue)}"
+                )
+
+    def handle_generate_request(
+        self,
+        recv_req: TokenizedGenerateReqInput,
+    ):
+        req = Req(recv_req.rid)
+        req.input_ids = recv_req.input_ids
+        req.pixel_values = recv_req.pixel_values
+        if req.pixel_values is not None:
+            pad_value = [
+                (recv_req.image_hash) % self.model_config.vocab_size,
+                (recv_req.image_hash >> 16) % self.model_config.vocab_size,
+                (recv_req.image_hash >> 32) % self.model_config.vocab_size,
+                (recv_req.image_hash >> 64) % self.model_config.vocab_size,
+            ]
+            req.input_ids, req.image_offset = self.model_runner.model.pad_input_ids(
+                req.input_ids, pad_value
+            )
+        req.sampling_params = recv_req.sampling_params
+        req.return_normalized_logprob = recv_req.return_normalized_logprob
+        req.normalized_logprob_start_len = recv_req.normalized_logprob_start_len
+        req.stream = recv_req.stream
+        req.tokenizer = self.tokenizer
+
+        # init the regex fsm
+        if req.sampling_params.regex is not None:
+            req.regex_fsm_state = 0
+            req.regex_fsm = self.regex_fsm_cache.get_fsm(req.sampling_params.regex)
+
+        # Truncate long prompts
+        req.input_ids = req.input_ids[: self.model_config.context_len - 1]
+        req.sampling_params.max_new_tokens = min(
+            req.sampling_params.max_new_tokens,
+            self.model_config.context_len - 1 - len(req.input_ids),
+        )
+        self.forward_queue.append(req)
+
+    def get_new_fill_batch(self):
+        if (
+            self.running_batch is not None
+            and len(self.running_batch.reqs) > self.max_num_running_seq
+        ):
+            return None
+
+        for req in self.forward_queue:
+            prefix_indices, last_node = self.tree_cache.match_prefix(req.input_ids)
+            if req.return_normalized_logprob:
+                prefix_indices = prefix_indices[: req.normalized_logprob_start_len]
+            req.adjust_input_len = len(req.input_ids) - len(prefix_indices)
+            req.prefix_indices = prefix_indices
+            req.last_node = last_node
+
+        # Get priority queue
+        self.forward_queue = self.scheduler.get_priority_queue(self.forward_queue)
+
+        # Add requests if there is available space
+        can_run_list = []
+        new_batch_total_tokens = 0
+        new_batch_input_tokens = 0
+        new_batch_prefix_tokens = 0
+
+        available_size = (
+            self.token_to_kv_pool.available_size() + self.tree_cache.evictable_size()
+        )
+        new_ratio = self.scheduler.new_token_estimation_ratio()
+        if self.running_batch:
+            available_size -= sum(
+                [
+                    (r.max_new_tokens() - len(r.output_ids)) * new_ratio
+                    for r in self.running_batch.reqs
+                ]
+            )
+
+        for req in self.forward_queue:
+            if req.return_normalized_logprob:
+                # Need at least two tokens to compute normalized logprob
+                if req.adjust_input_len < 2:
+                    delta = 2 - req.adjust_input_len
+                    req.adjust_input_len += delta
+                    req.prefix_indices = req.prefix_indices[:-delta]
+                    if req.image_offset is not None:
+                        req.image_offset += delta
+            if req.adjust_input_len == 0 and req.max_new_tokens() > 0:
+                # Need at least one token to compute logits
+                req.adjust_input_len = 1
+                req.prefix_indices = req.prefix_indices[:-1]
+                if req.image_offset is not None:
+                    req.image_offset += 1
+
+            if (
+                req.adjust_input_len + req.max_new_tokens() + new_batch_total_tokens
+                < available_size
+                and req.adjust_input_len + new_batch_input_tokens
+                < self.max_prefill_num_token
+            ):
+                delta = self.tree_cache.inc_ref_counter(req.last_node)
+                available_size += delta
+
+                if not (
+                    req.adjust_input_len + req.max_new_tokens() + new_batch_total_tokens
+                    < available_size
+                ):
+                    delta = self.tree_cache.dec_ref_counter(req.last_node)
+                    available_size += delta
+                else:
+                    self.token_to_kv_pool.add_refs(req.prefix_indices)
+                    can_run_list.append(req)
+                    new_batch_total_tokens += (
+                        req.adjust_input_len + req.max_new_tokens()
+                    )
+                    new_batch_input_tokens += req.adjust_input_len
+
+        if len(can_run_list) == 0:
+            return None
+
+        if self.tp_rank == 0:
+            logger.info(
+                f"new fill batch. #seq: {len(can_run_list)}. "
+                f"#cached_token: {sum(len(x.prefix_indices) for x in can_run_list)}. "
+                f"#new_token: {new_batch_input_tokens}. "
+                f"#remaining_req: {len(self.forward_queue) - len(can_run_list)}. "
+                f"#running_req: {0 if self.running_batch is None else len(self.running_batch.reqs)}"
+            )
+
+        new_batch = Batch(
+            can_run_list,
+            self.req_to_token_pool,
+            self.token_to_kv_pool,
+            self.tree_cache,
+        )
+        self.forward_queue = [x for x in self.forward_queue if x not in can_run_list]
+        return new_batch
+
+    def forward_fill_batch(self, batch: Batch):
+        # Build batch tensors
+        batch.init_extend_batch(self.model_config.vocab_size, self.int_token_logit_bias)
+        if batch.extend_num_tokens != 0:
+            # Forward
+            logits, normalized_logprobs = self.model_runner.forward(
+                batch, ForwardMode.EXTEND, batch.return_normalized_logprob
+            )
+            # print("extend logits", logits)
+            if normalized_logprobs is not None:
+                normalized_logprobs = normalized_logprobs.cpu().tolist()
+
+            next_token_ids, next_token_probs = batch.sample(logits)
+            next_token_ids = next_token_ids.cpu().tolist()
+        else:
+            next_token_ids = [self.tokenizer.eos_token_id] * len(batch.reqs)
+            normalized_logprobs = None
+
+        # Check finish condition
+        reqs = batch.reqs
+        for i in range(len(reqs)):
+            reqs[i].output_ids = [next_token_ids[i]]
+            reqs[i].check_finished()
+
+            if normalized_logprobs is not None:
+                reqs[i].normalized_logprob = normalized_logprobs[i]
+
+        self.handle_finished_requests(batch)
+
+    def forward_decode_batch(self, batch: Batch):
+        # Update batch tensors
+        self.decode_forward_ct += 1
+        batch.update_for_decode()
+
+        # Forward
+        logits = self.model_runner.forward(batch, ForwardMode.DECODE)
+        next_token_ids, next_token_probs = batch.sample(logits)
+        next_token_ids = next_token_ids.cpu().tolist()
+
+        # Check finish condition
+        reqs = batch.reqs
+        for i in range(len(reqs)):
+            reqs[i].output_ids.append(next_token_ids[i])
+            reqs[i].check_finished()
+
+        self.handle_finished_requests(batch)
+
+    def handle_finished_requests(self, batch: Batch):
+        output_rids = []
+        output_tokens = []
+        output_hit_stop_str = []
+        output_skip_special_tokens = []
+        output_meta_info = []
+        output_finished = []
+        finished_indices = []
+        unfinished_indices = []
+        for i, req in enumerate(batch.reqs):
+            if req.finished:
+                finished_indices.append(i)
+            else:
+                unfinished_indices.append(i)
+
+            if req.finished or (
+                req.stream and self.decode_forward_ct % self.stream_interval == 0
+            ):
+                output_rids.append(req.rid)
+                output_tokens.append(req.output_ids)
+                output_hit_stop_str.append(req.hit_stop_str)
+                output_skip_special_tokens.append(
+                    req.sampling_params.skip_special_tokens
+                )
+                meta_info = {
+                    "prompt_tokens": len(req.input_ids),
+                    "completion_tokens": len(req.output_ids),
+                }
+                if req.return_normalized_logprob:
+                    meta_info["normalized_logprob"] = req.normalized_logprob
+                output_meta_info.append(meta_info)
+                output_finished.append(req.finished)
+
+        # Send to detokenizer
+        if output_rids:
+            self.out_pyobjs.append(
+                BatchTokenIDOut(
+                    output_rids,
+                    output_tokens,
+                    output_hit_stop_str,
+                    output_skip_special_tokens,
+                    output_meta_info,
+                    output_finished,
+                )
+            )
+
+        # Remove finished reqs
+        if finished_indices:
+            # Update radix cache
+            req_pool_indices_cpu = batch.req_pool_indices.cpu().tolist()
+            for i in finished_indices:
+                req = batch.reqs[i]
+                req_pool_idx = req_pool_indices_cpu[i]
+                token_ids = tuple(req.input_ids + req.output_ids)
+                seq_len = len(token_ids) - 1
+                indices = self.req_to_token_pool.req_to_token[req_pool_idx, :seq_len]
+                prefix_len = self.tree_cache.insert(token_ids, indices.clone())
+
+                self.token_to_kv_pool.free(indices[:prefix_len])
+                self.req_to_token_pool.free(req_pool_idx)
+                self.tree_cache.dec_ref_counter(req.last_node)
+
+            # Update batch tensors
+            if unfinished_indices:
+                batch.filter_batch(unfinished_indices)
+            else:
+                batch.reqs = []
+
+
+class ModelRpcClient:
+    def __init__(self, server_args: ServerArgs, port_args: PortArgs):
+        tp_size = server_args.tp_size
+
+        if tp_size == 1:
+            # Init model
+            self.model_server = ModelRpcServer()
+            self.model_server.exposed_init_model(0, server_args, port_args)
+
+            # Wrap functions
+            def async_wrap(f):
+                async def _func(*args, **kwargs):
+                    return f(*args, **kwargs)
+
+                return _func
+
+            self.step = async_wrap(self.model_server.exposed_step)
+        else:
+            with ThreadPoolExecutor(tp_size) as executor:
+                # Launch model processes
+                rets = executor.map(start_model_process, port_args.model_rpc_ports)
+                self.model_servers = [x[0] for x in rets]
+                self.procs = [x[1] for x in rets]
+
+                # Init model
+                def init_model(i):
+                    return self.model_servers[i].init_model(i, server_args, port_args)
+
+                rets = [obtain(x) for x in executor.map(init_model, range(tp_size))]
+
+            # Wrap functions
+            def async_wrap(func_name):
+                fs = [rpyc.async_(getattr(m, func_name)) for m in self.model_servers]
+
+                async def _func(*args, **kwargs):
+                    tasks = [f(*args, **kwargs) for f in fs]
+                    await asyncio.gather(*[asyncio.to_thread(t.wait) for t in tasks])
+                    return obtain(tasks[0].value)
+
+                return _func
+
+            self.step = async_wrap("step")
+
+
+def start_model_process(port):
+    def _init_service(port):
+        t = ThreadedServer(
+            ModelRpcServer(),
+            port=port,
+            protocol_config={"allow_pickle": True, "sync_request_timeout": 600},
+        )
+        t.start()
+
+    proc = multiprocessing.Process(target=_init_service, args=(port,))
+    proc.start()
+    time.sleep(1)
+
+    repeat_count = 0
+    while repeat_count < 20:
+        try:
+            con = rpyc.connect(
+                "localhost",
+                port,
+                config={"allow_pickle": True, "sync_request_timeout": 600},
+            )
+            break
+        except ConnectionRefusedError:
+            time.sleep(1)
+        repeat_count += 1
+    if repeat_count == 20:
+        raise RuntimeError("init rpc env error!")
+
+    assert proc.is_alive()
+    return con.root, proc

python/sglang/srt/managers/router/model_runner.py

@@ -0,0 +1,458 @@
+from dataclasses import dataclass
+from enum import Enum, auto
+from typing import List
+
+import numpy as np
+import torch
+from sglang.srt.managers.router.infer_batch import Batch, ForwardMode
+from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
+from sglang.srt.utils import is_multimodal_model
+from sglang.utils import get_available_gpu_memory
+from vllm.model_executor.layers.quantization.awq import AWQConfig
+from vllm.model_executor.model_loader import _set_default_torch_dtype
+from vllm.model_executor.parallel_utils.parallel_state import initialize_model_parallel
+
+# for model_mode
+global_model_mode: List[str] = []
+
+
+@dataclass
+class InputMetadata:
+    model_runner: "ModelRunner"
+    forward_mode: ForwardMode
+    batch_size: int
+    total_num_tokens: int
+    max_seq_len: int
+    req_pool_indices: torch.Tensor
+    start_loc: torch.Tensor
+    seq_lens: torch.Tensor
+    prefix_lens: torch.Tensor
+    positions: torch.Tensor
+    req_to_token_pool: ReqToTokenPool
+    token_to_kv_pool: TokenToKVPool
+
+    # for extend
+    extend_seq_lens: torch.Tensor = None
+    extend_start_loc: torch.Tensor = None
+    max_extend_len: int = 0
+
+    out_cache_loc: torch.Tensor = None
+    out_cache_cont_start: torch.Tensor = None
+    out_cache_cont_end: torch.Tensor = None
+
+    other_kv_index: torch.Tensor = None
+    return_normalized_logprob: bool = False
+
+    # for flashinfer
+    use_flashinfer: bool = False
+    qo_indptr: torch.Tensor = None
+    kv_indptr: torch.Tensor = None
+    kv_indices: torch.Tensor = None
+    kv_last_page_len: torch.Tensor = None
+    prefill_wrapper = None
+    decode_wrapper = None
+
+    def init_flashinfer_args(self, tp_size):
+        self.kv_indptr = torch.zeros(
+            (self.batch_size + 1,), dtype=torch.int32, device="cuda"
+        )
+        self.kv_indptr[1:] = torch.cumsum(self.seq_lens, dim=0)
+        self.kv_indices = torch.cat(
+            [
+                self.req_to_token_pool.req_to_token[
+                    self.req_pool_indices[i].item(), : self.seq_lens[i].item()
+                ]
+                for i in range(self.batch_size)
+            ],
+            dim=0,
+        ).contiguous()
+        self.kv_last_page_len = torch.ones(
+            (self.batch_size,), dtype=torch.int32, device="cuda"
+        )
+
+        from flashinfer.ops import (
+            BatchDecodeWithPagedKVCacheWrapper,
+            BatchPrefillWithPagedKVCacheWrapper,
+        )
+
+        if (
+            self.forward_mode == ForwardMode.PREFILL
+            or self.forward_mode == ForwardMode.EXTEND
+        ):
+            self.qo_indptr = torch.zeros(
+                (self.batch_size + 1,), dtype=torch.int32, device="cuda"
+            )
+            self.qo_indptr[1:] = torch.cumsum(self.extend_seq_lens, dim=0)
+            self.prefill_wrapper = BatchPrefillWithPagedKVCacheWrapper()
+            self.prefill_wrapper.begin_forward(
+                self.qo_indptr,
+                self.batch_size,
+                self.model_runner.model_config.num_attention_heads // tp_size,
+                self.model_runner.model_config.num_key_value_heads // tp_size,
+            )
+        else:
+            self.decode_wrapper = BatchDecodeWithPagedKVCacheWrapper()
+            self.decode_wrapper.begin_forward(
+                self.kv_indptr,
+                self.kv_last_page_len,
+                self.batch_size,
+                self.model_runner.model_config.num_attention_heads // tp_size,
+                self.model_runner.model_config.num_key_value_heads // tp_size,
+                self.model_runner.model_config.head_dim,
+                1,
+                "NONE",
+                "float16",
+            )
+
+    def init_extend_args(self):
+        self.extend_seq_lens = self.seq_lens - self.prefix_lens
+        self.extend_start_loc = torch.zeros_like(self.seq_lens)
+        self.extend_start_loc[1:] = torch.cumsum(self.extend_seq_lens[:-1], 0)
+        self.max_extend_len = int(torch.max(self.extend_seq_lens))
+
+    @classmethod
+    def create(
+        cls,
+        model_runner,
+        tp_size,
+        forward_mode,
+        req_pool_indices,
+        seq_lens,
+        prefix_lens,
+        position_ids_offsets,
+        out_cache_loc,
+        out_cache_cont_start=None,
+        out_cache_cont_end=None,
+        return_normalized_logprob=False,
+    ):
+        batch_size = len(req_pool_indices)
+        start_loc = torch.zeros((batch_size,), dtype=torch.int32, device="cuda")
+        start_loc[1:] = torch.cumsum(seq_lens[:-1], dim=0)
+        total_num_tokens = int(torch.sum(seq_lens))
+        max_seq_len = int(torch.max(seq_lens))
+
+        if forward_mode == ForwardMode.DECODE:
+            positions = ((seq_lens - 1) + position_ids_offsets).to(torch.int64)
+            other_kv_index = model_runner.req_to_token_pool.req_to_token[
+                req_pool_indices[0], seq_lens[0] - 1
+            ].item()
+        else:
+            seq_lens_np = seq_lens.cpu().numpy()
+            prefix_lens_np = prefix_lens.cpu().numpy()
+            position_ids_offsets_np = position_ids_offsets.cpu().numpy()
+            positions = torch.tensor(
+                np.concatenate(
+                    [
+                        np.arange(
+                            prefix_lens_np[i] + position_ids_offsets_np[i],
+                            seq_lens_np[i] + position_ids_offsets_np[i],
+                        )
+                        for i in range(batch_size)
+                    ],
+                    axis=0,
+                ),
+                device="cuda",
+            )
+            other_kv_index = None
+
+        ret = cls(
+            model_runner=model_runner,
+            forward_mode=forward_mode,
+            batch_size=batch_size,
+            total_num_tokens=total_num_tokens,
+            max_seq_len=max_seq_len,
+            req_pool_indices=req_pool_indices,
+            start_loc=start_loc,
+            seq_lens=seq_lens,
+            prefix_lens=prefix_lens,
+            positions=positions,
+            req_to_token_pool=model_runner.req_to_token_pool,
+            token_to_kv_pool=model_runner.token_to_kv_pool,
+            out_cache_loc=out_cache_loc,
+            out_cache_cont_start=out_cache_cont_start,
+            out_cache_cont_end=out_cache_cont_end,
+            return_normalized_logprob=return_normalized_logprob,
+            other_kv_index=other_kv_index,
+        )
+
+        if forward_mode == ForwardMode.EXTEND:
+            ret.init_extend_args()
+
+        ret.use_flashinfer = "flashinfer" in model_runner.model_mode
+        if ret.use_flashinfer:
+            ret.init_flashinfer_args(tp_size)
+
+        return ret
+
+
+class ModelRunner:
+    def __init__(
+        self,
+        model_config,
+        mem_fraction_static,
+        tp_rank,
+        tp_size,
+        nccl_port,
+        load_format="auto",
+        trust_remote_code=True,
+        model_mode: List[str] = (),
+    ):
+        self.model_config = model_config
+        self.mem_fraction_static = mem_fraction_static
+        self.tp_rank = tp_rank
+        self.tp_size = tp_size
+        self.nccl_port = nccl_port
+        self.load_format = load_format
+        self.trust_remote_code = trust_remote_code
+        self.model_mode = model_mode
+
+        global global_model_mode
+        global_model_mode = model_mode
+
+        # Init torch distributed
+        torch.cuda.set_device(self.tp_rank)
+        torch.distributed.init_process_group(
+            backend="nccl",
+            world_size=self.tp_size,
+            rank=self.tp_rank,
+            init_method=f"tcp://127.0.0.1:{self.nccl_port}",
+        )
+
+        # A small all_reduce for warmup.
+        if self.tp_size > 1:
+            torch.distributed.all_reduce(torch.zeros(1).cuda())
+        initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
+
+        total_gpu_memory = get_available_gpu_memory(
+            self.tp_rank, distributed=self.tp_size > 1
+        ) * (1 << 30)
+        self.load_model()
+        self.init_memory_pool(total_gpu_memory)
+
+        self.is_multimodal_model = is_multimodal_model(self.model_config)
+
+    def load_model(self):
+        """See also vllm/model_executor/model_loader.py::get_model"""
+        from sglang.srt.models.llama2 import LlamaForCausalLM
+        from sglang.srt.models.llava import LlavaLlamaForCausalLM
+        from sglang.srt.models.mixtral import MixtralForCausalLM
+
+        # Select model class
+        architectures = getattr(self.model_config.hf_config, "architectures", [])
+
+        model_class = None
+        for arch in architectures:
+            if arch == "LlamaForCausalLM":
+                model_class = LlamaForCausalLM
+                break
+            if arch == "MistralForCausalLM":
+                model_class = LlamaForCausalLM
+                break
+            if arch == "LlavaLlamaForCausalLM":
+                model_class = LlavaLlamaForCausalLM
+                break
+            if arch == "MixtralForCausalLM":
+                model_class = MixtralForCausalLM
+                break
+        if model_class is None:
+            raise ValueError(f"Unsupported architectures: {architectures}")
+
+        # Load weights
+        linear_method = None
+        with _set_default_torch_dtype(torch.float16):
+            with torch.device("cuda"):
+                hf_quant_config = getattr(
+                    self.model_config.hf_config, "quantization_config", None
+                )
+                if hf_quant_config is not None:
+                    # TODO: config quantization awq etc
+                    quant_config = AWQConfig.from_config(hf_quant_config)
+                    print(f"quant_config: {quant_config}")
+                    linear_method = quant_config.get_linear_method()
+                model = model_class(
+                    config=self.model_config.hf_config, linear_method=linear_method
+                )
+            model.load_weights(
+                self.model_config.path,
+                cache_dir=None,
+                load_format=self.load_format,
+                revision=None,
+            )
+        self.model = model
+
+    def profile_max_num_token(self, total_gpu_memory):
+        available_gpu_memory = get_available_gpu_memory(
+            self.tp_rank, distributed=self.tp_size > 1
+        ) * (1 << 30)
+        head_dim = (
+            self.model_config.hidden_size // self.model_config.num_attention_heads
+        )
+        head_num = self.model_config.num_key_value_heads // self.tp_size
+        cell_size = head_num * head_dim * self.model_config.num_hidden_layers * 2 * 2
+        rest_memory = available_gpu_memory - total_gpu_memory * (
+            1 - self.mem_fraction_static
+        )
+        max_num_token = int(rest_memory // cell_size)
+        return max_num_token
+
+    def init_memory_pool(self, total_gpu_memory):
+        self.max_total_num_token = self.profile_max_num_token(total_gpu_memory)
+        self.req_to_token_pool = ReqToTokenPool(
+            int(self.max_total_num_token / self.model_config.context_len * 256),
+            self.model_config.context_len + 8,
+        )
+        self.token_to_kv_pool = TokenToKVPool(
+            self.max_total_num_token,
+            dtype=torch.float16,
+            head_num=self.model_config.num_key_value_heads // self.tp_size,
+            head_dim=self.model_config.hidden_size
+            // self.model_config.num_attention_heads,
+            layer_num=self.model_config.num_hidden_layers,
+        )
+
+    @torch.inference_mode()
+    def forward_prefill(
+        self,
+        input_ids,
+        req_pool_indices,
+        seq_lens,
+        prefix_lens,
+        position_ids_offsets,
+        out_cache_loc,
+        return_normalized_logprob,
+    ):
+        input_metadata = InputMetadata.create(
+            self,
+            forward_mode=ForwardMode.PREFILL,
+            tp_size=self.tp_size,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            prefix_lens=prefix_lens,
+            position_ids_offsets=position_ids_offsets,
+            out_cache_loc=out_cache_loc,
+            return_normalized_logprob=return_normalized_logprob,
+        )
+        return self.model.forward(input_ids, input_metadata.positions, input_metadata)
+
+    @torch.inference_mode()
+    def forward_extend(
+        self,
+        input_ids,
+        req_pool_indices,
+        seq_lens,
+        prefix_lens,
+        position_ids_offsets,
+        out_cache_loc,
+        return_normalized_logprob,
+    ):
+        input_metadata = InputMetadata.create(
+            self,
+            forward_mode=ForwardMode.EXTEND,
+            tp_size=self.tp_size,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            prefix_lens=prefix_lens,
+            position_ids_offsets=position_ids_offsets,
+            out_cache_loc=out_cache_loc,
+            return_normalized_logprob=return_normalized_logprob,
+        )
+        return self.model.forward(input_ids, input_metadata.positions, input_metadata)
+
+    @torch.inference_mode()
+    def forward_decode(
+        self,
+        input_ids,
+        req_pool_indices,
+        seq_lens,
+        prefix_lens,
+        position_ids_offsets,
+        out_cache_loc,
+        out_cache_cont_start,
+        out_cache_cont_end,
+    ):
+        input_metadata = InputMetadata.create(
+            self,
+            forward_mode=ForwardMode.DECODE,
+            tp_size=self.tp_size,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            prefix_lens=prefix_lens,
+            position_ids_offsets=position_ids_offsets,
+            out_cache_loc=out_cache_loc,
+            out_cache_cont_start=out_cache_cont_start,
+            out_cache_cont_end=out_cache_cont_end,
+        )
+        return self.model.forward(input_ids, input_metadata.positions, input_metadata)[
+            0
+        ]
+
+    @torch.inference_mode()
+    def forward_extend_multi_modal(
+        self,
+        input_ids,
+        pixel_values,
+        image_offsets,
+        req_pool_indices,
+        seq_lens,
+        prefix_lens,
+        position_ids_offsets,
+        out_cache_loc,
+        return_normalized_logprob,
+    ):
+        input_metadata = InputMetadata.create(
+            self,
+            forward_mode=ForwardMode.EXTEND,
+            tp_size=self.tp_size,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            prefix_lens=prefix_lens,
+            position_ids_offsets=position_ids_offsets,
+            out_cache_loc=out_cache_loc,
+            return_normalized_logprob=return_normalized_logprob,
+        )
+        return self.model.forward(
+            input_ids,
+            input_metadata.positions,
+            input_metadata,
+            pixel_values,
+            image_offsets,
+        )
+
+    def forward(
+        self, batch: Batch, forward_mode: ForwardMode, return_normalized_logprob=False
+    ):
+        if self.is_multimodal_model and forward_mode == ForwardMode.EXTEND:
+            kwargs = {
+                "input_ids": batch.input_ids,
+                "pixel_values": batch.pixel_values,
+                "image_offsets": batch.image_offsets,
+                "req_pool_indices": batch.req_pool_indices,
+                "seq_lens": batch.seq_lens,
+                "prefix_lens": batch.prefix_lens,
+                "position_ids_offsets": batch.position_ids_offsets,
+                "out_cache_loc": batch.out_cache_loc,
+            }
+            kwargs["return_normalized_logprob"] = return_normalized_logprob
+            return self.forward_extend_multi_modal(**kwargs)
+        else:
+            kwargs = {
+                "input_ids": batch.input_ids,
+                "req_pool_indices": batch.req_pool_indices,
+                "seq_lens": batch.seq_lens,
+                "prefix_lens": batch.prefix_lens,
+                "position_ids_offsets": batch.position_ids_offsets,
+                "out_cache_loc": batch.out_cache_loc,
+            }
+
+        if forward_mode == ForwardMode.DECODE:
+            kwargs["out_cache_cont_start"] = batch.out_cache_cont_start
+            kwargs["out_cache_cont_end"] = batch.out_cache_cont_end
+            return self.forward_decode(**kwargs)
+        elif forward_mode == ForwardMode.EXTEND:
+            kwargs["return_normalized_logprob"] = return_normalized_logprob
+            return self.forward_extend(**kwargs)
+        elif forward_mode == ForwardMode.PREFILL:
+            kwargs["return_normalized_logprob"] = return_normalized_logprob
+            return self.forward_prefill(**kwargs)
+        else:
+            raise ValueError(f"Invaid forward mode: {forward_mode}")

python/sglang/srt/managers/router/radix_cache.py

@@ -0,0 +1,220 @@
+import heapq
+import time
+from collections import defaultdict
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+
+
+class TreeNode:
+    def __init__(self):
+        self.children = defaultdict(TreeNode)
+        self.parent = None
+        self.value = None
+        self.ref_counter = 0
+        self.last_access_time = time.time()
+
+    def __lt__(self, other):
+        return self.last_access_time < other.last_access_time
+
+
+def match(key, seq):
+    i = 0
+    for k, w in zip(key, seq):
+        if k != w:
+            break
+        i += 1
+    return i
+
+
+class RadixCache:
+    def __init__(self, disable=False):
+        self.root_node = TreeNode()
+        self.root_node.value = []
+        self.root_node.ref_counter = 1
+        self.evictable_size_ = 0
+
+        self.disable = disable
+
+    ##### Public API #####
+    def match_prefix(self, key):
+        if self.disable:
+            return [], self.root_node
+
+        value = []
+        last_node = [self.root_node]
+        self._match_prefix_helper(self.root_node, key, value, last_node)
+        if value:
+            value = torch.concat(value)
+        return value, last_node[0]
+
+    def insert(self, key, value=None):
+        if self.disable:
+            return len(key)
+
+        if value is None:
+            value = [x for x in key]
+        return self._insert_helper(self.root_node, key, value)
+
+    def pretty_print(self):
+        self._print_helper(self.root_node, 0)
+        print(f"#tokens: {self.total_size()}")
+
+    def total_size(self):
+        return self._total_size_helper(self.root_node)
+
+    def evict(self, num_tokens, evict_callback):
+        if self.disable:
+            raise RuntimeError()
+
+        leaves = self._collect_leaves()
+        heapq.heapify(leaves)
+
+        num_evicted = 0
+        while num_evicted < num_tokens and len(leaves):
+            x = heapq.heappop(leaves)
+
+            if x == self.root_node:
+                break
+            if x.ref_counter > 0:
+                continue
+
+            num_evicted += evict_callback(x.value)
+            self._delete_leaf(x)
+
+            if len(x.parent.children) == 0:
+                heapq.heappush(leaves, x.parent)
+
+    def inc_ref_counter(self, node):
+        delta = 0
+        while node != self.root_node:
+            if node.ref_counter == 0:
+                self.evictable_size_ -= len(node.value)
+                delta -= len(node.value)
+            node.ref_counter += 1
+            node = node.parent
+        return delta
+
+    def dec_ref_counter(self, node):
+        delta = 0
+        while node != self.root_node:
+            if node.ref_counter == 1:
+                self.evictable_size_ += len(node.value)
+                delta += len(node.value)
+            node.ref_counter -= 1
+            node = node.parent
+        return delta
+
+    def evictable_size(self):
+        return self.evictable_size_
+
+    ##### Internal Helper Functions #####
+    def _match_prefix_helper(self, node, key, value, last_node):
+        node.last_access_time = time.time()
+
+        for c_key, child in node.children.items():
+            prefix_len = match(c_key, key)
+            if prefix_len != 0:
+                if prefix_len == len(key) and prefix_len != len(c_key):
+                    new_node = self._split_node(c_key, child, prefix_len)
+                    value.append(new_node.value)
+                    last_node[0] = new_node
+                else:
+                    value.append(child.value[:prefix_len])
+                    last_node[0] = child
+                    self._match_prefix_helper(child, key[prefix_len:], value, last_node)
+                break
+
+    def _split_node(self, key, child, split_len):
+        # new_node -> child
+        new_node = TreeNode()
+        new_node.children = {key[split_len:]: child}
+        new_node.parent = child.parent
+        new_node.ref_counter = child.ref_counter
+        new_node.value = child.value[:split_len]
+        child.parent = new_node
+        child.value = child.value[split_len:]
+        new_node.parent.children[key[:split_len]] = new_node
+        del new_node.parent.children[key]
+        return new_node
+
+    def _insert_helper(self, node, key, value):
+        node.last_access_time = time.time()
+
+        for c_key, child in node.children.items():
+            prefix_len = match(c_key, key)
+
+            if prefix_len == len(c_key):
+                if prefix_len == len(key):
+                    return prefix_len
+                else:
+                    key = key[prefix_len:]
+                    value = value[prefix_len:]
+                    return prefix_len + self._insert_helper(child, key, value)
+
+            if prefix_len:
+                new_node = self._split_node(c_key, child, prefix_len)
+                return prefix_len + self._insert_helper(
+                    new_node, key[prefix_len:], value[prefix_len:]
+                )
+
+        if len(key):
+            new_node = TreeNode()
+            new_node.parent = node
+            new_node.value = value
+            node.children[key] = new_node
+            self.evictable_size_ += len(value)
+        return 0
+
+    def _print_helper(self, node, indent):
+        for key, child in node.children.items():
+            print(" " * indent, len(key), key[:10], f"r={child.ref_counter}")
+            self._print_helper(child, indent=indent + 2)
+
+    def _delete_leaf(self, node):
+        for k, v in node.parent.children.items():
+            if v == node:
+                break
+        del node.parent.children[k]
+        self.evictable_size_ -= len(k)
+
+    def _total_size_helper(self, node):
+        x = len(node.value)
+        for child in node.children.values():
+            x += self._total_size_helper(child)
+        return x
+
+    def _collect_leaves(self):
+        ret_list = []
+
+        def dfs_(cur_node):
+            if len(cur_node.children) == 0:
+                ret_list.append(cur_node)
+
+            for x in cur_node.children.values():
+                dfs_(x)
+
+        dfs_(self.root_node)
+        return ret_list
+
+
+if __name__ == "__main__":
+    tree = RadixCache(disable=False)
+
+    tree.insert("Hello")
+    tree.insert("Hello")
+    tree.insert("Hello_L.A.!")
+    # tree.insert("Hello_world! Happy")
+    # tree.insert("I love you!")
+    tree.pretty_print()
+
+    # print(tree.match_prefix("I love you! aha"))
+
+    # def evict_callback(x):
+    #    print("evict", x)
+    #    return len(x)
+
+    # tree.evict(5, evict_callback)
+    # tree.evict(10, evict_callback)
+    # tree.pretty_print()

python/sglang/srt/managers/router/scheduler.py

@@ -0,0 +1,73 @@
+import random
+from collections import defaultdict
+
+
+class Scheduler:
+    def __init__(
+        self,
+        schedule_heuristic,
+        max_running_seq,
+        max_prefill_num_token,
+        max_total_num_token,
+        tree_cache,
+    ):
+        self.schedule_heuristic = schedule_heuristic
+        self.max_running_seq = max_running_seq
+        self.max_prefill_num_token = max_prefill_num_token
+        self.max_total_num_token = max_total_num_token
+        self.tree_cache = tree_cache
+
+    def new_token_estimation_ratio(self):
+        return 0.4 if self.schedule_heuristic != "fcfs" else 0.5
+
+    def get_priority_queue(self, forward_queue):
+        if self.schedule_heuristic == "lpm":
+            # longest prefix match
+            forward_queue.sort(key=lambda x: -len(x.prefix_indices))
+            return forward_queue
+        elif self.schedule_heuristic == "random":
+            random.shuffle(forward_queue)
+            return forward_queue
+        elif self.schedule_heuristic == "fcfs":
+            return forward_queue
+        elif self.schedule_heuristic == "weight":
+            last_node_to_reqs = defaultdict(list)
+            for req in forward_queue:
+                last_node_to_reqs[req.last_node].append(req)
+            for node in last_node_to_reqs:
+                last_node_to_reqs[node].sort(key=lambda x: -len(x.prefix_indices))
+
+            node_to_weight = defaultdict(int)
+            self._calc_weight_recursive(
+                self.tree_cache.root_node, last_node_to_reqs, node_to_weight
+            )
+
+            tmp_queue = []
+            self._get_weight_priority_recursive(
+                self.tree_cache.root_node, node_to_weight, last_node_to_reqs, tmp_queue
+            )
+            assert len(tmp_queue) == len(forward_queue)
+            return tmp_queue
+        else:
+            raise ValueError(f"Unknown schedule_heuristic: {self.schedule_heuristic}")
+
+    def _calc_weight_recursive(self, cur_node, last_node_to_reqs, node_to_weight):
+        node_to_weight[cur_node] = 1
+        if cur_node in last_node_to_reqs:
+            node_to_weight[cur_node] += len(last_node_to_reqs[cur_node])
+        for child in cur_node.children.values():
+            self._calc_weight_recursive(child, last_node_to_reqs, node_to_weight)
+            node_to_weight[cur_node] += node_to_weight[child]
+
+    def _get_weight_priority_recursive(
+        self, cur_node, node_to_wight, last_node_to_reqs, tmp_queue
+    ):
+        visit_list = [child for child in cur_node.children.values()]
+        visit_list.sort(key=lambda x: -node_to_wight[x])
+        # for node in visit_list:
+        #     print(f"{node_to_wight[node]} {len(node.value) if node.value is not None else 0}")
+        for child in visit_list:
+            self._get_weight_priority_recursive(
+                child, node_to_wight, last_node_to_reqs, tmp_queue
+            )
+        tmp_queue.extend(last_node_to_reqs[cur_node])

python/sglang/srt/managers/tokenizer_manager.py

@@ -0,0 +1,219 @@
+import asyncio
+import concurrent.futures
+import dataclasses
+import os
+from typing import List
+
+import numpy as np
+import transformers
+import uvloop
+import zmq
+import zmq.asyncio
+from sglang.srt.hf_transformers_utils import (
+    get_config,
+    get_context_length,
+    get_processor,
+    get_tokenizer,
+)
+from sglang.srt.managers.io_struct import (
+    BatchStrOut,
+    GenerateReqInput,
+    TokenizedGenerateReqInput,
+)
+from sglang.srt.sampling_params import SamplingParams
+from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import get_exception_traceback, is_multimodal_model, load_image
+
+asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+
+
+@dataclasses.dataclass
+class ReqState:
+    out_list: List
+    finished: bool
+    event: asyncio.Event
+    lock: asyncio.Lock
+
+
+global global_processor
+
+
+def init_global_processor(server_args: ServerArgs):
+    global global_processor
+    transformers.logging.set_verbosity_error()
+    global_processor = get_processor(
+        server_args.tokenizer_path,
+        tokenizer_mode=server_args.tokenizer_mode,
+        trust_remote_code=server_args.trust_remote_code,
+    )
+
+
+def get_pixel_values(image_data, processor=None):
+    try:
+        processor = processor or global_processor
+        image = load_image(image_data)
+        image_hash = hash(image_data)
+        pixel_values = processor.image_processor(image)["pixel_values"][0]
+        pixel_values = pixel_values.astype(np.float16)
+        return pixel_values, image_hash
+    except Exception:
+        print("Exception in TokenizerManager:\n" + get_exception_traceback())
+
+
+class TokenizerManager:
+    def __init__(
+        self,
+        server_args: ServerArgs,
+        port_args: PortArgs,
+    ):
+        context = zmq.asyncio.Context(2)
+        self.recv_from_detokenizer = context.socket(zmq.PULL)
+        self.recv_from_detokenizer.bind(f"tcp://127.0.0.1:{port_args.tokenizer_port}")
+
+        self.send_to_router = context.socket(zmq.PUSH)
+        self.send_to_router.connect(f"tcp://127.0.0.1:{port_args.router_port}")
+
+        self.model_path = server_args.model_path
+        self.hf_config = get_config(
+            self.model_path, trust_remote_code=server_args.trust_remote_code
+        )
+        self.context_len = get_context_length(self.hf_config)
+
+        if is_multimodal_model(self.model_path):
+            self.processor = get_processor(
+                server_args.tokenizer_path,
+                tokenizer_mode=server_args.tokenizer_mode,
+                trust_remote_code=server_args.trust_remote_code,
+            )
+            self.tokenizer = self.processor.tokenizer
+            os.environ["TOKENIZERS_PARALLELISM"] = "false"
+            self.executor = concurrent.futures.ProcessPoolExecutor(
+                initializer=init_global_processor, initargs=(server_args,)
+            )
+        else:
+            self.tokenizer = get_tokenizer(
+                server_args.tokenizer_path,
+                tokenizer_mode=server_args.tokenizer_mode,
+                trust_remote_code=server_args.trust_remote_code,
+            )
+
+        self.to_create_loop = True
+        self.rid_to_state = {}  # Dict[str -> ReqState]
+
+    async def get_pixel_values(self, image_data):
+        if self.executor is not None:
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(
+                self.executor, get_pixel_values, image_data
+            )
+        else:
+            return get_pixel_values(image_data, self.processor)
+
+    async def generate_request(self, obj: GenerateReqInput):
+        if self.to_create_loop:
+            await self.create_handle_loop()
+
+        is_single = isinstance(obj.text, str)
+
+        if is_single:
+            rid = obj.rid
+            input_ids = self.tokenizer.encode(obj.text)
+            sampling_params = SamplingParams(**obj.sampling_params)
+            if sampling_params.max_new_tokens != 0:
+                sampling_params.normalize(self.tokenizer)
+                sampling_params.verify()
+            if obj.image_data is None:
+                pixel_values, image_hash = None, None
+            else:
+                pixel_values, image_hash = await self.get_pixel_values(obj.image_data)
+            tokenized_obj = TokenizedGenerateReqInput(
+                rid=rid,
+                input_ids=input_ids,
+                pixel_values=pixel_values,
+                image_hash=image_hash,
+                sampling_params=sampling_params,
+                return_normalized_logprob=obj.return_normalized_logprob,
+                normalized_logprob_start_len=obj.normalized_logprob_start_len,
+                stream=obj.stream,
+            )
+            self.send_to_router.send_pyobj(tokenized_obj)
+
+            lock = asyncio.Lock()
+            event = asyncio.Event()
+            state = ReqState([], False, event, lock)
+            self.rid_to_state[rid] = state
+
+            while True:
+                await event.wait()
+                yield state.out_list[-1]
+                state.out_list = []
+                if state.finished:
+                    del self.rid_to_state[rid]
+                    break
+                event.clear()
+        else:
+            assert obj.stream is False
+            bs = len(obj.text)
+            for i in range(bs):
+                rid = obj.rid[i]
+                input_ids = self.tokenizer.encode(obj.text[i])
+                sampling_params = SamplingParams(**obj.sampling_params[i])
+                if sampling_params.max_new_tokens != 0:
+                    sampling_params.normalize(self.tokenizer)
+                    sampling_params.verify()
+                if obj.image_data[i] is None:
+                    pixel_values, image_hash = None, None
+                else:
+                    pixel_values, image_hash = await self.get_pixel_values(
+                        obj.image_data[i]
+                    )
+                tokenized_obj = TokenizedGenerateReqInput(
+                    rid=rid,
+                    input_ids=input_ids,
+                    pixel_values=pixel_values,
+                    image_hash=image_hash,
+                    sampling_params=sampling_params,
+                    return_normalized_logprob=obj.return_normalized_logprob[i],
+                    normalized_logprob_start_len=obj.normalized_logprob_start_len[i],
+                    stream=obj.stream,
+                )
+                self.send_to_router.send_pyobj(tokenized_obj)
+
+                lock = asyncio.Lock()
+                event = asyncio.Event()
+                state = ReqState([], False, event, lock)
+                self.rid_to_state[rid] = state
+
+            output_list = []
+            for i in range(bs):
+                rid = obj.rid[i]
+                state = self.rid_to_state[rid]
+                await state.event.wait()
+                output_list.append(state.out_list[-1])
+                assert state.finished
+                del self.rid_to_state[rid]
+
+            yield output_list
+
+    async def create_handle_loop(self):
+        self.to_create_loop = False
+        loop = asyncio.get_event_loop()
+        loop.create_task(self.handle_loop())
+
+    async def handle_loop(self):
+        while True:
+            recv_obj = await self.recv_from_detokenizer.recv_pyobj()
+
+            if isinstance(recv_obj, BatchStrOut):
+                for i, rid in enumerate(recv_obj.rids):
+                    recv_obj.meta_info[i]["id"] = rid
+                    out_dict = {
+                        "text": recv_obj.output_str[i],
+                        "meta_info": recv_obj.meta_info[i],
+                    }
+                    state = self.rid_to_state[rid]
+                    state.out_list.append(out_dict)
+                    state.finished = recv_obj.finished[i]
+                    state.event.set()
+            else:
+                raise ValueError(f"Invalid object: {recv_obj}")