chore: import upstream snapshot with attribution

python/mlc_llm/serve/__init__.py

@@ -0,0 +1,11 @@
+"""Subdirectory of serving."""
+
+# Load MLC LLM library by importing base
+from .. import base
+from .config import EngineConfig
+from .data import Data, ImageData, RequestStreamOutput, TextData, TokenData
+from .embedding_engine import AsyncEmbeddingEngine
+from .engine import AsyncMLCEngine, MLCEngine
+from .radix_tree import PagedRadixTree
+from .request import Request
+from .server import PopenServer

python/mlc_llm/serve/_ffi_api.py

@@ -0,0 +1,7 @@
+"""FFI APIs for mlc_llm.serve"""
+
+import tvm_ffi
+
+# Exports functions registered via TVM_FFI_REGISTER_GLOBAL with the "mlc.serve" prefix.
+# e.g. TVM_FFI_REGISTER_GLOBAL("mlc.serve.TextData")
+tvm_ffi.init_ffi_api("mlc.serve", __name__)

python/mlc_llm/serve/config.py

@@ -0,0 +1,169 @@
+"""Configuration dataclasses used in MLC LLM serving"""
+
+import json
+from dataclasses import asdict, dataclass, field
+from typing import List, Literal, Optional, Tuple, Union  # noqa: UP035
+
+
+@dataclass
+class EngineConfig:
+    """The class of MLCEngine execution configuration.
+
+    Parameters
+    ----------
+    model : str
+        The path to the model directory.
+
+    model_lib : str
+        The path to the model library.
+
+    additional_models : List[Union[str, Tuple[str, str]]]
+        The paths to the additional models' directories (and model libraries).
+        Each element is a single string (denoting the model directory)
+        or a tuple of two strings (denoting the model directory and model lib path).
+
+    mode : Literal["local", "interactive", "server"]
+        The engine mode in MLC LLM.
+        We provide three preset modes: "local", "interactive" and "server".
+        The default mode is "local".
+        The choice of mode decides the values of "max_num_sequence", "max_total_sequence_length"
+        and "prefill_chunk_size" when they are not explicitly specified.
+        1. Mode "local" refers to the local server deployment which has low
+        request concurrency. So the max batch size will be set to 4, and max
+        total sequence length and prefill chunk size are set to the context
+        window size (or sliding window size) of the model.
+        2. Mode "interactive" refers to the interactive use of server, which
+        has at most 1 concurrent request. So the max batch size will be set to 1,
+        and max total sequence length and prefill chunk size are set to the context
+        window size (or sliding window size) of the model.
+        3. Mode "server" refers to the large server use case which may handle
+        many concurrent request and want to use GPU memory as much as possible.
+        In this mode, we will automatically infer the largest possible max batch
+        size and max total sequence length.
+
+        You can manually specify arguments "max_num_sequence", "max_total_sequence_length" and
+        "prefill_chunk_size" to override the automatic inferred values.
+
+    tensor_parallel_shards : Optional[int]
+        Number of shards to split the model into in tensor parallelism multi-gpu inference.
+        When "model_lib" is given, this field will be ignored, and the tensor_parallel_shards
+        in the model_lib metadata will be used.
+
+    pipeline_parallel_stages : Optional[int]
+        Number of pipeline stages to split the model layers for pipeline parallelism.
+        When "model_lib" is given, this field will be ignored, and the pipeline_parallel_stages
+        in the model_lib metadata will be used.
+
+    opt : Optional[str]
+        The optimization flags for JIT compilation.
+        When "model_lib" is given, this field will be ignored.
+        MLC LLM maintains a predefined set of optimization flags,
+        denoted as O0, O1, O2, O3, where O0 means no optimization, O2 means majority of them,
+        and O3 represents extreme optimization that could potentially break the system.
+        Meanwhile, optimization flags could be explicitly specified via details knobs, e.g.
+        "cublas_gemm=1;cudagraph=0".
+
+    gpu_memory_utilization : Optional[float]
+        A number in (0, 1) denoting the fraction of GPU memory used by the server in total.
+        It is used to infer to maximum possible KV cache capacity.
+        When it is unspecified, it defaults to 0.85.
+        Under mode "local" or "interactive", the actual memory usage may be
+        significantly smaller than this number. Under mode "server", the actual
+        memory usage may be slightly larger than this number.
+
+    kv_cache_page_size : int
+        The number of consecutive tokens handled in each page in paged KV cache.
+
+    max_num_sequence : Optional[int]
+        The maximum number of sequences that are allowed to be
+        processed by the KV cache at any time.
+
+    max_total_sequence_length : Optional[int]
+        The maximum total number of tokens whose KV data are allowed
+        to exist in the KV cache at any time.
+
+    max_single_sequence_length : Optional[int]
+        The maximum length allowed for a single sequence in the engine.
+
+    prefill_chunk_size : Optional[int]
+        The maximum total sequence length in a prefill.
+
+    sliding_window_size : Optional[int]
+        The sliding window size in sliding window attention (SWA).
+
+    attention_sink_size : Optional[int]
+        The number of attention sinks when sliding window is enabled..
+
+    max_history_size: Optional[int]
+        The maximum history size for RNN state to roll back.
+
+    kv_state_kind: Optional[Literal["kv_cache", "rnn_state"]]
+        The kind of cache.
+
+    speculative_mode : Literal["disable", "small_draft", "eagle", "medusa"]
+        The speculative mode.
+        "disable" means speculative decoding is disabled.
+        "small_draft" means the normal speculative decoding (small draft) mode.
+        "eagle" means the eagle-style speculative decoding.
+        "medusa" means the medusa-style speculative decoding.
+
+    spec_draft_length : int
+        The number of tokens to generate in speculative proposal (draft).
+        Being 0 means to enable adaptive speculative mode, where the draft length
+        will be automatically adjusted based on engine state.
+
+    spec_tree_width : int
+        The width of the speculative decoding tree.
+
+    prefix_cache_mode : Literal["disable", "radix"]
+        The prefix cache mode.
+        "disable" means no prefix cache is disabled.
+        "radix" means the paged radix tree based prefix cache mode.
+
+    prefix_cache_max_num_recycling_seqs: Optional[int]
+        The maximum number of recycling sequences in prefix cache, default as max_num_sequence.
+        And set 0 to disable prefix cache, set -1 to have infinite capacity prefix cache.
+
+    prefill_mode : Literal["chunked", "hybrid"]
+        The prefill mode.
+        "chunked" means the basic prefill with chunked input enabled.
+        "hybrid" means the hybrid prefill or split-fuse,
+        so that decode step will be converted into prefill.
+
+    verbose : bool
+        A boolean indicating whether to print logging info in engine.
+    """
+
+    model: Optional[str] = None
+    model_lib: Optional[str] = None
+    additional_models: List[Union[str, Tuple[str, str]]] = field(default_factory=list)  # noqa: UP006
+    mode: Optional[Literal["local", "interactive", "server"]] = None
+    tensor_parallel_shards: Optional[int] = None
+    pipeline_parallel_stages: Optional[int] = None
+    opt: Optional[str] = None
+    gpu_memory_utilization: Optional[float] = None
+    kv_cache_page_size: int = 16
+    max_num_sequence: Optional[int] = None
+    max_total_sequence_length: Optional[int] = None
+    max_single_sequence_length: Optional[int] = None
+    prefill_chunk_size: Optional[int] = None
+    sliding_window_size: Optional[int] = None
+    attention_sink_size: Optional[int] = None
+    max_history_size: Optional[int] = None
+    kv_state_kind: Optional[Literal["kv_cache", "rnn_state"]] = None
+    speculative_mode: Literal["disable", "small_draft", "eagle", "medusa"] = "disable"
+    spec_draft_length: int = 0
+    spec_tree_width: int = 1
+    prefix_cache_mode: Literal["disable", "radix"] = "radix"
+    prefix_cache_max_num_recycling_seqs: Optional[int] = None
+    prefill_mode: Literal["chunked", "hybrid"] = "hybrid"
+    verbose: bool = True
+
+    def asjson(self) -> str:
+        """Return the config in string of JSON format."""
+        return json.dumps(asdict(self))
+
+    @staticmethod
+    def from_json(json_str: str) -> "EngineConfig":
+        """Construct a config from JSON string."""
+        return EngineConfig(**json.loads(json_str))

python/mlc_llm/serve/data.py

