chore: import upstream snapshot with attribution

tests/python/serve/evaluate_engine.py

@@ -0,0 +1,73 @@
+import argparse
+import os
+import random
+from typing import List, Tuple  # noqa: UP035
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine
+
+
+def _parse_args():
+    args = argparse.ArgumentParser()
+    args.add_argument("--model-lib", type=str)
+    args.add_argument("--device", type=str, default="auto")
+    args.add_argument("--batch-size", type=int, default=80)
+    args.add_argument("--max-total-seq-length", type=int)
+    args.add_argument("--seed", type=int, default=0)
+
+    parsed = args.parse_args()
+    parsed.model = os.path.dirname(parsed.model_lib)
+    assert parsed.batch_size % 16 == 0
+    return parsed
+
+
+def generate_requests(
+    num_requests: int, input_length: int, output_length: int
+) -> Tuple[List[List[int]], List[GenerationConfig]]:  # noqa: UP006
+    prompt_ids = []
+    for _ in range(num_requests):
+        token_ids = []
+        for _ in range(input_length):
+            token_ids.append(random.randint(0, 30000))
+        prompt_ids.append(token_ids)
+    generation_config_list = [
+        GenerationConfig(temperature=1.0, top_p=1.0, max_tokens=output_length)
+    ] * num_requests
+    return prompt_ids, generation_config_list
+
+
+def benchmark(args: argparse.Namespace):
+    random.seed(args.seed)
+
+    # Create engine
+    engine = SyncMLCEngine(
+        model=args.model,
+        device=args.device,
+        model_lib=args.model_lib,
+        mode="server",
+        engine_config=EngineConfig(
+            max_num_sequence=args.batch_size,
+            max_total_sequence_length=args.max_total_seq_length,
+        ),
+    )
+
+    print(args)
+    for num_requests in [1, 2, 4, 8, 16, 32, 64]:
+        if num_requests > args.batch_size:
+            continue
+        for input_length in [64, 128, 256, 512, 1024]:
+            if num_requests * input_length >= 16384:
+                continue
+            for output_length in [4]:
+                print(f"nreq={num_requests}\tin={input_length}\tout={output_length}")
+                prompt_ids, generation_config = generate_requests(
+                    num_requests, input_length, output_length
+                )
+                engine.reset()
+                engine.generate(prompt_ids, generation_config)
+                print()
+
+
+if __name__ == "__main__":
+    ARGS = _parse_args()
+    benchmark(ARGS)

tests/python/serve/server/conftest.py

@@ -0,0 +1,34 @@
+import os
+from typing import Tuple  # noqa: UP035
+
+import pytest
+
+from mlc_llm.serve import PopenServer
+
+
+@pytest.fixture(scope="session")
+def served_model() -> Tuple[str, str]:  # noqa: UP006
+    model_lib = os.environ.get("MLC_SERVE_MODEL_LIB")
+    if model_lib is None:
+        raise ValueError(
+            'Environment variable "MLC_SERVE_MODEL_LIB" not found. '
+            "Please set it to model lib compiled by MLC LLM "
+            "(e.g., `dist/Llama-2-7b-chat-hf-q0f16-MLC/Llama-2-7b-chat-hf-q0f16-MLC-cuda.so`)."
+        )
+    model = os.path.dirname(model_lib)
+    return model, model_lib
+
+
+@pytest.fixture(scope="session")
+def launch_server(served_model):
+    """A pytest session-level fixture which launches the server in a subprocess."""
+    server = PopenServer(
+        model=served_model[0],
+        model_lib=served_model[1],
+        enable_tracing=True,
+        enable_debug=True,
+        port=8000,
+    )
+
+    with server:
+        yield

tests/python/serve/server/test_embedding_server.py

@@ -0,0 +1,369 @@
+"""Embedding server endpoint tests in MLC LLM.
+
+Tests the /v1/embeddings endpoint via HTTP using the OpenAI client,
+following the same patterns as test_server.py.
+
+Reuses MLC LLM test infrastructure:
+  - Pytest markers (endpoint)
+  - expect_error() response validation pattern from test_server.py
+  - OpenAI client usage pattern from test_server.py
+  - Session-scoped server fixture pattern from conftest.py
+
+Run (launches its own embedding-only server):
+  MLC_SERVE_EMBEDDING_MODEL_LIB="path/to/model.dylib" \
+    pytest -m endpoint tests/python/serve/server/test_embedding_server.py -v
+
+Environment variables:
+  MLC_SERVE_EMBEDDING_MODEL_LIB  Path to compiled embedding model library (required)
+  MLC_SERVE_EMBEDDING_MODEL      Path to embedding model weight directory
+                                  (optional, defaults to dirname of model lib)
+"""
+
+import json
+import os
+import signal
+import subprocess
+import sys
+import time
+from pathlib import Path
+from typing import Dict, Optional  # noqa: UP035
+
+import numpy as np
+import pytest
+import requests
+from openai import OpenAI
+
+# Reuse MLC LLM marker system
+pytestmark = [pytest.mark.endpoint]
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+EMBEDDING_MODEL_LIB = os.environ.get("MLC_SERVE_EMBEDDING_MODEL_LIB")
+EMBEDDING_MODEL_DIR = os.environ.get(
+    "MLC_SERVE_EMBEDDING_MODEL",
+    os.path.dirname(EMBEDDING_MODEL_LIB) if EMBEDDING_MODEL_LIB else None,
+)
+EMBEDDING_SERVER_HOST = "127.0.0.1"
+EMBEDDING_SERVER_PORT = 8321
+EMBEDDING_BASE_URL = f"http://{EMBEDDING_SERVER_HOST}:{EMBEDDING_SERVER_PORT}/v1"
+EMBEDDING_MODEL_NAME = "embedding"
+
+
+def _skip_if_no_model():
+    if EMBEDDING_MODEL_LIB is None:
+        pytest.skip(
+            'Environment variable "MLC_SERVE_EMBEDDING_MODEL_LIB" not found. '
+            "Set it to a compiled embedding model library."
+        )
+    if not os.path.isfile(EMBEDDING_MODEL_LIB):
+        pytest.skip(f"Embedding model library not found at: {EMBEDDING_MODEL_LIB}")
+    if EMBEDDING_MODEL_DIR is None or not os.path.isdir(EMBEDDING_MODEL_DIR):
+        pytest.skip(f"Embedding model directory not found at: {EMBEDDING_MODEL_DIR}")
+
+
+# ---------------------------------------------------------------------------
+# Response validation helpers — adapted from test_server.py patterns
+# ---------------------------------------------------------------------------
+
+
+def check_embedding_response(
+    response: Dict,  # noqa: UP006
+    *,
+    model: str,
+    num_embeddings: int,
+    expected_dim: Optional[int] = None,
+    check_unit_norm: bool = True,
+):
+    """Validate an OpenAI-compatible embedding response.
+
+    Adapted from check_openai_nonstream_response() in test_server.py,
+    specialized for embedding responses.
+    """
+    assert response["object"] == "list"
+    assert response["model"] == model
+
+    data = response["data"]
+    assert isinstance(data, list)
+    assert len(data) == num_embeddings
+
+    for item in data:
+        assert item["object"] == "embedding"
+        assert isinstance(item["index"], int)
+        emb = item["embedding"]
+        assert isinstance(emb, list)
+        assert len(emb) > 0
+
+        if expected_dim is not None:
+            assert len(emb) == expected_dim, f"Expected dim={expected_dim}, got {len(emb)}"
+
+        if check_unit_norm:
+            norm = float(np.linalg.norm(emb))
+            assert abs(norm - 1.0) < 1e-3, f"Expected unit norm, got {norm}"
+
+    # Usage validation — same pattern as test_server.py
+    usage = response["usage"]
+    assert isinstance(usage, dict)
+    assert usage["prompt_tokens"] > 0
+    assert usage["total_tokens"] == usage["prompt_tokens"]
+
+
+def expect_error(response_str: str, msg_prefix: Optional[str] = None):
+    """Validate error response — reused directly from test_server.py."""
+    response = json.loads(response_str)
+    assert response["object"] == "error"
+    assert isinstance(response["message"], str)
+    if msg_prefix is not None:
+        assert response["message"].startswith(msg_prefix)
+
+
+# ---------------------------------------------------------------------------
+# Server fixture — follows PopenServer/launch_server pattern from conftest.py
+# ---------------------------------------------------------------------------
+
+
+@pytest.fixture(scope="module")
+def launch_embedding_server():
+    """Launch an embedding-only server as a subprocess.
+
+    Follows the same lifecycle pattern as the launch_server fixture
+    in serve/server/conftest.py, but uses a lightweight embedding-only
+    server since PopenServer doesn't support --embedding-model yet.
+    """
+    _skip_if_no_model()
+
+    mlc_llm_path = str(Path(__file__).resolve().parents[4] / "python")
+    server_code = f"""
+import sys
+sys.path.insert(0, "{mlc_llm_path}")
+
+import fastapi
+import uvicorn
+from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine
+from mlc_llm.serve.server import ServerContext
+from mlc_llm.serve.entrypoints import openai_entrypoints
+
+app = fastapi.FastAPI()
+app.include_router(openai_entrypoints.app)
+
+engine = AsyncEmbeddingEngine(
+    model="{EMBEDDING_MODEL_DIR}",
+    model_lib="{EMBEDDING_MODEL_LIB}",
+    device="auto",
+)
+ctx = ServerContext()
+ServerContext.server_context = ctx
+ctx.add_embedding_engine("{EMBEDDING_MODEL_NAME}", engine)
+
+uvicorn.run(app, host="{EMBEDDING_SERVER_HOST}", port={EMBEDDING_SERVER_PORT}, log_level="info")
+"""
+    with subprocess.Popen(
+        [sys.executable, "-c", server_code],
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+    ) as proc:
+        # Wait for server readiness — same polling pattern as PopenServer.start()
+        timeout = 120
+        attempts = 0.0
+        ready = False
+        while attempts < timeout:
+            try:
+                response = requests.get(f"{EMBEDDING_BASE_URL}/models", timeout=2)
+                if response.status_code == 200:
+                    ready = True
+                    break
+            except requests.RequestException:
+                pass
+            attempts += 0.5
+            time.sleep(0.5)
+
+        if not ready:
+            stderr = proc.stderr.read().decode() if proc.stderr else ""
+            proc.kill()
+            raise RuntimeError(f"Embedding server failed to start in {timeout}s.\nStderr: {stderr}")
+
+        yield proc
+
+        # Cleanup — same pattern as PopenServer.terminate()
+        proc.send_signal(signal.SIGINT)
+        try:
+            proc.wait(timeout=10)
+        except subprocess.TimeoutExpired:
+            proc.kill()
+
+
+@pytest.fixture(scope="module")
+def client(launch_embedding_server):
+    """OpenAI client connected to the embedding server."""
+    assert launch_embedding_server is not None
+    return OpenAI(base_url=EMBEDDING_BASE_URL, api_key="none")
+
+
+# ===================================================================
+# /v1/models
+# ===================================================================
+
+
+@pytest.mark.usefixtures("client")
+def test_models_endpoint():
+    """The /v1/models endpoint lists the embedding model."""
+    resp = requests.get(f"{EMBEDDING_BASE_URL}/models", timeout=5)
+    assert resp.status_code == 200
+    data = resp.json()
+    assert isinstance(data["data"], list)
+
+
+# ===================================================================
+# Single input
+# ===================================================================
+
+
+def test_single_string_input(client):
+    """Single string input returns one embedding."""
+    resp = client.embeddings.create(input="What is machine learning?", model=EMBEDDING_MODEL_NAME)
+    raw = resp.model_dump()
+    check_embedding_response(raw, model=EMBEDDING_MODEL_NAME, num_embeddings=1)
+
+
+# ===================================================================
+# Batch input
+# ===================================================================
+
+BATCH_INPUTS = [
+    "What is machine learning?",
+    "How to brew coffee?",
+    "ML is a subset of AI.",
+]
+
+
+def test_batch_string_input(client):
+    """List of strings returns one embedding per input."""
+    resp = client.embeddings.create(input=BATCH_INPUTS, model=EMBEDDING_MODEL_NAME)
+    raw = resp.model_dump()
+    check_embedding_response(raw, model=EMBEDDING_MODEL_NAME, num_embeddings=len(BATCH_INPUTS))
+
+
+def test_batch_index_ordering(client):
+    """Embedding indices are sequential."""
+    resp = client.embeddings.create(input=BATCH_INPUTS, model=EMBEDDING_MODEL_NAME)
+    indices = [d.index for d in resp.data]
+    assert indices == list(range(len(BATCH_INPUTS)))
+
+
+# ===================================================================
+# Cosine similarity — semantic quality via endpoint
+# ===================================================================
+
+
+def test_cosine_similarity_via_endpoint(client):
+    """Related texts have higher similarity than unrelated (end-to-end)."""
+    resp = client.embeddings.create(
+        input=[
+            "What is machine learning?",
+            "Explain deep learning",
+            "Order a pizza",
+        ],
+        model=EMBEDDING_MODEL_NAME,
+    )
+    e0, e1, e2 = [np.array(d.embedding) for d in resp.data]
+    sim_related = float(np.dot(e0, e1))
+    sim_unrelated = float(np.dot(e0, e2))
+    assert sim_related > sim_unrelated, (
+        f"Related ({sim_related:.4f}) should > unrelated ({sim_unrelated:.4f})"
+    )
+
+
+# ===================================================================
+# Dimension truncation (Matryoshka)
+# ===================================================================
+
+
+def test_dimension_truncation(client):
+    """dimensions parameter truncates and re-normalizes output."""
+    target_dim = 256
+    resp = client.embeddings.create(
+        input="Hello world", model=EMBEDDING_MODEL_NAME, dimensions=target_dim
+    )
+    raw = resp.model_dump()
+    check_embedding_response(
+        raw,
+        model=EMBEDDING_MODEL_NAME,
+        num_embeddings=1,
+        expected_dim=target_dim,
+    )
+
+
+# ===================================================================
+# Encoding format
+# ===================================================================
+
+
+@pytest.mark.usefixtures("launch_embedding_server")
+def test_base64_encoding():
+    """base64 encoding format returns base64-encoded embeddings."""
+    resp = requests.post(
+        f"{EMBEDDING_BASE_URL}/embeddings",
+        json={
+            "input": "Hello world",
+            "model": EMBEDDING_MODEL_NAME,
+            "encoding_format": "base64",
+        },
+        timeout=5,
+    )
+    assert resp.status_code == 200
+    data = resp.json()
+    assert data["data"][0]["object"] == "embedding"
+    # base64 string should be a non-empty string (not a list)
+    emb = data["data"][0]["embedding"]
+    assert isinstance(emb, str) and len(emb) > 0
+
+
+# ===================================================================
+# Error handling — reuses expect_error() pattern from test_server.py
+# ===================================================================
+
+
+@pytest.mark.usefixtures("launch_embedding_server")
+def test_any_model_name_works_with_single_engine():
+    """When only one embedding engine is served, any model name works.
+
+    This mirrors ServerContext.get_engine() behavior: a single served
+    model is returned regardless of the requested model name.
+    """
+    resp = requests.post(
+        f"{EMBEDDING_BASE_URL}/embeddings",
+        json={"input": "test", "model": "any-name-works"},
+        timeout=5,
+    )
+    assert resp.status_code == 200
+    data = resp.json()
+    assert len(data["data"]) == 1
+
+
+# ===================================================================
+# Standalone runner (same pattern as test_server.py __main__)
+# ===================================================================
+
+if __name__ == "__main__":
+    _skip_if_no_model()
+
+    print(f"Using model: {EMBEDDING_MODEL_DIR}")
+    print(f"Using model lib: {EMBEDDING_MODEL_LIB}")
+    print(f"Server URL: {EMBEDDING_BASE_URL}")
+    print(
+        "\nMake sure the embedding server is running, or set env vars "
+        "and use pytest to auto-launch."
+    )
+
+    # Allow running against an already-running server
+    c = OpenAI(base_url=EMBEDDING_BASE_URL, api_key="none")
+    test_models_endpoint()
+    test_single_string_input(c)
+    test_batch_string_input(c)
+    test_batch_index_ordering(c)
+    test_cosine_similarity_via_endpoint(c)
+    test_dimension_truncation(c)
+    test_base64_encoding()
+    test_any_model_name_works_with_single_engine()
+    print("\nAll embedding server tests passed!")

