chore: import upstream snapshot with attribution

tests/python/compiler_pass/test_fuse_ft_dequantize_matmul_epilogue.py

@@ -0,0 +1,341 @@
+import tvm
+from tvm.ir import assert_structural_equal
+from tvm.script import ir as I
+from tvm.script import relax as R
+
+from mlc_llm.compiler_pass.fuse_ft_dequantize_matmul_epilogue import (
+    FuseFTDequantizeEpilogue,
+)
+
+
+def test_fuse_bias():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+            bias: R.Tensor((1, 1, 1024), "float16"),
+        ):
+            with R.dataflow():
+                lv1 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        "identity",
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                lv2 = R.add(lv1, bias)
+                R.output(lv2)
+            return lv2
+
+    @I.ir_module
+    class After:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+            bias: R.Tensor((1, 1, 1024), "float16"),
+        ) -> R.Tensor((1, 1, 1024), "float16"):
+            with R.dataflow():
+                lv2 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int_bias",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        bias,
+                        R.str("identity"),
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                        R.prim_value(0),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                R.output(lv2)
+            return lv2
+
+    seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()])
+    mod = seq(Before)
+    assert_structural_equal(mod, After)
+
+
+def test_fuse_activation():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+        ):
+            with R.dataflow():
+                lv1 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        "identity",
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                lv2 = R.nn.silu(lv1)
+                R.output(lv2)
+            return lv2
+
+    @I.ir_module
+    class After:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+        ) -> R.Tensor((1, 1, 1024), "float16"):
+            with R.dataflow():
+                lv2 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        R.str("silu"),
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                R.output(lv2)
+            return lv2
+
+    seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()])
+    mod = seq(Before)
+    assert_structural_equal(mod, After)
+
+
+def test_fuse_bias_activation():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+            bias: R.Tensor((1, 1, 1024), "float16"),
+        ):
+            with R.dataflow():
+                lv1 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        "identity",
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                lv2 = R.add(lv1, bias)
+                lv3 = R.nn.relu(lv2)
+                R.output(lv3)
+            return lv3
+
+    @I.ir_module
+    class After:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+            bias: R.Tensor((1, 1, 1024), "float16"),
+        ) -> R.Tensor((1, 1, 1024), "float16"):
+            with R.dataflow():
+                lv2 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int_bias",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        bias,
+                        R.str("relu"),
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                        R.prim_value(0),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                R.output(lv2)
+            return lv2
+
+    seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()])
+    mod = seq(Before)
+    assert_structural_equal(mod, After)
+
+
+def test_fuse_residual_binary():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+            bias: R.Tensor((1, 1, 1024), "float16"),
+            residual: R.Tensor((1, 1, 1024), "float16"),
+        ):
+            with R.dataflow():
+                lv1 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        "identity",
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                lv2 = R.add(lv1, bias)
+                lv3 = R.nn.relu(lv2)
+                lv4 = R.multiply(lv3, residual)
+                R.output(lv4)
+            return lv4
+
+    @I.ir_module
+    class After:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+            bias: R.Tensor((1, 1, 1024), "float16"),
+            residual: R.Tensor((1, 1, 1024), "float16"),
+        ) -> R.Tensor((1, 1, 1024), "float16"):
+            with R.dataflow():
+                lv2 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int_bias_residual",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        bias,
+                        residual,
+                        R.str("relu"),
+                        R.str("multiply"),
+                        R.str("identity"),
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                R.output(lv2)
+            return lv2
+
+    seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()])
+    mod = seq(Before)
+    assert_structural_equal(mod, After)
+
+
+def test_fuse_residual_unary():
+    @I.ir_module
+    class Before:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+            bias: R.Tensor((1, 1, 1024), "float16"),
+            residual: R.Tensor((1, 1, 1024), "float16"),
+        ):
+            with R.dataflow():
+                lv1 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        "identity",
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                lv2 = R.add(lv1, bias)
+                lv3 = R.nn.relu(lv2)
+                lv4 = R.add(lv3, residual)
+                lv5 = R.nn.gelu(lv4)
+                R.output(lv5)
+            return lv5
+
+    @I.ir_module
+    class After:
+        @R.function
+        def main(
+            x: R.Tensor((1, 1, 4096), "float16"),
+            weight: R.Tensor((4096, 512), "int8"),
+            scale: R.Tensor((1, 1024), "float16"),
+            bias: R.Tensor((1, 1, 1024), "float16"),
+            residual: R.Tensor((1, 1, 1024), "float16"),
+        ) -> R.Tensor((1, 1, 1024), "float16"):
+            with R.dataflow():
+                lv2 = R.call_dps_packed(
+                    "fastertransformer.gemm_fp16_int_bias_residual",
+                    (
+                        x,
+                        weight,
+                        scale,
+                        bias,
+                        residual,
+                        R.str("relu"),
+                        R.str("plus"),
+                        R.str("gelu"),
+                        R.prim_value(1),
+                        R.prim_value(1024),
+                        R.prim_value(4096),
+                        R.prim_value(4096),
+                    ),
+                    out_sinfo=R.Tensor((1, 1, 1024), "float16"),
+                )
+                R.output(lv2)
+            return lv2
+
+    seq = tvm.transform.Sequential([FuseFTDequantizeEpilogue()])
+    mod = seq(Before)
+    assert_structural_equal(mod, After)
+
+
+if __name__ == "__main__":
+    test_fuse_bias()
+    test_fuse_activation()
+    test_fuse_bias_activation()
+    test_fuse_residual_binary()
+    test_fuse_residual_unary()

tests/python/conftest.py

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#   http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+
+def pytest_configure(config):
+    """Register markers"""
+    config.addinivalue_line(
+        "markers",
+        "unittest: unittests for modules, do not require GPU, usually run fast",
+    )
+    config.addinivalue_line("markers", "op_correctness: unittest for op corectness, requires GPU")
+    config.addinivalue_line(
+        "markers",
+        (
+            "engine: testing engine feature functionalities, requires model and GPU, "
+            "note: for most request related tests, use endpoint test instead."
+        ),
+    )
+    config.addinivalue_line(
+        "markers",
+        (
+            "endpoint: sending requests to a global endpoint fixture(can be an rest or API), "
+            "tests compatibilities of API behaviors"
+        ),
+    )
+    config.addinivalue_line(
+        "markers",
+        "uncategorized: this test is not yet categorized, team should work to categorize it",
+    )

tests/python/conversation_template/test_conversation_protocol.py

@@ -0,0 +1,83 @@
+import pytest
+
+from mlc_llm.conversation_template import ConvTemplateRegistry
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+
+def get_conv_templates():
+    return [
+        "llama-3",
+        "llama-2",
+        "mistral_default",
+        "gorilla",
+        "gorilla-openfunctions-v2",
+        "chatml",
+        "phi-2",
+        "codellama_completion",
+        "codellama_instruct",
+        "rwkv_world",
+    ]
+
+
+@pytest.mark.parametrize("conv_template_name", get_conv_templates())
+def test_json(conv_template_name):
+    template = ConvTemplateRegistry.get_conv_template(conv_template_name)
+    j = template.to_json_dict()
+    template_parsed = Conversation.from_json_dict(j)
+    assert template == template_parsed
+
+
+@pytest.mark.parametrize("conv_template_name", get_conv_templates())
+def test_prompt(conv_template_name):
+    conversation = ConvTemplateRegistry.get_conv_template(conv_template_name)
+    user_msg = "test1"
+    assistant_msg = "test2"
+    prompt = "test3"
+
+    expected_user_msg = (
+        conversation.role_templates["user"]
+        .replace(MessagePlaceholders.USER.value, user_msg)
+        .replace(MessagePlaceholders.FUNCTION.value, "")
+    )
+
+    expected_prompt = (
+        conversation.role_templates["user"]
+        .replace(MessagePlaceholders.USER.value, prompt)
+        .replace(MessagePlaceholders.FUNCTION.value, "")
+    )
+
+    conversation.messages.append(("user", user_msg))
+    conversation.messages.append(("assistant", assistant_msg))
+    conversation.messages.append(("user", prompt))
+    conversation.messages.append(("assistant", None))
+    res = conversation.as_prompt()
+
+    system_msg = conversation.system_template.replace(
+        MessagePlaceholders.SYSTEM.value, conversation.system_message
+    )
+    expected_final_prompt = (
+        system_msg
+        + (conversation.seps[0] if system_msg != "" else "")
+        + (
+            conversation.roles["user"] + conversation.role_content_sep
+            if conversation.add_role_after_system_message
+            else ""
+        )
+        + expected_user_msg
+        + conversation.seps[0 % len(conversation.seps)]
+        + conversation.roles["assistant"]
+        + conversation.role_content_sep
+        + assistant_msg
+        + conversation.seps[1 % len(conversation.seps)]
+        + conversation.roles["user"]
+        + conversation.role_content_sep
+        + expected_prompt
+        + conversation.seps[0 % len(conversation.seps)]
+        + conversation.roles["assistant"]
+        + conversation.role_empty_sep
+    )
+    assert res == expected_final_prompt
+
+
+if __name__ == "__main__":
+    test_json("llama-3")

tests/python/conversation_template/test_llama_template.py

@@ -0,0 +1,43 @@
+import pytest
+
+from mlc_llm.conversation_template import ConvTemplateRegistry
+
+pytestmark = [pytest.mark.runtime_unittest]
+
+
+# From the official Llama-3 example:
+# https://llama.meta.com/docs/model-cards-and-prompt-formats/meta-llama-3/
+def test_llama3_prompt():
+    conversation = ConvTemplateRegistry.get_conv_template("llama-3")
+    system_msg = "You are a helpful AI assistant for travel tips and recommendations"
+    user_msg1 = "What is France's capital?"
+    assistant_msg1 = "Bonjour! The capital of France is Paris!"
+    user_msg2 = "What can I do there?"
+    assistant_msg2 = "Paris, the City of Light, offers a romantic getaway with must-see attractions like the Eiffel Tower and Louvre Museum, romantic experiences like river cruises and charming neighborhoods, and delicious food and drink options, with helpful tips for making the most of your trip."  # noqa: E501
+    prompt = "Give me a detailed list of the attractions I should visit, and time it takes in each one, to plan my trip accordingly."  # noqa: E501
+
+    conversation.system_message = system_msg
+    conversation.messages.append(("user", user_msg1))
+    conversation.messages.append(("assistant", assistant_msg1))
+    conversation.messages.append(("user", user_msg2))
+    conversation.messages.append(("assistant", assistant_msg2))
+    conversation.messages.append(("user", prompt))
+    conversation.messages.append(("assistant", None))
+    res = conversation.as_prompt()
+
+    expected = (
+        "<|start_header_id|>system<|end_header_id|>\n\n"
+        "You are a helpful AI assistant for travel tips and recommendations<|eot_id|>\n"
+        "<|start_header_id|>user<|end_header_id|>\n\n"
+        "What is France's capital?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        "Bonjour! The capital of France is Paris!<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n"  # noqa: E501
+        "What can I do there?<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+        "Paris, the City of Light, offers a romantic getaway with must-see attractions like the Eiffel Tower and Louvre Museum, romantic experiences like river cruises and charming neighborhoods, and delicious food and drink options, with helpful tips for making the most of your trip.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n"  # noqa: E501
+        "Give me a detailed list of the attractions I should visit, and time it takes in each one, to plan my trip accordingly.<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"  # noqa: E501
+    )
+
+    assert res[0] == expected
+
+
+if __name__ == "__main__":
+    test_llama3_prompt()

tests/python/integration/test_model_compile.py

@@ -0,0 +1,169 @@
+import concurrent.futures as cf
+import os
+import shlex
+import subprocess
+import sys
+import tempfile
+from itertools import product
+
+import tvm
+
+from mlc_llm.model import MODEL_PRESETS
+from mlc_llm.model import MODELS as SUPPORTED_MODELS
+from mlc_llm.quantization import QUANTIZATION as SUPPORTED_QUANTS
+from mlc_llm.support.constants import MLC_TEMP_DIR
+
+OPT_LEVEL = "O2"
+DEVICE2TARGET = {
+    "cuda": {
+        "kind": "cuda",
+        "arch": "sm_86",
+        "max_threads_per_block": 1024,
+        "max_num_threads": 1024,
+        "max_shared_memory_per_block": 49152,
+        "thread_warp_size": 32,
+    },
+    "rocm": {
+        "kind": "rocm",
+        "mtriple": "amdgcn-amd-amdhsa-hcc",
+        "mcpu": "gfx1100",
+        "thread_warp_size": 32,
+        "max_threads_per_block": 1024,
+        "max_num_threads": 256,
+        "max_shared_memory_per_block": 65536,
+    },
+    "vulkan": {
+        "kind": "vulkan",
+        "max_threads_per_block": 1024,
+        "max_num_threads": 256,
+        "max_shared_memory_per_block": 32768,
+        "thread_warp_size": 1,
+        "supports_float32": 1,
+        "supports_float16": 1,
+        "supports_int64": 1,
+        "supports_int32": 1,
+        "supports_int16": 1,
+        "supports_int8": 1,
+        "supports_16bit_buffer": 1,
+    },
+    "metal": "metal",
+    "wasm": "webgpu",
+    "android": "android",
+    "ios": "iphone",
+}
+DEVICE2SUFFIX = {
+    "cuda": "so",
+    "rocm": "so",
+    "vulkan": "so",
+    "metal": "dylib",
+    "wasm": "wasm",
+    "android": "tar",
+    "ios": "tar",
+}
+MODELS = list(MODEL_PRESETS.keys())
+QUANTS = [  # TODO(@junrushao): use `list(mlc_llm.quantization.QUANTIZATION.keys())`
+    "q0f16",
+    "q0f32",
+    "q3f16_1",
+    "q4f16_1",
+    "q4f32_1",
+    "q4f16_ft",
+]
+TENSOR_PARALLEL_SHARDS = [
+    1,
+]
+
+
+def run_command(log_file, cmd):
+    with open(log_file, "w", encoding="utf-8") as file:
+        subprocess.check_call(
+            cmd,
+            stdout=file,
+            stderr=subprocess.STDOUT,
+        )
+
+
+def test_model_compile():
+    device = sys.argv[1]
+    num_workers = int(sys.argv[2])
+    target = DEVICE2TARGET[device]
+    if not isinstance(target, str):
+        target = str(tvm.target.Target(target))
+    suffix = DEVICE2SUFFIX[device]
+
+    passed_cmds = []
+    failed_cmds = []
+    with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir:
+        with cf.ProcessPoolExecutor(max_workers=num_workers) as executor:
+            log_files = []
+            cmds = []
+            futures = []
+            for idx, (model, quant, tp_shard) in enumerate(
+                product(
+                    MODELS,
+                    QUANTS,
+                    TENSOR_PARALLEL_SHARDS,
+                )
+            ):
+                if (
+                    SUPPORTED_QUANTS[quant].kind
+                    not in SUPPORTED_MODELS[MODEL_PRESETS[model]["model_type"]].quantize
+                ):
+                    continue
+                if not target.startswith("cuda") and quant == "q4f16_ft":
+                    # FasterTransformer only works with cuda
+                    continue
+                if "deepseek_v2" in model and "32" in quant:
+                    # Skip f32 for deepseek v2 model for now.
+                    continue
+                log_file = os.path.join(tmp_dir, f"lib{idx}.log")
+                cmd = [
+                    sys.executable,
+                    "-m",
+                    "mlc_llm",
+                    "compile",
+                    model,
+                    "--quantization",
+                    quant,
+                    "--overrides",
+                    f"tensor_parallel_shards={tp_shard}",
+                    "--device",
+                    target,
+                    "--opt",
+                    OPT_LEVEL,
+                    "-o",
+                    os.path.join(tmp_dir, f"lib{idx}.{suffix}"),
+                ]
+                future = executor.submit(run_command, log_file, cmd)
+                log_files.append(log_file)
+                cmds.append(cmd)
+                futures.append(future)
+            for log_file, cmd, future in zip(log_files, cmds, futures):
+                cmd = shlex.join(cmd)
+                try:
+                    future.result()
+                    passed_cmds.append(cmd)
+                    print(f"[PASS] {cmd}")
+                except Exception:
+                    failed_cmds.append(cmd)
+                    print("-------------------------------")
+                    print(f"[FAIL] {cmd}")
+                    with open(log_file, encoding="utf-8") as file:
+                        print(file.read())
+                    print("-------------------------------")
+    print("-------------------------------")
+    print(f"Total {len(passed_cmds)} passed, {len(failed_cmds)} failed.")
+    print("-------------------------------")
+    print("Passed commands:")
+    for cmd in passed_cmds:
+        print(cmd)
+    if failed_cmds:
+        print("-------------------------------")
+        print("Failed commands:")
+        for cmd in failed_cmds:
+            print(cmd)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    test_model_compile()