@@ -0,0 +1,214 @@
+"""Classes denoting multi-modality data used in MLC LLM serving"""
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple  # noqa: UP035
+
+import tvm
+import tvm_ffi
+from tvm.runtime import Object, Tensor
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("mlc.serve.Data")
+class Data(Object):
+    """The base class of multi-modality data (text, tokens, embedding, etc)."""
+
+    def __init__(self):
+        pass
+
+
+@tvm_ffi.register_object("mlc.serve.TextData")
+class TextData(Data):
+    """The class of text data, containing a text string.
+
+    Parameters
+    ----------
+    text : str
+        The text string.
+    """
+
+    def __init__(self, text: str):
+        self.__init_handle_by_constructor__(_ffi_api.TextData, text)
+
+    @property
+    def text(self) -> str:
+        """The text data in `str`."""
+        return str(_ffi_api.TextDataGetTextString(self))
+
+    def __str__(self) -> str:
+        return self.text
+
+
+@tvm_ffi.register_object("mlc.serve.TokenData")
+class TokenData(Data):
+    """The class of token data, containing a list of token ids.
+
+    Parameters
+    ----------
+    token_ids : List[int]
+        The list of token ids.
+    """
+
+    def __init__(self, token_ids: List[int]):  # noqa: UP006
+        self.__init_handle_by_constructor__(_ffi_api.TokenData, *token_ids)
+
+    @property
+    def token_ids(self) -> List[int]:  # noqa: UP006
+        """Return the token ids of the TokenData."""
+        return list(_ffi_api.TokenDataGetTokenIds(self))
+
+
+# mypy: disable-error-code="attr-defined"
+@tvm_ffi.register_object("mlc.serve.ImageData")
+class ImageData(Data):
+    """The class of image data, containing the image as Tensor.
+
+    Parameters
+    ----------
+    image : tvm.runtime.Tensor
+        The image data.
+    """
+
+    def __init__(self, image: Tensor, embed_size: int):
+        self.embed_size = embed_size
+        self.__init_handle_by_constructor__(_ffi_api.ImageData, image, embed_size)
+
+    @property
+    def image(self) -> Tensor:
+        """Return the image data."""
+        return _ffi_api.ImageDataGetImage(self)
+
+    def __len__(self):
+        return self.embed_size
+
+    @staticmethod
+    def from_url(url: str, config: Dict) -> "ImageData":  # noqa: UP006
+        """Get the image from the given URL, process and return the image tensor as TVM Tensor."""
+
+        import base64
+        from io import BytesIO
+
+        import numpy as np
+        import requests
+        from PIL import Image
+
+        if url.startswith("data:image"):
+            # The image is encoded in base64 format
+            base64_image = url.split(",")[1]
+            image_data = base64.b64decode(base64_image)
+            image_tensor = Image.open(BytesIO(image_data)).convert("RGB")
+        elif url.startswith("http"):
+            response = requests.get(url, timeout=5)
+            image_tensor = Image.open(BytesIO(response.content)).convert("RGB")
+        else:
+            raise ValueError(f"Unsupported image URL format: {url}")
+
+        # image_embed_size = ImageData.get_embed_size(config)
+        # TODO: fix these hard-coded values for phi3.5-vision and llava
+        image_embed_size = 576
+        if config["model_type"] == "phi3_v":
+            image_embed_size = 1921
+        image_tensor = np.expand_dims(image_tensor, axis=0)  # HWC -> NHWC
+        image_features = tvm.runtime.tensor(image_tensor)
+        image_data = ImageData(image_features, image_embed_size)
+        return image_data
+
+    @staticmethod
+    def get_embed_size(config: Dict) -> int:  # noqa: UP006
+        """Get the image embedding size from the model config file."""
+        image_size = config["model_config"]["vision_config"]["image_size"]
+        patch_size = config["model_config"]["vision_config"]["patch_size"]
+        embed_size = (image_size // patch_size) ** 2
+        return embed_size
+
+    @staticmethod
+    def get_input_size(config: Dict) -> int:  # noqa: UP006
+        """Get the image input size from the model config file."""
+        image_size = config["model_config"]["vision_config"]["image_size"]
+        return image_size
+
+
+@dataclass
+class SingleRequestStreamOutput:
+    """The request stream output of a single request.
+
+    Attributes
+    ----------
+    delta_token_ids : List[int]
+        The new generated tokens since the last callback invocation
+        for the input request.
+
+    delta_logprob_json_strs : Optional[List[str]]
+        The logprobs JSON strings of the new generated tokens
+        since last invocation.
+
+    finish_reason : Optional[str]
+        The finish reason of the request when it is finished,
+        of None if the request has not finished yet.
+    """
+
+    delta_token_ids: List[int]  # noqa: UP006
+    delta_logprob_json_strs: Optional[List[str]]  # noqa: UP006
+    finish_reason: Optional[str]
+    request_final_usage_json_str: Optional[str]
+    extra_prefix_string: str
+
+
+@tvm_ffi.register_object("mlc.serve.RequestStreamOutput")
+class RequestStreamOutput(Object):
+    """The generated delta request output that is streamed back
+    through callback stream function.
+    It contains four fields (in order):
+
+    request_id : str
+        The id of the request that the function is invoked for.
+
+    stream_outputs : List[SingleRequestStreamOutput]
+        The output instances, one for a request.
+
+    Note
+    ----
+    We do not provide constructor, since in practice only C++ side
+    instantiates this class.
+    """
+
+    def unpack(self) -> Tuple[str, List[SingleRequestStreamOutput]]:  # noqa: UP006
+        """Return the fields of the delta output in a tuple.
+
+        Returns
+        -------
+        request_id : str
+            The id of the request that the function is invoked for.
+
+        stream_outputs : List[SingleRequestStreamOutput]
+            The output instances, one for a request.
+        """
+        fields = _ffi_api.RequestStreamOutputUnpack(self)
+        request_final_usage_json_str = fields[4]
+        request_id = str(fields[0])
+        if request_final_usage_json_str is not None:
+            return (
+                request_id,
+                [SingleRequestStreamOutput([], None, None, request_final_usage_json_str, "")],
+            )
+
+        stream_outputs = []
+        for i, (delta_token_ids, finish_reason, extra_prefix_string) in enumerate(
+            zip(fields[1], fields[3], fields[5])
+        ):
+            delta_logprob_json_strs = (
+                [str(logprob_json_str) for logprob_json_str in fields[2][i]]
+                if fields[2] is not None
+                else None
+            )
+            stream_outputs.append(
+                SingleRequestStreamOutput(
+                    delta_token_ids=list(delta_token_ids),
+                    delta_logprob_json_strs=delta_logprob_json_strs,
+                    finish_reason=str(finish_reason) if finish_reason is not None else None,
+                    request_final_usage_json_str=None,
+                    extra_prefix_string=str(extra_prefix_string),
+                )
+            )
+        return request_id, stream_outputs

python/mlc_llm/serve/embedding_engine.py

@@ -0,0 +1,490 @@
+"""Asynchronous embedding inference engine for encoder and decoder models."""
+
+import asyncio
+import concurrent.futures
+import json
+import os
+from typing import List, Literal, Optional, Tuple, Union  # noqa: UP035
+
+import numpy as np
+import tvm
+from tvm import relax
+from tvm.runtime import Device
+from tvm_ffi import Shape
+
+from mlc_llm.serve import engine_utils
+from mlc_llm.support.auto_device import detect_device
+from mlc_llm.tokenizers import Tokenizer
+
+
+class AsyncEmbeddingEngine:
+    """Asynchronous embedding inference engine.
+
+    Supports both encoder models (BERT-style) and decoder-only embedding models
+    (e.g. Qwen3-Embeddings). Uses a ThreadPoolExecutor for background inference
+    so that the asyncio event loop is not blocked.
+
+    Parameters
+    ----------
+    model : str
+        Path to the model weight directory.
+
+    model_lib : str
+        Path to the compiled model library (.so/.dylib file).
+
+    device : Union[str, Device]
+        Device string, e.g. "auto", "cuda:0", "metal".
+
+    pooling_strategy : Optional[str]
+        Pooling strategy: "cls" (first token), "mean" (masked average),
+        or "last" (last token). If None, auto-detected based on model type:
+        encoder -> "cls", decoder -> "last".
+    """
+
+    def __init__(
+        self,
+        model: str,
+        model_lib: str,
+        device: Union[str, Device] = "auto",
+        *,
+        pooling_strategy: Optional[str] = None,
+    ) -> None:
+        # Reuse existing utility: device detection
+        self.device = detect_device(device) if isinstance(device, str) else device
+        # Reuse existing utility: tokenizer
+        self.tokenizer = Tokenizer(model)
+
+        # Load TVM module, metadata, and params via engine_utils helpers
+        ex = tvm.runtime.load_module(model_lib)
+        vm = relax.VirtualMachine(ex, device=self.device)
+        self._mod = vm.module
+        self._metadata = json.loads(self._mod["_metadata"]())
+        self._params = engine_utils.load_embedding_params(model, self.device, self._metadata)
+
+        # Detect model type and set pooling strategy
+        self.embedding_metadata = engine_utils.get_embedding_metadata(self._metadata)
+        if self.embedding_metadata:
+            self.model_type = self.embedding_metadata["model_type"]
+            self.pooling_strategy = self.embedding_metadata["pooling_strategy"]
+            self.normalize = self.embedding_metadata["normalize"]
+        else:
+            self.model_type = engine_utils.detect_embedding_model_type(self._mod)
+            self.pooling_strategy = "cls" if self.model_type == "encoder" else "last"
+            self.normalize = True
+        # Allow caller to override pooling strategy
+        if pooling_strategy:
+            self.pooling_strategy = pooling_strategy
+
+        # Initialize model-type-specific functions
+        if self.model_type == "encoder":
+            self._init_encoder(model)
+        else:
+            self._init_decoder(model)
+
+        # Background thread pool (1 worker = serialized GPU inference)
+        self._executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="embedding"
+        )
+        self._terminated = False
+
+    def _init_encoder(self, model: str) -> None:
+        """Initialize encoder (BERT-style) model functions and special tokens."""
+        self._prefill_func = self._mod["prefill"]
+        self._cls_token_id: Optional[int] = None
+        self._sep_token_id: Optional[int] = None
+        tok_config_path = os.path.join(model, "tokenizer_config.json")
+        if os.path.exists(tok_config_path):
+            with open(tok_config_path, encoding="utf-8") as f:
+                tok_config = json.load(f)
+            # Try added_tokens_decoder first (newer HF format)
+            added = tok_config.get("added_tokens_decoder", {})
+            for tid, info in added.items():
+                if info.get("content") == tok_config.get("cls_token"):
+                    self._cls_token_id = int(tid)
+                if info.get("content") == tok_config.get("sep_token"):
+                    self._sep_token_id = int(tid)
+            # Fallback: encode the special token strings via tokenizer
+            if self._cls_token_id is None and tok_config.get("cls_token"):
+                ids = list(self.tokenizer.encode(tok_config["cls_token"]))
+                if len(ids) == 1:
+                    self._cls_token_id = ids[0]
+            if self._sep_token_id is None and tok_config.get("sep_token"):
+                ids = list(self.tokenizer.encode(tok_config["sep_token"]))
+                if len(ids) == 1:
+                    self._sep_token_id = ids[0]
+
+    def _init_decoder(self, model: str) -> None:
+        """Initialize decoder (Qwen3-Embeddings style) model functions."""
+        # Prefer tokenizer post-processing (HF-style) for terminal/pooling token handling.
+        # Only fall back to manual EOS append when tokenizer does not define a post-processor
+        # that actually appends a token at the end of the sequence.
+        self._decoder_tokenizer_appends_eos = False
+        tokenizer_json_path = os.path.join(model, "tokenizer.json")
+        if os.path.exists(tokenizer_json_path):
+            with open(tokenizer_json_path, encoding="utf-8") as f:
+                tokenizer_json = json.load(f)
+            post_proc = tokenizer_json.get("post_processor")
+            if post_proc is not None:
+                # Check if the post-processor actually appends a special token at the end
+                # (e.g. TemplateProcessing with "$A <|endoftext|>"). We verify by encoding
+                # a test string and checking if the last token is a known special token.
+                test_tokens = list(self.tokenizer.encode("test"))
+                if len(test_tokens) > 0:
+                    vocab = tokenizer_json.get("added_tokens", [])
+                    special_ids = {t["id"] for t in vocab if t.get("special", False)}
+                    if test_tokens[-1] in special_ids:
+                        self._decoder_tokenizer_appends_eos = True
+
+        # Read EOS token from config — fallback only when tokenizer does not auto-append.
+        self._decoder_eos_token_id: Optional[int] = None
+        config_path = os.path.join(model, "mlc-chat-config.json")
+        if os.path.exists(config_path):
+            with open(config_path, encoding="utf-8") as f:
+                chat_config = json.load(f)
+            eos = chat_config.get("eos_token_id")
+            if isinstance(eos, list):
+                self._decoder_eos_token_id = eos[0]
+            elif isinstance(eos, int):
+                self._decoder_eos_token_id = eos
+
+        self._embed_func = self._mod["embed"]
+        self._prefill_to_hidden_func = self._mod["prefill_to_last_hidden_states"]
+        self._batch_prefill_to_hidden_func = self._mod["batch_prefill_to_last_hidden_states"]
+        if self._mod.implements_function("create_tir_paged_kv_cache"):
+            self._create_kv_cache_func = self._mod["create_tir_paged_kv_cache"]
+        elif self._mod.implements_function("create_flashinfer_paged_kv_cache"):
+            self._create_kv_cache_func = self._mod["create_flashinfer_paged_kv_cache"]
+        else:
+            raise RuntimeError("Cannot find KV cache creation function in model library.")
+        self._kv_state_add_sequence = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
+        self._kv_state_remove_sequence = tvm.get_global_func("vm.builtin.kv_state_remove_sequence")
+        self._kv_state_begin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
+        self._kv_state_end_forward = tvm.get_global_func("vm.builtin.kv_state_end_forward")
+        self._nd_reshape = tvm.get_global_func("vm.builtin.reshape")
+
+    def embed(self, inputs: List[str]) -> Tuple[List[List[float]], int]:  # noqa: UP006
+        """Compute embeddings for a list of input strings (synchronous).
+
+        Parameters
+        ----------
+        inputs : List[str]
+            The input strings to embed.
+
+        Returns
+        -------
+        embeddings : List[List[float]]
+            The L2-normalized embedding vectors.
+        total_tokens : int
+            Total number of tokens processed.
+        """
+        if self.model_type == "encoder":
+            return self._embed_encoder(inputs)
+        return self._embed_decoder(inputs)
+
+    async def async_embed(self, inputs: List[str]) -> Tuple[List[List[float]], int]:  # noqa: UP006
+        """Compute embeddings asynchronously in a background thread.
+
+        This method does not block the asyncio event loop.
+
+        Parameters
+        ----------
+        inputs : List[str]
+            The input strings to embed.
+
+        Returns
+        -------
+        embeddings : List[List[float]]
+            The L2-normalized embedding vectors.
+        total_tokens : int
+            Total number of tokens processed.
+        """
+        loop = asyncio.get_running_loop()
+        return await loop.run_in_executor(self._executor, self.embed, inputs)
+
+    def _embed_encoder(
+        self,
+        inputs: List[str],  # noqa: UP006
+    ) -> Tuple[List[List[float]], int]:  # noqa: UP006
+        """Encoder model embedding (BERT-style).
+
+        Processes each input individually to avoid batch padding artifacts.
+
+        Encoder uses bidirectional attention, so chunked prefill is NOT possible
+        (each token must attend to all other tokens in the full sequence).
+        Inputs exceeding prefill_chunk_size are truncated.
+
+        (Additional Strategy)
+        TODO: For better long-text support, implement sliding window + mean pooling:
+          1. Split text into overlapping windows of prefill_chunk_size (stride=chunk/2)
+          2. Encode each window independently
+          3. Mean-pool all window embeddings → final embedding → L2 normalize
+          This preserves information from the full text at the cost of N× compute.
+        """  # noqa: RUF002
+        embeddings: List[List[float]] = []  # noqa: UP006
+        total_tokens = 0
+        prefill_chunk = self._metadata.get("prefill_chunk_size", 512)
+
+        for text in inputs:
+            tokens = list(self.tokenizer.encode(text))
+            # Add [CLS] and [SEP] if needed
+            if self._cls_token_id is not None and (
+                len(tokens) == 0 or tokens[0] != self._cls_token_id
+            ):
+                tokens = [self._cls_token_id, *tokens]
+            if self._sep_token_id is not None and (
+                len(tokens) == 0 or tokens[-1] != self._sep_token_id
+            ):
+                tokens = [*tokens, self._sep_token_id]
+
+            # Truncate to compiled buffer limit (keep [CLS] at start, [SEP] at end)
+            if len(tokens) > prefill_chunk:
+                tokens = tokens[:prefill_chunk]
+                if self._sep_token_id is not None:
+                    tokens[-1] = self._sep_token_id
+
+            seq_len = len(tokens)
+            total_tokens += seq_len
+
+            token_ids = np.array([tokens], dtype=np.int32)  # [1, seq_len]
+            attention_mask: np.ndarray = np.ones((1, seq_len), dtype=np.int32)  # [1, seq_len]
+
+            tokens_tvm = tvm.runtime.tensor(token_ids, device=self.device)
+            mask_tvm = tvm.runtime.tensor(attention_mask, device=self.device)
+
+            output = self._prefill_func(tokens_tvm, mask_tvm, self._params)
+            # .numpy() copies to CPU, escaping TVM workspace buffer reuse across calls.
+            output_np = output.numpy()  # [1, seq_len, hidden_size]
+
+            # Pooling
+            if self.pooling_strategy == "cls":
+                pooled = output_np[0, 0, :]
+            elif self.pooling_strategy == "mean":
+                pooled = output_np[0].mean(axis=0)
+            else:  # "last"
+                pooled = output_np[0, -1, :]
+
+            # L2 normalize
+            pooled = pooled.astype(np.float32)
+            if self.normalize:
+                norm = np.linalg.norm(pooled)
+                if norm > 1e-12:
+                    pooled = pooled / norm
+
+            embeddings.append(pooled.tolist())
+
+        return embeddings, total_tokens
+
+    def _embed_decoder(self, inputs: List[str]) -> Tuple[List[List[float]], int]:  # noqa: UP006
+        """Decoder model embedding with batch prefill optimization.
+
+        When total tokens fit within prefill_chunk_size, all inputs are processed
+        in a single batch forward pass using shared KV cache. Otherwise, falls back
+        to sequential chunked prefill per input.
+        """
+        # Read KV cache config from metadata
+        prefill_chunk = self._metadata.get("prefill_chunk_size", 2048)
+        max_seq_len = self._metadata.get("context_window_size", 32768)
+        if max_seq_len == -1:
+            max_seq_len = self._metadata.get("sliding_window_size", -1)
+        assert max_seq_len > 0, f"max_seq_len must be positive, got {max_seq_len}"
+        support_sliding = int(self._metadata.get("sliding_window_size", -1) != -1)
+
+        # Tokenize all inputs. Prefer tokenizer post-processor output. If absent (older models),
+        # fall back to appending eos_token_id when missing.
+        token_lists: List[List[int]] = []  # noqa: UP006
+        for text in inputs:
+            tokens = list(self.tokenizer.encode(text))
+            if (
+                not self._decoder_tokenizer_appends_eos
+                and self._decoder_eos_token_id is not None
+                and (len(tokens) == 0 or tokens[-1] != self._decoder_eos_token_id)
+            ):
+                tokens.append(self._decoder_eos_token_id)
+            if len(tokens) > max_seq_len:
+                tokens = tokens[:max_seq_len]
+            token_lists.append(tokens)
+
+        total_tokens = sum(len(t) for t in token_lists)
+
+        # Fast path: all tokens fit in one prefill chunk → batch forward
+        if total_tokens <= prefill_chunk and all(len(t) > 0 for t in token_lists):
+            return self._batch_embed_decoder(
+                token_lists, total_tokens, max_seq_len, prefill_chunk, support_sliding
+            )
+
+        # Greedy sub-batching: pack texts into sub-batches that fit within
+        # prefill_chunk, preserving input order. Oversize texts (single text
+        # exceeding prefill_chunk) fall back to sequential chunked prefill.
+        sub_batches = self._build_sub_batches(token_lists, prefill_chunk)
+        all_embeddings: List[List[float]] = []  # noqa: UP006
+        for batch_type, batch, batch_total in sub_batches:
+            if batch_type == "batch":
+                embs, _ = self._batch_embed_decoder(
+                    batch, batch_total, max_seq_len, prefill_chunk, support_sliding
+                )
+            else:
+                embs, _ = self._sequential_embed_decoder(
+                    batch, batch_total, max_seq_len, prefill_chunk, support_sliding
+                )
+            all_embeddings.extend(embs)
+
+        return all_embeddings, total_tokens
+
+    @staticmethod
+    def _build_sub_batches(
+        token_lists: List[List[int]],  # noqa: UP006
+        prefill_chunk: int,
+    ) -> List[Tuple[Literal["batch", "sequential"], List[List[int]], int]]:  # noqa: UP006
+        """Partition token lists into sub-batches that fit within prefill_chunk.
+
+        Each sub-batch is a tuple of (mode, token_lists, total_token_count).
+        Empty token lists are skipped to avoid invalid batch processing.
+        """
+        sub_batches: List[Tuple[Literal["batch", "sequential"], List[List[int]], int]] = []  # noqa: UP006
+        current_batch: List[List[int]] = []  # noqa: UP006
+        current_tokens = 0
+
+        for tokens in token_lists:
+            if not tokens:
+                continue
+            token_len = len(tokens)
+            is_oversized = token_len > prefill_chunk
+            if current_batch and (is_oversized or current_tokens + token_len > prefill_chunk):
+                sub_batches.append(("batch", current_batch, current_tokens))
+                current_batch, current_tokens = [], 0
+            if is_oversized:
+                sub_batches.append(("sequential", [tokens], token_len))
+            else:
+                current_batch.append(tokens)
+                current_tokens += token_len
+        if current_batch:
+            sub_batches.append(("batch", current_batch, current_tokens))
+
+        return sub_batches
+
+    def _batch_embed_decoder(
+        self,
+        token_lists: List[List[int]],  # noqa: UP006
+        total_tokens: int,
+        max_seq_len: int,
+        prefill_chunk: int,
+        support_sliding: int,
+    ) -> Tuple[List[List[float]], int]:  # noqa: UP006
+        """Batch prefill: process all inputs in a single forward pass."""
+        batch_size = len(token_lists)
+
+        # Create KV cache for the entire batch
+        kv_cache = self._create_kv_cache_func(
+            Shape([batch_size]),
+            Shape([max_seq_len]),
+            Shape([prefill_chunk]),
+            Shape([16]),
+            Shape([support_sliding]),
+        )
+
+        # Register all sequences
+        seq_ids = list(range(batch_size))
+        seq_lens = [len(t) for t in token_lists]
+        for sid in seq_ids:
+            self._kv_state_add_sequence(kv_cache, sid)
+
+        # Begin forward with all sequences at once
+        self._kv_state_begin_forward(kv_cache, Shape(seq_ids), Shape(seq_lens))
+
+        # Concatenate all tokens → embed → batch prefill
+        all_tokens = []
+        for tokens in token_lists:
+            all_tokens.extend(tokens)
+        token_ids = tvm.runtime.tensor(np.array(all_tokens, dtype=np.int32), device=self.device)
+        all_embed = self._embed_func(token_ids, self._params)
+        all_embed = self._nd_reshape(all_embed, Shape([1, total_tokens, all_embed.shape[-1]]))
+
+        hidden_states, _ = self._batch_prefill_to_hidden_func(all_embed, kv_cache, self._params)
+        # .numpy() copies to CPU, escaping TVM workspace buffer reuse across calls.
+        # (torch.from_dlpack is zero-copy and hits aliasing bugs on 2nd+ invocation.)
+        hidden_np = hidden_states.numpy()
+        self._kv_state_end_forward(kv_cache)
+        for sid in seq_ids:
+            self._kv_state_remove_sequence(kv_cache, sid)
+
+        # Extract last token hidden state per sequence
+        embeddings: List[List[float]] = []  # noqa: UP006
+        offset = 0
+        for tokens in token_lists:
+            last_pos = offset + len(tokens) - 1
+            pooled = hidden_np[0, last_pos, :].astype(np.float32)
+            if self.normalize:
+                norm = np.linalg.norm(pooled)
+                if norm > 1e-12:
+                    pooled = pooled / norm
+            embeddings.append(pooled.tolist())
+            offset += len(tokens)
+
+        return embeddings, total_tokens
+
+    def _sequential_embed_decoder(
+        self,
+        token_lists: List[List[int]],  # noqa: UP006
+        total_tokens: int,
+        max_seq_len: int,
+        prefill_chunk: int,
+        support_sliding: int,
+    ) -> Tuple[List[List[float]], int]:  # noqa: UP006
+        """Sequential chunked prefill: process each input independently."""
+        embeddings: List[List[float]] = []  # noqa: UP006
+
+        for tokens in token_lists:
+            if len(tokens) == 0:
+                continue
+
+            # Create KV cache for this single sequence
+            kv_cache = self._create_kv_cache_func(
+                Shape([1]),
+                Shape([max_seq_len]),
+                Shape([prefill_chunk]),
+                Shape([16]),
+                Shape([support_sliding]),
+            )
+            self._kv_state_add_sequence(kv_cache, 0)
+
+            # Process tokens in chunks
+            hidden = None
+            for chunk_start in range(0, len(tokens), prefill_chunk):
+                chunk_end = min(chunk_start + prefill_chunk, len(tokens))
+                chunk_tokens = tokens[chunk_start:chunk_end]
+                chunk_len = len(chunk_tokens)
+
+                token_ids = tvm.runtime.tensor(
+                    np.array(chunk_tokens, dtype=np.int32), device=self.device
+                )
+                chunk_embed = self._embed_func(token_ids, self._params)
+                chunk_embed = self._nd_reshape(
+                    chunk_embed, Shape([1, chunk_len, chunk_embed.shape[-1]])
+                )
+                self._kv_state_begin_forward(kv_cache, Shape([0]), Shape([chunk_len]))
+                hidden, kv_cache = self._prefill_to_hidden_func(chunk_embed, kv_cache, self._params)
+                # .numpy() copies to CPU, escaping TVM buffer aliasing.
+                hidden_np = hidden.numpy()
+                self._kv_state_end_forward(kv_cache)
+
+            self._kv_state_remove_sequence(kv_cache, 0)
+
+            pooled = hidden_np[0, -1, :] if hidden_np.ndim == 3 else hidden_np[-1, :]
+            pooled = pooled.astype(np.float32)
+            if self.normalize:
+                norm = np.linalg.norm(pooled)
+                if norm > 1e-12:
+                    pooled = pooled / norm
+            embeddings.append(pooled.tolist())
+
+        return embeddings, total_tokens
+
+    def terminate(self) -> None:
+        """Terminate the engine and clean up the thread pool."""
+        if getattr(self, "_terminated", True):
+            return
+        self._terminated = True
+        self._executor.shutdown(wait=False)
+
+    def __del__(self):
+        self.terminate()