tests/python/serve/server/test_server_function_call.py

@@ -0,0 +1,208 @@
+"""
+Test script for function call in chat completion. To run this script, use the following command:
+MLC_SERVE_MODEL_LIB=dist/gorilla-openfunctions-v1-q4f16_1_MLC/gorilla-openfunctions-v1-q4f16_1-cuda.so
+MLC_SERVE_MODEL_LIB=${MLC_SERVE_MODEL_LIB} python -m pytest -x tests/python/serve/server/test_server_function_call.py
+"""  # noqa: E501
+
+import json
+import os
+from typing import Dict, List, Optional, Tuple  # noqa: UP035
+
+import pytest
+import requests
+
+OPENAI_V1_CHAT_COMPLETION_URL = "http://127.0.0.1:8000/v1/chat/completions"
+
+
+def check_openai_nonstream_response(
+    response: Dict,  # noqa: UP006
+    *,
+    model: str,
+    object_str: str,
+    num_choices: int,
+    finish_reason: List[str],  # noqa: UP006
+    completion_tokens: Optional[int] = None,
+):
+    print(response)
+    assert response["model"] == model
+    assert response["object"] == object_str
+
+    choices = response["choices"]
+    assert isinstance(choices, list)
+    assert len(choices) == num_choices
+    for idx, choice in enumerate(choices):
+        assert choice["index"] == idx
+        assert choice["finish_reason"] in finish_reason
+
+        # text: str
+        message = choice["message"]
+        assert message["role"] == "assistant"
+        if choice["finish_reason"] == "tool_calls":
+            assert message["content"] is None
+            assert isinstance(message["tool_calls"], list)
+        else:
+            assert message["tool_calls"] is None
+            assert message["content"] is not None
+
+    usage = response["usage"]
+    assert isinstance(usage, dict)
+    assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]
+    assert usage["prompt_tokens"] > 0
+
+    if completion_tokens is not None:
+        assert usage["completion_tokens"] == completion_tokens
+
+
+def check_openai_stream_response(
+    responses: List[Dict],  # noqa: UP006
+    *,
+    model: str,
+    object_str: str,
+    num_choices: int,
+    finish_reason: str,
+    echo_prompt: Optional[str] = None,
+    suffix: Optional[str] = None,
+    stop: Optional[List[str]] = None,  # noqa: UP006
+    require_substr: Optional[List[str]] = None,  # noqa: UP006
+):
+    assert len(responses) > 0
+
+    finished = [False for _ in range(num_choices)]
+    outputs = ["" for _ in range(num_choices)]
+    for response in responses:
+        assert response["model"] == model
+        assert response["object"] == object_str
+
+        choices = response["choices"]
+        assert isinstance(choices, list)
+        assert len(choices) == num_choices
+        for idx, choice in enumerate(choices):
+            assert choice["index"] == idx
+
+            delta = choice["delta"]
+            assert delta["role"] == "assistant"
+            assert isinstance(delta["content"], str)
+            outputs[idx] += delta["content"]
+
+            if finished[idx]:
+                assert choice["finish_reason"] == finish_reason
+            elif choice["finish_reason"] is not None:
+                assert choice["finish_reason"] == finish_reason
+                finished[idx] = True
+
+    for output in outputs:
+        if echo_prompt is not None:
+            assert output.startswith(echo_prompt)
+        if suffix is not None:
+            assert output.endswith(suffix)
+        if stop is not None:
+            for stop_str in stop:
+                assert stop_str not in output
+        if require_substr is not None:
+            for substr in require_substr:
+                assert substr in output
+
+
+tools = [
+    {
+        "type": "function",
+        "function": {
+            "name": "get_current_weather",
+            "description": "Get the current weather in a given location",
+            "parameters": {
+                "type": "object",
+                "properties": {
+                    "location": {
+                        "type": "string",
+                        "description": "The city and state, e.g. San Francisco, CA",
+                    },
+                    "unit": {"type": "string", "enum": ["celsius", "fahrenheit"]},
+                },
+                "required": ["location"],
+            },
+        },
+    }
+]
+
+
+CHAT_COMPLETION_MESSAGES = [
+    # messages #0
+    [
+        {
+            "role": "user",
+            "content": "What is the current weather in Pittsburgh, PA?",
+        }
+    ],
+    # messages #1
+    [
+        {
+            "role": "user",
+            "content": "What is the current weather in Pittsburgh, PA and Tokyo, JP?",
+        }
+    ],
+    # messages #2
+    [
+        {
+            "role": "user",
+            "content": "What is the current weather in Pittsburgh, PA in fahrenheit?",
+        }
+    ],
+]
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("messages", CHAT_COMPLETION_MESSAGES)
+def test_openai_v1_chat_completion_function_call(
+    served_model: Tuple[str, str],  # noqa: UP006
+    launch_server,
+    stream: bool,
+    messages: List[Dict[str, str]],  # noqa: UP006
+):
+    # `served_model` and `launch_server` are pytest fixtures
+    # defined in conftest.py.
+
+    payload = {
+        "model": served_model[0],
+        "messages": messages,
+        "stream": stream,
+        "tools": tools,
+    }
+
+    response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=60)
+    if not stream:
+        check_openai_nonstream_response(
+            response.json(),
+            model=served_model[0],
+            object_str="chat.completion",
+            num_choices=1,
+            finish_reason=["tool_calls", "error"],
+        )
+    else:
+        responses = []
+        for chunk in response.iter_lines(chunk_size=512):
+            if not chunk or chunk == b"data: [DONE]":
+                continue
+            responses.append(json.loads(chunk.decode("utf-8")[6:]))
+        check_openai_stream_response(
+            responses,
+            model=served_model[0],
+            object_str="chat.completion.chunk",
+            num_choices=1,
+            finish_reason="tool_calls",
+        )
+
+
+if __name__ == "__main__":
+    model_lib = os.environ.get("MLC_SERVE_MODEL_LIB")
+    if model_lib is None:
+        raise ValueError(
+            'Environment variable "MLC_SERVE_MODEL_LIB" not found. '
+            "Please set it to model lib compiled by MLC LLM "
+            "(e.g., `./dist/gorilla-openfunctions-v1-q4f16_1_MLC/gorilla-openfunctions-v1-q4f16_1-cuda.so`) "  # noqa: E501
+            "which supports function calls."
+        )
+    MODEL = (os.path.dirname(model_lib), model_lib)
+
+    for msg in CHAT_COMPLETION_MESSAGES:
+        test_openai_v1_chat_completion_function_call(MODEL, None, stream=False, messages=msg)
+        test_openai_v1_chat_completion_function_call(MODEL, None, stream=True, messages=msg)

tests/python/serve/server/test_server_image.py

@@ -0,0 +1,257 @@
+import json
+import os
+from typing import Dict, List, Optional, Tuple  # noqa: UP035
+
+import pytest
+import regex
+import requests
+
+OPENAI_V1_CHAT_COMPLETION_URL = "http://127.0.0.1:8001/v1/chat/completions"
+
+JSON_TOKEN_PATTERN = (
+    r"((-?(?:0|[1-9]\d*))(\.\d+)?([eE][-+]?\d+)?)|null|true|false|"
+    r'("((\\["\\\/bfnrt])|(\\u[0-9a-fA-F]{4})|[^"\\\x00-\x1f])*")'
+)
+JSON_TOKEN_RE = regex.compile(JSON_TOKEN_PATTERN)
+
+
+def is_json_or_json_prefix(s: str) -> bool:
+    try:
+        json.loads(s)
+        return True
+    except json.JSONDecodeError as e:
+        # If the JSON decoder reaches the end of s, it is a prefix of a JSON string.
+        if e.pos == len(s):
+            return True
+        # Since json.loads is token-based instead of char-based, there may remain half a token after
+        # the matching position.
+        # If the left part is a prefix of a valid JSON token, the output is also valid
+        regex_match = JSON_TOKEN_RE.fullmatch(s[e.pos :], partial=True)
+        return regex_match is not None
+
+
+def check_openai_nonstream_response(
+    response: Dict,  # noqa: UP006
+    *,
+    is_chat_completion: bool,
+    model: str,
+    object_str: str,
+    num_choices: int,
+    finish_reasons: List[str],  # noqa: UP006
+    completion_tokens: Optional[int] = None,
+    echo_prompt: Optional[str] = None,
+    suffix: Optional[str] = None,
+    stop: Optional[List[str]] = None,  # noqa: UP006
+    require_substr: Optional[List[str]] = None,  # noqa: UP006
+    json_mode: bool = False,
+):
+    assert response["model"] == model
+    assert response["object"] == object_str
+
+    choices = response["choices"]
+    assert isinstance(choices, list)
+    assert len(choices) <= num_choices
+    texts: List[str] = ["" for _ in range(num_choices)]  # noqa: UP006
+    for choice in choices:
+        idx = choice["index"]
+        assert choice["finish_reason"] in finish_reasons
+
+        if not is_chat_completion:
+            assert isinstance(choice["text"], str)
+            texts[idx] = choice["text"]
+            if echo_prompt is not None:
+                assert texts[idx]
+            if suffix is not None:
+                assert texts[idx]
+        else:
+            message = choice["message"]
+            assert message["role"] == "assistant"
+            assert isinstance(message["content"], str)
+            texts[idx] = message["content"]
+
+        if stop is not None:
+            for stop_str in stop:
+                assert stop_str not in texts[idx]
+        if require_substr is not None:
+            for substr in require_substr:
+                assert substr in texts[idx]
+        if json_mode:
+            assert is_json_or_json_prefix(texts[idx])
+
+    usage = response["usage"]
+    assert isinstance(usage, dict)
+    assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]
+    assert usage["prompt_tokens"] > 0
+    if completion_tokens is not None:
+        assert usage["completion_tokens"] == completion_tokens
+
+
+def check_openai_stream_response(
+    responses: List[Dict],  # noqa: UP006
+    *,
+    is_chat_completion: bool,
+    model: str,
+    object_str: str,
+    num_choices: int,
+    finish_reasons: List[str],  # noqa: UP006
+    completion_tokens: Optional[int] = None,
+    echo_prompt: Optional[str] = None,
+    suffix: Optional[str] = None,
+    stop: Optional[List[str]] = None,  # noqa: UP006
+    require_substr: Optional[List[str]] = None,  # noqa: UP006
+    json_mode: bool = False,
+):
+    assert len(responses) > 0
+
+    finished = [False for _ in range(num_choices)]
+    outputs = ["" for _ in range(num_choices)]
+    for response in responses:
+        assert response["model"] == model
+        assert response["object"] == object_str
+
+        choices = response["choices"]
+        assert isinstance(choices, list)
+        assert len(choices) <= num_choices
+        for choice in choices:
+            idx = choice["index"]
+
+            if not is_chat_completion:
+                assert isinstance(choice["text"], str)
+                outputs[idx] += choice["text"]
+            else:
+                delta = choice["delta"]
+                assert delta["role"] == "assistant"
+                assert isinstance(delta["content"], str)
+                outputs[idx] += delta["content"]
+
+            if finished[idx]:
+                assert choice["finish_reason"] in finish_reasons
+            elif choice["finish_reason"] is not None:
+                assert choice["finish_reason"] in finish_reasons
+                finished[idx] = True
+
+        if not is_chat_completion:
+            usage = response["usage"]
+            assert isinstance(usage, dict)
+            assert usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"]
+            assert usage["prompt_tokens"] > 0
+            if completion_tokens is not None:
+                assert usage["completion_tokens"] <= completion_tokens
+
+    if not is_chat_completion:
+        if completion_tokens is not None:
+            assert responses[-1]["usage"]["completion_tokens"] == completion_tokens
+
+    for i, output in enumerate(outputs):
+        if echo_prompt is not None:
+            assert output.startswith(echo_prompt)
+        if suffix is not None:
+            assert output.endswith(suffix)
+        if stop is not None:
+            for stop_str in stop:
+                assert stop_str not in output
+        if require_substr is not None:
+            for substr in require_substr:
+                assert substr in output
+        if json_mode:
+            assert is_json_or_json_prefix(output)
+
+
+CHAT_COMPLETION_MESSAGES = [
+    # messages #0
+    [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": "https://llava-vl.github.io/static/images/view.jpg",
+                },
+                {"type": "text", "text": "What does this image represent?"},
+            ],
+        },
+    ],
+    # messages #1
+    [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": "https://llava-vl.github.io/static/images/view.jpg",
+                },
+                {"type": "text", "text": "What does this image represent?"},
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": "The image represents a serene and peaceful scene of a pier extending over a body of water, such as a lake or a river.er. The pier is made of wood and has a bench on it, providing a place for people to sit and enjoy the view. The pier is situated in a natural environment, surrounded by trees and mountains in the background. This setting creates a tranquil atmosphere, inviting visitors to relax and appreciate the beauty of the landscape.",  # noqa: E501
+        },
+        {
+            "role": "user",
+            "content": "What country is the image set in? Give me 10 ranked guesses and reasons why.",  # noqa: E501
+        },
+    ],
+]
+
+
+@pytest.mark.parametrize("stream", [False, True])
+@pytest.mark.parametrize("messages", CHAT_COMPLETION_MESSAGES)
+def test_openai_v1_chat_completions(
+    served_model: Tuple[str, str],  # noqa: UP006
+    launch_server,
+    stream: bool,
+    messages: List[Dict[str, str]],  # noqa: UP006
+):
+    # `served_model` and `launch_server` are pytest fixtures
+    # defined in conftest.py.
+
+    payload = {
+        "model": served_model[0],
+        "messages": messages,
+        "stream": stream,
+    }
+    response = requests.post(OPENAI_V1_CHAT_COMPLETION_URL, json=payload, timeout=180)
+    if not stream:
+        check_openai_nonstream_response(
+            response.json(),
+            is_chat_completion=True,
+            model=served_model[0],
+            object_str="chat.completion",
+            num_choices=1,
+            finish_reasons=["stop"],
+        )
+    else:
+        responses = []
+        for chunk in response.iter_lines(chunk_size=512):
+            if not chunk or chunk == b"data: [DONE]":
+                continue
+            responses.append(json.loads(chunk.decode("utf-8")[6:]))
+        check_openai_stream_response(
+            responses,
+            is_chat_completion=True,
+            model=served_model[0],
+            object_str="chat.completion.chunk",
+            num_choices=1,
+            finish_reasons=["stop"],
+        )
+
+
+if __name__ == "__main__":
+    model_lib = os.environ.get("MLC_SERVE_MODEL_LIB")
+    if model_lib is None:
+        raise ValueError(
+            'Environment variable "MLC_SERVE_MODEL_LIB" not found. '
+            "Please set it to model lib compiled by MLC LLM "
+            "(e.g., `dist/Llama-2-7b-chat-hf-q0f16-MLC/Llama-2-7b-chat-hf-q0f16-MLC-cuda.so`)."
+        )
+
+    model = os.environ.get("MLC_SERVE_MODEL")
+    if model is None:
+        MODEL = (os.path.dirname(model_lib), model_lib)
+    else:
+        MODEL = (model, model_lib)
+
+    for msg in CHAT_COMPLETION_MESSAGES:
+        test_openai_v1_chat_completions(MODEL, None, stream=False, messages=msg)
+        test_openai_v1_chat_completions(MODEL, None, stream=True, messages=msg)