tests/python/json_ffi/test_json_ffi_engine.py

@@ -0,0 +1,197 @@
+import json
+from typing import Dict, List, Optional  # noqa: UP035
+
+import pytest
+from pydantic import BaseModel
+
+from mlc_llm.json_ffi import JSONFFIEngine
+from mlc_llm.testing import require_test_model
+
+# test category "engine_feature"
+pytestmark = [pytest.mark.engine_feature]
+
+
+chat_completion_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
+]
+
+function_calling_prompts = [
+    "What is the temperature in Pittsburgh, PA?",
+    "What is the temperature in Tokyo, JP?",
+    "What is the temperature in Pittsburgh, PA and Tokyo, JP?",
+]
+
+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"],
+            },
+        },
+    }
+]
+
+
+def run_chat_completion(
+    engine: JSONFFIEngine,
+    model: str,
+    prompts: List[str] = chat_completion_prompts,  # noqa: UP006
+    tools: Optional[List[Dict]] = None,  # noqa: UP006
+):
+    num_requests = 2
+    max_tokens = 64
+    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"chat completion for request {rid}")
+        for response in engine.chat.completions.create(
+            messages=[{"role": "user", "content": [{"type": "text", "text": prompts[rid]}]}],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            tools=tools,
+        ):
+            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")
+
+
+def run_json_schema_function_calling(
+    engine: JSONFFIEngine,
+    model: str,
+    prompts: List[str] = function_calling_prompts,  # noqa: UP006
+    tools: Optional[List[Dict]] = None,  # noqa: UP006
+):
+    num_requests = 2
+    max_tokens = 64
+    n = 1
+    output_texts: List[List[str]] = [["" for _ in range(n)] for _ in range(num_requests)]  # noqa: UP006
+
+    class ToolCall(BaseModel):
+        name: str
+        arguments: Dict[str, str]  # noqa: UP006
+
+    class Schema(BaseModel):
+        tool_calls: List[ToolCall]  # noqa: UP006
+
+    schema_str = json.dumps(Schema.model_json_schema())
+    print("Schema str", schema_str)
+
+    for rid in range(num_requests):
+        print(f"chat completion for request {rid}")
+        for response in engine.chat.completions.create(
+            messages=[
+                {
+                    "role": "system",
+                    "content": "You are a function calling AI model. You are provided with function signatures within "  # noqa: E501
+                    "<tools></tools> XML tags. You may call one or more functions to assist with the user query. Don't make "  # noqa: E501
+                    f"assumptions about what values to plug into functions. Here are the available tools: <tools> {json.dumps(tools)} </tools> "  # noqa: E501
+                    "Do not stop calling functions until the task has been accomplished or you've reached max iteration of 10. "  # noqa: E501
+                    "Calling multiple functions at once can overload the system and increase cost so call one function at a time please. "  # noqa: E501
+                    "If you plan to continue with analysis, always call another function. Return a valid json object (using double "  # noqa: E501
+                    f"quotes) in the following schema: {schema_str}",
+                },
+                {"role": "user", "content": [{"type": "text", "text": prompts[rid]}]},
+            ],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            response_format={"type": "json_object", "schema": schema_str},
+        ):
+            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")
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_chat_completion(model):
+    # Create engine.
+    engine = JSONFFIEngine(model)
+
+    run_chat_completion(engine, model)
+
+    # Test malformed requests.
+    for response in engine._raw_chat_completion(
+        "malformed_string", include_usage=False, request_id="123"
+    ):
+        assert len(response.choices) == 1
+        assert response.choices[0].finish_reason == "error"
+
+    engine.terminate()
+
+
+@require_test_model("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_reload_reset_unload(model):
+    # Create engine.
+    engine = JSONFFIEngine(model)
+
+    # Run chat completion before and after reload/reset.
+    run_chat_completion(engine, model)
+    engine._test_reload()
+    run_chat_completion(engine, model)
+    engine._test_reset()
+    run_chat_completion(engine, model)
+    engine._test_unload()
+
+    engine.terminate()
+
+
+@require_test_model("Hermes-2-Pro-Mistral-7B-q4f16_1-MLC")
+def test_json_schema_with_system_prompt(model):
+    engine = JSONFFIEngine(model)
+
+    # run function calling
+    run_json_schema_function_calling(engine, model, function_calling_prompts, tools)
+
+    engine.terminate()
+
+
+if __name__ == "__main__":
+    test_chat_completion()
+    test_reload_reset_unload()
+    test_json_schema_with_system_prompt()

tests/python/json_ffi/test_json_ffi_engine_image.py

@@ -0,0 +1,92 @@
+import base64
+from typing import Dict, List, Optional  # noqa: UP035
+
+import requests
+
+from mlc_llm.json_ffi import JSONFFIEngine
+from mlc_llm.testing import require_test_model
+
+
+def base64_encode_image(url: str) -> str:
+    response = requests.get(url)
+    response.raise_for_status()  # Ensure we got a successful response
+    image_data = base64.b64encode(response.content)
+    image_data_str = image_data.decode("utf-8")
+    data_url = f"data:image/jpeg;base64,{image_data_str}"
+    return data_url
+
+
+image_prompts = [
+    [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "image_url",
+                    "image_url": f"{base64_encode_image('https://llava-vl.github.io/static/images/view.jpg')}",
+                },
+                {"type": "text", "text": "What does the image represent?"},
+            ],
+        }
+    ]
+]
+
+
+def run_chat_completion(
+    engine: JSONFFIEngine,
+    model: str,
+    prompts: List[List[Dict]] = image_prompts,  # noqa: UP006
+    tools: Optional[List[Dict]] = None,  # noqa: UP006
+):
+    num_requests = 1
+    max_tokens = 64
+    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"chat completion for request {rid}")
+        for response in engine.chat.completions.create(
+            messages=prompts[rid],
+            model=model,
+            max_tokens=max_tokens,
+            n=n,
+            request_id=str(rid),
+            tools=tools,
+        ):
+            for choice in response.choices:
+                assert choice.delta.role == "assistant"
+                assert isinstance(choice.delta.content[0], Dict)  # noqa: UP006
+                assert choice.delta.content[0]["type"] == "text"
+                output_texts[rid][choice.index] += choice.delta.content[0]["text"]
+
+    # 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")
+
+
+@require_test_model("llava-1.5-7b-hf-q4f16_1-MLC")
+def test_chat_completion():
+    # Create engine.
+    engine = JSONFFIEngine(
+        model,  # noqa: F821
+        max_total_sequence_length=1024,
+    )
+
+    run_chat_completion(engine, model)  # noqa: F821
+
+    # Test malformed requests.
+    for response in engine._raw_chat_completion("malformed_string", n=1, request_id="123"):
+        assert len(response.choices) == 1
+        assert response.choices[0].finish_reason == "error"
+
+    engine.terminate()
+
+
+if __name__ == "__main__":
+    test_chat_completion()

tests/python/json_ffi/test_json_ffi_engine_mock.py

@@ -0,0 +1,105 @@
+import json
+
+import pytest
+import tvm
+
+from mlc_llm.json_ffi import JSONFFIEngine
+from mlc_llm.testing import require_test_model
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+def check_error_handling(engine, expect_str, **params):
+    """Check error handling in raw completion API"""
+    body = {
+        "messages": [{"role": "user", "content": "hello"}],
+        "stream_options": {"include_usage": True},
+    }
+    body.update(params)
+
+    for response in engine._raw_chat_completion(
+        json.dumps(body), include_usage=False, request_id="123"
+    ):
+        if response.choices[0].finish_reason is not None:
+            break
+    if response.choices[0].finish_reason != "error":
+        raise RuntimeError(f"expect the request {params} to hit an error")
+
+    if expect_str not in response.choices[0].delta.content:
+        raise RuntimeError(
+            f"expect '{expect_str}' in error msg, but get '{response.choices[0].delta.content}'"
+        )
+
+
+# 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_chat_completion_misuse(model: str):
+    engine = JSONFFIEngine(model, tvm.cpu(), model_lib="mock://echo")
+    # Test malformed requests.
+    for response in engine._raw_chat_completion(
+        "malformed_string", include_usage=False, request_id="123"
+    ):
+        assert len(response.choices) == 1
+        assert response.choices[0].finish_reason == "error"
+    # check parameters
+    check_error_handling(engine, "should be non-negative", temperature=-1)
+    check_error_handling(engine, "in range [0, 1]", top_p=100)
+    check_error_handling(engine, "frequency_penalty", frequency_penalty=100)
+
+
+def check_normal_param_passing(engine):
+    json_schema = """
+    {"properties": {"result": {"items": {"type": "Integer"}, "title": "Result", "type": "array"}},
+      "required": ["result"], "title": "Output", "type": "object"}
+    """
+    param_dict = {
+        "top_p": 0.6,
+        "temperature": 0.8,
+        "frequency_penalty": 0.1,
+        "presence_penalty": 0.1,
+    }
+    usage = None
+    for response in engine.chat.completions.create(
+        messages=[{"role": "user", "content": "hello"}],
+        stream=True,
+        stream_options={"include_usage": True},
+        response_format={"type": "json_object", "schema": json_schema},
+        **param_dict,
+    ):
+        if response.usage is not None:
+            usage = response.usage
+
+    # echo mock will echo back the generation config
+    for k, v in param_dict.items():
+        assert usage.extra[k] == v, f"{k} mismatch"
+    assert "response_format" in usage.extra
+    assert usage.extra["response_format"]["type"] == "json_object"
+    assert "schema" in usage.extra["response_format"]
+
+
+def check_n_generation(engine):
+    hit_set = set()
+    for response in engine.chat.completions.create(
+        messages=[{"role": "user", "content": "hello"}],
+        stream=True,
+        stream_options={"include_usage": True},
+        n=3,
+    ):
+        for choice in response.choices:
+            hit_set.add(choice.index)
+    for i in range(3):
+        assert i in hit_set, f"{i} not in n generation"
+
+
+@require_test_model("Llama-3-8B-Instruct-q4f16_1-MLC")
+def test_chat_completion_api(model: str):
+    engine = JSONFFIEngine(model, tvm.cpu(), model_lib="mock://echo")
+    check_normal_param_passing(engine)
+    check_n_generation(engine)
+
+
+if __name__ == "__main__":
+    test_chat_completion_api()
+    test_chat_completion_misuse()

tests/python/loader/test_awq.py

@@ -0,0 +1,39 @@
+from pathlib import Path
+from typing import Union
+
+import pytest
+import tvm
+
+from mlc_llm.loader import HuggingFaceLoader
+from mlc_llm.model import MODEL_PRESETS, MODELS
+from mlc_llm.quantization import QUANTIZATION
+from mlc_llm.support import logging, tqdm
+
+logging.enable_logging()
+
+
+@pytest.mark.parametrize(
+    "param_path",
+    [
+        "./dist/models/llama-2-7b-w4-g128-awq.pt",
+        "./dist/models/Llama-2-7B-AWQ/model.safetensors",
+    ],
+)
+def test_load_llama(param_path: Union[str, Path]):
+    path_params = Path(param_path)
+
+    model = MODELS["llama"]
+    quantization = QUANTIZATION["q4f16_awq"]
+    config = model.config.from_dict(MODEL_PRESETS["llama2_7b"])
+    loader = HuggingFaceLoader(
+        path=path_params,
+        extern_param_map=model.source["awq"](config, quantization),
+    )
+    with tqdm.redirect():
+        for _name, _param in loader.load(tvm.device("cpu")):
+            ...
+
+
+if __name__ == "__main__":
+    test_load_llama(param_path="./dist/models/llama-2-7b-w4-g128-awq.pt")
+    test_load_llama(param_path="./dist/models/Llama-2-7B-AWQ/model.safetensors")

tests/python/loader/test_huggingface.py

@@ -0,0 +1,68 @@
+from pathlib import Path
+from typing import Union
+
+import pytest
+import tvm
+
+from mlc_llm.loader import HuggingFaceLoader
+from mlc_llm.model import MODELS
+from mlc_llm.support import logging, tqdm
+
+logging.enable_logging()
+
+
+@pytest.mark.parametrize(
+    "base_path",
+    [
+        "./dist/models/Llama-2-7b-hf",
+        "./dist/models/Llama-2-13b-hf",
+        "./dist/models/Llama-2-70b-hf",
+    ],
+)
+def test_load_torch_llama(base_path: Union[str, Path]):
+    base_path = Path(base_path)
+    path_config = base_path / "config.json"
+    path_params = base_path / "pytorch_model.bin.index.json"
+
+    model = MODELS["llama"]
+    config = model.config.from_file(path_config)
+    loader = HuggingFaceLoader(
+        path=path_params,
+        extern_param_map=model.source["huggingface-torch"](config, None),
+    )
+    with tqdm.redirect():
+        for _name, _param in loader.load(device=tvm.device("cpu")):
+            return  # To reduce the time of the test
+
+
+@pytest.mark.parametrize(
+    "base_path",
+    [
+        "./dist/models/Llama-2-7b-hf",
+        "./dist/models/Llama-2-13b-hf",
+        "./dist/models/Llama-2-70b-hf",
+    ],
+)
+def test_load_safetensor_llama(base_path: Union[str, Path]):
+    base_path = Path(base_path)
+    path_config = base_path / "config.json"
+    path_params = base_path / "model.safetensors.index.json"
+
+    model = MODELS["llama"]
+    config = model.config.from_file(path_config)
+    loader = HuggingFaceLoader(
+        path=path_params,
+        extern_param_map=model.source["huggingface-safetensor"](config, None),
+    )
+    with tqdm.redirect():
+        for _name, _param in loader.load(device=tvm.device("cpu")):
+            return  # To reduce the time of the test
+
+
+if __name__ == "__main__":
+    test_load_torch_llama(base_path="./dist/models/Llama-2-7b-hf")
+    test_load_torch_llama(base_path="./dist/models/Llama-2-13b-hf")
+    test_load_torch_llama(base_path="./dist/models/Llama-2-70b-hf")
+    test_load_safetensor_llama(base_path="./dist/models/Llama-2-7b-hf")
+    test_load_safetensor_llama(base_path="./dist/models/Llama-2-13b-hf")
+    test_load_safetensor_llama(base_path="./dist/models/Llama-2-70b-hf")

tests/python/model/test_gemma3.py