python/mlc_llm/serve/engine_utils.py

@@ -0,0 +1,337 @@
+"""Utility functions for MLC Serve engine"""
+
+import uuid
+from typing import Any, Callable, Dict, List, Literal, Optional, Union  # noqa: UP035
+
+from mlc_llm.protocol import error_protocol, openai_api_protocol
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import data
+
+RequestProtocol = Union[
+    openai_api_protocol.CompletionRequest, openai_api_protocol.ChatCompletionRequest
+]
+
+
+def get_unsupported_fields(request: RequestProtocol) -> List[str]:  # noqa: UP006
+    """Get the unsupported fields of the request.
+    Return the list of unsupported field names.
+    """
+    if isinstance(
+        request,
+        (
+            openai_api_protocol.CompletionRequest,
+            openai_api_protocol.ChatCompletionRequest,
+        ),
+    ):
+        return openai_api_protocol.openai_api_get_unsupported_fields(request)
+    raise RuntimeError("Cannot reach here")
+
+
+def openai_api_get_generation_config(request: RequestProtocol) -> Dict[str, Any]:  # noqa: UP006
+    """Create the generation config from the given request."""
+    kwargs: Dict[str, Any] = {}  # noqa: UP006
+    arg_names = [
+        "n",
+        "temperature",
+        "top_p",
+        "max_tokens",
+        "frequency_penalty",
+        "presence_penalty",
+        "logit_bias",
+        "seed",
+        "response_format",
+        "debug_config",
+    ]
+    for arg_name in arg_names:
+        kwargs[arg_name] = getattr(request, arg_name)
+    if kwargs["max_tokens"] is None:
+        # Setting to -1 means the generation will not stop until
+        # exceeding model capability or hit any stop criteria.
+        kwargs["max_tokens"] = -1
+    if request.stop is not None:
+        kwargs["stop_strs"] = [request.stop] if isinstance(request.stop, str) else request.stop
+    if isinstance(request, openai_api_protocol.ChatCompletionRequest):
+        kwargs["logprobs"] = request.logprobs
+        kwargs["top_logprobs"] = request.top_logprobs
+    else:
+        logprobs = request.logprobs is not None
+        kwargs["logprobs"] = logprobs
+        kwargs["top_logprobs"] = request.logprobs if logprobs else 0
+    return kwargs
+
+
+def get_generation_config(
+    request: RequestProtocol,
+    extra_stop_token_ids: Optional[List[int]] = None,  # noqa: UP006
+    extra_stop_str: Optional[List[str]] = None,  # noqa: UP006
+) -> GenerationConfig:
+    """Create the generation config in MLC LLM out from the input request protocol."""
+    kwargs: Dict[str, Any]  # noqa: UP006
+    if isinstance(
+        request,
+        (
+            openai_api_protocol.CompletionRequest,
+            openai_api_protocol.ChatCompletionRequest,
+        ),
+    ):
+        kwargs = openai_api_get_generation_config(request)
+    else:
+        raise RuntimeError("Cannot reach here")
+
+    if extra_stop_token_ids is not None:
+        stop_token_ids = kwargs.get("stop_token_ids", [])
+        assert isinstance(stop_token_ids, list)
+        stop_token_ids += extra_stop_token_ids
+        kwargs["stop_token_ids"] = stop_token_ids
+
+    if extra_stop_str is not None:
+        stop_strs = kwargs.get("stop_strs", [])
+        assert isinstance(stop_strs, list)
+        stop_strs += extra_stop_str
+        kwargs["stop_strs"] = stop_strs
+
+    return GenerationConfig(**kwargs)
+
+
+def random_uuid() -> str:
+    """Generate a random id in hexadecimal string."""
+    return uuid.uuid4().hex
+
+
+def check_unsupported_fields(request: RequestProtocol) -> None:
+    """Check if the request has unsupported fields. Raise BadRequestError if so."""
+    unsupported_fields = get_unsupported_fields(request)
+    if len(unsupported_fields) != 0:
+        unsupported_fields = [f'"{field}"' for field in unsupported_fields]
+        raise error_protocol.BadRequestError(
+            f"Request fields {', '.join(unsupported_fields)} are not supported right now.",
+        )
+
+
+def check_and_get_prompts_length(
+    prompts: List[Union[List[int], data.ImageData]],  # noqa: UP006
+    max_input_sequence_length: int,
+) -> int:
+    """Check if the total prompt length exceeds the max single sequence
+    sequence length allowed by the served model. Raise BadRequestError if so.
+    Return the total prompt length.
+    """
+    total_length: int = 0
+    for prompt in prompts:
+        total_length += len(prompt)
+    if total_length > max_input_sequence_length:
+        raise error_protocol.BadRequestError(
+            f"Request prompt has {total_length} tokens in total,"
+            f" larger than the model input length limit {max_input_sequence_length}.",
+        )
+    return total_length
+
+
+def process_prompts(
+    input_prompts: Union[str, List[int], List[Union[str, List[int], data.ImageData]]],  # noqa: UP006
+    ftokenize: Callable[[str], List[int]],  # noqa: UP006
+) -> List[Union[List[int], data.ImageData]]:  # noqa: UP006
+    """Convert all input tokens to list of token ids with regard to the
+    given tokenization function.
+    For each input prompt, return the list of token ids after tokenization.
+    """
+    error_msg = f"Invalid request prompt {input_prompts}"
+
+    # Case 1. The prompt is a single string.
+    if isinstance(input_prompts, str):
+        return [ftokenize(input_prompts)]
+
+    assert isinstance(input_prompts, list)
+    if len(input_prompts) == 0:
+        raise error_protocol.BadRequestError(error_msg)
+
+    # Case 2. The prompt is a list of token ids.
+    if isinstance(input_prompts[0], int):
+        assert isinstance(input_prompts, list)
+        if not all(isinstance(token_id, int) for token_id in input_prompts):
+            raise error_protocol.BadRequestError(error_msg)
+        return [input_prompts]
+
+    # Case 3. A list of prompts.
+    output_prompts: List[Union[List[int], data.ImageData]] = []  # noqa: UP006
+    for input_prompt in input_prompts:
+        if isinstance(input_prompt, str):
+            output_prompts.append(ftokenize(input_prompt))
+        elif isinstance(input_prompt, list) and all(
+            isinstance(token_id, int) for token_id in input_prompt
+        ):
+            output_prompts.append(input_prompt)
+        elif isinstance(input_prompt, data.ImageData):
+            output_prompts.append(input_prompt)
+        else:
+            raise error_protocol.BadRequestError(error_msg)
+    return output_prompts
+
+
+def convert_prompts_to_data(
+    prompts: Union[str, List[int], List[Union[str, List[int], data.Data]]],  # noqa: UP006
+) -> List[data.Data]:  # noqa: UP006
+    """Convert the given prompts in the combination of token id lists
+    and/or data to all data."""
+    if isinstance(prompts, data.Data):
+        return [prompts]
+    if isinstance(prompts, str):
+        return [data.TextData(prompts)]
+    if isinstance(prompts[0], int):
+        assert isinstance(prompts, list) and all(isinstance(token_id, int) for token_id in prompts)
+        return [data.TokenData(prompts)]
+    return [convert_prompts_to_data(x)[0] for x in prompts]
+
+
+class ErrorCleanupScope:
+    """Scope to call cleanup when an error is thrown.
+
+    This class provides an important pattern properly cleanup
+    when async scope CancelledError or other exception happens.
+
+    Parameters
+    ----------
+    cleanup : Callable
+        A callable function to trigger at scope exit during an exception.
+
+    Note
+    ----
+    This helper is motivated by the need to properly
+    abort an async generator and trigger corresponding
+    cleanup functions. Naively use the try except
+    pattern will results in bug when we chain up
+    async generators.
+
+    .. code:: python
+
+        class EngineNotSafe:
+            async def _inner_gen(self, request):
+                request_id = self.get_request_id()
+                self.add_request(request)
+                try:
+                    async for res in await producer_stream:
+                        yield res
+                except asyncio.CancelledError:
+                    self.abort(request_id)
+
+            async def generate(self, request):
+                async for res in await self._inner_gen(request):
+                    # async error can he raised in here
+                    # this will cause
+                    res = await process(res)
+                    yield res
+
+    The above except pattern is not safe.
+    This is because CancelledError may also be raised
+    outside _inner_gen during the process of generate
+    function in between iterations.
+
+    Instead, we use ErrorCleanupScope to safeguard the
+    generation process. The scope will always properly
+    cleanup in exit function when the exception is raised
+
+     .. code:: python
+
+        class EngineSafe:
+            async def _inner_gen(self, request):
+                request_id = self.get_request_id()
+                self.add_request(request)
+                with ErrorCleanupScope(lambda: self.abort(request_id))
+                    async for res in await producer_stream:
+                        yield res
+
+            async def generate(self, request):
+                async for res in await self._inner_gen(request):
+                    # even if async error is raised here
+                    # it will cleanup the ErrorCleanupScope
+                    # properly during function exit
+                    res = await process(res)
+                    yield res
+    """
+
+    cleanup: Callable
+
+    def __init__(self, cleanup: Callable):
+        self.cleanup = cleanup
+
+    def __enter__(self):
+        pass
+
+    def __exit__(self, exc_type, exc_value, traceback) -> None:
+        # only cleanup when exc type is not none
+        if exc_type is not None:
+            self.cleanup()
+
+
+# ====== Embedding Engine Utilities ======
+
+
+def load_embedding_params(model_weight_path, device, model_metadata) -> list:
+    """Load embedding model parameters from weight directory.
+
+    Parameters
+    ----------
+    model_weight_path : str
+        Path to the model weight directory.
+    device : tvm.runtime.Device
+        The target device.
+    model_metadata : dict
+        The model metadata dictionary containing param info.
+
+    Returns
+    -------
+    params : list
+        List of tvm.runtime.Tensor parameters in metadata order.
+    """
+    from tvm.contrib import tvmjs
+
+    params, meta = tvmjs.load_tensor_cache(model_weight_path, device)
+    param_names = [param["name"] for param in model_metadata["params"]]
+    assert len(param_names) == meta["ParamSize"]
+    return [params[name] for name in param_names]
+
+
+def get_embedding_metadata(config: Dict[str, Any]) -> Optional[Dict[str, Any]]:  # noqa: UP006
+    """Read emedding metadata from mlc-chat-config or model lib metadata.
+
+    Parameters
+    ----------
+    config : Dict[str, Any]
+        The configuration dictionary containing model metadata.
+
+    Returns
+    -------
+    embedding_metadata : Optional[Dict[str, Any]] = None if it's not an embedding model.
+        The embedding metadata dictionary.
+    """
+    if config.get("model_task") == "embedding":
+        return config.get("embedding_metadata")
+    return None
+
+
+def detect_embedding_model_type(mod) -> Literal["encoder", "decoder"]:
+    """Detect embedding model type from compiled TVM module functions.
+
+    Parameters
+    ----------
+    mod : tvm.runtime.Module
+        The VM module with model functions.
+
+    Returns
+    -------
+    model_type : str
+        "encoder" for BERT-style models, "decoder" for Qwen3-Embeddings style.
+    """
+    has_embed = mod.implements_function("embed")
+    has_prefill_to_hidden = mod.implements_function("prefill_to_last_hidden_states")
+    has_prefill = mod.implements_function("prefill")
+
+    if has_embed and has_prefill_to_hidden:
+        return "decoder"
+    if has_prefill:
+        return "encoder"
+    raise ValueError(
+        "Model does not support embedding inference. "
+        "Expected 'embed' + 'prefill_to_last_hidden_states' (decoder) "
+        "or 'prefill' (encoder)."
+    )