tests/python/serve/test_embedding_engine.py

@@ -0,0 +1,365 @@
+"""Embedding engine tests in MLC LLM.
+
+Tests AsyncEmbeddingEngine for both direct (sync) and async embedding inference.
+Reuses MLC LLM test infrastructure: markers, require_test_model pattern,
+and conventions from test_serve_engine.py.
+
+Run with real model (requires GPU + compiled embedding model):
+  MLC_SERVE_EMBEDDING_MODEL_LIB="path/to/model.dylib" \
+    pytest -m engine tests/python/serve/test_embedding_engine.py -v
+
+Environment variables:
+  MLC_SERVE_EMBEDDING_MODEL_LIB  Path to compiled embedding model library (required)
+  MLC_SERVE_EMBEDDING_MODEL      Path to embedding model weight directory
+                                  (optional, defaults to dirname of model lib)
+"""
+
+import asyncio
+import os
+
+import numpy as np
+import pytest
+
+# Reuse MLC LLM marker system (registered in tests/python/conftest.py)
+pytestmark = [pytest.mark.engine]
+
+# ---------------------------------------------------------------------------
+# Fixtures — follows pattern from serve/server/conftest.py (served_model)
+# ---------------------------------------------------------------------------
+
+EMBEDDING_MODEL_LIB = os.environ.get("MLC_SERVE_EMBEDDING_MODEL_LIB")
+EMBEDDING_MODEL_DIR = os.environ.get(
+    "MLC_SERVE_EMBEDDING_MODEL",
+    os.path.dirname(EMBEDDING_MODEL_LIB) if EMBEDDING_MODEL_LIB else None,
+)
+
+
+def _skip_if_no_model():
+    if EMBEDDING_MODEL_LIB is None:
+        pytest.skip(
+            'Environment variable "MLC_SERVE_EMBEDDING_MODEL_LIB" not found. '
+            "Set it to a compiled embedding model library "
+            "(e.g., Qwen3-Embedding-0.6B-q0f32-MLC.dylib)."
+        )
+    if not os.path.isfile(EMBEDDING_MODEL_LIB):
+        pytest.skip(f"Embedding model library not found at: {EMBEDDING_MODEL_LIB}")
+    if EMBEDDING_MODEL_DIR is None or not os.path.isdir(EMBEDDING_MODEL_DIR):
+        pytest.skip(f"Embedding model directory not found at: {EMBEDDING_MODEL_DIR}")
+
+
+@pytest.fixture(scope="module")
+def embedding_engine():
+    """Module-scoped AsyncEmbeddingEngine — loaded once, shared across tests."""
+    _skip_if_no_model()
+    from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine
+
+    engine = AsyncEmbeddingEngine(
+        model=EMBEDDING_MODEL_DIR,
+        model_lib=EMBEDDING_MODEL_LIB,
+        device="auto",
+    )
+    yield engine
+    engine.terminate()
+
+
+# ---------------------------------------------------------------------------
+# Helpers — reuse cosine_similarity pattern from test_serve_engine.py
+# ---------------------------------------------------------------------------
+
+
+def cosine_similarity(a, b):
+    """Return cosine similarity between two vectors."""
+    a, b = np.array(a), np.array(b)
+    return float(np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b)))
+
+
+# ===================================================================
+# Engine initialization tests
+# ===================================================================
+
+
+def test_engine_model_type(embedding_engine):
+    """Engine reports a valid model type."""
+    assert embedding_engine.model_type in ("encoder", "decoder")
+
+
+def test_engine_pooling_strategy(embedding_engine):
+    """Engine selects appropriate default pooling strategy."""
+    if embedding_engine.model_type == "encoder":
+        assert embedding_engine.pooling_strategy == "cls"
+    else:
+        assert embedding_engine.pooling_strategy == "last"
+
+
+# ===================================================================
+# Single-text embedding
+# ===================================================================
+
+
+def test_single_text_shape(embedding_engine):
+    """Single text returns exactly one embedding vector."""
+    embeddings, tokens = embedding_engine.embed(["Hello world"])
+    assert len(embeddings) == 1
+    assert len(embeddings[0]) > 0
+    assert tokens > 0
+
+
+def test_single_text_unit_norm(embedding_engine):
+    """Embedding output is L2-normalized."""
+    embeddings, _ = embedding_engine.embed(["Hello world"])
+    norm = float(np.linalg.norm(embeddings[0]))
+    assert abs(norm - 1.0) < 1e-4, f"Expected unit norm, got {norm}"
+
+
+# ===================================================================
+# Batch embedding
+# ===================================================================
+
+BATCH_TEXTS = [
+    "Machine learning is fascinating",
+    "I love pizza",
+    "Deep learning uses neural networks",
+]
+
+
+def test_batch_count(embedding_engine):
+    """Batch embedding returns one vector per input."""
+    embeddings, tokens = embedding_engine.embed(BATCH_TEXTS)
+    assert len(embeddings) == len(BATCH_TEXTS)
+    assert tokens > 0
+
+
+def test_batch_all_normalized(embedding_engine):
+    """Every vector in a batch is L2-normalized."""
+    embeddings, _ = embedding_engine.embed(BATCH_TEXTS)
+    for i, emb in enumerate(embeddings):
+        norm = float(np.linalg.norm(emb))
+        assert abs(norm - 1.0) < 1e-4, f"Embedding [{i}] norm={norm}"
+
+
+def test_batch_consistent_dimension(embedding_engine):
+    """All embeddings in a batch have the same dimension."""
+    embeddings, _ = embedding_engine.embed(BATCH_TEXTS)
+    dims = {len(emb) for emb in embeddings}
+    assert len(dims) == 1, f"Inconsistent dimensions: {dims}"
+
+
+# ===================================================================
+# Semantic quality — cosine similarity ranking
+# ===================================================================
+
+SIMILARITY_TEXTS = [
+    "What is machine learning?",
+    "Explain deep learning algorithms",
+    "I want to order pizza",
+]
+
+
+def test_cosine_similarity_ranking(embedding_engine):
+    """Related texts have higher cosine similarity than unrelated texts."""
+    embeddings, _ = embedding_engine.embed(SIMILARITY_TEXTS)
+    e_ml, e_dl, e_pizza = [np.array(e) for e in embeddings]
+    sim_related = float(np.dot(e_ml, e_dl))
+    sim_unrelated = float(np.dot(e_ml, e_pizza))
+    assert sim_related > sim_unrelated, (
+        f"Related sim ({sim_related:.4f}) should > unrelated sim ({sim_unrelated:.4f})"
+    )
+
+
+# ===================================================================
+# Determinism
+# ===================================================================
+
+
+def test_deterministic_output(embedding_engine):
+    """Same input produces identical output across calls."""
+    text = ["Deterministic test"]
+    emb1, _ = embedding_engine.embed(text)
+    emb2, _ = embedding_engine.embed(text)
+    cos = cosine_similarity(emb1[0], emb2[0])
+    assert cos > 0.9999, f"Expected deterministic output, cosine={cos}"
+
+
+# ===================================================================
+# Async embedding
+# ===================================================================
+
+
+def test_async_embed(embedding_engine):
+    """async_embed produces same result as sync embed."""
+    text = ["Async test"]
+    sync_emb, sync_tokens = embedding_engine.embed(text)
+
+    loop = asyncio.new_event_loop()
+    try:
+        async_emb, async_tokens = loop.run_until_complete(embedding_engine.async_embed(text))
+    finally:
+        loop.close()
+
+    assert sync_tokens == async_tokens
+    cos = cosine_similarity(sync_emb[0], async_emb[0])
+    assert cos > 0.9999, f"Async vs sync mismatch, cosine={cos}"
+
+
+# ===================================================================
+# Edge cases
+# ===================================================================
+
+
+def test_empty_string(embedding_engine):
+    """Empty string should still produce a valid embedding for supported models."""
+    embeddings, tokens = embedding_engine.embed([""])
+    if embedding_engine.model_type == "encoder":
+        assert len(embeddings) == 1
+        assert len(embeddings[0]) > 0
+        assert tokens > 0
+    else:
+        assert len(embeddings) == 1
+        assert len(embeddings[0]) > 0
+        assert tokens > 0
+
+
+# ===================================================================
+# Long text handling (model-type dependent)
+# ===================================================================
+
+
+def test_long_text_decoder_chunked_prefill(embedding_engine):
+    """[Decoder only] Text >prefill_chunk_size triggers chunked prefill.
+    ~5000 tokens processed in 3 chunks. Result is unit-norm embedding."""
+    if embedding_engine.model_type != "decoder":
+        pytest.skip("Chunked prefill is decoder-only")
+    long_text = "word " * 5000
+    embeddings, tokens = embedding_engine.embed([long_text])
+    assert tokens > 2048, f"Expected >2048 tokens to trigger chunking, got {tokens}"
+    norm = float(np.linalg.norm(embeddings[0]))
+    assert abs(norm - 1.0) < 1e-3
+
+
+def _get_encoder_tokens(embedding_engine, text):
+    """Replicate encoder preprocessing: tokenize and add [CLS]/[SEP]."""
+    tokens = list(embedding_engine.tokenizer.encode(text))
+    if embedding_engine._cls_token_id is not None and (
+        len(tokens) == 0 or tokens[0] != embedding_engine._cls_token_id
+    ):
+        tokens = [embedding_engine._cls_token_id, *tokens]
+    if embedding_engine._sep_token_id is not None and (
+        len(tokens) == 0 or tokens[-1] != embedding_engine._sep_token_id
+    ):
+        tokens = [*tokens, embedding_engine._sep_token_id]
+    return tokens
+
+
+def test_long_text_encoder_truncation(embedding_engine):
+    """[Encoder only] Text exceeding prefill_chunk_size is truncated.
+    Two texts with the same shared prefix but different suffixes beyond the
+    limit should produce identical embeddings, since the suffix is truncated
+    and the retained token prefixes are verified to be identical."""
+    if embedding_engine.model_type != "encoder":
+        pytest.skip("Truncation test is encoder-only")
+    prefill_chunk = embedding_engine._metadata.get("prefill_chunk_size", 512)
+
+    # Dynamically construct input that exceeds prefill_chunk_size.
+    unit = "machine learning is great "
+    suffix_a = " alpha beta gamma " * 200
+    suffix_b = " totally different ending " * 200
+    unit_tokens = len(list(embedding_engine.tokenizer.encode(unit)))
+    repeats = max(1, prefill_chunk // max(unit_tokens, 1) + 64)
+
+    # Increase prefix length until both inputs exceed prefill_chunk_size
+    # and their truncated token prefixes are identical.
+    while True:
+        shared_prefix = unit * repeats
+        full_tokens_a = _get_encoder_tokens(embedding_engine, shared_prefix + suffix_a)
+        full_tokens_b = _get_encoder_tokens(embedding_engine, shared_prefix + suffix_b)
+        if (
+            len(full_tokens_a) > prefill_chunk
+            and len(full_tokens_b) > prefill_chunk
+            and full_tokens_a[:prefill_chunk] == full_tokens_b[:prefill_chunk]
+        ):
+            break
+        repeats += 64
+        assert repeats < 200000, "Failed to construct truncation test inputs"
+
+    text_a = shared_prefix + suffix_a
+    text_b = shared_prefix + suffix_b
+
+    emb_a, tokens_a = embedding_engine.embed([text_a])
+    emb_b, tokens_b = embedding_engine.embed([text_b])
+
+    # Verify truncation happened
+    assert tokens_a <= prefill_chunk, (
+        f"Encoder should truncate to {prefill_chunk}, got {tokens_a} tokens"
+    )
+    assert tokens_b <= prefill_chunk
+    # Both should be valid unit-norm embeddings
+    assert abs(float(np.linalg.norm(emb_a[0])) - 1.0) < 1e-3
+    assert abs(float(np.linalg.norm(emb_b[0])) - 1.0) < 1e-3
+
+    # Both truncated to identical token sequences → embeddings must match
+    cos = cosine_similarity(emb_a[0], emb_b[0])
+    assert cos > 0.999, f"Same truncated tokens should match, cosine={cos:.6f}"
+
+
+def test_long_vs_short_semantic_quality(embedding_engine):
+    """Long text should still capture semantic meaning correctly.
+    Decoder: chunked prefill preserves full context.
+    Encoder: truncation keeps most relevant prefix."""
+    short_ml = "Machine learning enables systems to learn from data"
+    long_ml = (
+        "Machine learning is a fascinating field of study. " * 200
+        + "It enables systems to learn from data."
+    )
+    pizza = "I want to order a pepperoni pizza for dinner"
+
+    embs, _ = embedding_engine.embed([short_ml, long_ml, pizza])
+    e_short, e_long, e_pizza = [np.array(e) for e in embs]
+
+    sim_same_topic = float(np.dot(e_short, e_long))
+    sim_different = float(np.dot(e_short, e_pizza))
+    assert sim_same_topic > sim_different, (
+        f"Same topic ({sim_same_topic:.4f}) should > different ({sim_different:.4f})"
+    )
+
+
+def test_unicode_text(embedding_engine):
+    """Unicode input is handled correctly."""
+    texts = ["Привет мир", "你好世界", "こんにちは世界"]
+    embeddings, _ = embedding_engine.embed(texts)
+    assert len(embeddings) == 3
+    for emb in embeddings:
+        assert abs(float(np.linalg.norm(emb)) - 1.0) < 1e-4
+
+
+# ===================================================================
+# Standalone runner (like test_serve_engine.py)
+# ===================================================================
+
+if __name__ == "__main__":
+    _skip_if_no_model()
+    from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine
+
+    engine = AsyncEmbeddingEngine(
+        model=EMBEDDING_MODEL_DIR,
+        model_lib=EMBEDDING_MODEL_LIB,
+        device="auto",
+    )
+    try:
+        test_engine_model_type(engine)
+        test_engine_pooling_strategy(engine)
+        test_single_text_shape(engine)
+        test_single_text_unit_norm(engine)
+        test_batch_count(engine)
+        test_batch_all_normalized(engine)
+        test_batch_consistent_dimension(engine)
+        test_cosine_similarity_ranking(engine)
+        test_deterministic_output(engine)
+        test_async_embed(engine)
+        test_empty_string(engine)
+        test_long_text_decoder_chunked_prefill(engine)
+        test_long_text_encoder_truncation(engine)
+        test_long_vs_short_semantic_quality(engine)
+        test_unicode_text(engine)
+        print("\nAll embedding engine tests passed!")
+    finally:
+        engine.terminate()

tests/python/serve/test_event_trace_recorder.py

@@ -0,0 +1,48 @@
+import json
+
+import pytest
+
+from mlc_llm.serve.event_trace_recorder import EventTraceRecorder
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+def test_event_trace_recorder():
+    trace_recorder = EventTraceRecorder()
+    request_ids = ["x", "y"]
+    num_decode = 5
+
+    for request_id in request_ids:
+        trace_recorder.add_event(request_id, event="start tokenization")
+        trace_recorder.add_event(request_id, event="finish tokenization")
+        trace_recorder.add_event(request_id, event="add request")
+        trace_recorder.add_event(request_id, event="start embed")
+        trace_recorder.add_event(request_id, event="finish embed")
+        trace_recorder.add_event(request_id, event="start prefill")
+        trace_recorder.add_event(request_id, event="finish prefill")
+
+    for _ in range(num_decode):
+        for request_id in request_ids:
+            trace_recorder.add_event(request_id, event="start decode")
+            trace_recorder.add_event(request_id, event="finish decode")
+    for request_id in request_ids:
+        trace_recorder.add_event(request_id, event="start detokenization")
+        trace_recorder.add_event(request_id, event="finish detokenization")
+
+    events = json.loads(trace_recorder.dump_json())
+    decode_count = {}
+    for event in events:
+        request_id = event["tid"]
+        if event["name"].startswith("decode"):
+            if request_id not in decode_count:
+                decode_count[request_id] = 1
+            else:
+                decode_count[request_id] += 1
+
+    for _, decode_cnt in decode_count.items():
+        assert decode_cnt == num_decode * 2, decode_cnt
+
+
+if __name__ == "__main__":
+    test_event_trace_recorder()

tests/python/serve/test_radix_tree.py

@@ -0,0 +1,138 @@
+import pytest
+
+from mlc_llm.serve import PagedRadixTree
+
+# category "runtime_module"
+pytestmark = [pytest.mark.unittest]
+
+
+def test_add():
+    prt = PagedRadixTree()
+    prt.add(0)
+    assert list(prt.get(0)) == []
+    prt.add(1)
+    assert list(prt.get(1)) == []
+
+
+def test_remove():
+    prt = PagedRadixTree()
+    capacity = prt.free_capacity()
+    prt.add(0)
+    prt.remove(0)
+    prt.add(0)
+    prt.extend(0, [1 for _ in range(200)])
+    prt.remove(0)
+    assert prt.free_capacity() == capacity
+
+    prt.add(1)
+    prt.extend(1, [1 for _ in range(200)])
+    capacity = prt.free_capacity()
+    prt.add(2)
+    prt.extend(2, [1 for _ in range(100)] + [2 for _ in range(100)])
+    prt.remove(2)
+    assert prt.free_capacity() == capacity
+
+    prt.add(3)
+    prt.extend(3, [1 for _ in range(200)])
+    prt.remove(3)
+    assert prt.free_capacity() == capacity
+
+    prt.add(4)
+    prt.add(5)
+    prt.add(6)
+    assert prt.free_capacity() == capacity
+    prt.remove(4)
+    assert prt.free_capacity() == capacity
+    prt.remove(5)
+    assert prt.free_capacity() == capacity
+    prt.remove(6)
+    assert prt.free_capacity() == capacity
+
+
+def test_extend():
+    prt = PagedRadixTree()
+    L = prt.free_capacity() // 64
+    H = L // 2
+    Q = L // 4
+    seq_id = 0
+    for start_pos in [0, H, L, L + H]:
+        for length in [Q, L - H, L, 2 * L - H, 2 * L]:
+            prt.add(seq_id)
+            if start_pos:
+                tokens_1 = [seq_id for _ in range(start_pos)]
+                prt.extend(seq_id, tokens_1)
+                assert list(prt.get(seq_id)) == tokens_1
+            else:
+                tokens_1 = []
+            tokens_2 = [seq_id for _ in range(length)]
+            prt.extend(seq_id, tokens_2)
+            assert list(prt.get(seq_id)) == tokens_1 + tokens_2
+            seq_id += 1
+
+
+def test_fork():
+    prt = PagedRadixTree()
+    L = prt.free_capacity() // 64
+    H = L // 2
+    Q = L // 4
+    seq_id = 0
+    length_list = [Q, H, L, L + Q, L + H, L * 2]
+    for p_idx in range(1, len(length_list)):
+        for c_idx in range(0, p_idx + 1):
+            prt.add(seq_id)
+            tokens = [seq_id for _ in range(length_list[p_idx])]
+            prt.extend(seq_id, tokens)
+            prt.fork(seq_id + 1, seq_id, length_list[c_idx])
+            assert list(prt.get(seq_id + 1)) == tokens[: length_list[c_idx]]
+            seq_id += 2
+
+
+def test_fork_2():
+    prt = PagedRadixTree()
+    prt.add(0)
+    prt.extend(0, [0, 1, 2, 3])
+    prt.fork(1, 0, 3)
+    prt.extend(1, [4])
+    prt.fork(2, 0, 3)
+    prt.extend(2, [5])
+    assert prt.match([0, 1, 2, 4]) == (4, (1,))
+    assert prt.match([0, 1, 2, 5]) == (4, (2,))
+
+
+def test_rollback():
+    prt = PagedRadixTree()
+    L = prt.free_capacity() // 64
+    H = L // 2
+    Q = L // 4
+    seq_id = 0
+    for start_pos in [H, L, L + H, 2 * L, 3 * L + H]:
+        for length in [Q, H, L + Q, 2 * L, 2 * L + Q]:
+            if length > start_pos:
+                continue
+            prt.add(seq_id)
+            tokens = [seq_id for _ in range(start_pos)]
+            prt.extend(seq_id, tokens)
+            prt.rollback(seq_id, length)
+            assert list(prt.get(seq_id)) == tokens[:-length]
+            seq_id += 1
+
+    for start_pos in [H, L, L + H, 2 * L, 3 * L + H]:
+        for length in [Q, H, L + Q, 2 * L, 2 * L + Q]:
+            if length > start_pos:
+                continue
+            prt.add(seq_id)
+            tokens = [seq_id for _ in range(start_pos)]
+            prt.extend(seq_id, tokens)
+            prt.fork(seq_id + 1, seq_id, start_pos)
+            prt.rollback(seq_id + 1, length)
+            assert list(prt.get(seq_id + 1)) == tokens[:-length]
+            seq_id += 2
+
+
+if __name__ == "__main__":
+    test_add()
+    test_remove()
+    test_extend()
+    test_fork()
+    test_fork_2()
+    test_rollback()

tests/python/serve/test_serve_async_engine.py

@@ -0,0 +1,285 @@
+import asyncio
+from typing import List  # noqa: UP035
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import AsyncMLCEngine, EngineConfig
+from mlc_llm.testing import require_test_model
+
+prompts = [
+    "What is the meaning of life?",
+    "Introduce the history of Pittsburgh to me. Please elaborate in detail.",
+    "Write a three-day Seattle travel plan. Please elaborate in detail.",
+    "What is Alaska famous of? Please elaborate in detail.",
+    "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.",  # noqa: E501
+    "What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
+    "Why is Vitamin D important to human beings? Please elaborate in detail.",
+    "Where is milk tea originated from? Please elaborate in detail.",
+    "Where is the southernmost place in United States? Please elaborate in detail.",
+    "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.",  # noqa: E501
+]
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+async def test_engine_generate(model: str):
+    # Create engine
+    async_engine = AsyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+
+    num_requests = 10
+    max_tokens = 256
+    generation_cfg = GenerationConfig(max_tokens=max_tokens, n=7)
+
+    output_texts: List[List[str]] = [  # noqa: UP006
+        ["" for _ in range(generation_cfg.n)] for _ in range(num_requests)
+    ]
+
+    async def generate_task(
+        async_engine: AsyncMLCEngine,
+        prompt: str,
+        generation_cfg: GenerationConfig,
+        request_id: str,
+    ):
+        print(f"generate task for request {request_id}")
+        rid = int(request_id)
+        async for delta_outputs in async_engine._generate(
+            prompt, generation_cfg, request_id=request_id
+        ):
+            if len(delta_outputs) == generation_cfg.n:
+                for i, delta_output in enumerate(delta_outputs):
+                    output_texts[rid][i] += delta_output.delta_text
+            else:
+                assert len(delta_outputs) == 1
+                assert len(delta_outputs[0].request_final_usage_json_str) != 0
+
+    tasks = [
+        asyncio.create_task(
+            generate_task(async_engine, prompts[i], generation_cfg, request_id=str(i))
+        )
+        for i in range(num_requests)
+    ]
+
+    await asyncio.gather(*tasks)
+
+    # Print output.
+    print("All finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    async_engine.terminate()
+    del async_engine
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+async def test_chat_completion(model: str):
+    # Create engine
+    async_engine = AsyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+
+    num_requests = 2
+    max_tokens = 32
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    async def generate_task(prompt: str, request_id: str):
+        print(f"generate chat completion task for request {request_id}")
+        rid = int(request_id)
+        async for response in await async_engine.chat.completions.create(  # noqa: F821
+            messages=[{"role": "user", "content": prompt}],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=request_id,
+            stream=True,
+        ):
+            for choice in response.choices:
+                assert choice.delta.role == "assistant"
+                assert isinstance(choice.delta.content, str)
+                output_texts[rid][choice.index] += choice.delta.content
+
+    tasks = [
+        asyncio.create_task(generate_task(prompts[i], request_id=str(i)))
+        for i in range(num_requests)
+    ]
+
+    await asyncio.gather(*tasks)
+
+    # Print output.
+    print("Chat completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    async_engine.terminate()
+    del async_engine
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+async def test_chat_completion_non_stream(model: str):
+    # Create engine
+    async_engine = AsyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+
+    num_requests = 2
+    max_tokens = 32
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    async def generate_task(prompt: str, request_id: str):
+        print(f"generate chat completion task for request {request_id}")
+        rid = int(request_id)
+        response = await async_engine.chat.completions.create(  # noqa: F821
+            messages=[{"role": "user", "content": prompt}],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=request_id,
+        )
+        for choice in response.choices:
+            assert choice.message.role == "assistant"
+            assert isinstance(choice.message.content, str)
+            output_texts[rid][choice.index] += choice.message.content
+
+    tasks = [
+        asyncio.create_task(generate_task(prompts[i], request_id=str(i)))
+        for i in range(num_requests)
+    ]
+
+    await asyncio.gather(*tasks)
+
+    # Print output.
+    print("Chat completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    async_engine.terminate()
+    del async_engine
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+async def test_completion(model: str):
+    # Create engine
+    async_engine = AsyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+
+    num_requests = 2
+    max_tokens = 128
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    async def generate_task(prompt: str, request_id: str):
+        print(f"generate completion task for request {request_id}")
+        rid = int(request_id)
+        async for response in await async_engine.completions.create(  # noqa: F821
+            prompt=prompt,
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=request_id,
+            stream=True,
+            extra_body={"debug_config": {"ignore_eos": True}},
+        ):
+            for choice in response.choices:
+                output_texts[rid][choice.index] += choice.text
+
+    tasks = [
+        asyncio.create_task(generate_task(prompts[i], request_id=str(i)))
+        for i in range(num_requests)
+    ]
+
+    await asyncio.gather(*tasks)
+
+    # Print output.
+    print("Completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    async_engine.terminate()
+    del async_engine
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+async def test_completion_non_stream(model: str):
+    # Create engine
+    async_engine = AsyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+
+    num_requests = 2
+    max_tokens = 128
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    async def generate_task(prompt: str, request_id: str):
+        print(f"generate completion task for request {request_id}")
+        rid = int(request_id)
+        response = await async_engine.completions.create(  # noqa: F821
+            prompt=prompt,
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=request_id,
+            extra_body={"debug_config": {"ignore_eos": True}},
+        )
+        for choice in response.choices:
+            output_texts[rid][choice.index] += choice.text
+
+    tasks = [
+        asyncio.create_task(generate_task(prompts[i], request_id=str(i)))
+        for i in range(num_requests)
+    ]
+
+    await asyncio.gather(*tasks)
+
+    # Print output.
+    print("Completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    async_engine.terminate()
+    del async_engine
+
+
+if __name__ == "__main__":
+    asyncio.run(test_engine_generate())
+    asyncio.run(test_chat_completion())
+    asyncio.run(test_chat_completion_non_stream())
+    asyncio.run(test_completion())
+    asyncio.run(test_completion_non_stream())

tests/python/serve/test_serve_async_engine_spec.py

@@ -0,0 +1,84 @@
+import asyncio
+from typing import List  # noqa: UP035
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import AsyncMLCEngine, EngineConfig
+from mlc_llm.testing import require_test_model
+
+prompts = [
+    "What is the meaning of life?",
+    "Introduce the history of Pittsburgh to me. Please elaborate in detail.",
+    "Write a three-day Seattle travel plan. Please elaborate in detail.",
+    "What is Alaska famous of? Please elaborate in detail.",
+    "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.",  # noqa: E501
+    "What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
+    "Why is Vitamin D important to human beings? Please elaborate in detail.",
+    "Where is milk tea originated from? Please elaborate in detail.",
+    "Where is the southernmost place in United States? Please elaborate in detail.",
+    "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.",  # noqa: E501
+]
+
+
+@require_test_model(
+    "Llama-2-7b-chat-hf-q0f16-MLC",
+    "Llama-2-7b-chat-hf-q4f16_1-MLC",
+)
+async def test_engine_generate(model: str, small_model: str):
+    # Create engine
+    async_engine = AsyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            additional_models=[small_model],
+            speculative_mode="small_draft",
+        ),
+    )
+
+    num_requests = 10
+    max_tokens = 256
+    generation_cfg = GenerationConfig(max_tokens=max_tokens)
+
+    output_texts: List[List[str]] = [  # noqa: UP006
+        ["" for _ in range(generation_cfg.n)] for _ in range(num_requests)
+    ]
+
+    async def generate_task(
+        async_engine: AsyncMLCEngine,
+        prompt: str,
+        generation_cfg: GenerationConfig,
+        request_id: str,
+    ):
+        print(f"generate task for request {request_id}")
+        rid = int(request_id)
+        async for delta_outputs in async_engine._generate(
+            prompt, generation_cfg, request_id=request_id
+        ):
+            assert len(delta_outputs) == generation_cfg.n
+            for i, delta_output in enumerate(delta_outputs):
+                output_texts[rid][i] += delta_output.delta_text
+
+    tasks = [
+        asyncio.create_task(
+            generate_task(async_engine, prompts[i], generation_cfg, request_id=str(i))
+        )
+        for i in range(num_requests)
+    ]
+
+    await asyncio.gather(*tasks)
+
+    # Print output.
+    print("All finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    async_engine.terminate()
+    del async_engine
+
+
+if __name__ == "__main__":
+    asyncio.run(test_engine_generate())

tests/python/serve/test_serve_engine.py

@@ -0,0 +1,241 @@
+from typing import List  # noqa: UP035
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import EngineConfig, MLCEngine
+from mlc_llm.testing import require_test_model
+
+prompts = [
+    "What is the meaning of life?",
+    "Introduce the history of Pittsburgh to me. Please elaborate in detail.",
+    "Write a three-day Seattle travel plan. Please elaborate in detail.",
+    "What is Alaska famous of? Please elaborate in detail.",
+    "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.",  # noqa: E501
+    "What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
+    "Why is Vitamin D important to human beings? Please elaborate in detail.",
+    "Where is milk tea originated from? Please elaborate in detail.",
+    "Where is the southernmost place in United States? Please elaborate in detail.",
+    "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.",  # noqa: E501
+]
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_generate(model: str):
+    # Create engine
+    engine = MLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+        ),
+    )
+
+    num_requests = 10
+    max_tokens = 256
+    generation_cfg = GenerationConfig(max_tokens=max_tokens, n=7)
+
+    output_texts: List[List[str]] = [  # noqa: UP006
+        ["" for _ in range(generation_cfg.n)] for _ in range(num_requests)
+    ]
+    for rid in range(num_requests):
+        print(f"generating for request {rid}")
+        for delta_outputs in engine._generate(prompts[rid], generation_cfg, request_id=str(rid)):
+            assert len(delta_outputs) == generation_cfg.n
+            for i, delta_output in enumerate(delta_outputs):
+                output_texts[rid][i] += delta_output.delta_text
+
+    # Print output.
+    print("All finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    engine.terminate()
+    del engine
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_chat_completion(model: str):
+    # Create engine
+    engine = MLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+        ),
+    )
+
+    num_requests = 2
+    max_tokens = 64
+    n = 2
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    for rid in range(num_requests):
+        print(f"chat completion for request {rid}")
+        for response in engine.chat.completions.create(
+            messages=[{"role": "user", "content": prompts[rid]}],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            stream=True,
+        ):
+            for choice in response.choices:
+                assert choice.delta.role == "assistant"
+                assert isinstance(choice.delta.content, str)
+                output_texts[rid][choice.index] += choice.delta.content
+
+    # Print output.
+    print("Chat completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    engine.terminate()
+    del engine
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_chat_completion_non_stream(model: str):
+    # Create engine
+    engine = MLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+        ),
+    )
+
+    num_requests = 2
+    max_tokens = 64
+    n = 2
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    for rid in range(num_requests):
+        print(f"chat completion for request {rid}")
+        response = engine.chat.completions.create(
+            messages=[{"role": "user", "content": prompts[rid]}],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+        )
+        for choice in response.choices:
+            assert choice.message.role == "assistant"
+            assert isinstance(choice.message.content, str)
+            output_texts[rid][choice.index] += choice.message.content
+
+    # Print output.
+    print("Chat completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    engine.terminate()
+    del engine
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_completion(model: str):
+    # Create engine
+    engine = MLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+        ),
+    )
+
+    num_requests = 2
+    max_tokens = 128
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    for rid in range(num_requests):
+        print(f"completion for request {rid}")
+        for response in engine.completions.create(
+            prompt=prompts[rid],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            stream=True,
+            extra_body={"debug_config": {"ignore_eos": True}},
+        ):
+            for choice in response.choices:
+                output_texts[rid][choice.index] += choice.text
+
+    # Print output.
+    print("Completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    engine.terminate()
+    del engine
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_completion_non_stream(model: str):
+    # Create engine
+    engine = MLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+        ),
+    )
+
+    num_requests = 2
+    max_tokens = 128
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    for rid in range(num_requests):
+        print(f"completion for request {rid}")
+        response = engine.completions.create(
+            prompt=prompts[rid],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            extra_body={"debug_config": {"ignore_eos": True}},
+        )
+        for choice in response.choices:
+            output_texts[rid][choice.index] += choice.text
+
+    # Print output.
+    print("Completion all finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    engine.terminate()
+    del engine
+
+
+if __name__ == "__main__":
+    test_engine_generate()
+    test_chat_completion()
+    test_chat_completion_non_stream()
+    test_completion()
+    test_completion_non_stream()

tests/python/serve/test_serve_engine_grammar.py

@@ -0,0 +1,356 @@
+import asyncio
+import json
+import random
+from typing import Dict, List, Literal  # noqa: UP035
+
+from pydantic import BaseModel
+
+from mlc_llm.protocol.debug_protocol import DebugConfig
+from mlc_llm.protocol.openai_api_protocol import ChatCompletionResponse
+from mlc_llm.serve import AsyncMLCEngine, MLCEngine
+from mlc_llm.testing import require_test_model
+
+LLAMA_2_MODEL = "Llama-2-7b-chat-hf-q4f16_1-MLC"
+LLAMA_3_MODEL = "Meta-Llama-3-8B-Instruct-q4f16_1-MLC"
+
+
+@require_test_model(LLAMA_3_MODEL)
+def test_batch_generation_with_grammar(model: str):
+    # Engine
+    engine = MLCEngine(model=model, mode="server")
+
+    # Inputs
+    system_prompt = "You are a helpful assistant. Always respond only with json."
+    prompts_list = [
+        "Generate a JSON string containing 20 objects:",
+        "Generate a JSON containing a non-empty list:",
+        "Generate a JSON with 5 elements:",
+        "Generate a JSON with a number list, counting from 1 to 20:",
+    ]
+
+    repeat = 3
+    top_p = 0.9
+    temperature = 0.6
+    max_tokens = 4096
+
+    # non-json output
+    responses_text: List[ChatCompletionResponse] = []  # noqa: UP006
+    for _ in range(repeat):
+        for p in prompts_list:
+            print(f"Start generation task for request {len(responses_text)}")
+            responses_text.append(
+                engine.chat.completions.create(
+                    messages=[
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": p},
+                    ],
+                    response_format={"type": "text"},
+                    top_p=top_p,
+                    temperature=temperature,
+                    max_tokens=max_tokens,
+                    seed=random.randint(0, 1 << 30),
+                    extra_body={"debug_config": DebugConfig(grammar_execution_mode="constraint")},
+                )
+            )
+
+    print("Text output")
+    for req_id, response in enumerate(responses_text):
+        prompt = prompts_list[req_id % len(prompts_list)]
+        output = response.choices[0].message.content
+        print(f"Prompt {req_id}: {prompt}")
+        print(f"Output {req_id}: {output}\n")
+
+    # json output
+    responses_json: List[ChatCompletionResponse] = []  # noqa: UP006
+    for _ in range(repeat):
+        for p in prompts_list:
+            print(f"Start generation task for request {len(responses_json)}")
+            responses_json.append(
+                engine.chat.completions.create(
+                    messages=[
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": p},
+                    ],
+                    response_format={"type": "json_object"},
+                    top_p=top_p,
+                    temperature=temperature,
+                    seed=random.randint(0, 1 << 30),
+                )
+            )
+
+    print("JSON output")
+    for req_id, response in enumerate(responses_json):
+        prompt = prompts_list[req_id % len(prompts_list)]
+        output = str(response.choices[0].message.content)
+        print(f"Prompt {req_id}: {prompt}")
+        print(f"Output {req_id}: {output}\n")
+        json.loads(output)
+
+    print("Engine metrics:", engine.metrics())
+
+    engine.terminate()
+
+
+@require_test_model(LLAMA_3_MODEL)
+def test_batch_generation_with_schema(model: str):
+    # Create engine
+    engine = MLCEngine(model=model, mode="server")
+
+    class Product(BaseModel):
+        product_id: int
+        is_available: bool
+        price: float
+        is_featured: Literal[True]
+        category: Literal["Electronics", "Clothing", "Food"]
+        tags: List[str]  # noqa: UP006
+        stock: Dict[str, int]  # noqa: UP006
+
+    schema_str = json.dumps(Product.model_json_schema())
+
+    system_prompt = (
+        "You are a helpful assistant. Always respond only with JSON based on the "
+        f"following JSON schema: {schema_str}."
+    )
+    prompt = "Generate a JSON that describes the product according to the given JSON schema."
+
+    repeat = 8
+    top_p = 0.9
+    temperature = 0.6
+    max_tokens = 4096
+
+    # non-json output
+    responses_text: List[ChatCompletionResponse] = []  # noqa: UP006
+    for i in range(repeat):
+        print(f"Start generation task for request {i}")
+        responses_text.append(
+            engine.chat.completions.create(
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt},
+                ],
+                response_format={"type": "text"},
+                top_p=top_p,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                seed=random.randint(0, 1 << 30),
+                extra_body={"debug_config": DebugConfig(grammar_execution_mode="constraint")},
+            )
+        )
+
+    print("Text output")
+    for req_id, response in enumerate(responses_text):
+        output = response.choices[0].message.content
+        print(f"Prompt {req_id}: {prompt}")
+        print(f"Output {req_id}: {output}\n")
+
+    # json output without schema
+    responses_json: List[ChatCompletionResponse] = []  # noqa: UP006
+    for i in range(repeat):
+        print(f"Start generation task for request {i}")
+        responses_json.append(
+            engine.chat.completions.create(
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt},
+                ],
+                response_format={"type": "json_object"},
+                top_p=top_p,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                seed=random.randint(0, 1 << 30),
+                extra_body={"debug_config": DebugConfig(grammar_execution_mode="constraint")},
+            )
+        )
+
+    print("JSON output")
+    for req_id, response in enumerate(responses_json):
+        output = response.choices[0].message.content
+        print(f"Prompt {req_id}: {prompt}")
+        print(f"Output {req_id}: {output}\n")
+
+    # json output with schema
+    responses_schema: List[ChatCompletionResponse] = []  # noqa: UP006
+    for i in range(repeat):
+        print(f"Start generation task for request {i}")
+        responses_schema.append(
+            engine.chat.completions.create(
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt},
+                ],
+                response_format={"type": "json_object", "schema": schema_str},
+                top_p=top_p,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                seed=random.randint(0, 1 << 30),
+                extra_body={"debug_config": DebugConfig(grammar_execution_mode="constraint")},
+            )
+        )
+
+    print("JSON Schema output")
+    for req_id, response in enumerate(responses_schema):
+        output = response.choices[0].message.content
+        print(f"Prompt {req_id}: {prompt}")
+        print(f"Output {req_id}: {output}\n")
+
+    print("Engine metrics:", engine.metrics())
+
+    engine.terminate()
+
+
+@require_test_model(LLAMA_3_MODEL)
+def test_batch_generation_jump_forward(model: str, jump_forward: bool = True, repeat: int = 1):
+    # Create engine
+    engine = MLCEngine(model=model, mode="server")
+
+    class Product(BaseModel):
+        product_id: int
+        is_available: bool
+        price: float
+        is_featured: Literal[True]
+        category: Literal["Electronics", "Clothing", "Food"]
+        tags: List[str]  # noqa: UP006
+        stock: Dict[str, int]  # noqa: UP006
+
+    schema_str = json.dumps(Product.model_json_schema())
+
+    system_prompt = (
+        "You are a helpful assistant. Always respond only with JSON based on the "
+        f"following JSON schema: {schema_str}."
+    )
+    prompt = "Generate a JSON that describes the product according to the given JSON schema."
+
+    top_p = 0.9
+    temperature = 0.6
+    max_tokens = 4096
+    grammar_execution_mode = "jump_forward" if jump_forward else "constraint"
+
+    # json output with schema
+    responses: List[ChatCompletionResponse] = []  # noqa: UP006
+    for i in range(repeat):
+        print(f"Start generation task for request {i}")
+        responses.append(
+            engine.chat.completions.create(
+                messages=[
+                    {"role": "system", "content": system_prompt},
+                    {"role": "user", "content": prompt},
+                ],
+                response_format={"type": "json_object", "schema": schema_str},
+                top_p=top_p,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                seed=random.randint(0, 1 << 30),
+                extra_body={
+                    "debug_config": DebugConfig(grammar_execution_mode=grammar_execution_mode)
+                },
+            )
+        )
+
+    print(f"Jump forward: {jump_forward}, Repeat: {repeat}")
+    for req_id, response in enumerate(responses):
+        output = response.choices[0].message.content
+        print(f"Prompt {req_id}: {prompt}")
+        print(f"Output {req_id}: {output}\n")
+
+    print("Engine metrics:", engine.metrics())
+
+    engine.terminate()
+
+
+@require_test_model(LLAMA_3_MODEL)
+async def run_async_engine(
+    model: str,
+    mode: Literal["text", "json", "schema"] = "schema",
+    jump_forward: bool = True,
+    num_requests: int = 8,
+):
+    # Create engine
+    async_engine = AsyncMLCEngine(model=model, mode="server")
+
+    class Product(BaseModel):
+        product_id: int
+        is_available: bool
+        price: float
+        is_featured: Literal[True]
+        category: Literal["Electronics", "Clothing", "Food"]
+        tags: List[str]  # noqa: UP006
+        stock: Dict[str, int]  # noqa: UP006
+
+    schema_str = json.dumps(Product.model_json_schema())
+
+    if mode == "text":
+        response_format = {"type": "text"}
+    elif mode == "json":
+        response_format = {"type": "json_object"}
+    elif mode == "schema":
+        response_format = {"type": "json_object", "schema": schema_str}
+
+    system_prompt = (
+        "You are a helpful assistant. Always respond only with JSON based on the "
+        f"following JSON schema: {schema_str}."
+    )
+    prompt = "Generate a JSON that describes the product according to the given JSON schema."
+
+    top_p = 0.9
+    temperature = 0.6
+    max_tokens = 4096
+    grammar_execution_mode = "jump_forward" if jump_forward else "constraint"
+
+    responses = ["" for _ in range(num_requests)]
+
+    async def generate_task(prompt: str, request_id: str):
+        print(f"Start generation task for request {request_id}")
+        rid = int(request_id)
+        async for response in await async_engine.chat.completions.create(  # noqa: F821
+            messages=[
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": prompt},
+            ],
+            response_format=response_format,
+            top_p=top_p,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            seed=random.randint(0, 1 << 30),
+            stream=True,
+            extra_body={"debug_config": DebugConfig(grammar_execution_mode=grammar_execution_mode)},
+        ):
+            assert len(response.choices) == 1
+            choice = response.choices[0]
+            assert choice.delta.role == "assistant"
+            assert isinstance(choice.delta.content, str)
+            responses[rid] += choice.delta.content
+
+    tasks = [
+        asyncio.create_task(generate_task(prompt, request_id=str(i))) for i in range(num_requests)
+    ]
+
+    await asyncio.gather(*tasks)
+
+    print(f"Mode: {mode}, Jump forward: {jump_forward}, Num requests: {num_requests}")
+    for req_id, output in enumerate(responses):
+        print(f"Prompt {req_id}: {prompt}")
+        print(f"Output {req_id}: {output}\n")
+
+    print("Engine metrics:", await async_engine.metrics())
+
+    async_engine.terminate()
+    del async_engine
+
+
+def test_async_engine(
+    mode: Literal["text", "json", "schema"] = "schema",
+    jump_forward: bool = True,
+    num_requests: int = 8,
+):
+    asyncio.run(run_async_engine(mode, jump_forward, num_requests))
+
+
+if __name__ == "__main__":
+    test_batch_generation_with_grammar()
+    test_batch_generation_with_schema()
+    test_batch_generation_jump_forward(False)
+    test_batch_generation_jump_forward(True)
+    test_async_engine("schema", False, 1)
+    test_async_engine("schema", True, 1)
+    test_async_engine("schema", False, 8)
+    test_async_engine("schema", True, 8)