@@ -0,0 +1,70 @@
+"""Unit tests for Gemma3 model architecture."""
+
+import pytest
+
+from mlc_llm.model import MODEL_PRESETS, MODELS
+
+
+def test_gemma3_model_registered():
+    """Verify Gemma3 model is in the registry."""
+    assert "gemma3" in MODELS, "gemma3 should be registered in MODELS"
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    [
+        "gemma3_2b",
+        "gemma3_9b",
+    ],
+)
+def test_gemma3_creation(model_name: str):
+    """Test Gemma3 model creation and export to TVM IR.
+
+    Verifies:
+    - Config can be loaded from preset
+    - Model instance can be created
+    - Model exports to TVM IR successfully
+    - Named parameters are extracted
+    """
+    model_info = MODELS["gemma3"]
+    config = model_info.config.from_dict(MODEL_PRESETS[model_name])
+    model = model_info.model(config)
+    mod, named_params = model.export_tvm(
+        spec=model.get_default_spec(),
+    )
+
+    # Verify export succeeded
+    assert mod is not None
+    assert len(named_params) > 0
+
+    # Optional: show module structure
+    mod.show(black_format=False)
+
+    # Print parameters for debugging
+    for name, param in named_params:
+        print(name, param.shape, param.dtype)
+
+
+def test_gemma3_config_validation():
+    """Test Gemma3 configuration has required fields."""
+    model_info = MODELS["gemma3"]
+    config = model_info.config.from_dict(MODEL_PRESETS["gemma3_2b"])
+
+    # Check required config parameters
+    assert hasattr(config, "hidden_size") and config.hidden_size > 0
+    assert hasattr(config, "num_hidden_layers") and config.num_hidden_layers > 0
+    assert hasattr(config, "num_attention_heads") and config.num_attention_heads > 0
+    assert hasattr(config, "vocab_size") and config.vocab_size > 0
+
+    print(
+        f"Gemma3 Config: hidden_size={config.hidden_size}, "
+        f"layers={config.num_hidden_layers}, "
+        f"heads={config.num_attention_heads}, "
+        f"vocab={config.vocab_size}"
+    )
+
+
+if __name__ == "__main__":
+    # Allow running tests directly
+    test_gemma3_creation("gemma3_2b")
+    test_gemma3_creation("gemma3_9b")

tests/python/model/test_gpt2.py

@@ -0,0 +1,20 @@
+import pytest
+
+from mlc_llm.model import MODEL_PRESETS, MODELS
+
+
+@pytest.mark.parametrize("model_name", ["gpt2"])
+def test_gpt2_creation(model_name: str):
+    model_info = MODELS["gpt2"]
+    config = model_info.config.from_dict(MODEL_PRESETS[model_name])
+    model = model_info.model(config)
+    mod, named_params = model.export_tvm(
+        spec=model.get_default_spec(),
+    )
+    mod.show(black_format=False)
+    for name, param in named_params:
+        print(name, param.shape, param.dtype)
+
+
+if __name__ == "__main__":
+    test_gpt2_creation("gpt2")

tests/python/model/test_gptNeox.py

@@ -0,0 +1,20 @@
+import pytest
+
+from mlc_llm.model import MODEL_PRESETS, MODELS
+
+
+@pytest.mark.parametrize("model_name", ["redpajama_3b_v1"])
+def test_mistral_creation(model_name: str):
+    model_info = MODELS["gpt_neox"]
+    config = model_info.config.from_dict(MODEL_PRESETS[model_name])
+    model = model_info.model(config)
+    mod, named_params = model.export_tvm(
+        spec=model.get_default_spec(),
+    )
+    mod.show(black_format=False)
+    for name, param in named_params:
+        print(name, param.shape, param.dtype)
+
+
+if __name__ == "__main__":
+    test_mistral_creation("redpajama_3b_v1")

tests/python/model/test_kv_cache.py

@@ -0,0 +1,114 @@
+import tvm
+from tvm import tirx
+from tvm.relax.frontend.nn import core, modules, spec
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tirx as T
+
+from mlc_llm.nn.kv_cache import PagedKVCache, RopeMode
+
+# mypy: disable-error-code="attr-defined"
+
+
+def test_nn_module_paged_kv_cache():
+    # fmt: off
+    @I.ir_module
+    class Module:
+        @R.function
+        def create_paged_kv_cache(
+            max_batch_size: R.Shape(["max_batch_size_1"]),
+            max_total_seq_len: R.Shape(["max_total_seq_len_1"]),
+            prefill_chunk_size: R.Shape(["prefill_chunk_size_1"]),
+            page_size: R.Shape(["page_size_1"]),
+            support_sliding_window: R.Shape(["support_sliding_window_1"]),
+        ) -> R.Object:
+            max_batch_size_1 = T.int64()
+            max_total_seq_len_1 = T.int64()
+            prefill_chunk_size_1 = T.int64()
+            page_size_1 = T.int64()
+            support_sliding_window_1 = T.int64()
+            R.func_attr({"num_input": 5})
+            with R.dataflow():
+                paged_kv_cache: R.Object = R.call_pure_packed("mlc.create_paged_kv_cache_generic", R.shape([max_batch_size_1, max_total_seq_len_1, prefill_chunk_size_1, page_size_1, support_sliding_window_1]), R.prim_value(32), R.prim_value(32), R.prim_value(32), R.prim_value(128), R.prim_value(1), R.prim_value(1), R.prim_value(10000), R.prim_value(128), R.dtype("float16"), sinfo_args=(R.Object,))  # noqa: E501
+                gv1: R.Object = paged_kv_cache
+                R.output(gv1)
+            return gv1
+
+        @R.function
+        def forward(
+            cache: R.Object, qkv: R.Tensor((1, 100, 96, 128), dtype="float16")
+        ) -> R.Tensor((1, 100, 32, 128), dtype="float16"):
+            R.func_attr({"num_input": 2})
+            with R.dataflow():
+                reshape: R.Tensor((100, 96, 128), dtype="float16") = R.reshape(
+                    qkv, R.shape([100, 96, 128])
+                )
+                lv = R.call_dps_packed(
+                    "vm.builtin.attention_kv_cache_attention_with_fused_qkv",
+                    (cache, R.prim_value(0), R.prim_value(T.float32(1)), reshape),
+                    out_sinfo=R.Tensor((100, 32, 128), dtype="float16"),
+                )
+                reshape1: R.Tensor((1, 100, 32, 128), dtype="float16") = R.reshape(
+                    lv, R.shape([1, 100, 32, 128])
+                )
+                gv: R.Tensor((1, 100, 32, 128), dtype="float16") = reshape1
+                R.output(gv)
+            return gv
+    # fmt: on
+
+    class PagedKVCacheTest(modules.Module):
+        def forward(
+            self,
+            cache: PagedKVCache,
+            qkv: core.Tensor,
+        ) -> core.Tensor:
+            return cache.attention_with_fused_qkv(0, qkv, num_qo_heads=32, sm_scale=128**-0.5)
+
+        def create_paged_kv_cache(
+            self,
+            max_batch_size: tirx.Var,
+            max_total_seq_len: tirx.Var,
+            prefill_chunk_size: tirx.Var,
+            page_size: tirx.Var,
+            support_sliding_window: tirx.Var,
+        ) -> PagedKVCache:
+            return PagedKVCache.create_generic(
+                attn_kind="mha",
+                max_batch_size=max_batch_size,
+                max_total_seq_len=max_total_seq_len,
+                prefill_chunk_size=prefill_chunk_size,
+                page_size=page_size,
+                support_sliding_window=support_sliding_window,
+                num_hidden_layers=32,
+                num_attention_heads=32,
+                num_key_value_heads=32,
+                qk_head_dim=128,
+                v_head_dim=128,
+                rope_mode=RopeMode.NORMAL,
+                rope_scale=1,
+                rope_theta=10000,
+                rotary_dim=128,
+                dtype="float16",
+            )
+
+    export_results = PagedKVCacheTest().export_tvm(
+        spec={
+            "forward": {
+                "cache": spec.Object(object_type=PagedKVCache),
+                "qkv": spec.Tensor((1, 100, 96, 128), "float16"),
+            },
+            "create_paged_kv_cache": {
+                "max_batch_size": int,
+                "max_total_seq_len": int,
+                "prefill_chunk_size": int,
+                "page_size": int,
+                "support_sliding_window": int,
+            },
+        },
+    )
+    tvm_mod = export_results[0]
+    tvm.ir.assert_structural_equal(tvm_mod, Module, True)
+
+
+if __name__ == "__main__":
+    test_nn_module_paged_kv_cache()

tests/python/model/test_llama.py

@@ -0,0 +1,25 @@
+import pytest
+
+from mlc_llm.model import MODEL_PRESETS, MODELS
+
+
+@pytest.mark.parametrize(
+    "model_name", ["llama2_7b", "llama2_13b", "llama2_70b", "tinyllama_1b_chat_v1.0"]
+)
+def test_llama2_creation(model_name: str):
+    model_info = MODELS["llama"]
+    config = model_info.config.from_dict(MODEL_PRESETS[model_name])
+    model = model_info.model(config)
+    mod, named_params = model.export_tvm(
+        spec=model.get_default_spec(),
+    )
+    mod.show(black_format=False)
+    for name, param in named_params:
+        print(name, param.shape, param.dtype)
+
+
+if __name__ == "__main__":
+    test_llama2_creation("llama2_7b")
+    test_llama2_creation("llama2_13b")
+    test_llama2_creation("llama2_70b")
+    test_llama2_creation("tinyllama_1b_chat_v1")

tests/python/model/test_llama_quantization.py

@@ -0,0 +1,72 @@
+import pytest
+
+from mlc_llm.model import MODEL_PRESETS, MODELS
+from mlc_llm.quantization import QUANTIZATION
+from mlc_llm.quantization.group_quantization import (
+    GroupQuantizeEmbedding,
+    GroupQuantizeLinear,
+)
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    ["llama2_7b", "llama2_13b", "llama2_70b"],
+)
+@pytest.mark.parametrize(
+    "quant_name",
+    ["q3f16_1", "q4f16_1", "q4f32_1"],
+)
+def test_llama2_group_quantization(model_name: str, quant_name: str):
+    model_info = MODELS["llama"]
+    config = model_info.config.from_dict(MODEL_PRESETS[model_name])
+    model, quant_map = model_info.quantize["group-quant"](config, QUANTIZATION[quant_name])
+    assert "model.embed_tokens.weight" in quant_map.param_map
+    assert isinstance(
+        model.model.embed_tokens,
+        GroupQuantizeEmbedding,
+    )
+    assert "lm_head.weight" in quant_map.param_map
+    assert isinstance(model.lm_head, GroupQuantizeLinear)
+    for i in range(config.num_hidden_layers):
+        assert f"model.layers.{i}.self_attn.qkv_proj.weight" in quant_map.param_map
+        assert isinstance(
+            model.model.layers[i].self_attn.qkv_proj,
+            GroupQuantizeLinear,
+        )
+        assert f"model.layers.{i}.self_attn.o_proj.weight" in quant_map.param_map
+        assert isinstance(
+            model.model.layers[i].self_attn.o_proj,
+            GroupQuantizeLinear,
+        )
+        assert f"model.layers.{i}.mlp.gate_up_proj.weight" in quant_map.param_map
+        assert isinstance(
+            model.model.layers[i].mlp.gate_up_proj,
+            GroupQuantizeLinear,
+        )
+        assert f"model.layers.{i}.mlp.down_proj.weight" in quant_map.param_map
+        assert isinstance(
+            model.model.layers[i].mlp.down_proj,
+            GroupQuantizeLinear,
+        )
+
+
+@pytest.mark.parametrize(
+    "model_name",
+    ["llama2_7b", "llama2_13b", "llama2_70b"],
+)
+@pytest.mark.parametrize(
+    "quant_name",
+    ["q0f16", "q0f32"],
+)
+def test_llama2_no_quantization(model_name: str, quant_name: str):
+    model_info = MODELS["llama"]
+    config = model_info.config.from_dict(MODEL_PRESETS[model_name])
+    _, quant_map = model_info.quantize["no-quant"](config, QUANTIZATION[quant_name])
+    assert len(quant_map.param_map) == 0
+    assert len(quant_map.map_func) == 0
+
+
+if __name__ == "__main__":
+    test_llama2_group_quantization("llama2_7b", "q4f16_1")
+    test_llama2_group_quantization("llama2_13b", "q4f16_1")
+    test_llama2_group_quantization("llama2_70b", "q4f16_1")

tests/python/model/test_mistral.py

@@ -0,0 +1,20 @@
+import pytest
+
+from mlc_llm.model import MODEL_PRESETS, MODELS
+
+
+@pytest.mark.parametrize("model_name", ["mistral_7b"])
+def test_mistral_creation(model_name: str):
+    model_info = MODELS["mistral"]
+    config = model_info.config.from_dict(MODEL_PRESETS[model_name])
+    model = model_info.model(config)
+    mod, named_params = model.export_tvm(
+        spec=model.get_default_spec(),
+    )
+    mod.show(black_format=False)
+    for name, param in named_params:
+        print(name, param.shape, param.dtype)
+
+
+if __name__ == "__main__":
+    test_mistral_creation("mistral_7b")

tests/python/model/test_phi.py

@@ -0,0 +1,21 @@
+import pytest
+
+from mlc_llm.model import MODEL_PRESETS, MODELS
+
+
+@pytest.mark.parametrize("model_name", ["phi-1_5", "phi-2"])
+def test_phi_creation(model_name: str):
+    model_info = MODELS["phi-msft"]
+    config = model_info.config.from_dict(MODEL_PRESETS[model_name])
+    model = model_info.model(config)
+    mod, named_params = model.export_tvm(
+        spec=model.get_default_spec(),
+    )
+    mod.show(black_format=False)
+    for name, param in named_params:
+        print(name, param.shape, param.dtype)
+
+
+if __name__ == "__main__":
+    test_phi_creation("phi-1_5")
+    test_phi_creation("phi-2")

tests/python/model/test_qwen3_embedding.py