python/mlc_llm/serve/entrypoints/__init__.py

@@ -0,0 +1,8 @@
+"""The entrypoints for MLC LLM server."""
+
+from . import (
+    debug_entrypoints,
+    metrics_entrypoints,
+    microserving_entrypoints,
+    openai_entrypoints,
+)

python/mlc_llm/serve/entrypoints/debug_entrypoints.py

@@ -0,0 +1,128 @@
+"""MLC LLM server debug entrypoints"""
+
+import json
+from http import HTTPStatus
+
+import fastapi
+
+from mlc_llm.protocol import error_protocol
+from mlc_llm.serve.server import ServerContext
+
+app = fastapi.APIRouter()
+
+################ /debug/dump_event_trace ################
+
+
+@app.post("/debug/dump_event_trace")
+async def debug_dump_event_trace(request: fastapi.Request):
+    """Return the recorded events in Chrome Trace Event Format in JSON string.
+    The input request payload should have only one field, specifying the
+    model to query. For example: `{"model": "Llama-2-7b-chat-hf-q0f16"}`.
+    """
+    # Get the raw request body as bytes
+    request_raw_data = await request.body()
+    request_json_str = request_raw_data.decode("utf-8")
+    try:
+        # Parse the JSON string
+        request_dict = json.loads(request_json_str)
+    except json.JSONDecodeError:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+    if "model" not in request_dict:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+
+    # Check the requested model.
+    model = request_dict["model"]
+
+    server_context: ServerContext = ServerContext.current()
+    async_engine = server_context.get_engine(model)
+
+    if async_engine is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{model}" is not served.',
+        )
+    if async_engine.state.trace_recorder is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{model}" does not enable tracing',
+        )
+
+    return json.loads(async_engine.state.trace_recorder.dump_json())
+
+
+################ /debug/cuda_profiler_start/end ################
+
+
+@app.post("/debug/cuda_profiler_start")
+async def debug_cuda_profiler_start(_request: fastapi.Request):
+    """Start the cuda profiler for the engine. Only for debug purpose."""
+    server_context: ServerContext = ServerContext.current()
+    # Since the CUDA profiler is process-wise, call the function for one model is sufficient.
+    for model in server_context.get_model_list():
+        async_engine = server_context.get_engine(model)
+        async_engine._debug_call_func_on_all_worker("mlc.debug_cuda_profiler_start")
+        break
+
+
+@app.post("/debug/cuda_profiler_stop")
+async def debug_cuda_profiler_stop(_request: fastapi.Request):
+    """Stop the cuda profiler for the engine. Only for debug purpose."""
+    server_context: ServerContext = ServerContext.current()
+    # Since the CUDA profiler is process-wise, call the function for one model is sufficient.
+    for model in server_context.get_model_list():
+        async_engine = server_context.get_engine(model)
+        async_engine._debug_call_func_on_all_worker("mlc.debug_cuda_profiler_stop")
+        break
+
+
+@app.post("/debug/dump_engine_metrics")
+async def debug_dump_engine_metrics(request: fastapi.Request):
+    """Dump the engine metrics for the engine. Only for debug purpose."""
+    # Get the raw request body as bytes
+    request_raw_data = await request.body()
+    request_json_str = request_raw_data.decode("utf-8")
+    try:
+        # Parse the JSON string
+        request_dict = json.loads(request_json_str)
+    except json.JSONDecodeError:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+
+    # Check the requested model.
+    model = request_dict.get("model", None)
+
+    server_context: ServerContext = ServerContext.current()
+    async_engine = server_context.get_engine(model)
+    res = await async_engine.metrics()
+    return res
+
+
+@app.post("/debug/reset_engine")
+async def debug_reset_engine_stats(request: fastapi.Request):
+    """Reset the engine, clean up all running data and metrics."""
+    # Get the raw request body as bytes
+    request_raw_data = await request.body()
+    request_json_str = request_raw_data.decode("utf-8")
+    try:
+        # Parse the JSON string
+        request_dict = json.loads(request_json_str)
+    except json.JSONDecodeError:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+    if "model" not in request_dict:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST, message=f"Invalid request {request_json_str}"
+        )
+
+    # Check the requested model.
+    model = request_dict["model"]
+
+    server_context: ServerContext = ServerContext.current()
+    async_engine = server_context.get_engine(model)
+    async_engine.reset()