tests/python/serve/test_serve_engine_image.py

@@ -0,0 +1,56 @@
+import json
+from pathlib import Path
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import data
+from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine
+
+
+def get_test_image(config) -> data.ImageData:
+    return data.ImageData.from_url("https://llava-vl.github.io/static/images/view.jpg", config)
+
+
+def test_engine_generate():
+    # Create engine
+    model = "dist/llava-1.5-7b-hf-q4f16_1-MLC/params"
+    model_lib = "dist/llava-1.5-7b-hf-q4f16_1-MLC/llava-1.5-7b-hf-q4f16_1-MLC.so"
+    engine = SyncMLCEngine(
+        model=model,
+        model_lib=model_lib,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+    max_tokens = 256
+
+    with open(Path(model) / "mlc-chat-config.json", encoding="utf-8") as file:
+        model_config = json.load(file)
+
+    prompts = [
+        [
+            data.TextData("USER: "),
+            get_test_image(model_config),
+            data.TextData("\nWhat does this image represent? ASSISTANT:"),
+        ],
+        [
+            data.TextData("USER: "),
+            get_test_image(model_config),
+            data.TextData("\nIs there a dog in this image? ASSISTANT:"),
+        ],
+        [data.TextData("USER: What is the meaning of life? ASSISTANT:")],
+    ]
+
+    output_texts, _ = engine.generate(
+        prompts, GenerationConfig(max_tokens=max_tokens, stop_token_ids=[2])
+    )
+
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+
+if __name__ == "__main__":
+    test_engine_generate()

tests/python/serve/test_serve_engine_mock.py

@@ -0,0 +1,39 @@
+"""Mock testing engine I/O conventions
+
+Mock test only can help checking the overall input
+output processing options are passed correctly
+"""
+
+import pytest
+import tvm
+
+from mlc_llm.serve import MLCEngine
+from mlc_llm.testing import require_test_model
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+# NOTE: we only need tokenizers in folder
+# launch time of mock test is fast so we can put it in unittest
+@require_test_model("Llama-3-8B-Instruct-q4f16_1-MLC")
+def test_completion_api(model: str):
+    engine = MLCEngine(model, tvm.cpu(), model_lib="mock://echo")
+    param_dict = {
+        "top_p": 0.6,
+        "temperature": 0.9,
+        "frequency_penalty": 0.1,
+        "presence_penalty": 0.1,
+        "n": 2,
+    }
+    response = engine.chat.completions.create(
+        messages=[{"role": "user", "content": "hello"}],
+        **param_dict,
+    )
+    # echo mock will echo back the generation config
+    for k, v in param_dict.items():
+        assert response.usage.extra[k] == v
+
+
+if __name__ == "__main__":
+    test_completion_api()

tests/python/serve/test_serve_engine_prefix_cache.py

@@ -0,0 +1,141 @@
+from mlc_llm.protocol.debug_protocol import DebugConfig
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine
+from mlc_llm.testing import require_test_model
+
+prompts = [
+    "The meaning of life is",
+    "According to the history of Pittsburgh,",
+    "I have a three-day Seattle travel plan. On the first day,",
+    "Undoubtedly, Alaska is one of the most beautiful places on Earth,",
+    "Explain difference between Lambda calculus and Turing machine is",
+    "To assemble a desktop computer, we need the necessary components of",
+    "Vitamin D is important to human beings, because",
+    "Refer to history, the milk tea is originated from",
+    "In the southernmost place in United States,",
+    "AlphaGo has the capabilities of",
+]
+
+
+def test_engine_system_prompt(engine):
+    system_prompt = "This is a system prompt"
+    system_prompt_tokens = len(engine.tokenizer.encode(system_prompt))
+    max_tokens = 8
+    _, _ = engine.generate(
+        system_prompt,
+        GenerationConfig(
+            temperature=0,
+            max_tokens=max_tokens,
+            debug_config=DebugConfig(pinned_system_prompt=True),
+        ),
+    )
+    metrics = engine.metrics()
+    assert metrics["prefill_tokens_sum"] == system_prompt_tokens
+    sum_prefill_tokens = system_prompt_tokens
+
+    input_token_lens = [len(engine.tokenizer.encode(prompt)) for prompt in prompts]
+
+    generation_config = GenerationConfig(temperature=0, max_tokens=max_tokens)
+    _, _ = engine.generate(prompts, generation_config)
+    metrics = engine.metrics()
+    assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + sum(input_token_lens)
+    sum_prefill_tokens = metrics["prefill_tokens_sum"]
+
+    _, _ = engine.generate(system_prompt + " and why ?", generation_config)
+    metrics = engine.metrics()
+    # system prompt is reused entirely
+    assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + 3
+    sum_prefill_tokens = metrics["prefill_tokens_sum"]
+
+    _, _ = engine.generate(prompts[:4], generation_config)
+    metrics = engine.metrics()
+    # first 4 prompts are removed and need to prefill again
+    assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + sum(input_token_lens[:4])
+
+
+def test_engine_multi_round(engine):
+    num_requests = 10
+    max_tokens = 8
+    generation_config = GenerationConfig(temperature=0, max_tokens=max_tokens)
+    input_token_lens = [len(engine.tokenizer.encode(prompt)) for prompt in prompts[:num_requests]]
+
+    output_texts, _ = engine.generate(prompts[:num_requests], generation_config)
+    metrics = engine.metrics()
+    assert metrics["prefill_tokens_sum"] == sum(input_token_lens)
+    sum_prefill_tokens = metrics["prefill_tokens_sum"]
+    concat_prompt = []
+    for i, output in enumerate(output_texts):
+        concat_prompt.append(prompts[i] + " " + output[0] + " ?")
+    output_texts, _ = engine.generate(concat_prompt[:num_requests], generation_config)
+    metrics = engine.metrics()
+    assert metrics["prefill_tokens_sum"] == sum_prefill_tokens + 2 * num_requests
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_basic_engine_system_prompt(model: str):
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="local",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            prefix_cache_max_num_recycling_seqs=5,
+        ),
+    )
+    test_engine_system_prompt(engine)
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_basic_engine_multi_round(model: str):
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+    test_engine_multi_round(engine)
+
+
+@require_test_model(
+    "Llama-2-7b-chat-hf-q0f16-MLC",
+    "Llama-2-7b-chat-hf-q4f16_1-MLC",
+)
+def test_engine_spec_multi_round(model: str, small_model: str):
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[small_model],
+            speculative_mode="small_draft",
+        ),
+    )
+
+    test_engine_multi_round(engine)
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_eagle_multi_round(model: str):
+    # Create engine
+    small_model = "dist/Eagle-llama2-7b-chat-q0f16-MLC"
+    small_model_lib = "dist/Eagle-llama2-7b-chat-q0f16-MLC/Eagle-llama2-7b-chat-q0f16-MLC-cuda.so"
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[(small_model, small_model_lib)],
+            speculative_mode="eagle",
+            max_num_sequence=80,
+        ),
+    )
+
+    test_engine_multi_round(engine)
+
+
+if __name__ == "__main__":
+    test_basic_engine_system_prompt()
+    test_basic_engine_multi_round()
+    test_engine_spec_multi_round()
+    test_engine_eagle_multi_round()