@@ -0,0 +1,147 @@
+import json
+import os
+
+import numpy as np
+import pytest
+import torch
+import tvm
+from safetensors import safe_open
+from transformers import AutoModel, AutoTokenizer
+from tvm import relax
+from tvm.contrib import tvmjs
+from tvm.runtime import ShapeTuple
+from tvm.runtime.vm import VirtualMachine
+
+MLC_QWEN3_EMB_HF_DIR = os.environ.get("MLC_QWEN3_EMB_HF_DIR")
+MLC_QWEN3_EMB_MODEL_DIR = os.environ.get("MLC_QWEN3_EMB_MODEL_DIR")
+MLC_QWEN3_EMB_MODEL_LIB = os.environ.get("MLC_QWEN3_EMB_MODEL_LIB")
+MLC_QWEN3_EMB_DEVICE = os.environ.get("MLC_QWEN3_EMB_DEVICE", "cuda")
+
+_skip = not all([MLC_QWEN3_EMB_HF_DIR, MLC_QWEN3_EMB_MODEL_DIR, MLC_QWEN3_EMB_MODEL_LIB])
+_skip_reason = (
+    "Set MLC_QWEN3_EMB_HF_DIR, MLC_QWEN3_EMB_MODEL_DIR, MLC_QWEN3_EMB_MODEL_LIB to run this test"
+)
+
+TEST_TEXTS = [
+    "What is machine learning?",
+    "CMU is Carnegie Mellon University",
+    "机器学习是人工智能的一个分支",
+    "量子コンピュータの基本原理を説明してください",
+    "머신러닝은 인공지능의 한 분야입니다.",
+    (
+        "Instruct: Given a web search query, retrieve relevant passages "
+        "that answer the query\nQuery: What is the capital of China?"
+    ),
+    (
+        "The Transformer architecture, introduced in the paper Attention Is All You Need, "
+        "revolutionized natural language processing by replacing recurrent layers with "
+        "self-attention mechanisms. This allows the model to process all positions in a "
+        "sequence simultaneously rather than sequentially, leading to significant improvements "
+        "in both training efficiency and the ability to capture long-range dependencies. "
+        "The key innovation is the multi-head attention mechanism, which allows the model "
+        "to jointly attend to information from different representation subspaces at "
+        "different positions."
+    ),
+    "Hello",
+    "def fibonacci(n): return n if n <= 1 else fibonacci(n-1) + fibonacci(n-2)",
+]
+
+
+def _load_embed_weight(hf_dir):
+    safetensor_files = [f for f in os.listdir(hf_dir) if f.endswith(".safetensors")]
+    for sf in safetensor_files:
+        with safe_open(os.path.join(hf_dir, sf), framework="pt", device="cpu") as f:
+            if "embed_tokens.weight" in f.keys():
+                return f.get_tensor("embed_tokens.weight")
+    raise FileNotFoundError(f"embed_tokens.weight not found in {hf_dir}")
+
+
+def _hf_logits(text, tokenizer, hf_model, embed_weight):
+    inputs = tokenizer(text, return_tensors="pt")
+    with torch.no_grad():
+        hidden = hf_model(**inputs).last_hidden_state.float()
+        logits = hidden @ embed_weight.float().T
+    return logits[0, -1, :].numpy()
+
+
+def _mlc_logits(text, tokenizer, mlc_module, params, metadata, dev, embed_weight):
+    input_ids = tokenizer(text, return_tensors="pt")["input_ids"][0].numpy().astype(np.int32)
+    seq_len = len(input_ids)
+
+    embed_func = mlc_module["embed"]
+    prefill_func = mlc_module["prefill_to_last_hidden_states"]
+
+    if mlc_module.implements_function("create_flashinfer_paged_kv_cache"):
+        create_kv = mlc_module["create_flashinfer_paged_kv_cache"]
+    elif mlc_module.implements_function("create_tir_paged_kv_cache"):
+        create_kv = mlc_module["create_tir_paged_kv_cache"]
+    else:
+        raise RuntimeError("Cannot find KV cache creation function")
+
+    sliding_window = metadata.get("sliding_window_size", -1)
+    context_window = metadata.get("context_window_size", 32768)
+    prefill_chunk = metadata.get("prefill_chunk_size", 2048)
+    max_seq_len = sliding_window if context_window == -1 else context_window
+
+    kv_cache = create_kv(
+        ShapeTuple([1]),
+        ShapeTuple([max_seq_len]),
+        ShapeTuple([prefill_chunk]),
+        ShapeTuple([16]),
+        ShapeTuple([int(sliding_window != -1)]),
+    )
+
+    nd_view = tvm.get_global_func("vm.builtin.reshape")
+    add_sequence = tvm.get_global_func("vm.builtin.kv_state_add_sequence")
+    begin_forward = tvm.get_global_func("vm.builtin.kv_state_begin_forward")
+    end_forward = tvm.get_global_func("vm.builtin.kv_state_end_forward")
+
+    tokens_tvm = tvm.runtime.tensor(input_ids, device=dev)
+    embedding = embed_func(tokens_tvm, params)
+    embedding = nd_view(embedding, ShapeTuple([1, seq_len, embedding.shape[-1]]))
+
+    add_sequence(kv_cache, 0)
+    begin_forward(kv_cache, ShapeTuple([0]), ShapeTuple([seq_len]))
+    hidden_states, _ = prefill_func(embedding, kv_cache, params)
+    end_forward(kv_cache)
+
+    # Compute logits from hidden states using embed_tokens weight (tie_word_embeddings)
+    hidden = hidden_states.numpy().astype(np.float32)
+    logits = hidden @ embed_weight.float().numpy().T
+    return logits[0, -1, :]
+
+
+@pytest.mark.skipif(_skip, reason=_skip_reason)
+def test_mlc_hf_logit_match():
+    tokenizer = AutoTokenizer.from_pretrained(MLC_QWEN3_EMB_HF_DIR, padding_side="left")
+    hf_model = AutoModel.from_pretrained(MLC_QWEN3_EMB_HF_DIR)
+    embed_weight = _load_embed_weight(MLC_QWEN3_EMB_HF_DIR)
+
+    dev = tvm.runtime.device(MLC_QWEN3_EMB_DEVICE, 0)
+    ex = tvm.runtime.load_module(MLC_QWEN3_EMB_MODEL_LIB)
+    vm = relax.VirtualMachine(ex, dev)
+    mlc_module = vm.module
+
+    metadata = json.loads(VirtualMachine(ex, tvm.runtime.device("cpu"))["_metadata"]())
+    params_dict, _ = tvmjs.load_tensor_cache(MLC_QWEN3_EMB_MODEL_DIR, dev)
+    param_names = [p["name"] for p in metadata["params"]]
+    params = [params_dict[name] for name in param_names]
+
+    for text in TEST_TEXTS:
+        hf = _hf_logits(text, tokenizer, hf_model, embed_weight)
+        mlc = _mlc_logits(text, tokenizer, mlc_module, params, metadata, dev, embed_weight)
+
+        cos_sim = np.dot(hf, mlc) / (np.linalg.norm(hf) * np.linalg.norm(mlc))
+        assert cos_sim > 0.99, f"[{text[:30]}] Cosine similarity {cos_sim:.6f} below 0.99"
+
+        max_diff = np.max(np.abs(hf - mlc))
+        assert max_diff < 1.0, f"[{text[:30]}] Max absolute diff {max_diff:.6e} exceeds 1.0"
+
+        hf_top10 = set(np.argsort(hf)[-10:])
+        mlc_top10 = set(np.argsort(mlc)[-10:])
+        overlap = len(hf_top10 & mlc_top10)
+        assert overlap >= 7, f"[{text[:30]}] Top-10 overlap {overlap}/10 below 7"
+
+
+if __name__ == "__main__":
+    test_mlc_hf_logit_match()

tests/python/op/test_batch_spec_verify.py