python/mlc_llm/serve/entrypoints/metrics_entrypoints.py

@@ -0,0 +1,23 @@
+"""MLC LLM server metrics entrypoints"""
+
+import fastapi
+from fastapi.responses import PlainTextResponse
+
+from mlc_llm.serve.server import ServerContext
+
+app = fastapi.APIRouter()
+
+################ /metrics ################
+
+
+@app.get("/metrics", response_class=PlainTextResponse)
+async def metrics(_request: fastapi.Request):
+    """Start the cuda profiler for the engine. Only for debug purpose."""
+    server_context: ServerContext = ServerContext.current()
+    # Use the metrics from first engine for now
+    # TODO(mlc-team): consider refactor server context to
+    # single engine since multiple AsyncMLCEngine do not work well with each other
+    # We need to work within the internal engine instead.
+    for model in server_context.get_model_list():
+        async_engine = server_context.get_engine(model)
+        return (await async_engine.metrics()).prometheus_text()

python/mlc_llm/serve/entrypoints/microserving_entrypoints.py

@@ -0,0 +1,73 @@
+"""MicroServing server entrypoints in MLC LLM"""
+
+import fastapi
+
+from mlc_llm.protocol.debug_protocol import DisaggConfig
+from mlc_llm.protocol.microserving_protocol import (
+    PrepRecvRequest,
+    PrepRecvResponse,
+    RemoteSendRequest,
+    StartGenerateRequest,
+)
+from mlc_llm.protocol.openai_api_protocol import StreamOptions
+
+from .openai_entrypoints import request_completion
+
+app = fastapi.APIRouter()
+
+
+################ MicroServing Endpoints ################
+
+
+@app.post("/microserving/prep_recv")
+async def prep_recv(request: PrepRecvRequest, raw_request: fastapi.Request) -> PrepRecvResponse:
+    """Handle the microserving request for receive preparation.
+    Match the prompt in the prefix cache (when enabled),
+    allocate entries in the KV cache to prepare receiving the KV data of the prompt.
+    Return the matched prefix length and the allocated KV entry metadata.
+    """
+    request.debug_config.disagg_config = DisaggConfig(
+        kind="prepare_receive",
+        kv_window_begin=0,  # always zero for prepare_receive
+        kv_window_end=request.end,
+    )
+    request.stream_options = StreamOptions(include_usage=True)
+    request.stream = False
+
+    response = await request_completion(request=request, raw_request=raw_request)
+    assert response.usage is not None
+    assert response.usage.extra is not None
+    assert "prefix_matched_length" in response.usage.extra
+    assert "kv_append_metadata" in response.usage.extra
+    return PrepRecvResponse(
+        prefix_matched_length=response.usage.extra["prefix_matched_length"],
+        kv_append_metadata=response.usage.extra["kv_append_metadata"],
+    )
+
+
+@app.post("/microserving/remote_send")
+async def remote_send(request: RemoteSendRequest, raw_request: fastapi.Request):
+    """Compute and generate the KV data of the prompt in the specified KV window.
+    Send the KV data to the destination server."""
+    request.debug_config.disagg_config = DisaggConfig(
+        kind="remote_send",
+        kv_window_begin=request.begin,
+        kv_window_end=request.end,
+        kv_append_metadata=request.kv_addr_info,
+        dst_group_offset=request.recv_rank,
+    )
+    request.stream_options = StreamOptions(include_usage=True)
+    request.stream = False
+
+    await request_completion(request=request, raw_request=raw_request)
+    return {}
+
+
+@app.post("/microserving/start_generate")
+async def start_generate(request: StartGenerateRequest, raw_request: fastapi.Request):
+    """Prefill the prompt in the specified KV window, and start decode."""
+    request.debug_config.disagg_config = DisaggConfig(
+        kind="start_generation",
+        kv_window_begin=request.begin,
+    )
+    return await request_completion(request=request, raw_request=raw_request)

python/mlc_llm/serve/entrypoints/openai_entrypoints.py