tests/python/serve/test_serve_engine_rnn.py

@@ -0,0 +1,60 @@
+from typing import List  # noqa: UP035
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import EngineConfig, MLCEngine
+
+prompts = [
+    "What is the meaning of life?",
+    "Introduce the history of Pittsburgh to me. Please elaborate in detail.",
+    "Write a three-day Seattle travel plan. Please elaborate in detail.",
+    "What is Alaska famous of? Please elaborate in detail.",
+    "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.",  # noqa: E501
+    "What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
+    "Why is Vitamin D important to human beings? Please elaborate in detail.",
+    "Where is milk tea originated from? Please elaborate in detail.",
+    "Where is the southernmost place in United States? Please elaborate in detail.",
+    "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.",  # noqa: E501
+]
+
+
+def test_engine_generate() -> None:
+    engine = MLCEngine(
+        model="dist/rwkv-6-world-1b6-q0f16-MLC",
+        model_lib="dist/rwkv-6-world-1b6-q0f16-MLC/rwkv-6-world-1b6-q0f16-MLC-cuda.so",
+        mode="server",
+        engine_config=EngineConfig(
+            max_num_sequence=8,
+            max_history_size=1,
+        ),
+    )
+
+    num_requests = 10
+    max_tokens = 256
+    generation_cfg = GenerationConfig(max_tokens=max_tokens, n=7)
+
+    output_texts: List[List[str]] = [  # noqa: UP006
+        ["" for _ in range(generation_cfg.n)] for _ in range(num_requests)
+    ]
+    for rid in range(num_requests):
+        print(f"generating for request {rid}")
+        for delta_outputs in engine._generate(prompts[rid], generation_cfg, request_id=str(rid)):
+            assert len(delta_outputs) == generation_cfg.n
+            for i, delta_output in enumerate(delta_outputs):
+                output_texts[rid][i] += delta_output.delta_text
+
+    # Print output.
+    print("All finished")
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+    engine.terminate()
+    del engine
+
+
+if __name__ == "__main__":
+    test_engine_generate()