@@ -0,0 +1,163 @@
+import numpy as np
+import pytest
+import tvm
+import tvm.testing
+
+from mlc_llm.op.batch_spec_verify import batch_spec_verify
+
+# test category "op_correctness"
+pytestmark = [pytest.mark.op_correctness]
+
+
+@pytest.mark.parametrize("nbatch", [32, 64])
+@pytest.mark.parametrize("vocab", [3, 32, 64, 32000, 33, 65, 32001, 128000])
+@pytest.mark.parametrize("plist", [[0.5, 0.5], [1, 0], [0, 1]])
+def test_batch_spec_verify(nbatch, vocab, plist):
+    def numpy_reference(
+        draft_probs,
+        draft_tokens,
+        model_probs,
+        token_tree_first_child,
+        token_tree_next_sibling,
+        uniform_samples,
+        token_tree_parent_ptr,
+    ):
+        nbatch = token_tree_parent_ptr.shape[0]
+        for b in range(nbatch):
+            parent_ptr = token_tree_parent_ptr[b]
+            child_ptr = token_tree_first_child[parent_ptr]
+            while child_ptr != -1:
+                child_token = draft_tokens[child_ptr]
+                p_child = model_probs[parent_ptr, child_token]
+                q_child = draft_probs[child_ptr, child_token]
+                uniform_sample = uniform_samples[child_ptr]
+                if p_child / q_child >= uniform_sample:
+                    parent_ptr = child_ptr
+                    child_ptr = token_tree_first_child[child_ptr]
+                else:
+                    model_probs[parent_ptr, :] = np.maximum(
+                        model_probs[parent_ptr, :] - draft_probs[child_ptr, :], 0.0
+                    )
+                    psum = np.sum(model_probs[parent_ptr, :])
+                    model_probs[parent_ptr, :] /= psum
+                    child_ptr = token_tree_next_sibling[child_ptr]
+            token_tree_parent_ptr[b] = parent_ptr
+
+    np.random.seed(0)
+
+    def gen_chain(num_nodes, base):
+        token_tree_first_child = list()
+        token_tree_next_sibling = list()
+        for i in range(num_nodes):
+            token_tree_first_child.append(base + i + 1 if i + 1 < num_nodes else -1)
+            token_tree_next_sibling.append(-1)
+        return token_tree_first_child, token_tree_next_sibling, base, base + 1
+
+    def gen_full_binary_tree(height, base):
+        token_tree_first_child = list()
+        token_tree_next_sibling = list()
+        num_nodes = 2**height - 1
+        for i in range(num_nodes):
+            token_tree_first_child.append(base + i * 2 + 1 if i * 2 + 1 < num_nodes else -1)
+            token_tree_next_sibling.append(base + i * 2 + 2 if i * 2 + 2 < num_nodes else -1)
+        return token_tree_first_child, token_tree_next_sibling, base, base + 1
+
+    ### Inputs
+    num_nodes = 0
+    token_tree_first_child = list()
+    token_tree_next_sibling = list()
+    token_tree_parent_ptr = list()
+
+    for _ in range(nbatch):
+        choice = np.random.choice(2, 1, p=plist)
+        if choice == 0:
+            nodes_batch = np.random.randint(3, 32)
+            res = gen_chain(nodes_batch, num_nodes)
+            num_nodes += nodes_batch
+        else:
+            height = np.random.randint(3, 5)
+            res = gen_full_binary_tree(height, num_nodes)
+            num_nodes += 2**height - 1
+        token_tree_first_child.extend(res[0])
+        token_tree_next_sibling.extend(res[1])
+        token_tree_parent_ptr.append(res[2])
+
+    token_tree_first_child = np.array(token_tree_first_child).astype("int32")
+    token_tree_next_sibling = np.array(token_tree_next_sibling).astype("int32")
+    token_tree_parent_ptr = np.array(token_tree_parent_ptr).astype("int32")
+
+    draft_probs = np.random.rand(num_nodes, vocab).astype("float32")
+    draft_probs /= np.sum(draft_probs, axis=1, keepdims=True)
+    draft_tokens = np.random.randint(0, vocab, num_nodes).astype("int32")
+    model_probs = np.random.rand(num_nodes, vocab).astype("float32")
+    model_probs /= np.sum(model_probs, axis=1, keepdims=True)
+    uniform_samples = np.random.rand(num_nodes).astype("float32")
+
+    ### TVM Inputs
+    dev = tvm.cuda(0)
+    draft_probs_tvm = tvm.runtime.tensor(draft_probs, dev)
+    draft_tokens_tvm = tvm.runtime.tensor(draft_tokens, dev)
+    model_probs_tvm = tvm.runtime.tensor(model_probs, dev)
+    token_tree_first_child_tvm = tvm.runtime.tensor(token_tree_first_child, dev)
+    token_tree_next_sibling_tvm = tvm.runtime.tensor(token_tree_next_sibling, dev)
+    uniform_samples_tvm = tvm.runtime.tensor(uniform_samples, dev)
+    token_tree_parent_ptr_tvm = tvm.runtime.tensor(token_tree_parent_ptr, dev)
+
+    # print("draft_probs", draft_probs)
+    # print("draft_tokens", draft_tokens)
+    # print("model_probs", model_probs)
+    # print("token_tree_first_child", token_tree_first_child)
+    # print("token_tree_next_sibling", token_tree_next_sibling)
+    # print("uniform_samples", uniform_samples)
+    # print("token_tree_parent_ptr", token_tree_parent_ptr)
+
+    ### Numpy reference
+    numpy_reference(
+        draft_probs,
+        draft_tokens,
+        model_probs,
+        token_tree_first_child,
+        token_tree_next_sibling,
+        uniform_samples,
+        token_tree_parent_ptr,
+    )
+    # print("model_probs", model_probs)
+    # print("token_tree_parent_ptr", token_tree_parent_ptr)
+
+    ### TVM
+    kernel = batch_spec_verify(vocab)
+    mod = tvm.build(kernel, target="cuda")
+    mod(
+        draft_probs_tvm,
+        draft_tokens_tvm,
+        model_probs_tvm,
+        token_tree_first_child_tvm,
+        token_tree_next_sibling_tvm,
+        uniform_samples_tvm,
+        token_tree_parent_ptr_tvm,
+    )
+    # print("model_probs", model_probs_tvm.asnumpy())
+    # print("token_tree_parent_ptr", token_tree_parent_ptr_tvm.asnumpy())
+
+    tvm.testing.assert_allclose(model_probs, model_probs_tvm.asnumpy())
+    tvm.testing.assert_allclose(
+        token_tree_parent_ptr, token_tree_parent_ptr_tvm.asnumpy(), rtol=0, atol=0
+    )
+
+    time_evaluator = mod.time_evaluator(mod.entry_name, dev, number=10, repeat=3)
+    print(f"batch_size: {nbatch}, vocab_size: {vocab}, tree_structure: {plist}")
+    print(
+        time_evaluator(
+            draft_probs_tvm,
+            draft_tokens_tvm,
+            model_probs_tvm,
+            token_tree_first_child_tvm,
+            token_tree_next_sibling_tvm,
+            uniform_samples_tvm,
+            token_tree_parent_ptr_tvm,
+        )
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/op/test_mrope.py

@@ -0,0 +1,251 @@
+import numpy as np
+import pytest
+
+tvm = pytest.importorskip("tvm")
+from tvm import relax  # noqa: E402
+from tvm.relax.frontend import nn  # noqa: E402
+from tvm.relax.frontend.nn import spec  # noqa: E402
+from tvm.runtime import tensor as tvm_tensor  # noqa: E402
+
+from mlc_llm.op import (  # noqa: E402
+    MultimodalRotaryEmbedding,
+    VisionPositionMetadata,
+    apply_multimodal_rotary_pos_emb,
+    get_mrope_position_ids,
+)
+
+
+def _numpy_rotate_half(x: np.ndarray) -> np.ndarray:
+    x1, x2 = np.split(x, 2, axis=-1)
+    return np.concatenate([-x2, x1], axis=-1)
+
+
+def _numpy_apply_mrope(
+    q: np.ndarray,
+    k: np.ndarray,
+    position_ids: np.ndarray,
+    theta: float,
+    mrope_section: tuple[int, ...],
+) -> tuple[np.ndarray, np.ndarray]:
+    if position_ids.ndim != 3:
+        raise ValueError(f"position_ids must be rank-3, got shape {position_ids.shape}")
+    if position_ids.shape[0] == 3:
+        position_ids = np.transpose(position_ids, (1, 2, 0))
+    elif position_ids.shape[-1] != 3:
+        raise ValueError(
+            "position_ids must have shape (batch, seq, 3) or (3, batch, seq), "
+            f"got {position_ids.shape}"
+        )
+
+    head_dim = q.shape[-1]
+    inv_freq = 1.0 / (theta ** (np.arange(0, head_dim, 2, dtype=np.float32) / float(head_dim)))
+    pos = np.transpose(position_ids, (2, 0, 1))
+    inv = inv_freq.reshape(1, 1, -1, 1).astype(np.float32)
+    inv = np.broadcast_to(inv, (3, pos.shape[1], inv_freq.size, 1))
+    pos = pos.reshape(3, pos.shape[1], 1, pos.shape[2]).astype(np.float32)
+    freqs = np.matmul(inv, pos)
+    freqs = np.transpose(freqs, (0, 1, 3, 2))
+    emb = np.concatenate([freqs, freqs], axis=-1)
+    cos = np.cos(emb)
+    sin = np.sin(emb)
+    split_sizes = list(mrope_section) * 2
+    split_points = np.cumsum(split_sizes)[:-1]
+    cos_chunks = np.split(cos, split_points, axis=-1)
+    sin_chunks = np.split(sin, split_points, axis=-1)
+    cos = np.concatenate([chunk[idx % 3] for idx, chunk in enumerate(cos_chunks)], axis=-1)
+    sin = np.concatenate([chunk[idx % 3] for idx, chunk in enumerate(sin_chunks)], axis=-1)
+    cos = np.expand_dims(cos, axis=2)
+    sin = np.expand_dims(sin, axis=2)
+    q_out = q * cos + _numpy_rotate_half(q) * sin
+    k_out = k * cos + _numpy_rotate_half(k) * sin
+    return q_out, k_out
+
+
+def _evaluate_tensor(expr):
+    mod = tvm.IRModule.from_expr(expr)
+    target = tvm.target.Target("llvm")
+    ex = tvm.relax.build(mod, target)
+    vm = tvm.relax.VirtualMachine(ex, tvm.cpu())
+    return vm["main"]().numpy()
+
+
+def _run_mlc_mrope(
+    q_np: np.ndarray,
+    k_np: np.ndarray,
+    position_ids_np: np.ndarray,
+    theta: float,
+    mrope_section: tuple[int, ...],
+) -> tuple[np.ndarray, np.ndarray]:
+    class RopeModule(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.rotary = MultimodalRotaryEmbedding(q_np.shape[-1], theta, mrope_section)
+
+        def forward(
+            self,
+            q: nn.Tensor,
+            k: nn.Tensor,
+            pos: nn.Tensor,
+        ):
+            """Run MRoPE on test tensors and return rotated query/key outputs."""
+            cos, sin = self.rotary(q, pos)
+            return apply_multimodal_rotary_pos_emb(q, k, cos, sin, mrope_section)
+
+    module = RopeModule()
+    mod, _, _ = module.export_tvm(
+        spec={
+            "forward": {
+                "q": spec.Tensor(q_np.shape, "float32"),
+                "k": spec.Tensor(k_np.shape, "float32"),
+                "pos": spec.Tensor(position_ids_np.shape, "int64"),
+            }
+        },
+        allow_extern=True,
+    )
+    target = tvm.target.Target("llvm")
+    exec_mod = relax.build(mod, target=target)
+    vm = relax.VirtualMachine(exec_mod, tvm.cpu())
+    device = tvm.cpu()
+    q_nd = tvm_tensor(q_np.astype("float32"), device=device)
+    k_nd = tvm_tensor(k_np.astype("float32"), device=device)
+    pos_nd = tvm_tensor(position_ids_np.astype("int64"), device=device)
+    out_q, out_k = vm["forward"](q_nd, k_nd, pos_nd)
+    return out_q.numpy(), out_k.numpy()
+
+
+def test_apply_mrope_matches_numpy_reference():
+    theta = 10000.0
+    mrope_section = (2, 2, 2)
+    batch, seq_len, heads, head_dim = 1, 4, 2, 12
+    rng = np.random.default_rng(0)
+    q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+    k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+    position_ids = np.zeros((batch, seq_len, 3), dtype=np.int64)
+    position_ids[0, :, 0] = np.arange(seq_len)
+    position_ids[0, :, 1] = np.arange(seq_len) * 2
+    position_ids[0, :, 2] = np.arange(seq_len) * 3
+
+    mlc_q, mlc_k = _run_mlc_mrope(q_np, k_np, position_ids, theta, mrope_section)
+    ref_q, ref_k = _numpy_apply_mrope(q_np, k_np, position_ids, theta, mrope_section)
+
+    np.testing.assert_allclose(mlc_q, ref_q, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_k, ref_k, rtol=1e-5, atol=1e-5)
+
+
+def test_get_mrope_position_ids_text_only():
+    input_ids = np.array([[1, 2, 3, 0, 0]], dtype=np.int64)
+    attention_mask = np.array([[1, 1, 1, 0, 0]], dtype=np.int64)
+    meta = VisionPositionMetadata(
+        vision_start_token_id=1000,
+        image_token_id=1001,
+        video_token_id=1002,
+        spatial_merge_size=2,
+        tokens_per_second=4.0,
+    )
+    position_ids, deltas = get_mrope_position_ids(
+        input_ids,
+        meta,
+        attention_mask=attention_mask,
+        image_grid_thw=None,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+    )
+    expected = attention_mask.cumsum(axis=-1) - 1
+    expected = np.where(attention_mask == 0, 1, expected)
+    expected = np.expand_dims(expected, axis=0).repeat(3, axis=0)
+    np.testing.assert_array_equal(position_ids, expected)
+    np.testing.assert_array_equal(deltas, np.array([[-2]], dtype=np.int64))
+
+
+def test_get_mrope_position_ids_single_image_block():
+    meta = VisionPositionMetadata(
+        vision_start_token_id=5000,
+        image_token_id=5001,
+        video_token_id=6000,
+        spatial_merge_size=2,
+        tokens_per_second=4.0,
+    )
+    input_ids = np.array(
+        [[11, 12, 5000, 5001, 21, 22, 23, 24, 31, 32]],
+        dtype=np.int64,
+    )
+    attention_mask = np.ones_like(input_ids, dtype=np.int64)
+    image_grid_thw = np.array([[1, 4, 4]], dtype=np.int64)
+    position_ids, deltas = get_mrope_position_ids(
+        input_ids,
+        meta,
+        attention_mask=attention_mask,
+        image_grid_thw=image_grid_thw,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+    )
+    expected = np.array(
+        [
+            [0, 1, 2, 3, 3, 3, 3, 5, 6, 7],
+            [0, 1, 2, 3, 3, 4, 4, 5, 6, 7],
+            [0, 1, 2, 3, 4, 3, 4, 5, 6, 7],
+        ],
+        dtype=np.int64,
+    ).reshape(3, 1, -1)
+    np.testing.assert_array_equal(position_ids, expected)
+    np.testing.assert_array_equal(deltas, np.array([[-2]], dtype=np.int64))
+
+
+def test_apply_mrope_accepts_3_batch_seq_layout():
+    theta = 10000.0
+    mrope_section = (2, 2, 2)
+    batch, seq_len, heads, head_dim = 1, 4, 2, 12
+    rng = np.random.default_rng(1)
+    q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+    k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+
+    position_ids_bsc = np.zeros((batch, seq_len, 3), dtype=np.int64)
+    position_ids_bsc[0, :, 0] = np.arange(seq_len)
+    position_ids_bsc[0, :, 1] = np.arange(seq_len) * 2
+    position_ids_bsc[0, :, 2] = np.arange(seq_len) * 3
+    position_ids_3bs = np.transpose(position_ids_bsc, (2, 0, 1))
+
+    mlc_q_bsc, mlc_k_bsc = _run_mlc_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
+    mlc_q_3bs, mlc_k_3bs = _run_mlc_mrope(q_np, k_np, position_ids_3bs, theta, mrope_section)
+    ref_q, ref_k = _numpy_apply_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
+
+    np.testing.assert_allclose(mlc_q_bsc, ref_q, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_k_bsc, ref_k, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_q_3bs, ref_q, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_k_3bs, ref_k, rtol=1e-5, atol=1e-5)
+
+
+def test_get_mrope_position_ids_output_is_directly_usable():
+    theta = 10000.0
+    mrope_section = (2, 2, 2)
+    meta = VisionPositionMetadata(
+        vision_start_token_id=7000,
+        image_token_id=7001,
+        video_token_id=7002,
+        spatial_merge_size=2,
+        tokens_per_second=4.0,
+    )
+    input_ids = np.array([[11, 12, 7000, 7001, 21, 22, 23, 24, 31, 32]], dtype=np.int64)
+    attention_mask = np.ones_like(input_ids, dtype=np.int64)
+    image_grid_thw = np.array([[1, 4, 4]], dtype=np.int64)
+    position_ids_3bs, _ = get_mrope_position_ids(
+        input_ids,
+        meta,
+        attention_mask=attention_mask,
+        image_grid_thw=image_grid_thw,
+        video_grid_thw=None,
+        second_per_grid_ts=None,
+    )
+    position_ids_bsc = np.transpose(position_ids_3bs, (1, 2, 0))
+
+    batch, seq_len = input_ids.shape
+    heads, head_dim = 2, 12
+    rng = np.random.default_rng(2)
+    q_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+    k_np = rng.standard_normal((batch, seq_len, heads, head_dim), dtype=np.float32)
+
+    mlc_q_3bs, mlc_k_3bs = _run_mlc_mrope(q_np, k_np, position_ids_3bs, theta, mrope_section)
+    mlc_q_bsc, mlc_k_bsc = _run_mlc_mrope(q_np, k_np, position_ids_bsc, theta, mrope_section)
+
+    np.testing.assert_allclose(mlc_q_3bs, mlc_q_bsc, rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(mlc_k_3bs, mlc_k_bsc, rtol=1e-5, atol=1e-5)

tests/python/op/test_top_p_pivot.py

@@ -0,0 +1,86 @@
+import numpy as np
+import pytest
+import tvm
+import tvm.testing
+
+from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm
+
+# mypy: disable-error-code="var-annotated"
+
+# test category "op_correctness"
+pytestmark = [pytest.mark.op_correctness]
+
+
+@pytest.mark.parametrize("batch_size", [32, 64])
+@pytest.mark.parametrize("vocab", [3, 32, 64, 128])
+def test_top_p_renorm(batch_size, vocab):
+    top_p = 0.95
+    init_pivots_np = np.array([1 - top_p, 0.02, 0.01]).astype(np.float32)
+    top_p_np = np.array([top_p]).astype(np.float32)
+
+    p_np = np.random.exponential(3, size=(batch_size, vocab)).astype(np.float32)
+    p_np /= np.sum(p_np, axis=-1, keepdims=True)
+    final_pivot_np = np.zeros(batch_size).astype(np.float32)
+    final_lsum_np = np.zeros(batch_size).astype(np.float32)
+
+    dev = tvm.cuda(0)
+    var_prob = tvm.runtime.tensor(p_np, dev)
+    var_init_pivots = tvm.runtime.tensor(init_pivots_np, dev)
+    top_p_global = tvm.runtime.tensor(top_p_np, dev)
+    var_final_pivot = tvm.runtime.tensor(final_pivot_np, dev)
+    var_final_lsum = tvm.runtime.tensor(final_lsum_np, dev)
+
+    kernel = top_p_pivot(init_pivots_np.shape[0])
+    mod = tvm.build(kernel, target="cuda")
+    mod(var_prob, top_p_global, var_init_pivots, var_final_pivot, var_final_lsum)
+
+    final_pivot = var_final_pivot.asnumpy()
+    final_lsum = var_final_lsum.asnumpy()
+
+    renorm_np = p_np.copy()
+    var_renorm = tvm.runtime.tensor(renorm_np, dev)
+
+    kernel_renorm = top_p_renorm()
+    mod_renorm = tvm.build(kernel_renorm, target="cuda")
+    mod_renorm(var_prob, var_final_pivot, var_final_lsum, var_renorm)
+
+    renorm = var_renorm.asnumpy()
+
+    def verify_pivot(probs: np.ndarray, pivot: float, lsum: float, renorm: np.ndarray):
+        sorted_probs = np.sort(probs, axis=-1)[::-1]
+        num_larger_than_pivot = np.sum(sorted_probs >= pivot)
+        filtered_sorted_probs = sorted_probs[:num_larger_than_pivot]
+        min_larger_than_pivot = min(filtered_sorted_probs)
+
+        sum_larger_than_pivot = np.sum(np.where(sorted_probs >= pivot, sorted_probs, 0))
+        sum_larger_than_pivot_exclude_min = np.sum(
+            np.where(filtered_sorted_probs != min_larger_than_pivot, filtered_sorted_probs, 0)
+        )
+
+        probs[probs < pivot] = 0
+        renorm_prob = probs / np.sum(probs, axis=-1, keepdims=True)
+        try:
+            assert sum_larger_than_pivot >= top_p
+            assert sum_larger_than_pivot_exclude_min < top_p
+            assert abs(lsum - sum_larger_than_pivot) < 1e-6
+            assert np.allclose(renorm, renorm_prob, atol=1e-6, rtol=1e-6)
+        except AssertionError:
+            print("Failed")
+            print("probs:", repr(probs))
+            print("pivot:", pivot)
+            print("sorted_probs:", sorted_probs)
+            print("num_larger_than_pivot:", num_larger_than_pivot)
+            print("filtered_sorted_probs:", filtered_sorted_probs)
+            print("min_larger_than_pivot:", min_larger_than_pivot)
+            print("sum_larger_than_pivot:", sum_larger_than_pivot)
+            print("sum_larger_than_pivot_exclude_min:", sum_larger_than_pivot_exclude_min)
+            print("renom_prob:", renorm_prob)
+            print("renorm:", renorm)
+            raise
+
+    for i in range(batch_size):
+        verify_pivot(p_np[i], final_pivot[i], final_lsum[i], renorm[i])
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/op/test_tree_attn.py

@@ -0,0 +1,243 @@
+import math
+
+import numpy as np
+import pytest
+import tvm
+import tvm.testing
+from tvm.relax.frontend.nn.llm import tree_attn
+
+# test category "op_correctness"
+pytestmark = [pytest.mark.op_correctness]
+
+
+@pytest.mark.parametrize("nbatch", [1, 4, 32])
+@pytest.mark.parametrize("h_q", [8, 16])
+@pytest.mark.parametrize("h_kv", [4, 8])
+@pytest.mark.parametrize("d", [128])
+@pytest.mark.parametrize("rotary_mode", [0, 1])
+def test_tree_attn(nbatch, h_q, h_kv, d, rotary_mode):
+    np.random.seed(0)
+    np.set_printoptions(linewidth=10000)
+
+    def gen_chain(num_nodes):
+        mask = np.tril(np.ones((num_nodes, num_nodes)))
+        return num_nodes, list(mask.flatten()), np.arange(num_nodes)
+
+    def gen_full_binary_tree(height):
+        mask = list()
+        pos = list()
+        num_nodes = 2**height - 1
+        for i in range(num_nodes):
+            if i == 0:
+                mask_0 = [0] * num_nodes
+                mask_0[0] = 1
+                mask.append(mask_0)
+                pos.append(0)
+            else:
+                mask_i = mask[(i + 1) // 2 - 1].copy()
+                mask_i[i] = 1
+                mask.append(mask_i)
+                pos.append(pos[(i + 1) // 2 - 1] + 1)
+        return num_nodes, list(np.array(mask).flatten()), pos
+
+    ### Inputs
+    num_nodes = 0
+    m_list = list()
+    mn_list = list()
+    mask_list = list()
+    q_pos_list = list()
+
+    mn_list.append(0)
+
+    for _ in range(nbatch):
+        choice = np.random.choice(2, 1, p=[1, 0])
+        if choice == 0:
+            nodes_batch = np.random.randint(3, 32)
+            res = gen_chain(nodes_batch)
+            num_nodes += nodes_batch
+        else:
+            height = np.random.randint(2, 6)
+            res = gen_full_binary_tree(height)
+            num_nodes += 2**height - 1
+        m_list.append(res[0])
+        mn_list.append(res[0] ** 2)
+        mask_list.extend(res[1])
+        q_pos_list.extend(res[2])
+
+    qkv_indptr = np.array(np.cumsum([0, *m_list])).astype(np.int32)
+    m_list = np.array(m_list).astype(np.int32)
+    mn_list = np.array(mn_list).astype(np.int32)
+    mn_list = np.cumsum(mn_list).astype(np.int32)
+    mask_list = np.array(mask_list).astype(np.int32)
+    q_pos_list = np.array(q_pos_list).astype(np.int32)
+
+    # print("qkv_indptr:", qkv_indptr)
+    # print("m_list:", m_list)
+    # print("mn_list:", mn_list)
+    # for num_nodes, base in zip(m_list, mn_list):
+    #     print("num_nodes:", num_nodes)
+    #     print("indptr:", base)
+    #     print(
+    #         "mask:",
+    #         mask_list[base : base + num_nodes * num_nodes].reshape(num_nodes, num_nodes),
+    #     )
+    #     print("q_pos:", q_pos_list[base : base + num_nodes])
+
+    q = np.random.rand(num_nodes, h_q, d).astype(np.float16)
+    q_indptr = qkv_indptr
+    k = np.random.rand(num_nodes, h_kv, d).astype(np.float16)
+    v = np.random.rand(num_nodes, h_kv, d).astype(np.float16)
+    kv_indptr = qkv_indptr
+    q_rope_position = q_pos_list
+    m_arr = m_list
+    mn_indptr = mn_list
+    mask = mask_list
+    output = np.zeros((num_nodes, h_q, d), dtype=np.float16)
+    lse = np.zeros((num_nodes, h_q), dtype=np.float32)
+    rotary_scale = 1.0
+    rotary_theta = 10000.0
+    attn_score_scaling_factor = 1.0
+
+    ### TVM Inputs
+    dev = tvm.cuda(0)
+    q_tvm = tvm.runtime.tensor(q, dev)
+    q_indptr_tvm = tvm.runtime.tensor(q_indptr, dev)
+    k_tvm = tvm.runtime.tensor(k, dev)
+    v_tvm = tvm.runtime.tensor(v, dev)
+    kv_indptr_tvm = tvm.runtime.tensor(kv_indptr, dev)
+    q_rope_position_tvm = tvm.runtime.tensor(q_rope_position, dev)
+    # m_arr_tvm = tvm.runtime.tensor(m_arr, dev)
+    mn_indptr_tvm = tvm.runtime.tensor(mn_indptr, dev)
+    mask_tvm = tvm.runtime.tensor(mask, dev)
+    output_tvm = tvm.runtime.tensor(output, dev)
+    lse_tvm = tvm.runtime.tensor(lse, dev)
+
+    target = tvm.target.Target("cuda")
+    kernel = tree_attn(h_kv=h_kv, h_q=h_q, d=d, dtype="float16", rope_scaling={}, target=target)
+    mod = tvm.build(kernel, target=target)
+    mod(
+        q_tvm,
+        q_indptr_tvm,
+        k_tvm,
+        v_tvm,
+        kv_indptr_tvm,
+        q_rope_position_tvm,
+        # m_arr_tvm,
+        mn_indptr_tvm,
+        mask_tvm,
+        output_tvm,
+        lse_tvm,
+        rotary_mode,
+        rotary_scale,
+        rotary_theta,
+        attn_score_scaling_factor,
+        nbatch,
+    )
+
+    ### Numpy reference
+    def numpy_reference(
+        q,
+        q_indptr,
+        k,
+        v,
+        kv_indptr,
+        q_rope_position,
+        m_arr,
+        mn_indptr,
+        mask,
+        rotary_mode,
+        rotary_scale,
+        rotary_theta,
+        attn_score_scaling_factor,
+        output_tvm,
+    ):
+        def rope_freq(s, d, d_range, theta, dtype):
+            freq = s / math.pow(theta, (d * 2 % d_range) / float(d_range))
+            cos_freq = np.cos(freq).astype(dtype)
+            sin_freq = np.sin(freq).astype(dtype)
+            return cos_freq, sin_freq
+
+        def rope(buffer, offset, rotary_dim, theta, scale, dtype):
+            result = buffer.copy()
+            for pos, h, d in np.ndindex(buffer.shape):
+                cos_freq, sin_freq = rope_freq(offset[pos] * scale, d, rotary_dim, theta, dtype)
+                cos = cos_freq * buffer[pos, h, d]
+                sin = sin_freq * (
+                    -buffer[pos, h, d + rotary_dim // 2]
+                    if d < rotary_dim // 2
+                    else buffer[pos, h, d - rotary_dim // 2]
+                )
+                result[pos, h, d] = cos + sin
+            return result
+
+        for i in range(len(m_arr)):
+            num_nodes = m_arr[i]
+            base = mn_indptr[i]
+            q_base = q_indptr[i]
+            kv_base = kv_indptr[i]
+            q_pos = q_rope_position[q_base : q_base + num_nodes]  # (num_nodes,)
+            q_i = q[q_base : q_base + num_nodes]  # (num_nodes, h_q, d)
+            k_i = k[kv_base : kv_base + num_nodes]  # (num_nodes, h_kv, d)
+            v_i = v[kv_base : kv_base + num_nodes]  # (num_nodes, h_kv, d)
+            mask_i = mask[base : base + num_nodes * num_nodes].reshape(num_nodes, num_nodes)
+
+            if rotary_mode == 1:
+                q_i = rope(q_i, q_pos, d, rotary_theta, rotary_scale, q_i.dtype)
+                k_i = rope(k_i, q_pos, d, rotary_theta, rotary_scale, k_i.dtype)
+
+            # group attention
+            # q: (num_nodes, h_q, d)
+            # k: (num_nodes, h_kv, d)
+            # v: (num_nodes, h_kv, d)
+            group_size = h_q // h_kv
+            q_reshape = q_i.transpose(1, 0, 2)  # (h_q, num_nodes, d)
+            k_reshape = k_i.transpose(1, 2, 0)  # (h_kv, d, num_nodes)
+            v_reshape = v_i.transpose(1, 0, 2)  # (h_kv, num_nodes, d)
+            # expand k_reshape
+            k_reshape = k_reshape.reshape(h_kv, 1, d, num_nodes)
+            k_reshape = np.repeat(k_reshape, group_size, axis=1)
+            k_reshape = k_reshape.reshape(h_q, d, num_nodes)
+            # expand v_reshape
+            v_reshape = v_reshape.reshape(h_kv, 1, num_nodes, d)
+            v_reshape = np.repeat(v_reshape, group_size, axis=1)
+            v_reshape = v_reshape.reshape(h_q, num_nodes, d)
+            # print("q_reshape:", q_reshape.shape)
+            # print("k_reshape:", k_reshape.shape)
+            # print("v_reshape:", v_reshape.shape)
+
+            # qk: (h_q, num_nodes, num_nodes)
+            qk = np.matmul(q_reshape, k_reshape) * attn_score_scaling_factor / math.sqrt(float(d))
+            # softmax(qk, axis=-1), numerical stability
+            qk[:, mask_i == 0] = -np.inf
+            qk_max = np.max(qk, axis=-1, keepdims=True)
+            qk = np.exp(qk - qk_max)
+            qk = qk / np.sum(qk, axis=-1, keepdims=True)
+
+            # attention
+            output_i = np.matmul(qk, v_reshape).transpose(1, 0, 2)  # (num_nodes, h_q, d)
+            # print(output_i)
+
+            tvm.testing.assert_allclose(
+                output_i, output_tvm[q_base : q_base + num_nodes], rtol=1e-3, atol=1e-3
+            )
+
+    numpy_reference(
+        q,
+        q_indptr,
+        k,
+        v,
+        kv_indptr,
+        q_rope_position,
+        m_arr,
+        mn_indptr,
+        mask,
+        rotary_mode,
+        rotary_scale,
+        rotary_theta,
+        attn_score_scaling_factor,
+        output_tvm.numpy(),
+    )
+
+
+if __name__ == "__main__":
+    tvm.testing.main()

tests/python/op/test_two_stage_softmax.py

@@ -0,0 +1,51 @@
+import numpy as np
+import pytest
+import scipy.special
+import tvm
+from tvm.s_tir import dlight
+
+# test category "op_correctness"
+pytestmark = [pytest.mark.op_correctness]
+
+
+def test_two_stage_softmax():
+    from mlc_llm.compiler_pass.rewrite_softmax import _get_lse_and_softmax_func
+
+    chunk_size = 4096
+    target = tvm.target.Target("cuda")
+    f_chunk_lse, f_softmax_with_lse = _get_lse_and_softmax_func(target, chunk_size)
+    mod = tvm.IRModule({"chunk_lse": f_chunk_lse, "softmax_with_chunked_lse": f_softmax_with_lse})
+    with target:
+        mod = dlight.ApplyDefaultSchedule(dlight.gpu.GeneralReduction())(mod)
+
+    runtime_mod = tvm.build(mod, target=target)
+    device = tvm.cuda()
+
+    num_runs = 5
+    vocab_size = 128256
+    for batch_size in [1, 2, 4, 8, 16, 32, 64, 128]:
+        for _ in range(num_runs):
+            x_np = np.random.uniform(low=-10, high=10, size=(batch_size, vocab_size)).astype(
+                "float32"
+            )
+            y_np = scipy.special.softmax(x_np, axis=-1)
+
+            x_nd = tvm.runtime.tensor(x_np, device=device)
+            r_nd = tvm.runtime.empty(
+                (batch_size, (vocab_size + chunk_size - 1) // chunk_size),
+                x_np.dtype,
+                device=device,
+            )
+            y_nd = tvm.runtime.empty(x_np.shape, x_np.dtype, device=device)
+
+            runtime_mod["chunk_lse"](x_nd, r_nd)
+            runtime_mod["softmax_with_chunked_lse"](x_nd, r_nd, y_nd)
+
+            y_nd_arr = y_nd.numpy()
+            np.testing.assert_allclose(y_nd_arr, y_np, atol=1e-6, rtol=1e-6)
+
+        print(f"pass batch size {batch_size}")
+
+
+if __name__ == "__main__":
+    test_two_stage_softmax()

tests/python/quantization/test_awq_quantization.py

@@ -0,0 +1,86 @@
+from typing import List  # noqa: UP035
+
+import numpy as np
+import pytest
+import torch
+import tvm
+import tvm.testing
+from tvm import DataType
+from tvm.relax.frontend import nn
+
+from mlc_llm.loader import QuantizeMapping
+from mlc_llm.quantization import QUANTIZATION, AWQQuantize
+
+
+def dequantize_np(
+    config: AWQQuantize,
+    weight: np.ndarray,
+    zeros: np.ndarray,
+    scale: np.ndarray,
+) -> np.ndarray:
+    def decode_int_arr(int_arr: np.ndarray, num_elem_per_storage: int, bits: int):
+        bin_mask = (1 << bits) - 1
+        int_arr_repeated = np.repeat(int_arr, num_elem_per_storage, axis=-1)
+        indice_j = np.indices(int_arr_repeated.shape)[1]
+        arr_bin = np.bitwise_and(
+            np.right_shift(
+                int_arr_repeated,
+                (indice_j % num_elem_per_storage) * bits,
+            ),
+            bin_mask,
+        )
+        return arr_bin
+
+    weight_bin = decode_int_arr(
+        weight, config.num_elem_per_storage, DataType(config.quantize_dtype).bits
+    )
+    zero_bin = decode_int_arr(
+        zeros, config.num_elem_per_storage, DataType(config.quantize_dtype).bits
+    )
+    scale_repeated = np.repeat(scale, config.group_size, axis=-1)
+    zero_bin_repeated = np.repeat(zero_bin, config.group_size, axis=-1)
+    return (weight_bin - zero_bin_repeated) * scale_repeated
+
+
+@pytest.mark.parametrize(
+    "quant_name, shape, dtype",
+    [
+        ("q4f16_awq", [2, 4096], "float16"),
+    ],
+)
+def test_dequantize_weight(quant_name: str, shape: List[int], dtype: str):  # noqa: UP006
+    class Test(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.linear = nn.Linear(shape[1], shape[0], bias=False, dtype=dtype)
+
+        def forward(self, x: nn.Tensor):
+            return self.linear(x)
+
+    config = QUANTIZATION[quant_name]
+    assert isinstance(config, AWQQuantize)
+    weight_np = np.random.randint(
+        np.iinfo(config.storage_dtype).min,
+        np.iinfo(config.storage_dtype).max,
+        (shape[0], shape[1] // config.num_elem_per_storage),
+    ).astype(config.storage_dtype)
+    zeros_np = np.random.randint(
+        np.iinfo(config.storage_dtype).min,
+        np.iinfo(config.storage_dtype).max,
+        (shape[0], shape[1] // config.num_elem_per_storage // config.group_size),
+    ).astype(config.storage_dtype)
+    scale_np = np.random.random((shape[0], shape[1] // config.group_size)).astype(
+        config.model_dtype
+    )
+    mod = config.quantize_model(Test(), QuantizeMapping({}, {}), "")
+    mod.linear.qweight.data = weight_np
+    mod.linear.qzeros.data = zeros_np
+    mod.linear.scales.data = scale_np
+    model = mod.jit(spec={"forward": {"x": nn.spec.Tensor((shape[1], shape[1]), dtype)}})
+    out = model["forward"](torch.from_numpy(np.diag(np.ones(shape[1]).astype(dtype))))
+    ref = dequantize_np(config, weight_np, zeros_np, scale_np).T
+    tvm.testing.assert_allclose(out, ref, rtol=1e-3, atol=1e-3)
+
+
+if __name__ == "__main__":
+    test_dequantize_weight("q4f16_awq", [2, 4096], "float16")

tests/python/quantization/test_group_quantization.py

@@ -0,0 +1,188 @@
+from typing import List, Optional  # noqa: UP035
+
+import numpy as np
+import pytest
+import torch
+import tvm
+import tvm.testing
+from tvm import DataType
+from tvm.relax.frontend import nn
+
+from mlc_llm.loader import QuantizeMapping
+from mlc_llm.quantization import QUANTIZATION
+from mlc_llm.quantization.group_quantization import (
+    GroupQuantize,
+    GroupQuantizeEmbedding,
+    GroupQuantizeLinear,
+)
+
+
+def quantize_np(config: GroupQuantize, weight: np.ndarray):
+    n, k = weight.shape
+    weight_padded = np.pad(
+        weight,
+        ((0, 0), (0, (config.group_size - k % config.group_size) % config.group_size)),
+    )
+    n, k = weight_padded.shape
+    weight_reshaped = np.reshape(weight_padded, (n, k // config.group_size, config.group_size))
+    max_abs = np.maximum(np.max(np.abs(weight_reshaped), axis=-1), 1e-4)
+    scale = np.divide(max_abs, config.max_int_value)
+    scale_reshaped = np.reshape(scale, (*scale.shape, 1))
+    weight_scaled_reshaped = np.clip(
+        np.add(
+            np.round(np.divide(weight_reshaped, scale_reshaped)),
+            config.max_int_value,
+        ),
+        0,
+        config.max_int_value * 2,
+    ).astype(config.storage_dtype)
+    weight_filtered = np.reshape(weight_scaled_reshaped, (n, k))
+    weight_filtered[..., weight.shape[1] :] = 0
+    weight_scaled = np.reshape(
+        weight_filtered,
+        (n, k // config.num_elem_per_storage, config.num_elem_per_storage),
+    )
+    indice_k = np.indices(weight_scaled.shape, dtype=config.storage_dtype)[-1]
+    quantized_weight = np.sum(
+        np.left_shift(weight_scaled, indice_k * DataType(config.quantize_dtype).bits),
+        axis=-1,
+        dtype=config.storage_dtype,
+    )
+    return quantized_weight, scale
+
+
+def dequantize_np(
+    config: GroupQuantize,
+    weight: np.ndarray,
+    scale: np.ndarray,
+    out_shape: Optional[List[int]] = None,  # noqa: UP006
+):
+    assert weight.shape[0] == scale.shape[0]
+    bin_mask = (1 << DataType(config.quantize_dtype).bits) - 1
+    max_int = config.max_int_value
+    out_shape = (
+        [weight.shape[0], weight.shape[1] * config.num_elem_per_storage]
+        if out_shape is None
+        else out_shape
+    )
+    weight_repeated = np.repeat(weight, config.num_elem_per_storage, axis=-1)
+    scale_repeated = np.repeat(scale, config.group_size, axis=-1)
+    indice_j = np.indices(weight_repeated.shape)[1]
+    weight_bin = np.bitwise_and(
+        np.right_shift(
+            weight_repeated,
+            (indice_j % config.num_elem_per_storage) * DataType(config.quantize_dtype).bits,
+        ),
+        bin_mask,
+    )
+    assert weight_bin.shape[1] <= scale_repeated.shape[1]
+    return ((weight_bin - max_int) * scale_repeated[..., : weight_bin.shape[1]])[
+        : out_shape[0], : out_shape[1]
+    ]
+
+
+@pytest.mark.parametrize(
+    "quant_name, shape, dtype, device",
+    [
+        ("q3f16_1", [2, 13], "float16", "cpu"),
+        ("q3f16_1", [16, 120], "float16", "cpu"),
+        ("q4f16_1", [2, 13], "float16", "cpu"),
+        ("q4f16_1", [16, 128], "float16", "cpu"),
+        ("q4f32_1", [2, 13], "float32", "cpu"),
+        ("q4f32_1", [16, 128], "float32", "cpu"),
+    ],
+)
+def test_quantize_weight(quant_name: str, shape: List[int], dtype: str, device: str):  # noqa: UP006
+    config = QUANTIZATION[quant_name]
+    assert isinstance(config, GroupQuantize)
+    weight_np = np.random.random(shape).astype(dtype)
+    output = config.quantize_weight(tvm.runtime.tensor(weight_np, device=tvm.device(device)))
+    quantized_weight, scale = output[0].numpy(), output[1].numpy()
+    quantized_weight_ref, scale_ref = quantize_np(config, weight_np)
+    tvm.testing.assert_allclose(scale, scale_ref, rtol=1e-3, atol=1e-3)
+    tvm.testing.assert_allclose(
+        dequantize_np(config, quantized_weight, scale, shape),
+        dequantize_np(config, quantized_weight_ref, scale_ref, shape),
+        rtol=1e-2 if quant_name.startswith("q3") else 1e-3,
+        atol=0.4 if quant_name.startswith("q3") else 0.2,
+    )
+
+
+@pytest.mark.parametrize(
+    "quant_name, shape, dtype",
+    [
+        ("q3f16_1", [2, 13], "float16"),
+        ("q3f16_1", [16, 120], "float16"),
+        ("q4f16_1", [2, 13], "float16"),
+        ("q4f16_1", [16, 128], "float16"),
+        ("q4f32_1", [2, 13], "float32"),
+        ("q4f32_1", [16, 128], "float32"),
+    ],
+)
+def test_dequantize_weight(quant_name: str, shape: List[int], dtype: str):  # noqa: UP006
+    class Test(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.linear = nn.Linear(shape[1], shape[0], bias=False, dtype=dtype)
+
+        def forward(self, x: nn.Tensor):
+            return self.linear(x)
+
+    config = QUANTIZATION[quant_name]
+    assert isinstance(config, GroupQuantize)
+    num_group = -(shape[1] // -config.group_size)
+    weight_np = np.random.randint(
+        np.iinfo(config.storage_dtype).min,
+        np.iinfo(config.storage_dtype).max,
+        (shape[0], config.num_storage_per_group * num_group),
+    ).astype(config.storage_dtype)
+    scale_np = np.random.random((shape[0], num_group)).astype(config.model_dtype)
+    mod = config.quantize_model(Test(), QuantizeMapping({}, {}), "")
+    mod.linear.q_weight.data = weight_np
+    mod.linear.q_scale.data = scale_np
+    model = mod.jit(spec={"forward": {"x": nn.spec.Tensor((shape[1], shape[1]), dtype)}})
+    out = model["forward"](torch.from_numpy(np.diag(np.ones(shape[1]).astype(dtype))))
+    ref = dequantize_np(config, weight_np, scale_np, shape).T
+    tvm.testing.assert_allclose(out, ref, rtol=1e-3, atol=1e-3)
+
+
+@pytest.mark.parametrize(
+    "quant_name, shape, dtype",
+    [
+        ("q3f16_1", [16, 128], "float16"),
+        ("q4f16_1", [16, 128], "float16"),
+        ("q4f32_1", [16, 128], "float32"),
+    ],
+)
+def test_quantize_model(quant_name: str, shape: List[int], dtype: str):  # noqa: UP006
+    class Test(nn.Module):
+        def __init__(self) -> None:
+            super().__init__()
+            self.linear = nn.Linear(shape[0], shape[1], dtype=dtype)
+            self.embedding = nn.Embedding(shape[0], shape[1], dtype=dtype)
+
+        def forward(self, x: nn.Tensor):
+            return self.linear(x)
+
+    config = QUANTIZATION[quant_name]
+    assert isinstance(config, GroupQuantize)
+    quant_map = QuantizeMapping({}, {})
+    mod = config.quantize_model(Test(), quant_map, "model")
+    assert quant_map.param_map["model.linear.weight"] == [
+        "model.linear.q_weight",
+        "model.linear.q_scale",
+    ]
+    assert quant_map.map_func["model.linear.weight"] == config.quantize_weight
+    assert isinstance(mod.linear, GroupQuantizeLinear)
+    assert quant_map.param_map["model.embedding.weight"] == [
+        "model.embedding.q_weight",
+        "model.embedding.q_scale",
+    ]
+    assert quant_map.map_func["model.embedding.weight"] == config.quantize_weight
+    assert isinstance(mod.embedding, GroupQuantizeEmbedding)
+
+
+if __name__ == "__main__":
+    test_quantize_weight("q4f16_1", [16, 128], "float16", "llvm")
+    test_quantize_model("q4f16_1", [16, 128], "float16")
+    test_dequantize_weight("q4f16_1", [16, 128], "float16")

tests/python/router/test_router.py

@@ -0,0 +1,101 @@
+import asyncio
+
+from mlc_llm.protocol import openai_api_protocol
+from mlc_llm.router import Router
+
+model_tp1 = "./dist/Llama-3.2-1B-Instruct-q0f16-MLC/"
+model_lib_tp1 = "./dist/lib/Llama-3.2-1B-q0f16-cuda.so"
+# model_lib_tp1 = None
+
+model_tp2 = "./dist/Llama-3.2-1B-Instruct-q0f16-MLC-tp2/"
+model_lib_tp2 = "./dist/lib/Llama-3.2-1B-q0f16-cuda-tp2.so"
+# model_lib_tp2 = None
+
+
+def get_router_1tp1():
+    return (
+        Router(
+            model_tp1,
+            model_lib=model_lib_tp1,
+            hosts=["127.0.0.1"],
+            ports=[8080],
+        ),
+        model_tp1,
+    )
+
+
+def get_router_2tp1():
+    return (
+        Router(
+            model_tp1,
+            model_lib=model_lib_tp1,
+            hosts=["127.0.0.1", "127.0.0.1"],
+            ports=[8080, 8081],
+            device_id_starts=[0, 1],
+            npes=2,
+        ),
+        model_tp1,
+    )
+
+
+def get_router_1tp2():
+    return (
+        Router(
+            model_tp2,
+            model_lib=model_lib_tp2,
+            hosts=["127.0.0.1"],
+            ports=[8080],
+            npes=2,
+        ),
+        model_tp2,
+    )
+
+
+def get_router_2tp2():
+    return (
+        Router(
+            model_tp2,
+            model_lib=model_lib_tp2,
+            hosts=["127.0.0.1", "127.0.0.1"],
+            ports=[8080, 8081],
+            device_id_starts=[0, 2],
+            npes=4,
+        ),
+        model_tp2,
+    )
+
+
+CONFIG_TO_ROUTER = {
+    "1tp1": get_router_1tp1,
+    "2tp1": get_router_2tp1,
+    "1tp2": get_router_1tp2,
+    "2tp2": get_router_2tp2,
+}
+
+
+async def test_router(schedule: str = "round_robin", endpoints_config: str = "1tp1"):
+    router, model_id = CONFIG_TO_ROUTER[endpoints_config]()
+
+    request = openai_api_protocol.CompletionRequest(
+        prompt="The meaning of life ",
+        model=model_id,
+        stream=True,
+        max_tokens=64,
+        stream_options=openai_api_protocol.StreamOptions(include_usage=True),
+    )
+    if schedule == "round_robin":
+        async for chunk in router._handle_completion_round_robin(request, "1"):
+            print(chunk)
+    elif schedule == "disagg":
+        async for chunk in router._handle_completion_disagg(request, "1"):
+            print(chunk)
+    else:
+        raise ValueError(f"Unknown scheduling method: {schedule}")
+    router.terminate()
+
+
+if __name__ == "__main__":
+    # asyncio.run(test_router("round_robin", endpoints_config="1tp1"))
+    # asyncio.run(test_router("round_robin", endpoints_config="1tp2"))
+    # asyncio.run(test_router("round_robin", endpoints_config="2tp1"))
+    asyncio.run(test_router("round_robin", endpoints_config="2tp2"))

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()

tests/python/support/test_auto_config.py

@@ -0,0 +1,43 @@
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from mlc_llm.support import logging
+from mlc_llm.support.auto_config import detect_config
+
+logging.enable_logging()
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+def _create_json_file(json_path, data):
+    with open(json_path, "w", encoding="utf-8") as i_f:
+        json.dump(data, i_f)
+
+
+def test_detect_config():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir)
+        config_json_path = base_path / "config.json"
+        _create_json_file(config_json_path, {})
+
+        assert detect_config(base_path) == config_json_path
+        assert detect_config(config_json_path) == config_json_path
+
+
+def test_detect_config_fail():
+    with pytest.raises(ValueError):
+        detect_config(Path("do/not/exist"))
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir)
+        with pytest.raises(ValueError):
+            assert detect_config(base_path)
+
+
+if __name__ == "__main__":
+    test_detect_config()
+    test_detect_config_fail()

tests/python/support/test_auto_target.py

@@ -0,0 +1,37 @@
+import pytest
+from tvm.target import Target
+
+from mlc_llm.support.auto_target import _apply_webgpu_subgroups
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+def test_apply_webgpu_subgroups_enables_webgpu_target():
+    target = Target("webgpu")
+
+    updated = _apply_webgpu_subgroups(target, True)
+
+    assert updated is not target
+    assert dict(target.export())["supports_subgroups"] is False
+    assert dict(updated.export())["supports_subgroups"] is True
+
+
+def test_apply_webgpu_subgroups_non_webgpu_target_is_unchanged():
+    target = Target("llvm")
+
+    updated = _apply_webgpu_subgroups(target, True)
+
+    assert updated is target
+    assert dict(updated.export()) == dict(target.export())
+
+
+@pytest.mark.parametrize("target_kind", ["webgpu", "llvm"])
+@pytest.mark.parametrize("enable_subgroups", [False, None])
+def test_apply_webgpu_subgroups_disabled_is_unchanged(target_kind, enable_subgroups):
+    target = Target(target_kind)
+
+    updated = _apply_webgpu_subgroups(target, enable_subgroups)
+
+    assert updated is target
+    assert dict(updated.export()) == dict(target.export())

tests/python/support/test_auto_weight.py

@@ -0,0 +1,150 @@
+import json
+import os
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from mlc_llm.support import logging
+from mlc_llm.support.auto_weight import detect_weight
+
+logging.enable_logging()
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+def _create_json_file(json_path, data):
+    with open(json_path, "w", encoding="utf-8") as i_f:
+        json.dump(data, i_f)
+
+
+@pytest.mark.parametrize(
+    "weight_format, index_filename, result",
+    [
+        ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"),
+        (
+            "huggingface-safetensor",
+            "model.safetensors.index.json",
+            "huggingface-safetensor",
+        ),
+        ("auto", "pytorch_model.bin.index.json", "huggingface-torch"),
+        ("auto", "model.safetensors.index.json", "huggingface-safetensor"),
+    ],
+)
+def test_detect_weight(weight_format, index_filename, result):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir)
+        if index_filename is not None:
+            weight_index_file = base_path / index_filename
+            _create_json_file(weight_index_file, {})
+        assert detect_weight(base_path, None, weight_format) == (
+            weight_index_file,
+            result,
+        )
+
+
+@pytest.mark.parametrize(
+    "weight_format, index_filename, result",
+    [
+        ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"),
+        (
+            "huggingface-safetensor",
+            "model.safetensors.index.json",
+            "huggingface-safetensor",
+        ),
+        ("auto", "pytorch_model.bin.index.json", "huggingface-torch"),
+        ("auto", "model.safetensors.index.json", "huggingface-safetensor"),
+    ],
+)
+def test_detect_weight_in_config_json(weight_format, index_filename, result):
+    with (
+        tempfile.TemporaryDirectory() as config_dir,
+        tempfile.TemporaryDirectory() as weight_dir,
+    ):
+        config_path = Path(config_dir)
+        weight_path = Path(weight_dir)
+        config_json_path = config_path / "config.json"
+        _create_json_file(config_json_path, {"weight_path": weight_dir})
+        if index_filename is not None:
+            weight_index_file = weight_path / index_filename
+            _create_json_file(weight_index_file, {})
+
+        assert detect_weight(None, config_json_path, weight_format) == (
+            weight_index_file,
+            result,
+        )
+
+
+@pytest.mark.parametrize(
+    "weight_format, index_filename, result",
+    [
+        ("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"),
+        (
+            "huggingface-safetensor",
+            "model.safetensors.index.json",
+            "huggingface-safetensor",
+        ),
+        ("auto", "pytorch_model.bin.index.json", "huggingface-torch"),
+        ("auto", "model.safetensors.index.json", "huggingface-safetensor"),
+    ],
+)
+def test_detect_weight_same_dir_config_json(weight_format, index_filename, result):
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir)
+        config_json_path = base_path / "config.json"
+        _create_json_file(config_json_path, {})
+        if index_filename is not None:
+            weight_index_file = Path(os.path.join(tmpdir, index_filename))
+            _create_json_file(weight_index_file, {})
+        assert detect_weight(None, config_json_path, weight_format) == (
+            weight_index_file,
+            result,
+        )
+
+
+def test_find_weight_fail():
+    with tempfile.TemporaryDirectory() as tmpdir:
+        base_path = Path(tmpdir)
+        with pytest.raises(ValueError):
+            detect_weight(Path("do/not/exist"), base_path, "awq")
+    with pytest.raises(AssertionError):
+        detect_weight(None, Path("do/not/exist"), "awq")
+
+
+if __name__ == "__main__":
+    test_detect_weight("huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch")
+    test_detect_weight(
+        "huggingface-safetensor",
+        "model.safetensors.index.json",
+        "huggingface-safetensor",
+    )
+    test_detect_weight("auto", "pytorch_model.bin.index.json", "huggingface-torch")
+    test_detect_weight("auto", "model.safetensors.index.json", "huggingface-safetensor")
+    test_detect_weight_in_config_json(
+        "huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"
+    )
+    test_detect_weight_in_config_json(
+        "huggingface-safetensor",
+        "model.safetensors.index.json",
+        "huggingface-safetensor",
+    )
+    test_detect_weight_in_config_json("auto", "pytorch_model.bin.index.json", "huggingface-torch")
+    test_detect_weight_in_config_json(
+        "auto", "model.safetensors.index.json", "huggingface-safetensor"
+    )
+    test_detect_weight_same_dir_config_json(
+        "huggingface-torch", "pytorch_model.bin.index.json", "huggingface-torch"
+    )
+    test_detect_weight_same_dir_config_json(
+        "huggingface-safetensor",
+        "model.safetensors.index.json",
+        "huggingface-safetensor",
+    )
+    test_detect_weight_same_dir_config_json(
+        "auto", "pytorch_model.bin.index.json", "huggingface-torch"
+    )
+    test_detect_weight_same_dir_config_json(
+        "auto", "model.safetensors.index.json", "huggingface-safetensor"
+    )
+    test_find_weight_fail()

tests/python/support/test_cli_convert_weight.py

@@ -0,0 +1,69 @@
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from mlc_llm.cli import convert_weight as convert_weight_cli
+
+pytestmark = [pytest.mark.unittest]
+
+
+def test_convert_weight_cli_passes_lora_adapter(monkeypatch):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        temp_path = Path(tmp_dir)
+        config_path = temp_path / "config.json"
+        source_dir = temp_path / "source"
+        source_dir.mkdir(parents=True, exist_ok=True)
+        source_index = source_dir / "pytorch_model.bin.index.json"
+        adapter_dir = temp_path / "adapter"
+        adapter_dir.mkdir(parents=True, exist_ok=True)
+        output_dir = temp_path / "output"
+
+        config_path.write_text(json.dumps({}), encoding="utf-8")
+        source_index.write_text(json.dumps({"weight_map": {}}), encoding="utf-8")
+
+        def _fake_detect_device(device):
+            return device
+
+        def _fake_detect_weight(_weight_path, _config_json_path, _weight_format):
+            return source_index, "huggingface-torch"
+
+        def _fake_detect_model_type(_model_type, _config):
+            return "dummy"
+
+        monkeypatch.setattr(convert_weight_cli, "detect_config", Path)
+        monkeypatch.setattr(convert_weight_cli, "detect_device", _fake_detect_device)
+        monkeypatch.setattr(convert_weight_cli, "detect_weight", _fake_detect_weight)
+        monkeypatch.setattr(convert_weight_cli, "detect_model_type", _fake_detect_model_type)
+        monkeypatch.setattr(convert_weight_cli, "MODELS", {"dummy": object()})
+        monkeypatch.setattr(convert_weight_cli, "QUANTIZATION", {"q0f16": object()})
+
+        call_args = {}
+
+        def _fake_convert_weight(**kwargs):
+            call_args.update(kwargs)
+
+        monkeypatch.setattr(convert_weight_cli, "convert_weight", _fake_convert_weight)
+
+        convert_weight_cli.main(
+            [
+                str(config_path),
+                "--quantization",
+                "q0f16",
+                "--model-type",
+                "dummy",
+                "--source",
+                str(source_dir),
+                "--source-format",
+                "auto",
+                "--output",
+                str(output_dir),
+                "--lora-adapter",
+                str(adapter_dir),
+            ]
+        )
+
+        assert call_args["lora_adapter"] == adapter_dir
+        assert call_args["source"] == source_index
+        assert call_args["source_format"] == "huggingface-torch"

tests/python/support/test_convert_weight_lora_merge.py

@@ -0,0 +1,108 @@
+import contextlib
+import json
+import tempfile
+from pathlib import Path
+
+import pytest
+
+from mlc_llm.interface import convert_weight as convert_weight_interface
+
+pytestmark = [pytest.mark.unittest]
+
+
+def test_resolve_base_model_dir():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        temp_path = Path(tmp_dir)
+        model_dir = temp_path / "model"
+        model_dir.mkdir(parents=True, exist_ok=True)
+        source_file = model_dir / "pytorch_model.bin.index.json"
+        source_file.write_text(json.dumps({"weight_map": {}}), encoding="utf-8")
+
+        assert convert_weight_interface._resolve_base_model_dir(model_dir) == model_dir
+        assert convert_weight_interface._resolve_base_model_dir(source_file) == model_dir
+
+
+def test_convert_weight_with_lora_uses_merged_source(monkeypatch):
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        temp_path = Path(tmp_dir)
+        config_path = temp_path / "config.json"
+        config_path.write_text(json.dumps({}), encoding="utf-8")
+
+        source_dir = temp_path / "source"
+        source_dir.mkdir(parents=True, exist_ok=True)
+        source_file = source_dir / "pytorch_model.bin.index.json"
+        source_file.write_text(json.dumps({"weight_map": {}}), encoding="utf-8")
+
+        adapter_dir = temp_path / "adapter"
+        adapter_dir.mkdir(parents=True, exist_ok=True)
+
+        merged_dir = temp_path / "merged"
+        merged_dir.mkdir(parents=True, exist_ok=True)
+        merged_file = merged_dir / "pytorch_model.bin"
+        merged_file.write_bytes(b"")
+
+        captured = {}
+
+        @contextlib.contextmanager
+        def _fake_merge(base_source: Path, lora_adapter: Path):
+            captured["merge_base_source"] = base_source
+            captured["merge_lora_adapter"] = lora_adapter
+            yield merged_dir
+
+        def _fake_detect_weight(weight_path: Path, config_json_path: Path, weight_format: str):
+            captured["detect_weight_path"] = weight_path
+            captured["detect_weight_config"] = config_json_path
+            captured["detect_weight_format"] = weight_format
+            return merged_file, "huggingface-torch"
+
+        def _fake_convert_args(args):
+            captured["converted_args"] = args
+
+        monkeypatch.setattr(
+            convert_weight_interface, "_merge_lora_adapter_with_base_model", _fake_merge
+        )
+        monkeypatch.setattr(convert_weight_interface, "detect_weight", _fake_detect_weight)
+        monkeypatch.setattr(convert_weight_interface, "_convert_args", _fake_convert_args)
+        monkeypatch.setattr(convert_weight_interface.ConversionArgs, "display", lambda self: None)
+
+        convert_weight_interface.convert_weight(
+            config=config_path,
+            quantization=object(),
+            model=type("DummyModel", (), {"name": "dummy"})(),
+            device=object(),
+            source=source_file,
+            source_format="huggingface-safetensor",
+            output=temp_path / "output",
+            lora_adapter=adapter_dir,
+        )
+
+        converted_args = captured["converted_args"]
+        assert captured["merge_base_source"] == source_file
+        assert captured["merge_lora_adapter"] == adapter_dir
+        assert captured["detect_weight_path"] == merged_dir
+        assert captured["detect_weight_config"] == config_path
+        assert captured["detect_weight_format"] == "auto"
+        assert converted_args.source == merged_file
+        assert converted_args.source_format == "huggingface-torch"
+        assert converted_args.lora_adapter == adapter_dir
+
+
+def test_convert_weight_with_lora_rejects_awq():
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        temp_path = Path(tmp_dir)
+        config_path = temp_path / "config.json"
+        config_path.write_text(json.dumps({}), encoding="utf-8")
+        adapter_dir = temp_path / "adapter"
+        adapter_dir.mkdir(parents=True, exist_ok=True)
+
+        with pytest.raises(ValueError, match="only supports source formats"):
+            convert_weight_interface.convert_weight(
+                config=config_path,
+                quantization=object(),
+                model=type("DummyModel", (), {"name": "dummy"})(),
+                device=object(),
+                source=temp_path / "source",
+                source_format="awq",
+                output=temp_path / "output",
+                lora_adapter=adapter_dir,
+            )

tests/python/tokenizers/test_streamer.py

@@ -0,0 +1,192 @@
+"""Streamer tests in MLC LLM.
+
+Please specify the local path to llama2 tokenizer via environment
+variable before running this test.
+The recommended way to run the tests is to use the following command:
+  MLC_LLAMA_TOKENIZER_PATH="path/to/llama/tokenizer" \
+  pytest -vv tests/python/support/test_text_streamer_stop_handler.py
+
+Here "MLC_LLAMA_TOKENIZER_PATH" can be chosen from
+- a llama2 weight directory (e.g., "path/to/Llama-2-7b-chat-hf"),
+- a sentencepiece llama2 tokenizer path
+  (e.g., "path/to/Llama-2-7b-chat-hf/tokenizer.model").
+
+To directly run the Python file (a.k.a., not using pytest), you also need to
+specify the tokenizer path via environment variable.
+"""
+
+import time
+from typing import List, Tuple  # noqa: UP035
+
+import pytest
+
+from mlc_llm.testing import require_test_tokenizers
+from mlc_llm.tokenizers import StopStrHandler, TextStreamer, Tokenizer
+
+# test category "unittest"
+pytestmark = [pytest.mark.unittest]
+
+
+# fmt: off
+para_input_tokens = [18585, 29892, 1244, 29915, 29879, 263, 3273, 14880, 1048, 953, 29877, 2397,
+          29892, 988, 1269, 1734, 338, 5643, 491, 385, 953, 29877, 2397, 29901, 13, 13,
+          29950, 1032, 727, 29991, 29871, 243, 162, 148, 142, 306, 29915, 29885, 1244, 304,
+          1371, 1234, 738, 5155, 366, 505, 1048, 953, 29877, 2397, 29871, 243, 162, 167, 151,
+          29889, 7440, 366, 1073, 393, 953, 29877, 2397, 508, 367, 1304, 304, 27769, 23023,
+          1080, 322, 21737, 297, 263, 2090, 322, 1708, 1319, 982, 29973, 29871, 243, 162, 155,
+          135, 2688, 508, 884, 367, 1304, 304, 788, 263, 6023, 310, 2022, 2877, 304, 596, 7191,
+          322, 11803, 29889, 29871, 243, 162, 149, 152, 1126, 29892, 1258, 366, 1073, 393, 727,
+          526, 1584, 953, 29877, 2397, 8090, 322, 14188, 366, 508, 1708, 29973, 29871, 243, 162,
+          145, 177, 243, 162, 148, 131, 1105, 29892, 748, 14432, 322, 679, 907, 1230, 411, 953,
+          29877, 2397, 29991, 29871, 243, 162, 149, 168, 243, 162, 145, 171]
+
+DECODED_PARAGRAPH = (
+    "Sure, here's a short paragraph about emoji, "
+    "where each word is followed by an emoji:\n\n"
+    "Hey there! 👋 I'm here to help answer any questions you have about emoji 🤔. "
+    "Did you know that emoji can be used to convey emotions and feelings in a "
+    "fun and playful way? 😄 "
+    "They can also be used to add a touch of personality to your messages and posts. 💕 "
+    "And, did you know that there are even emoji games and activities you can play? 🎮👀 "
+    "So, go ahead and get creative with emoji! 💥🎨"
+)
+# fmt: on
+
+
+@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_text_streamer(llama_tokenizer_path: str):
+    text_streamer = TextStreamer(Tokenizer(llama_tokenizer_path))
+    total_text = ""
+    for token in para_input_tokens:
+        total_text += text_streamer.put([token])
+    total_text += text_streamer.finish()
+
+    assert total_text == DECODED_PARAGRAPH
+
+
+def stop_handler_process_tokens(
+    stop_handler: StopStrHandler,
+    tokens: List[int],  # noqa: UP006
+    tokenizer: Tokenizer,
+) -> str:
+    returned_tokens = []
+    for token in tokens:
+        returned_tokens += stop_handler.put(token)
+        if stop_handler.stop_triggered:
+            break
+
+    if not stop_handler.stop_triggered:
+        returned_tokens += stop_handler.finish()
+
+    return tokenizer.decode(returned_tokens)
+
+
+@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_stop_str_handler_stop(llama_tokenizer_path: str):
+    stop_strs = [" 🤔"]
+    tokenizer = Tokenizer(llama_tokenizer_path)
+    stop_handler = StopStrHandler(stop_strs, tokenizer)
+
+    total_text = stop_handler_process_tokens(stop_handler, para_input_tokens, tokenizer)
+    expected_text = (
+        "Sure, here's a short paragraph about emoji, "
+        "where each word is followed by an emoji:\n\n"
+        "Hey there! 👋 I'm here to help answer any questions you have about emoji"
+    )
+
+    assert total_text == expected_text
+
+
+@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_stop_str_handler_not_stop(
+    llama_tokenizer_path: str,
+):
+    stop_strs = ["^^"]
+    tokenizer = Tokenizer(llama_tokenizer_path)
+    stop_handler = StopStrHandler(stop_strs, tokenizer)
+
+    total_text = stop_handler_process_tokens(stop_handler, para_input_tokens, tokenizer)
+    assert total_text == DECODED_PARAGRAPH
+
+
+@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_stop_str_handler_return_cached_tokens(
+    llama_tokenizer_path: str,
+):
+    tokens = para_input_tokens[:26]  # until "\n\n"
+    stop_strs = ["\n\n\n"]
+    tokenizer = Tokenizer(llama_tokenizer_path)
+    stop_handler = StopStrHandler(stop_strs, tokenizer)
+
+    total_text = stop_handler_process_tokens(stop_handler, tokens, tokenizer)
+    expected_text = (
+        "Sure, here's a short paragraph about emoji, where each word is followed by an emoji:\n\n"
+    )
+
+    assert total_text == expected_text
+
+
+@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_stop_str_handler_throughput(
+    llama_tokenizer_path: str,
+):
+    stop_strs = ["[INST]"]
+    tokenizer = Tokenizer(llama_tokenizer_path)
+    stop_handler = StopStrHandler(stop_strs, tokenizer)
+
+    tokens = para_input_tokens * 20
+    returned_tokens = []
+
+    tbegin = time.perf_counter()
+    for token in tokens:
+        returned_tokens += stop_handler.put(token)
+        assert not stop_handler.stop_triggered
+    tend = time.perf_counter()
+
+    throughput = len(tokens) / (tend - tbegin)
+    print(
+        f"num tokens = {len(tokens)}, "
+        f"time elapsed = {tend - tbegin:.5f} sec, "
+        f"throughput = {throughput}"
+    )
+    assert throughput >= 100000
+
+
+emoji_tokens_expected_result = [
+    # HF: "�����", SentencePiece: "�👀"
+    ([177, 243, 162, 148, 131], ("�����", "�👀")),
+    # Both: "👀👀"
+    ([243, 162, 148, 131, 243, 162, 148, 131], ("👀👀",)),
+    # Both: "👀👀👀"
+    ([243, 162, 148, 131, 243, 162, 148, 131, 243, 162, 148, 131], ("👀👀👀",)),
+    # HF: "👀�������", SentencePiece: "👀���👀"
+    ([243, 162, 148, 131, 162, 148, 131, 243, 162, 148, 131], ("👀�������", "👀���👀")),
+    # Both: "👀��� have👀"
+    ([243, 162, 148, 131, 162, 148, 131, 505, 243, 162, 148, 131], ("👀��� have👀",)),
+]
+
+
+@pytest.mark.parametrize("tokens_and_results", emoji_tokens_expected_result)
+@require_test_tokenizers("Llama-2-7b-chat-hf-q4f16_1-MLC")
+def test_text_streamer_emojis(
+    llama_tokenizer_path: str,
+    tokens_and_results: Tuple[List[int], Tuple[str]],  # noqa: UP006
+):
+    text_streamer = TextStreamer(Tokenizer(llama_tokenizer_path))
+    total_text = ""
+    tokens, expected_results = tokens_and_results
+    for token in tokens:
+        total_text += text_streamer.put([token])
+    total_text += text_streamer.finish()
+    assert total_text in expected_results
+
+
+if __name__ == "__main__":
+    test_text_streamer()
+    test_stop_str_handler_stop()
+    test_stop_str_handler_not_stop()
+    test_stop_str_handler_return_cached_tokens()
+    test_stop_str_handler_throughput()
+
+    for tokens_and_res in emoji_tokens_expected_result:
+        test_text_streamer_emojis(tokens_and_res)