@@ -0,0 +1,340 @@
+"""OpenAI API-compatible server entrypoints in MLC LLM"""
+
+import base64
+import struct
+from collections.abc import AsyncGenerator
+from http import HTTPStatus
+from typing import List, Optional  # noqa: UP035
+
+import fastapi
+import numpy as np
+
+from mlc_llm.protocol import error_protocol
+from mlc_llm.protocol.openai_api_protocol import (
+    ChatCompletionRequest,
+    CompletionLogProbs,
+    CompletionRequest,
+    EmbeddingObject,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    EmbeddingUsage,
+    ListResponse,
+    LogProbsContent,
+    ModelResponse,
+)
+from mlc_llm.serve import engine_base, engine_utils
+from mlc_llm.serve.server import ServerContext
+
+
+def verify_api_key(request: fastapi.Request):
+    """Function to verify API key"""
+    server_context = ServerContext.current()
+    # Only perform verification when API key is configured
+    if server_context is not None and server_context.api_key is not None:
+        provided_key = request.headers.get("Authorization", "").replace("Bearer ", "")
+        if provided_key != server_context.api_key:
+            raise fastapi.HTTPException(status_code=401, detail="Invalid API Key")
+
+
+app = fastapi.APIRouter(dependencies=[fastapi.Depends(verify_api_key)])
+
+
+################ v1/embeddings ################
+
+
+@app.post("/v1/embeddings")
+async def request_embedding(request: EmbeddingRequest):
+    """OpenAI-compatible embedding API.
+    API reference: https://platform.openai.com/docs/api-reference/embeddings/create
+    """
+    server_context: ServerContext = ServerContext.current()
+    embedding_engine = server_context.get_embedding_engine(request.model)
+    if embedding_engine is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{request.model}" is not served as an embedding model.',
+        )
+
+    # Normalize input to List[str]
+    inputs: List[str]  # noqa: UP006
+    if isinstance(request.input, str):
+        inputs = [request.input]
+    elif (
+        isinstance(request.input, list)
+        and len(request.input) > 0
+        and isinstance(request.input[0], str)
+    ):
+        inputs = list(request.input)
+    else:
+        # Token ID inputs (List[int] or List[List[int]]) — decode back to strings
+        if isinstance(request.input[0], int):
+            inputs = [embedding_engine.tokenizer.decode(request.input)]
+        else:
+            inputs = [embedding_engine.tokenizer.decode(ids) for ids in request.input]
+
+    # Run embedding inference (async — does not block the event loop)
+    try:
+        embeddings, total_tokens = await embedding_engine.async_embed(inputs)
+    except Exception as exc:
+        return error_protocol.create_error_response(
+            HTTPStatus.INTERNAL_SERVER_ERROR,
+            message=f"Embedding inference failed: {exc}",
+        )
+
+    # Optional: truncate dimensions (Matryoshka-style).
+    # This is API-level renormalization after dimension truncation,
+    # independent of model metadata normalize. Always renormalize
+    # truncated vectors to maintain unit length per OpenAI API contract.
+    if request.dimensions is not None:
+        for i, emb in enumerate(embeddings):
+            vec = np.array(emb[: request.dimensions], dtype=np.float32)
+            norm = np.linalg.norm(vec)
+            if norm > 1e-12:
+                vec = vec / norm
+            embeddings[i] = vec.tolist()
+
+    # Build response data
+    resp_data = []
+    for i, emb in enumerate(embeddings):
+        if request.encoding_format == "base64":
+            binary = struct.pack(f"<{len(emb)}f", *emb)
+            resp_data.append(
+                EmbeddingObject(
+                    embedding=base64.b64encode(binary).decode("utf-8"),
+                    index=i,
+                )
+            )
+        else:
+            resp_data.append(EmbeddingObject(embedding=emb, index=i))
+
+    return EmbeddingResponse(
+        data=resp_data,
+        model=request.model,
+        usage=EmbeddingUsage(prompt_tokens=total_tokens, total_tokens=total_tokens),
+    )
+
+
+################ v1/models ################
+
+
+@app.get("/v1/models")
+async def request_models() -> ListResponse:
+    """OpenAI-compatible served model query API.
+    API reference: https://platform.openai.com/docs/api-reference/models
+    """
+    server_context: ServerContext = ServerContext.current()
+    return ListResponse(data=[ModelResponse(id=model) for model in server_context.get_model_list()])
+
+
+################ v1/completions ################
+
+
+@app.post("/v1/completions")
+async def request_completion(request: CompletionRequest, raw_request: fastapi.Request):
+    """OpenAI-compatible completion API.
+    API reference: https://platform.openai.com/docs/api-reference/completions/create
+    """
+    # - Check the requested model.
+    server_context: ServerContext = ServerContext.current()
+    request_final_usage_include_extra = server_context.enable_debug
+    request_include_debug_config = server_context.enable_debug
+
+    if not request_include_debug_config:
+        request.debug_config = None
+
+    async_engine = server_context.get_engine(request.model)
+    if async_engine is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{request.model}" is not served.',
+        )
+    # FIXME: This is a temporary solution to make sure
+    # prep_recv, remote_send and start_generation process the same request
+    request_id = request.user if request.user is not None else f"cmpl-{engine_utils.random_uuid()}"
+
+    # Streaming response.
+    if request.stream:
+        # We manually get the first response from generator to
+        # capture potential exceptions in this scope, rather then
+        # the StreamingResponse scope.
+        stream_generator = async_engine._handle_completion(
+            request,
+            request_id,
+            request_final_usage_include_extra=request_final_usage_include_extra,
+        )
+        first_response = await anext(  # noqa: F821
+            stream_generator
+        )
+
+        async def completion_stream_generator() -> AsyncGenerator[str, None]:
+            if isinstance(first_response, StopAsyncIteration):
+                yield "data: [DONE]\n\n"
+                return
+            yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
+            async for response in stream_generator:
+                yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
+            yield "data: [DONE]\n\n"
+
+        return fastapi.responses.StreamingResponse(
+            completion_stream_generator(), media_type="text/event-stream"
+        )
+
+    # Normal response.
+    request_final_usage = None
+    output_texts = [""] * request.n
+    finish_reasons: List[Optional[str]] = [None] * request.n  # noqa: UP006
+    logprob_results: List[Optional[CompletionLogProbs]] = [None] * request.n  # noqa: UP006
+
+    async for response in async_engine._handle_completion(
+        request,
+        request_id,
+        request_final_usage_include_extra=request_final_usage_include_extra,
+    ):
+        if await raw_request.is_disconnected():
+            # In non-streaming cases, the engine will not be notified
+            # when the request is disconnected.
+            # Therefore, we check if it is disconnected each time,
+            # and explicitly return.
+            # Note that requesta abort is triggered when the async for and funciton scope ends.
+            return error_protocol.create_error_response(
+                HTTPStatus.BAD_REQUEST, message="The request has disconnected"
+            )
+        # this is the final chunk
+        if response.usage is not None:
+            request_final_usage = response.usage
+            # remove extra information if debug is not enabled
+            if not server_context.enable_debug:
+                request_final_usage.extra = None
+            continue
+        for choice in response.choices:
+            output_texts[choice.index] += choice.text
+            if choice.finish_reason is not None and finish_reasons[choice.index] is None:
+                finish_reasons[choice.index] = choice.finish_reason
+            if choice.logprobs is not None:
+                if logprob_results[choice.index] is None:
+                    logprob_results[choice.index] = choice.logprobs
+                else:
+                    logprob_results[choice.index].token_logprobs.extend(
+                        choice.logprobs.token_logprobs
+                    )
+                    logprob_results[choice.index].tokens.extend(choice.logprobs.tokens)
+                    logprob_results[choice.index].top_logprobs.extend(choice.logprobs.top_logprobs)
+
+    return engine_base.wrap_completion_response(
+        request_id=request_id,
+        model=request.model,
+        output_texts=output_texts,
+        finish_reasons=finish_reasons,
+        logprob_results=logprob_results,
+        usage=request_final_usage,
+    )
+
+
+################ v1/chat/completions ################
+
+
+@app.post("/v1/chat/completions")
+async def request_chat_completion(request: ChatCompletionRequest, raw_request: fastapi.Request):
+    """OpenAI-compatible chat completion API.
+    API reference: https://platform.openai.com/docs/api-reference/chat
+    """
+    # - Check the requested model.
+    server_context: ServerContext = ServerContext.current()
+    request_final_usage_include_extra = server_context.enable_debug
+    request_include_debug_config = server_context.enable_debug
+
+    if not request_include_debug_config:
+        request.debug_config = None
+
+    async_engine = server_context.get_engine(request.model)
+    if async_engine is None:
+        return error_protocol.create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            message=f'The requested model "{request.model}" is not served.',
+        )
+    # FIXME: This is a temporary solution to make sure
+    # prep_recv, remote_send and start_generation process the same request
+    request_id = (
+        request.user if request.user is not None else f"chatcmpl-{engine_utils.random_uuid()}"
+    )
+
+    # Streaming response.
+    if request.stream:
+        # We manually get the first response from generator to
+        # capture potential exceptions in this scope, rather then
+        # the StreamingResponse scope.
+        stream_generator = async_engine._handle_chat_completion(
+            request,
+            request_id,
+            request_final_usage_include_extra=request_final_usage_include_extra,
+        )
+        first_response = await anext(  # noqa: F821
+            stream_generator
+        )
+
+        async def completion_stream_generator() -> AsyncGenerator[str, None]:
+            if isinstance(first_response, StopAsyncIteration):
+                yield "data: [DONE]\n\n"
+                return
+            yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
+            async for response in stream_generator:
+                yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
+            yield "data: [DONE]\n\n"
+
+        return fastapi.responses.StreamingResponse(
+            completion_stream_generator(), media_type="text/event-stream"
+        )
+
+    # Normal response.
+    request_final_usage = None
+    output_texts = ["" for _ in range(request.n)]
+    finish_reasons: List[Optional[str]] = [None for _ in range(request.n)]  # noqa: UP006
+    logprob_results: Optional[List[List[LogProbsContent]]] = (  # noqa: UP006
+        [[] for _ in range(request.n)] if request.logprobs else None
+    )
+
+    async for response in async_engine._handle_chat_completion(
+        request,
+        request_id,
+        request_final_usage_include_extra=request_final_usage_include_extra,
+    ):
+        if await raw_request.is_disconnected():
+            # In non-streaming cases, the engine will not be notified
+            # when the request is disconnected.
+            # Therefore, we check if it is disconnected each time,
+            # no need to explicitly abort, as the chat completion
+            # return will trigger abort call
+            return error_protocol.create_error_response(
+                HTTPStatus.BAD_REQUEST, message="The request has disconnected"
+            )
+        # usage is always the last chunk
+        if response.usage is not None:
+            request_final_usage = response.usage
+            # remove extra information if debug is not enabled
+            if not server_context.enable_debug:
+                request_final_usage.extra = None
+
+        for choice in response.choices:
+            assert isinstance(choice.delta.content, str)
+            output_texts[choice.index] += choice.delta.content
+            if choice.finish_reason is not None and finish_reasons[choice.index] is None:
+                finish_reasons[choice.index] = choice.finish_reason
+            if choice.logprobs is not None:
+                assert logprob_results is not None
+                logprob_results[choice.index] += choice.logprobs.content
+
+    assert all(finish_reason is not None for finish_reason in finish_reasons)
+    use_function_calling, tool_calls_list = engine_base.process_function_call_output(
+        output_texts, finish_reasons
+    )
+
+    return engine_base.wrap_chat_completion_response(
+        request_id=request_id,
+        model=request.model,
+        output_texts=output_texts,
+        finish_reasons=finish_reasons,
+        tool_calls_list=tool_calls_list,
+        logprob_results=logprob_results,
+        use_function_calling=use_function_calling,
+        usage=request_final_usage,
+    )

python/mlc_llm/serve/event_trace_recorder.py

@@ -0,0 +1,37 @@
+"""The event trace recorder in MLC LLM serving"""
+
+import tvm_ffi
+from tvm.runtime import Object
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("mlc.serve.EventTraceRecorder")
+class EventTraceRecorder(Object):
+    """The event trace recorder for requests."""
+
+    def __init__(self) -> None:
+        """Initialize a trace recorder."""
+        self.__init_handle_by_constructor__(_ffi_api.EventTraceRecorder)
+
+    def add_event(self, request_id: str, event: str) -> None:
+        """Record a event for the input request in the trace recorder.
+
+        Parameters
+        ----------
+        request_id : str
+            The subject request of the event.
+
+        event : str
+            The event in a string name.
+            It can have one of the following patterns:
+            - "start xxx", which marks the start of event "xxx",
+            - "finish xxx", which marks the finish of event "xxx",
+            - "yyy", which marks the instant event "yyy".
+            The "starts" and "finishes" will be automatically paired in the trace recorder.
+        """
+        return _ffi_api.EventTraceRecorderAddEvent(self, request_id, event)
+
+    def dump_json(self) -> str:
+        """Dump the logged events in Chrome Trace Event Format in JSON string."""
+        return _ffi_api.EventTraceRecorderDumpJSON(self)

python/mlc_llm/serve/radix_tree.py

@@ -0,0 +1,154 @@
+"""The Paged Radix Tree class."""
+
+from typing import List, Tuple, Union  # noqa: UP035
+
+import tvm_ffi
+from tvm.runtime import Object
+from tvm_ffi import Shape
+
+from . import _ffi_api
+
+
+@tvm_ffi.register_object("mlc.serve.PagedRadixTree")
+class PagedRadixTree(Object):
+    """The paged radix tree to manage prefix and sequence."""
+
+    def __init__(self):
+        """
+        Constructor of paged radix tree.
+        """
+        self.__init_handle_by_constructor__(_ffi_api.PagedRadixTree)
+
+    def match(self, tokens: Union[Shape, List, Tuple]) -> Tuple[int, Shape]:  # noqa: UP006
+        """
+        Get all sequences with longest common prefix with given prefix tokens.
+
+        Parameters
+        ----------
+        tokens : Union[Shape, List, Tuple]
+            The prefix tokens for reference.
+
+        Returns
+        ------
+        matched_offset : int
+            The matched prefix length.
+        seq_ids : Shape
+            The array of matched sequence indice.
+        """
+        if isinstance(tokens, (list, tuple)):
+            tokens = Shape(tokens)
+        output = _ffi_api.PagedRadixTreeMatchPrefix(self, tokens)
+        if len(output) == 1:
+            return output[0], []
+        return output[0], output[1:]
+
+    def add(self, seq_id: int) -> None:
+        """
+        Add an empty sequence.
+
+        Parameters
+        ----------
+        seq_id : int
+            The sequence ID for index.
+        """
+        _ffi_api.PagedRadixTreeAddSequence(self, seq_id)
+
+    def remove(self, seq_id: int) -> None:
+        """
+        Remove a sequence.
+
+        Parameters
+        ----------
+        seq_id : int
+            The sequence ID to remove.
+        """
+        _ffi_api.PagedRadixTreeRemoveSequence(self, seq_id)
+
+    def extend(self, seq_id: int, tokens: Union[Shape, List, Tuple]) -> None:  # noqa: UP006
+        """
+        Extend a sequence with given tokens.
+
+        Parameters
+        ----------
+        seq_id : int
+            The sequence ID for index.
+        tokens : Union[Shape, List, Tuple]
+            The given tokens to extend.
+        """
+        if isinstance(tokens, (list, tuple)):
+            tokens = Shape(tokens)
+        _ffi_api.PagedRadixTreeExtendSequence(self, seq_id, tokens)
+
+    def rollback(self, seq_id: int, num_tokens: int) -> None:
+        """
+        Roll back a sequence by number of tokens.
+
+        Parameters
+        ----------
+        seq_id : int
+            The sequence ID for index.
+        num_tokens : int
+            The number of tokens to be rolled back.
+        """
+        _ffi_api.PagedRadixTreeRollBackSequence(self, seq_id, num_tokens)
+
+    def fork(self, seq_id: int, parent_seq_id: int, forked_offset: int) -> None:
+        """
+        Fork a sequence from parent sequence at given position.
+
+        Parameters
+        ----------
+        seq_id : int
+            The new sequence ID.
+        parent_seq_id : int
+            The parent sequence ID to fork from.
+        forked_offset : int
+            The position of parent sequence to fork at.
+            The valid value is [1, length of forked sequence].
+            If the position equals the length of forked sequence,
+            the new sequence will copy the entire forked sequence.
+        """
+        _ffi_api.PagedRadixTreeForkSequence(self, seq_id, parent_seq_id, forked_offset)
+
+    def get(self, seq_id: int) -> Shape:
+        """
+        Get a sequence's all tokens.
+
+        Parameters
+        ----------
+        seq_id : int
+            The sequence ID for index.
+
+        Returns
+        ------
+        tokens : Shape
+            The sequence tokens.
+        """
+        return _ffi_api.PagedRadixTreeGetSequence(self, seq_id)
+
+    def get_length(self, seq_id: int) -> int:
+        """
+        Get a sequence's length.
+
+        Parameters
+        ----------
+        seq_id : int
+            The sequence ID for index.
+
+        Returns
+        ------
+        length : int
+            The sequence length.
+        """
+        return _ffi_api.PagedRadixTreeGetSequenceLength(self, seq_id)
+
+    def free_capacity(self) -> int:
+        """
+        Get the remaining token capacity of the paged radix tree.
+
+        Returns
+        ------
+        capacity : int
+            The remaining token capacity of the paged radix tree.
+        """
+        return _ffi_api.PagedRadixTreeFreeCapacity(self)