tests/python/serve/test_serve_engine_spec.py

@@ -0,0 +1,660 @@
+from typing import Callable, List, Optional  # noqa: UP035
+
+import numpy as np
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import Request, RequestStreamOutput, data
+from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine
+from mlc_llm.testing import require_test_model
+
+prompts = [
+    "What is the meaning of life?",
+    "Introduce the history of Pittsburgh to me. Please elaborate in detail.",
+    "Write a three-day Seattle travel plan. Please elaborate in detail.",
+    "What is Alaska famous of? Please elaborate in detail.",
+    "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.",  # noqa: E501
+    "What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
+    "Why is Vitamin D important to human beings? Please elaborate in detail.",
+    "Where is milk tea originated from? Please elaborate in detail.",
+    "Where is the southernmost place in United States? Please elaborate in detail.",
+    "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.",  # noqa: E501
+]
+
+
+def create_requests(
+    num_requests: int,
+    stop_token_id: Optional[int] = None,
+    temperature: float = 0.8,
+    repetition_penalty: float = 1.0,
+    max_tokens_low: int = 256,
+    max_tokens_high: int = 257,
+) -> List[Request]:  # noqa: UP006
+    assert num_requests >= 0 and num_requests <= len(prompts)
+
+    stop_token_ids = [stop_token_id] if stop_token_id is not None else []
+    requests = []
+    for req_id, prompt in zip(range(num_requests), prompts):
+        max_tokens = np.random.randint(max_tokens_low, max_tokens_high)
+        requests.append(
+            Request(
+                request_id=str(req_id),
+                inputs=data.TextData(prompt),
+                generation_config=GenerationConfig(
+                    temperature=temperature,
+                    repetition_penalty=repetition_penalty,
+                    max_tokens=max_tokens,
+                    stop_token_ids=stop_token_ids,
+                ),
+            )
+        )
+    return requests
+
+
+@require_test_model(
+    "Llama-2-7b-chat-hf-q0f16-MLC",
+    "Llama-2-7b-chat-hf-q4f16_1-MLC",
+)
+def test_engine_basic(model: str, small_model: str):
+    """Test engine **without continuous batching**.
+
+    - Add all requests to the engine altogether in the beginning.
+    - All requests have the same max_tokens. This means all requests
+    will end together.
+    - Engine keeps running `step` for estimated number of steps (number of
+    requests + max_tokens - 1). Then check the output of each request.
+    """
+
+    # Hyperparameters for tests (you can try different combinations).
+    num_requests = len(prompts)  # [4, 8, 10]
+    temperature = 0.9  # [0, 0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.0  # [1.0, 1.01]
+    max_tokens: int = 256  # [32, 128, 256]
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+
+    # Define the callback function for request generation results
+    def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+        for delta_output in delta_outputs:
+            request_id, stream_outputs = delta_output.unpack()
+            assert len(stream_outputs) == 1
+            outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[small_model],
+            speculative_mode="small_draft",
+        ),
+        request_stream_callback=fcallback,
+    )
+
+    # Create requests
+    requests = create_requests(
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens,
+        max_tokens_high=max_tokens + 1,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        engine.add_request(request)
+
+    num_steps = num_requests + max_tokens - 1
+    # Run steps
+    for step in range(num_steps):
+        engine.step()
+
+    for req_id, output in enumerate(outputs):
+        print(f"Prompt {req_id}: {requests[req_id].inputs[0]}")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_eagle_basic(model: str):
+    """Test engine **without continuous batching**.
+
+    - Add all requests to the engine altogether in the beginning.
+    - All requests have the same max_tokens. This means all requests
+    will end together.
+    - Engine keeps running `step` for estimated number of steps (number of
+    requests + max_tokens - 1). Then check the output of each request.
+    - Use Eagle model as speculative model
+    """
+
+    # Hyperparameters for tests (you can try different combinations).
+    num_requests = len(prompts)  # [4, 8, 10]
+    temperature = 0.9  # [0, 0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.0  # [1.0, 1.01]
+    max_tokens: int = 256  # [32, 128, 256]
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+
+    # Define the callback function for request generation results
+    def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+        for delta_output in delta_outputs:
+            request_id, stream_outputs = delta_output.unpack()
+            assert len(stream_outputs) == 1
+            outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+
+    # Create engine
+    small_model = "dist/Eagle-llama2-7b-chat-q0f16-MLC"
+    small_model_lib = "dist/Eagle-llama2-7b-chat-q0f16-MLC/Eagle-llama2-7b-chat-q0f16-MLC-cuda.so"
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[(small_model, small_model_lib)],
+            speculative_mode="eagle",
+            spec_draft_length=2,
+        ),
+        request_stream_callback=fcallback,
+    )
+
+    # Create requests
+    requests = create_requests(
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens,
+        max_tokens_high=max_tokens + 1,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        engine.add_request(request)
+
+    num_steps = num_requests + max_tokens - 1
+    # Run steps
+    for step in range(num_steps):
+        engine.step()
+
+    for req_id, output in enumerate(outputs):
+        print(f"Prompt {req_id}: {requests[req_id].inputs[0]}")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+
+
+@require_test_model(
+    "Llama-2-7b-chat-hf-q0f16-MLC",
+    "Llama-2-7b-chat-hf-q4f16_1-MLC",
+)
+def test_engine_continuous_batching_1(model: str, small_model: str):
+    """Test engine **with continuous batching**.
+
+    - Add all requests to the engine altogether in the beginning.
+    - All requests have a random maximum generation length. So each
+    request keeps generating until reaching the maximum length.
+    - Engine keeps running `step` for estimated number of steps (number of
+    requests + the maximum max_tokens - 1). Then check the output
+    of each request.
+    """
+
+    # Hyperparameters for tests (you can try different combinations)
+    num_requests = len(prompts)  # [4, 8, 10]
+    temperature = 0.9  # [0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.00  # [1.0, 1.01]
+    max_tokens_low = 128
+    max_tokens_high = 384
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+    finish_time: List[Optional[int]] = [None] * num_requests  # noqa: UP006
+
+    # Define the callback class for request generation results
+    class CallbackTimer:
+        timer: int = -1
+
+        def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]:  # noqa: UP006
+            def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+                for delta_output in delta_outputs:
+                    request_id, stream_outputs = delta_output.unpack()
+                    assert len(stream_outputs) == 1
+                    if stream_outputs[0].finish_reason is not None:
+                        print(f"Request {request_id} finished at step {self.timer}.")
+                    outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+                    finish_time[int(request_id)] = self.timer
+
+            return fcallback
+
+        def step(self) -> None:
+            self.timer += 1
+
+    # Create engine
+    timer = CallbackTimer()
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[small_model],
+            speculative_mode="small_draft",
+        ),
+        request_stream_callback=timer.callback_getter(),
+    )
+
+    # Create requests
+    requests = create_requests(
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens_low,
+        max_tokens_high=max_tokens_high,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        engine.add_request(request)
+
+    num_steps = num_requests + max(request.generation_config.max_tokens for request in requests) - 1
+    # Run steps
+    for step in range(num_steps):
+        timer.step()
+        assert timer.timer == step
+        engine.step()
+
+    for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)):
+        print(f"Prompt {req_id}: {request.inputs[0]}")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+        # assert fin_time == request.generation_config.max_tokens - 1
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_engine_eagle_continuous_batching_1(model: str):
+    """Test engine **with continuous batching**.
+
+    - Add all requests to the engine altogether in the beginning.
+    - All requests have a random maximum generation length. So each
+    request keeps generating until reaching the maximum length.
+    - Engine keeps running `step` for estimated number of steps (number of
+    requests + the maximum max_tokens - 1). Then check the output
+    of each request.
+    """
+
+    # Hyperparameters for tests (you can try different combinations)
+    num_requests = len(prompts)  # [4, 8, 10]
+    temperature = 0.9  # [0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.00  # [1.0, 1.01]
+    max_tokens_low = 128
+    max_tokens_high = 384
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+    finish_time: List[Optional[int]] = [None] * num_requests  # noqa: UP006
+
+    # Define the callback class for request generation results
+    class CallbackTimer:
+        timer: int = -1
+
+        def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]:  # noqa: UP006
+            def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+                for delta_output in delta_outputs:
+                    request_id, stream_outputs = delta_output.unpack()
+                    assert len(stream_outputs) == 1
+                    if stream_outputs[0].finish_reason is not None:
+                        print(f"Request {request_id} finished at step {self.timer}.")
+                    outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+                    finish_time[int(request_id)] = self.timer
+
+            return fcallback
+
+        def step(self) -> None:
+            self.timer += 1
+
+    # Create engine
+    small_model = "dist/Eagle-llama2-7b-chat-q4f16_1-MLC"
+    small_model_lib = (
+        "dist/Eagle-llama2-7b-chat-q4f16_1-MLC/Eagle-llama2-7b-chat-q4f16_1-MLC-cuda.so"
+    )
+    timer = CallbackTimer()
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[(small_model, small_model_lib)],
+            speculative_mode="eagle",
+        ),
+        request_stream_callback=timer.callback_getter(),
+    )
+
+    # Create requests
+    requests = create_requests(
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens_low,
+        max_tokens_high=max_tokens_high,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        engine.add_request(request)
+
+    num_steps = num_requests + max(request.generation_config.max_tokens for request in requests) - 1
+    # Run steps
+    for step in range(num_steps):
+        timer.step()
+        assert timer.timer == step
+        engine.step()
+
+    for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)):
+        print(f"Prompt {req_id}: {request.inputs[0]}")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+        # assert fin_time == request.generation_config.max_tokens - 1
+
+
+def compare_output_text(output_text1, output_text2):
+    if isinstance(output_text1, list) and isinstance(output_text2, list):
+        for item1, item2 in zip(output_text1, output_text2):
+            if not compare_output_text(item1, item2):
+                return False
+    elif output_text1 != output_text2:
+        print(output_text1)
+        print(output_text2)
+        return False
+    return True
+
+
+@require_test_model(
+    "Llama-2-7b-chat-hf-q0f16-MLC",
+    "Llama-2-7b-chat-hf-q4f16_1-MLC",
+)
+def test_engine_generate(model: str, small_model: str, compare_precision=False):
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[small_model],
+            speculative_mode="small_draft",
+        ),
+    )
+
+    num_requests = 10
+    max_tokens = 256
+
+    # Generate output.
+    if compare_precision:
+        print("compare precision")
+        generation_config = GenerationConfig(
+            temperature=0.0, top_p=0, max_tokens=1024, stop_token_ids=[2], n=1
+        )
+        engine_single_model = SyncMLCEngine(
+            model=model,
+            mode="server",
+            engine_config=EngineConfig(
+                max_total_sequence_length=4096,
+            ),
+        )
+        output_texts_single_model, _ = engine_single_model.generate(
+            prompts[:num_requests], generation_config
+        )
+        for req_id, outputs in enumerate(output_texts_single_model):
+            print(f"Prompt {req_id}: {prompts[req_id]}")
+            if len(outputs) == 1:
+                print(f"Output {req_id}:{outputs[0]}\n")
+            else:
+                for i, output in enumerate(outputs):
+                    print(f"Output {req_id}({i}):{output}\n")
+        # TODO: Add pytorch precision
+    else:
+        generation_config = GenerationConfig(max_tokens=max_tokens, n=3)
+    output_texts, _ = engine.generate(prompts[:num_requests], generation_config)
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+    if compare_precision:
+        precision_flag = compare_output_text(output_texts, output_texts_single_model)
+        if precision_flag:
+            print("Accuracy verification succeed\n")
+        else:
+            print("Accuracy verification failed\n")
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_eagle_generate(model: str):
+    # Create engine
+    small_model = "dist/Eagle-llama2-7b-chat-q4f16_1-MLC"
+    small_model_lib = (
+        "dist/Eagle-llama2-7b-chat-q4f16_1-MLC/Eagle-llama2-7b-chat-q4f16_1-MLC-cuda.so"
+    )
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[(small_model, small_model_lib)],
+            speculative_mode="eagle",
+        ),
+    )
+
+    num_requests = 10
+    max_tokens = 256
+
+    # Generate output.
+    output_texts, _ = engine.generate(
+        prompts[:num_requests], GenerationConfig(max_tokens=max_tokens, n=3)
+    )
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+
+@require_test_model("Llama-2-13b-chat-hf-q4f16_1-MLC")
+def test_engine_efficiency(model: str):
+    """Test engine speculative decoding efficiency."""
+
+    # Hyperparameters for tests (you can try different combinations).
+    num_requests = 1  # [4, 8, 10]
+    temperature = 0.9  # [0, 0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.0  # [1.0, 1.01]
+    max_tokens: int = 512
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+
+    # Define the callback function for request generation results
+    def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+        for delta_output in delta_outputs:
+            request_id, stream_outputs = delta_output.unpack()
+            assert len(stream_outputs) == 1
+            outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+        request_stream_callback=fcallback,
+    )
+
+    # Create requests
+    requests = create_requests(
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens,
+        max_tokens_high=max_tokens + 1,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        engine.add_request(request)
+
+    num_steps = num_requests + max_tokens - 1
+    # Run steps
+    for step in range(num_steps):
+        engine.step()
+
+    for eg, name in zip([engine], ["Normal Deconding"]):
+        metrics = eg.metrics()
+        print("engine name:", name)
+        if name == "Speculative Decoding":
+            print("spec decode metrics:", metrics["spec_decode"])
+        print("engine total decode time:", metrics["engine_decode_time_sum"])
+        print()
+
+
+@require_test_model(
+    "Llama-2-13b-chat-hf-q4f16_1-MLC",
+    "Llama-2-7b-chat-hf-q4f16_1-MLC",
+)
+def test_engine_spec_efficiency(model: str, small_model: str):
+    """Test engine speculative decoding efficiency."""
+
+    # Hyperparameters for tests (you can try different combinations).
+    num_requests = 1  # [4, 8, 10]
+    temperature = 0.9  # [0, 0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.0  # [1.0, 1.01]
+    max_tokens: int = 512
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+
+    # Define the callback function for request generation results
+    def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+        for delta_output in delta_outputs:
+            request_id, stream_outputs = delta_output.unpack()
+            assert len(stream_outputs) == 1
+            outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+
+    # Create engine
+    spec_engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[small_model],
+            spec_draft_length=6,
+            speculative_mode="small_draft",
+        ),
+        request_stream_callback=fcallback,
+    )
+
+    # Create requests
+    requests = create_requests(
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens,
+        max_tokens_high=max_tokens + 1,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        spec_engine.add_request(request)
+
+    num_steps = num_requests + max_tokens - 1
+    # Run steps
+    for step in range(num_steps):
+        spec_engine.step()
+
+    for eg, name in zip([spec_engine], ["Speculative Decoding"]):
+        metrics = eg.metrics()
+        print("engine name:", name)
+        if name == "Speculative Decoding":
+            print("total draft tokens:", metrics["sum_num_draft_tokens"])
+            print("total accepted tokens:", metrics["sum_num_accepted_tokens"])
+            print(
+                "Accept rate:",
+                metrics["sum_num_accepted_tokens"] / (1e-10 + metrics["sum_num_draft_tokens"]),
+            )
+        print("engine total decode time:", metrics["engine_decode_time_sum"])
+        print()
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_engine_eagle_spec_efficiency(model: str):
+    """Test engine speculative decoding efficiency."""
+
+    # Hyperparameters for tests (you can try different combinations).
+    num_requests = 1  # [4, 8, 10]
+    temperature = 0.9  # [0, 0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.0  # [1.0, 1.01]
+    max_tokens: int = 512
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+
+    # Define the callback function for request generation results
+    def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+        for delta_output in delta_outputs:
+            request_id, stream_outputs = delta_output.unpack()
+            assert len(stream_outputs) == 1
+            outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+
+    # Create engine
+    small_model = "dist/Eagle-llama2-7b-chat-q0f16-MLC"
+    small_model_lib = "dist/Eagle-llama2-7b-chat-q0f16-MLC/Eagle-llama2-7b-chat-q0f16-MLC-cuda.so"
+    spec_engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(
+            max_total_sequence_length=4096,
+            additional_models=[(small_model, small_model_lib)],
+            spec_draft_length=6,
+            speculative_mode="eagle",
+        ),
+        request_stream_callback=fcallback,
+    )
+
+    # Create requests
+    requests = create_requests(
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens,
+        max_tokens_high=max_tokens + 1,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        spec_engine.add_request(request)
+
+    num_steps = num_requests + max_tokens - 1
+    # Run steps
+    for step in range(num_steps):
+        spec_engine.step()
+
+    for eg, name in zip([spec_engine], ["Speculative Decoding"]):
+        metrics = eg.metrics()
+        print("engine name:", name)
+        if name == "Speculative Decoding":
+            print("spec decode:", metrics["spec_decode"])
+        print("engine total decode time:", metrics["engine_decode_time_sum"])
+        print()
+
+
+if __name__ == "__main__":
+    test_engine_basic()
+    test_engine_eagle_basic()
+    test_engine_continuous_batching_1()
+    test_engine_eagle_continuous_batching_1()
+    test_engine_generate(compare_precision=True)
+    test_engine_eagle_generate()
+    test_engine_efficiency()
+    test_engine_spec_efficiency()
+    test_engine_eagle_spec_efficiency()

tests/python/serve/test_serve_sync_engine.py

@@ -0,0 +1,474 @@
+from typing import Callable, List, Optional  # noqa: UP035
+
+import numpy as np
+
+from mlc_llm.protocol.generation_config import GenerationConfig
+from mlc_llm.serve import Request, RequestStreamOutput, data
+from mlc_llm.serve.sync_engine import EngineConfig, SyncMLCEngine
+from mlc_llm.testing import require_test_model
+
+prompts = [
+    "What is the meaning of life?",
+    "Introduce the history of Pittsburgh to me. Please elaborate in detail.",
+    "Write a three-day Seattle travel plan. Please elaborate in detail.",
+    "What is Alaska famous of? Please elaborate in detail.",
+    "What is the difference between Lambda calculus and Turing machine? Please elaborate in detail.",  # noqa: E501
+    "What are the necessary components to assemble a desktop computer? Please elaborate in detail.",
+    "Why is Vitamin D important to human beings? Please elaborate in detail.",
+    "Where is milk tea originated from? Please elaborate in detail.",
+    "Where is the southernmost place in United States? Please elaborate in detail.",
+    "Do you know AlphaGo? What capabilities does it have, and what achievements has it got? Please elaborate in detail.",  # noqa: E501
+]
+
+
+def create_requests(
+    engine: SyncMLCEngine,
+    num_requests: int,
+    stop_token_id: Optional[int] = None,
+    temperature: float = 0.8,
+    repetition_penalty: float = 1.0,
+    max_tokens_low: int = 256,
+    max_tokens_high: int = 257,
+) -> List[Request]:  # noqa: UP006
+    assert num_requests >= 0 and num_requests <= len(prompts)
+
+    stop_token_ids = [stop_token_id] if stop_token_id is not None else []
+    requests = []
+    for req_id, prompt in zip(range(num_requests), prompts):
+        max_tokens = np.random.randint(max_tokens_low, max_tokens_high)
+        requests.append(
+            engine.create_request(
+                request_id=str(req_id),
+                inputs=data.TextData(prompt),
+                generation_config=GenerationConfig(
+                    temperature=temperature,
+                    repetition_penalty=repetition_penalty,
+                    max_tokens=max_tokens,
+                    stop_token_ids=stop_token_ids,
+                ),
+            )
+        )
+    return requests
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_basic(model: str):
+    """Test engine **without continuous batching**.
+
+    - Add all requests to the engine altogether in the beginning.
+    - All requests have the same max_tokens. This means all requests
+    will end together.
+    - Engine keeps running `step` for estimated number of steps (number of
+    requests + max_tokens - 1). Then check the output of each request.
+    """
+
+    # Hyperparameters for tests (you can try different combinations).
+    num_requests = 10  # [4, 8, 10]
+    temperature = 0.9  # [0, 0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.0  # [1.0, 1.01]
+    max_tokens: int = 256  # [32, 128, 256]
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+
+    # Define the callback function for request generation results
+    def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+        for delta_output in delta_outputs:
+            request_id, stream_outputs = delta_output.unpack()
+            assert len(stream_outputs) == 1
+            outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        request_stream_callback=fcallback,
+    )
+
+    # Create requests
+    requests = create_requests(
+        engine,
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens,
+        max_tokens_high=max_tokens + 1,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        engine.add_request(request)
+
+    num_steps = num_requests + max_tokens - 1
+    # Run steps
+    for step in range(num_steps):
+        engine.step()
+
+    for req_id, output in enumerate(outputs):
+        print(f"Prompt {req_id}: {requests[req_id].inputs[0]}")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_continuous_batching_1(model: str):
+    """Test engine **with continuous batching**.
+
+    - Add all requests to the engine altogether in the beginning.
+    - All requests have a random maximum generation length. So each
+    request keeps generating until reaching the maximum length.
+    - Engine keeps running `step` for estimated number of steps (number of
+    requests + the maximum max_tokens - 1). Then check the output
+    of each request.
+    """
+
+    # Hyperparameters for tests (you can try different combinations)
+    num_requests = 10  # [4, 8, 10]
+    temperature = 0.9  # [0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.00  # [1.0, 1.01]
+    max_tokens_low = 128
+    max_tokens_high = 384
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+    finish_time: List[Optional[int]] = [None] * num_requests  # noqa: UP006
+
+    # Define the callback class for request generation results
+    class CallbackTimer:
+        timer: int = -1
+
+        def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]:  # noqa: UP006
+            def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+                for delta_output in delta_outputs:
+                    request_id, stream_outputs = delta_output.unpack()
+                    assert len(stream_outputs) == 1
+                    if stream_outputs[0].finish_reason is not None:
+                        print(f"Request {request_id} finished at step {self.timer}.")
+                    outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+                    finish_time[int(request_id)] = self.timer
+
+            return fcallback
+
+        def step(self) -> None:
+            self.timer += 1
+
+    # Create engine
+    timer = CallbackTimer()
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        request_stream_callback=timer.callback_getter(),
+    )
+
+    # Create requests
+    requests = create_requests(
+        engine,
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens_low,
+        max_tokens_high=max_tokens_high,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        engine.add_request(request)
+
+    num_steps = num_requests + max(request.generation_config.max_tokens for request in requests) - 1
+    # Run steps
+    for step in range(num_steps):
+        timer.step()
+        assert timer.timer == step
+        engine.step()
+
+    for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)):
+        print(f"Prompt {req_id}: {request.inputs[0]}")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+        assert fin_time == request.generation_config.max_tokens - 1, (
+            f"finish time = {fin_time}, max tokens = {request.generation_config.max_tokens - 1}"
+        )
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_continuous_batching_2(model: str):
+    """Test engine **with continuous batching**.
+
+    - Add all requests to the engine altogether in the beginning.
+    - All requests have the stop token. So each request keeps generating
+    until having the stop token or reaching the maximum length.
+    - Engine keeps running `step` for estimated number of steps (number of
+    requests + the maximum max_tokens - 1). Then check the output
+    of each request.
+    """
+
+    # Hyperparameters for tests (you can try different combinations)
+    num_requests = 10  # [4, 8, 10]
+    temperature = 0.9  # [0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.00  # [1.0, 1.01]
+    stop_token_id = 2
+    max_tokens = 512
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+    finish_time: List[Optional[int]] = [None] * num_requests  # noqa: UP006
+
+    # Define the callback class for request generation results
+    class CallbackTimer:
+        timer: int = -1
+
+        def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]:  # noqa: UP006
+            def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+                for delta_output in delta_outputs:
+                    request_id, stream_outputs = delta_output.unpack()
+                    assert len(stream_outputs) == 1
+                    if stream_outputs[0].finish_reason is not None:
+                        print(f"Request {request_id} finished at step {self.timer}.")
+                    outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+                    finish_time[int(request_id)] = self.timer
+
+            return fcallback
+
+        def step(self) -> None:
+            self.timer += 1
+
+    # Create engine
+    timer = CallbackTimer()
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        request_stream_callback=timer.callback_getter(),
+    )
+
+    # Create requests
+    requests = create_requests(
+        engine,
+        num_requests,
+        stop_token_id=stop_token_id,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens,
+        max_tokens_high=max_tokens + 1,
+    )
+
+    # Add all requests to engine
+    for request in requests:
+        engine.add_request(request)
+
+    num_steps = num_requests + max_tokens - 1
+    # Run steps
+    for step in range(num_steps):
+        timer.step()
+        assert timer.timer == step
+        engine.step()
+
+    for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)):
+        print(f"Prompt {req_id}: {request.inputs[0]}")
+        if fin_time < num_requests + max_tokens - 2:
+            print(f"Request {req_id} ends early on the stop token")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_continuous_batching_3(model: str):
+    """Test engine **with continuous batching**.
+
+    - Add requests randomly between time [0, 200).
+    - All requests have a random maximum generation length. So each
+    request keeps generating until reaching the maximum length.
+    - Engine keeps running `step` until all requests finish.
+    Then check the output of each request.
+    """
+
+    # Hyperparameters for tests (you can try different combinations)
+    num_requests = 10  # [4, 8, 10]
+    temperature = 0.9  # [0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.00  # [1.0, 1.01]
+    stop_token_id = 2
+    max_tokens_low = 64
+    max_tokens_high = 192
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+    finish_time: List[Optional[int]] = [None] * num_requests  # noqa: UP006
+
+    # Define the callback class for request generation results
+    class CallbackTimer:
+        timer: int = -1
+        finished_requests: int = 0
+
+        def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]:  # noqa: UP006
+            def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+                for delta_output in delta_outputs:
+                    request_id, stream_outputs = delta_output.unpack()
+                    assert len(stream_outputs) == 1
+                    if stream_outputs[0].finish_reason is not None:
+                        print(f"Request {request_id} finished at step {self.timer}.")
+                        self.finished_requests += 1
+                    outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+                    finish_time[int(request_id)] = self.timer
+
+            return fcallback
+
+        def step(self) -> None:
+            self.timer += 1
+
+        def all_finished(self) -> bool:
+            return self.finished_requests == num_requests
+
+    # Create engine
+    timer = CallbackTimer()
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        request_stream_callback=timer.callback_getter(),
+    )
+
+    # Create requests
+    requests = create_requests(
+        engine,
+        num_requests,
+        stop_token_id=stop_token_id,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens_low,
+        max_tokens_high=max_tokens_high,
+    )
+
+    # Assign the time to add requests to engine
+    request_add_time = [np.random.randint(0, 200) for _ in range(num_requests)]
+
+    # Run steps
+    while not timer.all_finished():
+        timer.step()
+
+        # Add requests to engine
+        for req_id, add_time in enumerate(request_add_time):
+            if add_time == timer.timer:
+                print(f"add request {req_id} at step {timer.timer}")
+                engine.add_request(requests[req_id])
+
+        engine.step()
+
+    for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)):
+        print(f"Prompt {req_id}: {request.inputs[0]}")
+        print(f"Finish time: {fin_time}")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_generate(model: str):
+    # Create engine
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        engine_config=EngineConfig(max_total_sequence_length=4096),
+    )
+
+    num_requests = 10
+    max_tokens = 256
+
+    # Generate output.
+    output_texts, _ = engine.generate(
+        prompts[:num_requests], GenerationConfig(max_tokens=max_tokens, n=7)
+    )
+    for req_id, outputs in enumerate(output_texts):
+        print(f"Prompt {req_id}: {prompts[req_id]}")
+        if len(outputs) == 1:
+            print(f"Output {req_id}:{outputs[0]}\n")
+        else:
+            for i, output in enumerate(outputs):
+                print(f"Output {req_id}({i}):{output}\n")
+
+
+@require_test_model("Llama-2-7b-chat-hf-q0f16-MLC")
+def test_engine_hybrid_prefill(model: str):
+    """Test engine **with hybrid prefill**.
+
+    - Add each single request step by step.
+    - All requests have the same generation length. But due to hybrid prefill,
+    the earlier request will decode with later request prefill, in single step.
+    So each request lasts the same steps, and stops generation step by step as well.
+    - Engine keeps running `step` for the generation length, to finish the last request.
+    Then check the output of each request.
+    """
+
+    # Hyperparameters for tests (you can try different combinations)
+    num_requests = 10  # [4, 8, 10]
+    temperature = 0.9  # [0.8, 0.9, 1.0, 1.1]
+    repetition_penalty = 1.00  # [1.0, 1.01]
+    max_tokens = 15
+    np.random.seed(0)
+
+    # Output list
+    outputs: List[List[int]] = [[] for _ in range(num_requests)]  # noqa: UP006
+    finish_time: List[Optional[int]] = [None] * num_requests  # noqa: UP006
+
+    # Define the callback class for request generation results
+    class CallbackTimer:
+        timer: int = -1
+
+        def callback_getter(self) -> Callable[[List[RequestStreamOutput]], None]:  # noqa: UP006
+            def fcallback(delta_outputs: List[RequestStreamOutput]):  # noqa: UP006
+                for delta_output in delta_outputs:
+                    request_id, stream_outputs = delta_output.unpack()
+                    assert len(stream_outputs) == 1
+                    if stream_outputs[0].finish_reason is not None:
+                        print(f"Request {request_id} finished at step {self.timer}.")
+                    outputs[int(request_id)] += stream_outputs[0].delta_token_ids
+                    finish_time[int(request_id)] = self.timer
+
+            return fcallback
+
+        def step(self) -> None:
+            self.timer += 1
+
+    # Create engine
+    timer = CallbackTimer()
+    engine = SyncMLCEngine(
+        model=model,
+        mode="server",
+        request_stream_callback=timer.callback_getter(),
+        engine_config=EngineConfig(prefill_mode="hybrid"),
+    )
+
+    # Create requests
+    requests = create_requests(
+        engine,
+        num_requests,
+        temperature=temperature,
+        repetition_penalty=repetition_penalty,
+        max_tokens_low=max_tokens,
+        max_tokens_high=max_tokens + 1,
+    )
+
+    # Add all requests to engine step by step
+    for step, request in enumerate(requests):
+        engine.add_request(request)
+        timer.step()
+        assert timer.timer == step
+        engine.step()
+
+    # Run steps
+    for step in range(max_tokens):
+        timer.step()
+        assert timer.timer == step + num_requests
+        engine.step()
+
+    for req_id, (request, output, fin_time) in enumerate(zip(requests, outputs, finish_time)):
+        print(f"Prompt {req_id}: {request.inputs[0]}")
+        print(f"Output {req_id}:{engine.tokenizer.decode(output)}\n")
+        assert fin_time == req_id + request.generation_config.max_tokens - 1, (
+            f"finish time = {fin_time}, max tokens = {req_id + request.generation_config.max_tokens - 1}"  # noqa: E501
+        )
+
+
+if __name__ == "__main__":
+    test_engine_basic()
+    test_engine_continuous_batching_1()
+    test_engine_continuous_batching_2()
+    test_engine_continuous_batching_3()
+    test_engine_generate()
+    test_engine_hybrid_prefill()