python/mlc_llm/serve/request.py

@@ -0,0 +1,34 @@
+"""The request class in MLC LLM serving"""
+
+from typing import List  # noqa: UP035
+
+import tvm_ffi
+from tvm.runtime import Object
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+
+from . import _ffi_api
+from .data import Data
+
+
+@tvm_ffi.register_object("mlc.serve.Request")
+class Request(Object):
+    """The user submitted text-generation request, which contains
+    a unique request id, a list of multi-modal inputs, a set of
+    generation configuration parameters.
+
+    Note
+    ----
+    Do not explicitly construct this class.
+    Construct this object via engine.create_request functions.
+    """
+
+    @property
+    def inputs(self) -> List[Data]:  # noqa: UP006
+        """The inputs of the request."""
+        return _ffi_api.RequestGetInputs(self)
+
+    @property
+    def generation_config(self) -> GenerationConfig:
+        """The generation config of the request."""
+        return GenerationConfig.model_validate_json(_ffi_api.RequestGetGenerationConfigJSON(self))

python/mlc_llm/serve/server/__init__.py

@@ -0,0 +1,4 @@
+"""The server related data structure and tools in MLC LLM serve."""
+
+from .popen_server import PopenServer
+from .server_context import ServerContext

python/mlc_llm/serve/server/popen_server.py

@@ -0,0 +1,200 @@
+"""The MLC LLM server launched in a subprocess."""
+
+import os
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Literal, Optional, Union
+
+import psutil
+import requests
+from tvm.runtime import Device
+
+from mlc_llm.serve.config import EngineConfig
+from mlc_llm.serve.engine_base import _check_engine_config
+
+
+class PopenServer:
+    """The wrapper of MLC LLM server, which runs the server in
+    a background subprocess.
+
+    This server can be used for debugging purposes.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        device: Union[str, Device] = "auto",
+        *,
+        model_lib: Optional[str] = None,
+        mode: Literal["local", "interactive", "server"] = "local",
+        engine_config: Optional[EngineConfig] = None,
+        enable_debug: bool = True,
+        enable_tracing: bool = False,
+        host: str = "127.0.0.1",
+        port: int = 8082,
+    ) -> None:
+        """Please check out `python/mlc_llm/cli/serve.py` for the server arguments."""
+        # - Check the fields fields of `engine_config`.
+        if engine_config is None:
+            engine_config = EngineConfig()
+        _check_engine_config(model, model_lib, mode, engine_config)
+
+        self.model = model
+        self.model_lib = model_lib
+        self.device = device
+        self.mode = mode
+        self.enable_debug = enable_debug
+        self.engine_config = engine_config
+        self.enable_tracing = enable_tracing
+        self.enable_debug = enable_debug
+        self.host = host
+        self.port = port
+        self._proc: Optional[subprocess.Popen] = None
+
+        self.base_url = ""
+        self.openai_v1_base_url = ""
+
+    def start(self, extra_env=None) -> None:
+        """Launch the server in a popen subprocess.
+        Wait until the server becomes ready before return.
+        """
+        extra_env = extra_env or {}
+        cmd = [sys.executable]
+        cmd += ["-m", "mlc_llm", "serve", self.model]
+        if self.model_lib is not None:
+            cmd += ["--model-lib", self.model_lib]
+        cmd += ["--device", self.device]
+
+        if self.enable_debug:
+            cmd += ["--enable-debug"]
+
+        if self.mode is not None:
+            cmd += ["--mode", self.mode]
+
+        if len(self.engine_config.additional_models) > 0:
+            args_additional_model = []
+            for additional_model in self.engine_config.additional_models:
+                if isinstance(additional_model, str):
+                    args_additional_model.append(additional_model)
+                else:
+                    args_additional_model.append(additional_model[0] + "," + additional_model[1])
+            cmd += ["--additional-models", *args_additional_model]
+        cmd += ["--speculative-mode", self.engine_config.speculative_mode]
+        cmd += ["--prefix-cache-mode", self.engine_config.prefix_cache_mode]
+
+        args_overrides = []
+        if self.engine_config.max_num_sequence is not None:
+            args_overrides.append(f"max_num_sequence={self.engine_config.max_num_sequence}")
+        if self.engine_config.max_total_sequence_length is not None:
+            args_overrides.append(
+                f"max_total_seq_length={self.engine_config.max_total_sequence_length}"
+            )
+        if self.engine_config.prefill_chunk_size is not None:
+            args_overrides.append(f"prefill_chunk_size={self.engine_config.prefill_chunk_size}")
+        if self.engine_config.max_history_size is not None:
+            args_overrides.append(f"max_history_size={self.engine_config.max_history_size}")
+        if self.engine_config.gpu_memory_utilization is not None:
+            args_overrides.append(
+                f"gpu_memory_utilization={self.engine_config.gpu_memory_utilization}"
+            )
+        if self.engine_config.spec_draft_length is not None:
+            args_overrides.append(f"spec_draft_length={self.engine_config.spec_draft_length}")
+        if self.engine_config.prefix_cache_max_num_recycling_seqs is not None:
+            args_overrides.append(
+                "prefix_cache_max_num_recycling_seqs="
+                + str(self.engine_config.prefix_cache_max_num_recycling_seqs)
+            )
+        if len(args_overrides) > 0:
+            cmd += ["--overrides", ";".join(args_overrides)]
+
+        if self.enable_tracing:
+            cmd += ["--enable-tracing"]
+        if self.enable_debug:
+            cmd += ["--enable-debug"]
+
+        cmd += ["--host", self.host]
+        cmd += ["--port", str(self.port)]
+        process_path = str(Path(__file__).resolve().parents[4])
+        final_env = os.environ.copy()
+        for key, value in extra_env.items():
+            final_env[key] = value
+        self._proc = subprocess.Popen(cmd, cwd=process_path, env=final_env)
+        # NOTE: DO NOT USE `stdout=subprocess.PIPE, stderr=subprocess.PIPE`
+        # in subprocess.Popen here. PIPE has a fixed-size buffer with may block
+        # and hang forever.
+
+        # Try to query the server until it is ready.
+        self.base_url = f"http://{self.host}:{str(self.port)}"
+        self.openai_v1_base_url = f"http://{self.host}:{str(self.port)}/v1"
+        openai_v1_models_url = f"{self.base_url}/v1/models"
+
+        query_result = None
+        timeout = 120
+        attempts = 0.0
+        while query_result is None and attempts < timeout:
+            try:
+                query_result = requests.get(openai_v1_models_url, timeout=60)
+                if query_result.status_code != 200:
+                    query_result = None
+                    attempts += 0.1
+                    time.sleep(0.1)
+            except Exception:
+                attempts += 0.1
+                time.sleep(0.1)
+
+        # Check if the subprocess terminates unexpectedly or
+        # the queries reach the timeout.
+        process_return_code = self._proc.poll()
+        if process_return_code is not None:
+            raise RuntimeError(
+                "The server fails to launch. "
+                f'Please check if "{self.model}" is a valid model compiled by MLC LLM.'
+            )
+        if attempts == timeout:
+            self.terminate()
+            raise RuntimeError(f"The server fails to launch in {timeout} seconds.")
+
+    def terminate(self) -> None:
+        """Terminate the server subprocess."""
+        if self._proc is None:
+            return
+
+        # Kill all the child processes.
+        def kill_child_processes():
+            try:
+                parent = psutil.Process(self._proc.pid)
+                children = parent.children(recursive=True)
+            except psutil.NoSuchProcess:
+                return
+
+            for process in children:
+                try:
+                    process.kill()
+                except psutil.NoSuchProcess:
+                    pass
+
+        kill_child_processes()
+
+        # Kill the process.
+        try:
+            self._proc.kill()
+        except OSError:
+            pass
+
+        # Join the process to avoid zombies.
+        try:
+            self._proc.wait(timeout=10.0)
+        except subprocess.TimeoutExpired:
+            pass
+        self._proc = None
+
+    def __enter__(self):
+        """Start the server."""
+        self.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Terminate the server."""
+        self.terminate()

python/mlc_llm/serve/server/server_context.py

@@ -0,0 +1,72 @@
+"""Server context that shared by multiple entrypoint files."""
+
+from typing import TYPE_CHECKING, Dict, List, Optional  # noqa: UP035
+
+from ..engine import AsyncMLCEngine
+
+if TYPE_CHECKING:
+    from ..embedding_engine import AsyncEmbeddingEngine
+
+
+class ServerContext:
+    """The global server context, including the running models
+    and corresponding async engines.
+    """
+
+    server_context: Optional["ServerContext"] = None
+    enable_debug: bool = False
+
+    def __init__(self) -> None:
+        self._models: Dict[str, AsyncMLCEngine] = {}  # noqa: UP006
+        self._embedding_engines: Dict[str, AsyncEmbeddingEngine] = {}  # noqa: UP006
+        self.api_key: Optional[str] = None
+
+    def __enter__(self):
+        if ServerContext.server_context is not None:
+            raise RuntimeError("Server context already exists.")
+        ServerContext.server_context = self
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        for model_engine in self._models.values():
+            model_engine.terminate()
+        for emb_engine in self._embedding_engines.values():
+            emb_engine.terminate()
+        self._models.clear()
+        self._embedding_engines.clear()
+        ServerContext.server_context = None
+
+    @staticmethod
+    def current():
+        """Returns the current ServerContext."""
+        return ServerContext.server_context
+
+    def add_model(self, hosted_model: str, engine: AsyncMLCEngine) -> None:
+        """Add a new model to the server context together with the engine."""
+        if hosted_model in self._models:
+            raise RuntimeError(f"Model {hosted_model} already running.")
+        self._models[hosted_model] = engine
+
+    def get_engine(self, model: Optional[str]) -> Optional[AsyncMLCEngine]:
+        """Get the async engine of the requested model, or the unique async engine
+        if only one engine is served."""
+        if len(self._models) == 1:
+            return next(iter(self._models.values()))
+        return self._models.get(model, None)
+
+    def get_model_list(self) -> List[str]:  # noqa: UP006
+        """Get the list of all models on serve, including embedding models."""
+        return list(self._models.keys()) + list(self._embedding_engines.keys())
+
+    def add_embedding_engine(self, hosted_model: str, engine: "AsyncEmbeddingEngine") -> None:
+        """Add a new embedding model to the server context."""
+        if hosted_model in self._embedding_engines:
+            raise RuntimeError(f"Embedding model {hosted_model} already running.")
+        self._embedding_engines[hosted_model] = engine
+
+    def get_embedding_engine(self, model: Optional[str]) -> Optional["AsyncEmbeddingEngine"]:
+        """Get the embedding engine of the requested model, or the unique
+        embedding engine if only one is served."""
+        if len(self._embedding_engines) == 1:
+            return next(iter(self._embedding_engines.values()))
+        return self._embedding_engines.get(model, None)

python/mlc_llm/serve/sync_engine.py

@@ -0,0 +1,363 @@
+"""The MLC LLM synchronized engine.
+
+NOTE: This engine defined in this file directly wraps the underlying
+Engine implementation in C++, is not optimized by multi-threading and
+does not offer standard OpenAI API interface.
+
+We do not expose it and use it by default. As of now it mainly serves
+the test and debug purpose because of its simplicity.
+"""
+
+import json
+from collections.abc import Sequence
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union  # noqa: UP035
+
+import tvm
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import data
+from mlc_llm.serve.config import EngineConfig
+from mlc_llm.serve.engine_base import (
+    EngineMetrics,
+    _check_engine_config,
+    _parse_models,
+    _print_engine_mode_logging_msg,
+    _process_model_args,
+    detect_device,
+)
+from mlc_llm.serve.event_trace_recorder import EventTraceRecorder
+from mlc_llm.serve.request import Request
+from mlc_llm.support import logging
+from mlc_llm.tokenizers import TextStreamer, Tokenizer
+
+logger = logging.getLogger(__name__)
+
+
+def _create_tvm_module(
+    creator: str,
+    ffi_funcs: Sequence[str],
+    creator_args: Optional[List[Any]] = None,  # noqa: UP006
+) -> Dict[str, Callable]:  # noqa: UP006
+    """Internal method to create a module."""
+    if creator_args is None:
+        creator_args = []
+    module = tvm.get_global_func(creator, allow_missing=False)(*creator_args)
+    return {key: module[key] for key in ffi_funcs}
+
+
+class SyncMLCEngine:
+    """The Python interface of synchronize request serving engine for MLC LLM.
+
+    The engine receives requests from the "add_request" method. For
+    an given request, the engine will keep generating new tokens for
+    the request until finish (under certain criterion). After finish,
+    the engine will return the generation result through the callback
+    function provided by the request.
+
+    NOTE: This engine directly wraps the underlying Engine implementation
+    in C++, is not optimized by multi-threading and does not offer standard
+    OpenAI API interface. We do not expose it and use it by default.
+    As of now it mainly serves the test and debug purpose because of its
+    simplicity.
+
+    Parameters
+    ----------
+    engine_config : Optional[EngineConfig]
+        Additional configurable arguments of MLC engine.
+        See class "EngineConfig" for more detail.
+
+    enable_tracing : bool
+        A boolean indicating if to enable event logging for requests.
+
+    request_stream_callback : Optional[Callable[[str, data.TokenData, Optional[str]], None]]
+        The provided callback function to handle the generation
+        output. It has the signature of `(str, data.TokenData, bool) -> None`,
+        where
+        - the first string is the request id,
+        - the TokenData contains the generated **delta** token ids since
+        the last invocation of the callback on the specific request,
+        - the optional string value denotes the finish reason if the
+        generation of the request is finished, or None if it has not finished.
+
+        The callback function is optional at construction, but it needs to
+        be set before the engine executing requests. This can be done via
+        the `set_request_stream_callback` method. Otherwise, the engine will raise
+        exception.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        device: Union[str, tvm.runtime.Device] = "auto",
+        *,
+        model_lib: Optional[str] = None,
+        mode: Literal["local", "interactive", "server"] = "local",
+        engine_config: Optional[EngineConfig] = None,
+        enable_tracing: bool = False,
+        request_stream_callback: Optional[Callable[[List[data.RequestStreamOutput]], None]] = None,  # noqa: UP006
+    ):
+        # - Check the fields fields of `engine_config`.
+        if engine_config is None:
+            engine_config = EngineConfig()
+        _check_engine_config(
+            model,
+            model_lib,
+            mode,
+            engine_config,
+        )
+
+        # - Initialize model loading info.
+        models = _parse_models(model, model_lib, engine_config.additional_models)
+        if isinstance(device, str):
+            device = detect_device(device)
+        assert isinstance(device, tvm.runtime.Device)
+        (
+            model_args,
+            model_config_paths,
+            self.conv_template,
+        ) = _process_model_args(models, device, engine_config)
+
+        # - Load the raw model config into dict
+        self.model_config_dicts = []
+        for i, model_info in enumerate(models):
+            model_info.model_lib = model_args[i][1]
+            with open(model_config_paths[i], encoding="utf-8") as file:
+                self.model_config_dicts.append(json.load(file))
+
+        # - Print logging info for regarding the mode selection.
+        if engine_config.verbose:
+            _print_engine_mode_logging_msg(mode)
+
+        self._ffi = _create_tvm_module(
+            "mlc.serve.create_engine",
+            ffi_funcs=[
+                "init",
+                "add_request",
+                "abort_request",
+                "step",
+                "reset",
+                "json_metrics",
+                "get_request_stream_callback",
+                "set_request_stream_callback",
+                "create_request",
+            ],
+        )
+        self.trace_recorder = EventTraceRecorder() if enable_tracing else None
+
+        engine_config.model = model_args[0][0]
+        engine_config.model_lib = model_args[0][1]
+        engine_config.additional_models = model_args[1:]
+        engine_config.mode = mode
+        self._ffi["init"](
+            engine_config.asjson(),
+            device,
+            request_stream_callback,
+            self.trace_recorder,
+        )
+        self.tokenizer = Tokenizer(model_args[0][0])
+
+    def generate(
+        self,
+        prompts: Union[str, List[str], List[int], List[List[int]], List[List[data.Data]]],  # noqa: UP006
+        generation_config: Union[GenerationConfig, List[GenerationConfig]],  # noqa: UP006
+    ) -> Tuple[List[List[str]], List[Optional[List[List[str]]]]]:  # noqa: UP006
+        """Generate texts for a list of input prompts.
+        Each prompt can be a string or a list of token ids.
+        The generation for each prompt is independent.
+        Return the generation results, one for each prompt.
+
+        Parameters
+        ----------
+        prompts : Union[str, List[str], List[int], List[List[int]]]
+            One or a list of input prompts for text generation.
+            Each prompt can be a string or a list of token ids.
+
+        generation_config : Union[GenerationConfig, List[GenerationConfig]]
+            The generation config for each requests.
+            If the it is a single GenerationConfig instance,
+            this config will be shared by all the prompts.
+            Otherwise, one generation config is required for every
+            prompt.
+
+        Returns
+        -------
+        output_text : List[List[str]]
+            The text generation results, one list of strings for each input prompt.
+            The length of each list is the parallel generation `n` in
+            generation config.
+
+        output_logprobs_str : List[Optional[List[List[str]]]]
+            The logprob strings of each token for each input prompt, or None
+            if an input prompt does not require logprobs.
+        """
+        if isinstance(prompts, str):
+            # `prompts` is a single string.
+            prompts = [prompts]
+        else:
+            assert isinstance(prompts, list), (
+                "Input `prompts` is expected to be a string, a list of "
+                "str, a list of token ids or multiple lists of token ids. "
+            )
+            if len(prompts) == 0:
+                return [], []
+            if isinstance(prompts[0], int):
+                # `prompts` is a list of token ids
+                prompts = [prompts]
+
+        num_requests = len(prompts)
+        if not isinstance(generation_config, list):
+            generation_config = [generation_config] * num_requests
+
+        assert len(generation_config) == num_requests, (
+            "Number of generation config and number of prompts mismatch"
+        )
+
+        num_finished_generations = 0
+        output_texts: List[List[str]] = []  # noqa: UP006
+        output_logprobs_str: List[Optional[List[List[str]]]] = []  # noqa: UP006
+        text_streamers: List[List[TextStreamer]] = []  # noqa: UP006
+        for i in range(num_requests):
+            output_texts.append([])
+            output_logprobs_str.append([] if generation_config[i].logprobs else None)
+            text_streamers.append([])
+            for _ in range(generation_config[i].n):
+                output_texts[i].append("")
+                text_streamers[i].append(TextStreamer(self.tokenizer))
+                if output_logprobs_str[i] is not None:
+                    output_logprobs_str[i].append([])
+
+        num_total_generations = sum(cfg.n for cfg in generation_config)
+
+        # Save a copy of the original function callback since `generate`
+        # overrides the callback function.
+        # The original callback will be set back later on.
+        original_callback = self._ffi["get_request_stream_callback"]()
+
+        # Define the callback function for request generation results
+        def request_stream_callback(delta_outputs: List[data.RequestStreamOutput]):  # noqa: UP006
+            nonlocal num_finished_generations
+            for delta_output in delta_outputs:
+                request_id, stream_outputs = delta_output.unpack()
+                rid = int(request_id)
+
+                assert len(stream_outputs) == generation_config[rid].n
+                for i, (stream_output, text_streamer) in enumerate(
+                    zip(stream_outputs, text_streamers[rid])
+                ):
+                    if output_logprobs_str[rid] is not None:
+                        assert stream_output.delta_logprob_json_strs is not None
+                        output_logprobs_str[rid][i] += stream_output.delta_logprob_json_strs
+
+                    delta_text = stream_output.extra_prefix_string + (
+                        text_streamer.put(stream_output.delta_token_ids)
+                        if len(stream_output.delta_token_ids) > 0
+                        else ""
+                    )
+                    if stream_output.finish_reason is not None:
+                        delta_text += text_streamer.finish()
+
+                    output_texts[rid][i] += delta_text
+                    if stream_output.finish_reason is not None:
+                        num_finished_generations += 1
+
+        # Override the callback function in engine.
+        self._ffi["set_request_stream_callback"](request_stream_callback)
+
+        def convert_to_data(
+            prompt: Union[str, List[int], List[data.Data]],  # noqa: UP006
+        ) -> List[data.Data]:  # noqa: UP006
+            if isinstance(prompt, str):
+                return [data.TextData(prompt)]
+            if isinstance(prompt[0], int):
+                return [data.TokenData(prompt)]
+            return prompt
+
+        # Add requests to engine.
+        for req_id, (prompt, generation_cfg) in enumerate(zip(prompts, generation_config)):
+            input_data = convert_to_data(prompt)
+            self.add_request(
+                self.create_request(
+                    request_id=str(req_id),
+                    inputs=input_data,
+                    generation_config=generation_cfg,
+                )
+            )
+
+        while num_finished_generations != num_total_generations:
+            self.step()
+
+        # Restore the callback function in engine.
+        self._ffi["set_request_stream_callback"](original_callback)
+        return output_texts, output_logprobs_str
+
+    def create_request(
+        self,
+        request_id: str,
+        inputs: Union[data.Data, List[data.Data]],  # noqa: UP006
+        generation_config: GenerationConfig,
+    ):
+        """Create a new request that can be added to engine.
+
+        Parameters
+        ----------
+        request_id : str
+            The unique identifier of the request.
+            Different requests should have different ids.
+
+        inputs : List[Data]
+            The user inputs of a request. Input may have multi-modality.
+
+        generation_config : GenerationConfig
+            The generation configuration of the request.
+
+        Note
+        ----
+        engine may fill in default generation config of the model.
+        """
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+        return self._ffi["create_request"](
+            request_id, inputs, generation_config.model_dump_json(by_alias=True)
+        )
+
+    def add_request(self, request: Request) -> None:
+        """Add a new request to the engine.
+
+        Parameters
+        ----------
+        request : Request
+            The request to add.
+        """
+        self._ffi["add_request"](request)
+
+    def abort_request(self, request_id: str) -> None:
+        """Abort the generation of the request corresponding to the input request id.
+
+        Parameters
+        ----------
+        request_id : str
+            The unique id of the request to abort.
+        """
+        self._ffi["abort_request"](request_id)
+
+    def step(self) -> None:
+        """The main function that the engine takes a step of action.
+
+        At each step, the engine may decide to
+        - run prefill for one (or more) requests,
+        - run one-step decode for the all existing requests
+        ...
+
+        In the end of certain actions (e.g., decode), the engine will
+        check if any request has finished, and will return the
+        generation results for those finished requests.
+        """
+        self._ffi["step"]()
+
+    def reset(self) -> None:
+        """Reset the engine, clean up all running data and metrics."""
+        self._ffi["reset"]()
+
+    def metrics(self) -> EngineMetrics:
+        """Reset the engine, clean up all running data and metrics."""
+        return EngineMetrics(json.loads(self._ffi["json_metrics"]()))