chore: import upstream snapshot with attribution

python/mlc_llm/__init__.py

@@ -0,0 +1,20 @@
+"""MLC Chat python package.
+
+MLC Chat is the app runtime of MLC LLM.
+"""
+
+from tvm import register_global_func
+
+from . import protocol, serve
+from .libinfo import __version__
+from .serve import AsyncMLCEngine, MLCEngine
+
+
+@register_global_func("runtime.disco.create_socket_session_local_workers", override=True)
+def _create_socket_session_local_workers(num_workers):
+    """Create the local session for each distributed node over socket session."""
+    from tvm.runtime.disco import (
+        ProcessSession,
+    )
+
+    return ProcessSession(num_workers, num_groups=1, entrypoint="mlc_llm.cli.worker")

python/mlc_llm/__main__.py

@@ -0,0 +1,67 @@
+"""Entrypoint of all CLI commands from MLC LLM"""
+
+import sys
+
+from mlc_llm.support import logging
+from mlc_llm.support.argparse import ArgumentParser
+
+logging.enable_logging()
+
+
+def main():
+    """Entrypoint of all CLI commands from MLC LLM"""
+    parser = ArgumentParser("MLC LLM Command Line Interface.")
+    parser.add_argument(
+        "subcommand",
+        type=str,
+        choices=[
+            "compile",
+            "convert_weight",
+            "gen_config",
+            "chat",
+            "serve",
+            "package",
+            "calibrate",
+            "router",
+        ],
+        help="Subcommand to to run. (choices: %(choices)s)",
+    )
+    parsed = parser.parse_args(sys.argv[1:2])
+    if parsed.subcommand == "compile":
+        from mlc_llm.cli import compile as cli
+
+        cli.main(sys.argv[2:])
+    elif parsed.subcommand == "convert_weight":
+        from mlc_llm.cli import convert_weight as cli
+
+        cli.main(sys.argv[2:])
+    elif parsed.subcommand == "gen_config":
+        from mlc_llm.cli import gen_config as cli
+
+        cli.main(sys.argv[2:])
+    elif parsed.subcommand == "chat":
+        from mlc_llm.cli import chat as cli
+
+        cli.main(sys.argv[2:])
+    elif parsed.subcommand == "serve":
+        from mlc_llm.cli import serve as cli
+
+        cli.main(sys.argv[2:])
+    elif parsed.subcommand == "package":
+        from mlc_llm.cli import package as cli
+
+        cli.main(sys.argv[2:])
+    elif parsed.subcommand == "calibrate":
+        from mlc_llm.cli import calibrate as cli
+
+        cli.main(sys.argv[2:])
+    elif parsed.subcommand == "router":
+        from mlc_llm.cli import router as cli
+
+        cli.main(sys.argv[2:])
+    else:
+        raise ValueError(f"Unknown subcommand {parsed.subcommand}")
+
+
+if __name__ == "__main__":
+    main()

python/mlc_llm/base.py

@@ -0,0 +1,45 @@
+"""Load MLC LLM library and _ffi_api functions."""
+
+import ctypes
+import os
+import sys
+
+import tvm
+import tvm.base
+
+from . import libinfo
+
+SKIP_LOADING_MLCLLM_SO = os.environ.get("SKIP_LOADING_MLCLLM_SO", "0")
+
+
+def _load_mlc_llm_lib():
+    """Load MLC LLM lib"""
+    if sys.platform.startswith("win32") and sys.version_info >= (3, 8):
+        for path in libinfo.get_dll_directories():
+            os.add_dll_directory(path)
+    lib_name = "mlc_llm" if tvm.base._RUNTIME_ONLY else "mlc_llm_module"
+    lib_path = libinfo.find_lib_path(lib_name, optional=False)
+    return ctypes.CDLL(lib_path[0]), lib_path[0]
+
+
+@tvm.register_global_func("mlc.debug_cuda_profiler_start")
+def _debug_cuda_profiler_start() -> None:
+    """Start cuda profiler."""
+    import cuda
+    import cuda.cudart
+
+    cuda.cudart.cudaProfilerStart()
+
+
+@tvm.register_global_func("mlc.debug_cuda_profiler_stop")
+def _debug_cuda_profiler_stop() -> None:
+    """Stop cuda profiler."""
+    import cuda
+    import cuda.cudart
+
+    cuda.cudart.cudaProfilerStop()
+
+
+# only load once here
+if SKIP_LOADING_MLCLLM_SO == "0":
+    _LIB, _LIB_PATH = _load_mlc_llm_lib()

python/mlc_llm/bench/__init__.py

@@ -0,0 +1 @@
+"""Subdirectory of bench."""

python/mlc_llm/bench/__main__.py

@@ -0,0 +1,403 @@
+"""MLC LLM benchmark main entrance"""
+
+import functools
+import json
+import random
+from typing import Any, Dict, List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+import requests
+from transformers import AutoTokenizer
+
+import mlc_llm
+from mlc_llm.bench.api_endpoint import SUPPORTED_BACKENDS, create_api_endpoint
+from mlc_llm.bench.dataset import SUPPORTED_DATASET, Dataset, create_dataset
+from mlc_llm.bench.request_processor import (
+    MetricAnalyzer,
+    RequestProcessor,
+    create_pipelines,
+)
+from mlc_llm.bench.request_record import (
+    RequestRecord,
+    convert_reports_to_df,
+    generate_metrics_summary,
+    pretty_print_report,
+)
+from mlc_llm.cli.serve import EngineConfigOverride
+from mlc_llm.serve import EngineConfig
+from mlc_llm.support import argparse, logging
+
+logging.enable_logging()
+logger = logging.getLogger(__name__)
+
+
+def _parse_num_concurrent_requests(num_str: Optional[str]) -> Optional[List[int]]:  # noqa: UP006
+    if num_str is None:
+        return None
+    numbers = num_str.split(",")
+    if any(not number.isdigit() for number in numbers):
+        raise ValueError(f"Unrecognized num_concurrent_requests list: {numbers}")
+    return list(int(number) for number in numbers)
+
+
+def _parse_request_rate(request_rate_str: Optional[str]) -> Optional[List[np.float32]]:  # noqa: UP006
+    if request_rate_str is None:
+        return None
+    request_rates = request_rate_str.split(",")
+    results = []
+    for rate_str in request_rates:
+        request_rate = float(rate_str)
+        if request_rate <= 0:
+            raise ValueError(f"Invalid request rate {request_rate}")
+        results.append(np.float32(request_rate))
+    return results
+
+
+def _parse_mlc_engine_config(config_str: Optional[str]) -> EngineConfig:
+    if config_str is None:
+        return None
+    engine_config_override = EngineConfigOverride.from_str(config_str)
+    return EngineConfig(
+        tensor_parallel_shards=engine_config_override.tensor_parallel_shards,
+        max_num_sequence=engine_config_override.max_num_sequence,
+        max_total_sequence_length=engine_config_override.max_total_seq_length,
+        prefill_chunk_size=engine_config_override.prefill_chunk_size,
+        sliding_window_size=engine_config_override.sliding_window_size,
+        attention_sink_size=engine_config_override.attention_sink_size,
+        max_history_size=engine_config_override.max_history_size,
+        gpu_memory_utilization=engine_config_override.gpu_memory_utilization,
+        spec_draft_length=engine_config_override.spec_draft_length,
+        prefill_mode=engine_config_override.prefill_mode,
+        prefix_cache_max_num_recycling_seqs=engine_config_override.prefix_cache_max_num_recycling_seqs,
+        prefix_cache_mode=engine_config_override.prefix_cache_mode,
+    )
+
+
+def _launch_mlc_server(args: argparse.argparse.Namespace):
+    return mlc_llm.serve.PopenServer(
+        model=args.tokenizer,
+        mode="server",
+        model_lib=args.mlc_model_lib,
+        enable_tracing=False,
+        host=args.host,
+        port=args.port,
+        engine_config=args.mlc_engine_config,
+    )
+
+
+def run_pipeline(
+    pipeline: RequestProcessor,
+    dataset: Dataset,
+    tokenizer: AutoTokenizer,
+    args: argparse.argparse.Namespace,
+) -> Tuple[Dict[str, Any], List[RequestRecord]]:  # noqa: UP006
+    """Run the pipeline with the given dataset and args. Return the benchmark report dict."""
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    request_records = dataset.generate_request_records(
+        args.input_len,
+        args.output_len,
+        args.input_len_std,
+        args.output_len_std,
+    )
+    request_records = pipeline(request_records)
+    num_total_requests = (
+        args.num_requests if not args.per_gpu_workload else args.num_requests * args.num_gpus
+    )
+    assert len(request_records) == num_total_requests
+    sorted_requests: List[RequestRecord] = [None] * num_total_requests  # noqa: UP006
+    for request_record in request_records:
+        assert request_record.request_id is not None
+        assert sorted_requests[request_record.request_id] is None
+        sorted_requests[request_record.request_id] = request_record
+
+    request_records = MetricAnalyzer(tokenizer)(request_records)
+    report = generate_metrics_summary(request_records, num_total_requests, args.num_gpus)
+    return report, sorted_requests
+
+
+def query_mlc_server_metrics(host: str, port: int):
+    """Try to get the MLC server metrics whenever it exists."""
+    try:
+        r = requests.post(f"http://{host}:{port}/debug/dump_engine_metrics", json={}, timeout=10)
+        if r.status_code == 200:
+            print(f"MLC server metrics: {r.json()}")
+    except Exception:
+        pass
+
+
+def main(args: argparse.argparse.Namespace):
+    """Main benchmark entrance."""
+    mlc_server = None
+    if args.mlc_model_lib:
+        mlc_server = _launch_mlc_server(args)
+    if args.num_requests <= 0:
+        raise ValueError("Number of requests to benchmark must be positive.")
+
+    def _main():
+        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
+        dataset = create_dataset(args, tokenizer)
+        f_create_api_endpoint = functools.partial(create_api_endpoint, args)
+        pipelines = create_pipelines(args, f_create_api_endpoint, dataset)
+        reports = []
+        alltime_records = {}
+        for i, pipeline in enumerate(pipelines):
+            report, request_records = run_pipeline(pipeline, dataset, tokenizer, args)
+            exec_feature = (
+                json.dumps(report["exec_feature"])
+                if report["exec_feature"] is not None
+                else f"pipeline{i}"
+            )
+            alltime_records[exec_feature] = [
+                request_record.model_dump() for request_record in request_records
+            ]
+            reports.append(report)
+            pretty_print_report(report)
+        query_mlc_server_metrics(args.host, args.port)
+
+        # Construct data frame
+        df = convert_reports_to_df(reports)
+        print(df)
+        df.to_csv(args.output, index=False)
+        logger.info("Benchmark results dumped to file %s", args.output)
+        if args.debug_dump:
+            debug_dump_filepath = (
+                args.output[:-4] if args.output.endswith(".csv") else args.output
+            ) + "_debug_dump.log"
+            with open(debug_dump_filepath, "w", encoding="utf-8") as file:
+                json.dump(alltime_records, file, indent=4)
+            logger.info("Debug log dumped to file %s", debug_dump_filepath)
+
+    if mlc_server is not None:
+        with mlc_server:
+            _main()
+    else:
+        _main()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser("MLC LLM benchmark")
+
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        choices=SUPPORTED_DATASET,
+        help=f"The benchmark dataset kind. Supporting {SUPPORTED_DATASET}",
+    )
+    parser.add_argument(
+        "--dataset-path",
+        type=str,
+        help="The dataset file path.",
+    )
+    parser.add_argument(
+        "--api-endpoint",
+        type=str,
+        choices=SUPPORTED_BACKENDS,
+        default="openai",
+        help="The API endpoint API for benchmarking.",
+    )
+    parser.add_argument(
+        "--tokenizer",
+        type=str,
+        required=True,
+        help="The path of the tokenizer directory.",
+    )
+    parser.add_argument(
+        "--num-gpus",
+        type=int,
+        required=True,
+        help="The number of GPUs used by the server. "
+        "We need this to better analyze the throughput per GPU.",
+    )
+    parser.add_argument(
+        "--num-requests",
+        type=int,
+        required=True,
+        help="The number of requests for benchmark.",
+    )
+    parser.add_argument(
+        "--num-warmup-requests",
+        type=int,
+        help="The number of requests for warmup. "
+        "It is optional when fixing the number of concurrent requests, and is required otherwise.",
+    )
+    parser.add_argument(
+        "--per-gpu-workload",
+        default=False,
+        action="store_true",
+        help='When set to True, the specified "num_concurrent_requests"/"request_rate" '
+        "denote the workload **per GPU**, which means that the real values of "
+        '"num_concurrent_requests"/"request_rate" used in benchmark'
+        'will be multiplied by "num_gpus".',
+    )
+    parser.add_argument(
+        "--num-concurrent-requests",
+        type=_parse_num_concurrent_requests,
+        help="The number(s) of concurrent requests to benchmark. "
+        'It can be either one integer or a list of integer separated by commas(","). '
+        "When specified, for each integer, the benchmark keeps these many consistent "
+        "number of concurrently running requests.",
+    )
+    parser.add_argument(
+        "--request-rate",
+        type=_parse_request_rate,
+        help="The request rate(s) denoting the number of new requests each second. "
+        'It can be either one float number (or "inf") or a list of numbers separated '
+        'by commas(","). '
+        "When specified, the benchmark sends these many new requests each second. "
+        'If it is "inf", all requests will be sent together at once.',
+    )
+    parser.add_argument(
+        "--replay-timestamp-scale",
+        type=float,
+        help="The timestamp scale when replaying the timestamps in a dataset. "
+        'The dataset replay mode is enabled when neither "--num-concurrent-requests" and '
+        '"--request-rate" is specified. '
+        "The scale is 1 by default in the replay mode.",
+    )
+    parser.add_argument(
+        "--input-len",
+        type=int,
+        help="The benchmark request average input length. Default to None, "
+        "which means the request input length depends on the dataset being used.",
+    )
+    parser.add_argument(
+        "--input-len-std",
+        type=float,
+        default=0,
+        help="The benchmark request input length standard deviation. Default to 0.",
+    )
+    parser.add_argument(
+        "--output-len",
+        type=int,
+        help="The benchmark request average output length. Default to None, "
+        "which means the request output length depends on the dataset being used.",
+    )
+    parser.add_argument(
+        "--output-len-std",
+        type=float,
+        default=0,
+        help="The benchmark request output length standard deviation. Default to 0.",
+    )
+    parser.add_argument(
+        "--stream",
+        type=bool,
+        default=True,
+        help="Whether to benchmark stream responses. "
+        "When not enabled, metrics such as time-to-first-token (TTFT) will not be available. "
+        "Default to True.",
+    )
+    parser.add_argument(
+        # NOTE: The current implementation of server metrics still has some issues that need fixes,
+        # which makes it not work to include server metrics.
+        "--include-server-metrics",
+        action="store_true",
+        help="Whether to also benchmark the server side request metrics. "
+        "This option is only available when benchmarking MLC server.",
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        required=True,
+        help="The host address of the backend API.",
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        required=True,
+        help="The port of the backend API.",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=3 * 60 * 60,
+        help="The timeout limit of each request.",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help="The random number seed. Default to 0.",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=1.0,
+        help="The temperature value for logit adjustment. Default to 1.",
+    )
+    parser.add_argument(
+        "--top-p",
+        type=float,
+        default=1.0,
+        help="The top-p value for sampling. Default to 1.",
+    )
+    parser.add_argument(
+        "--ignore-eos",
+        default=False,
+        action="store_true",
+        help='Whether to set the "ignore_eos" field.',
+    )
+    parser.add_argument(
+        "--apply-chat-template",
+        default=False,
+        action="store_true",
+        help="Whether to apply chat template to the request input text. "
+        'It is not supported when "--input-len" is specified.',
+    )
+    parser.add_argument(
+        "--num-process-workers",
+        type=int,
+        help="The number of parallel process workers to send the requests.",
+    )
+    parser.add_argument(
+        "--disable-tqdm",
+        action="store_true",
+        help="Whether to disable showing progress bar with tqdm during benchmarking.",
+    )
+    parser.add_argument(
+        "--max-schedule-gap",
+        type=float,
+        default=0.5,
+        help="The maximum allowed delay between the scheduled time in seconds.",
+    )
+    parser.add_argument(
+        "--mlc-model-lib",
+        type=str,
+        help="The model lib path when benchmarking MLC serve. "
+        "When specified, the server is automatic launched and no external server launch is needed.",
+    )
+    parser.add_argument(
+        "--mlc-engine-config",
+        type=_parse_mlc_engine_config,
+        help="The engine config used when launch MLC server.",
+    )
+    parser.add_argument(
+        "--cuda-profile",
+        default=False,
+        action="store_true",
+        help="Whether to enable cuda profile on server. "
+        "The --mlc-model-lib path should be provided when enabling this option.",
+    )
+    parser.add_argument(
+        "--debug-dump",
+        default=False,
+        action="store_true",
+        help="Whether to dump all request record raw data to file.",
+    )
+    parser.add_argument(
+        "--multi-round",
+        default=False,
+        action="store_true",
+        help="Whether to chat like multi round conversion with history log each request. "
+        "Only enabled when benchmarked with fixed concurrent request mode."
+        "The --num-concurrent-requests should be provided when enabling this option.",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        default="mlc_benchmark.csv",
+        help="The path of the output file where to dump the benchmark results.",
+    )
+
+    main(parser.parse_args())

python/mlc_llm/bench/api_endpoint.py

@@ -0,0 +1,452 @@
+"""MLC LLM bench backends"""
+
+import argparse
+import json
+import os
+import time
+import traceback
+from typing import Optional
+
+from typing_extensions import Self
+
+from mlc_llm.bench.request_record import Metrics, RequestRecord, ServerMetrics
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+class APIEndPoint:
+    """Manages the sending of requests to a specified API endpoint and gathers
+    inference statistics.
+    """
+
+    def __init__(self, include_server_metrics: bool = False) -> None:
+        self.include_server_metrics = include_server_metrics
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+        pass
+
+    async def __call__(self, request: RequestRecord) -> RequestRecord:
+        raise NotImplementedError()
+
+
+class OpenAIChatEndPoint(APIEndPoint):
+    """The backend of sending HTTP requests in OpenAI API through "v1/chat/completions"."""
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        timeout: Optional[float] = None,
+        include_server_metrics: bool = False,
+    ) -> None:
+        super().__init__(include_server_metrics=include_server_metrics)
+
+        import aiohttp
+
+        self.timeout = timeout
+        self.client: aiohttp.ClientSession = None
+        self.url = f"http://{host}:{port}/v1/chat/completions"
+        self.headers = {"Content-Type": "application/json"}
+        if os.getenv("MLC_LLM_API_KEY"):
+            self.headers["Authorization"] = f"Bearer {os.getenv('MLC_LLM_API_KEY')}"
+
+    async def __aenter__(self) -> Self:
+        import aiohttp
+
+        self.client = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(self.timeout))
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+        await self.client.close()
+
+    async def __call__(self, request_record: RequestRecord) -> RequestRecord:
+        payload = request_record.chat_cmpl.model_dump()
+        if self.timeout is not None and "timeout" not in payload:
+            payload["timeout"] = self.timeout
+        if self.include_server_metrics:
+            if "stream_options" not in payload or payload["stream_options"] is None:
+                payload["stream_options"] = {"include_usage": True}
+            else:
+                payload["stream_options"]["include_usage"] = True
+        if (
+            request_record.chat_cmpl.debug_config is not None
+            and request_record.chat_cmpl.debug_config.ignore_eos
+        ):
+            payload["ignore_eos"] = True
+
+        generated_text = ""
+        first_chunk_output_str = ""
+        time_to_first_token_s = None
+        start_time = time.monotonic()
+        server_metrics = None
+
+        try:
+            async with self.client.post(self.url, json=payload, headers=self.headers) as response:
+                assert response.status == 200, await response.text()
+                if payload["stream"]:
+                    async for chunk in response.content:
+                        chunk = chunk.strip()
+                        if not chunk or chunk == b"\n":
+                            continue
+                        # Get rid of the prefix "data: " and suffix "\n"
+                        raw_data = chunk[6:].strip()
+                        if raw_data == b"[DONE]":
+                            continue
+                        data = json.loads(raw_data)
+                        if not data["choices"]:
+                            continue
+                        delta = data["choices"][0]["delta"]
+                        content = delta.get("content", None)
+                        if content is not None and not time_to_first_token_s:
+                            time_to_first_token_s = time.monotonic() - start_time
+                            first_chunk_output_str = content
+                        if self.include_server_metrics and data["usage"] is not None:
+                            # fmt: off
+                            server_metrics = ServerMetrics(
+                                input_tokens=data["usage"]["extra"]["prompt_tokens"],
+                                prefill_tokens=data["usage"]["extra"]["prefill_tokens"],
+                                output_tokens=data["usage"]["extra"]["completion_tokens"],
+                                end_to_end_latency_s=data["usage"]["extra"]["end_to_end_latency_s"],
+                                prefill_tokens_per_s=data["usage"]["extra"]["prefill_tokens_per_s"],
+                                inter_token_latency_s=data["usage"]["extra"]["inter_token_latency_s"],
+                                time_per_output_token_s=1 / data["usage"]["extra"]["decode_tokens_per_s"],  # noqa: E501
+                                time_to_first_token_s=data["usage"]["extra"]["ttft_s"],
+                            )
+                            # fmt: on
+
+                        if content is not None:
+                            generated_text += content
+                else:
+                    data = await response.json()
+                    generated_text = data["choices"][0]["message"]["content"]
+                    if self.include_server_metrics and data["usage"] is not None:
+                        # fmt: off
+                        server_metrics = ServerMetrics(
+                            input_tokens=data["usage"]["extra"]["prompt_tokens"],
+                            prefill_tokens=data["usage"]["extra"]["prefill_tokens"],
+                            output_tokens=data["usage"]["extra"]["completion_tokens"],
+                            end_to_end_latency_s=data["usage"]["extra"]["end_to_end_latency_s"],
+                            prefill_tokens_per_s=data["usage"]["extra"]["prefill_tokens_per_s"],
+                            inter_token_latency_s=data["usage"]["extra"]["inter_token_latency_s"],
+                            time_per_output_token_s=1 / data["usage"]["extra"]["decode_tokens_per_s"],  # noqa: E501
+                            time_to_first_token_s=data["usage"]["extra"]["ttft_s"],
+                        )
+                        # fmt: on
+        except Exception:
+            error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
+            logger.info(error_msg)
+            finish_time = time.monotonic()
+            request_record.output_str = generated_text
+            request_record.first_chunk_output_str = first_chunk_output_str
+            request_record.metrics = Metrics(
+                success=False,
+                start_time=start_time,
+                finish_time=finish_time,
+                end_to_end_latency_s=finish_time - start_time,
+                input_tokens=request_record.metrics.input_tokens,
+                time_to_first_token_s=time_to_first_token_s,
+                server_metrics=server_metrics,
+                exec_feature=request_record.metrics.exec_feature,
+            )
+            request_record.error_msg = error_msg
+            return request_record
+
+        finish_time = time.monotonic()
+        request_record.output_str = generated_text
+        request_record.first_chunk_output_str = first_chunk_output_str
+        success = True
+        error_msg = None
+        if len(generated_text) == 0:
+            success = False
+            error_msg = "Empty generated text."
+        request_record.metrics = Metrics(
+            success=success,
+            start_time=start_time,
+            finish_time=finish_time,
+            end_to_end_latency_s=finish_time - start_time,
+            input_tokens=request_record.metrics.input_tokens,
+            time_to_first_token_s=time_to_first_token_s,
+            server_metrics=server_metrics,
+            exec_feature=request_record.metrics.exec_feature,
+        )
+        request_record.error_msg = error_msg
+        return request_record
+
+
+class OpenAIEndPoint(APIEndPoint):
+    """The backend of sending HTTP requests in OpenAI API through "v1/completions"."""
+
+    def __init__(
+        self,
+        host: str,
+        port: int,
+        timeout: Optional[float] = None,
+        include_server_metrics: bool = False,
+        no_debug_config: bool = False,
+    ) -> None:
+        super().__init__(include_server_metrics=include_server_metrics)
+
+        import aiohttp
+
+        self.timeout = timeout
+        self.client: aiohttp.ClientSession = None
+        self.url = f"http://{host}:{port}/v1/completions"
+        self.headers = {"Content-Type": "application/json"}
+        if os.getenv("MLC_LLM_API_KEY"):
+            self.headers["Authorization"] = f"Bearer {os.getenv('MLC_LLM_API_KEY')}"
+        assert not include_server_metrics, (
+            '"include_server_metrics" only works for "openai-chat" endpoint for now'
+        )
+        self.no_debug_config = no_debug_config
+
+    async def __aenter__(self) -> Self:
+        import aiohttp
+
+        self.client = aiohttp.ClientSession()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+        await self.client.close()
+
+    async def __call__(self, request_record: RequestRecord) -> RequestRecord:
+        assert len(request_record.chat_cmpl.messages) == 1, (
+            'Endpoint "openai" does not support system prompt and multi-round conversation.'
+        )
+        assert isinstance(request_record.chat_cmpl.messages[0].content, str)
+        payload = {
+            "model": request_record.chat_cmpl.model,
+            "prompt": request_record.chat_cmpl.messages[0].content,
+            "temperature": request_record.chat_cmpl.temperature,
+            "top_p": request_record.chat_cmpl.top_p,
+            "max_tokens": request_record.chat_cmpl.max_tokens,
+            "stream": True,
+        }
+        if self.timeout is not None and "timeout" not in payload:
+            payload["timeout"] = self.timeout
+        if (
+            request_record.chat_cmpl.debug_config is not None
+            and request_record.chat_cmpl.debug_config.ignore_eos
+        ):
+            payload["ignore_eos"] = True
+            if not self.no_debug_config:
+                payload["debug_config"] = {"ignore_eos": True}
+
+        generated_text = ""
+        first_chunk_output_str = ""
+        time_to_first_token_s = None
+        start_time = time.monotonic()
+
+        try:
+            async with self.client.post(
+                self.url, json=payload, headers=self.headers, timeout=3600
+            ) as response:
+                assert response.status == 200, await response.text()
+                if payload["stream"]:
+                    async for chunk in response.content:
+                        chunk = chunk.strip()
+                        if not chunk or chunk == b"\n":
+                            continue
+                        # Get rid of the prefix "data: " and suffix "\n"
+                        raw_data = chunk[6:].strip()
+                        if raw_data == b"[DONE]":
+                            continue
+                        data = json.loads(raw_data)
+                        if not data["choices"]:
+                            continue
+                        content = data["choices"][0]["text"]
+                        if content is not None and not time_to_first_token_s:
+                            time_to_first_token_s = time.monotonic() - start_time
+                            first_chunk_output_str = content
+                        if content is not None:
+                            generated_text += content
+                else:
+                    data = await response.json()
+                    generated_text = data["choices"][0]["message"]["content"]
+        except Exception:
+            error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
+            logger.info(error_msg)
+            finish_time = time.monotonic()
+            request_record.output_str = generated_text
+            request_record.first_chunk_output_str = first_chunk_output_str
+            request_record.metrics = Metrics(
+                success=False,
+                start_time=start_time,
+                finish_time=finish_time,
+                end_to_end_latency_s=finish_time - start_time,
+                input_tokens=request_record.metrics.input_tokens,
+                time_to_first_token_s=time_to_first_token_s,
+                server_metrics=None,
+                exec_feature=request_record.metrics.exec_feature,
+            )
+            request_record.error_msg = error_msg
+            return request_record
+
+        finish_time = time.monotonic()
+        request_record.output_str = generated_text
+        request_record.first_chunk_output_str = first_chunk_output_str
+        success = True
+        error_msg = None
+        if len(generated_text) == 0:
+            success = False
+            error_msg = "Empty generated text."
+        request_record.metrics = Metrics(
+            success=success,
+            start_time=start_time,
+            finish_time=finish_time,
+            end_to_end_latency_s=finish_time - start_time,
+            input_tokens=request_record.metrics.input_tokens,
+            time_to_first_token_s=time_to_first_token_s,
+            server_metrics=None,
+            exec_feature=request_record.metrics.exec_feature,
+        )
+        request_record.error_msg = error_msg
+        return request_record
+
+
+class TensorRTLLMEndPoint(APIEndPoint):
+    """The backend of sending HTTP requests in TensorRT-LLM API."""
+
+    def __init__(self, host: str, port: int, timeout: Optional[float] = None) -> None:
+        super().__init__(include_server_metrics=False)
+
+        import aiohttp
+
+        self.timeout = timeout
+        self.client: aiohttp.ClientSession = None
+        self.url_stream = f"http://{host}:{port}/v2/models/ensemble/generate_stream"
+        self.url_no_stream = f"http://{host}:{port}/v2/models/ensemble/generate"
+
+    async def __aenter__(self) -> Self:
+        import aiohttp
+
+        self.client = aiohttp.ClientSession()
+        return self
+
+    async def __aexit__(self, exc_type, exc_value, tb) -> None:
+        await self.client.close()
+
+    async def __call__(self, request_record: RequestRecord) -> RequestRecord:
+        assert len(request_record.chat_cmpl.messages) == 1
+        assert isinstance(request_record.chat_cmpl.messages[0].content, str)
+        payload = {
+            "accumulate_tokens": True,
+            "text_input": request_record.chat_cmpl.messages[0].content,
+            "temperature": (
+                max(request_record.chat_cmpl.temperature, 1e-5)
+                if request_record.chat_cmpl.temperature
+                else 1
+            ),
+            "top_p": request_record.chat_cmpl.top_p if request_record.chat_cmpl.top_p else 1,
+            "max_tokens": request_record.chat_cmpl.max_tokens,
+            "stream": request_record.chat_cmpl.stream,
+        }
+        if (
+            request_record.chat_cmpl.debug_config is not None
+            and request_record.chat_cmpl.debug_config.ignore_eos
+        ):
+            payload["min_length"] = payload["max_tokens"]
+        if self.timeout is not None and "timeout" not in payload:
+            payload["timeout"] = self.timeout
+
+        generated_text = ""
+        first_chunk_output_str = ""
+        url = self.url_stream if request_record.chat_cmpl.stream else self.url_no_stream
+        time_to_first_token_s = None
+        start_time = time.monotonic()
+
+        try:
+            async with self.client.post(url, json=payload) as response:
+                assert response.status == 200, await response.text()
+                if payload["stream"]:
+                    async for chunk in response.content:
+                        chunk = chunk.strip()
+                        if not chunk or chunk == b"\n":
+                            continue
+                        # Get rid of the prefix "data:" and suffix "\n"
+                        raw_data = chunk[5:].strip()
+                        data = json.loads(raw_data)
+                        delta = data["text_output"]
+                        if delta is None:
+                            continue
+
+                        if not time_to_first_token_s:
+                            time_to_first_token_s = time.monotonic() - start_time
+                            first_chunk_output_str = delta
+                        generated_text += delta
+                else:
+                    data = await response.json()
+                    generated_text = data["text_output"]
+        except Exception:
+            error_msg = "API endpoint errored when sending request: " + traceback.format_exc()
+            logger.info(error_msg)
+            finish_time = time.monotonic()
+            request_record.output_str = generated_text
+            request_record.first_chunk_output_str = first_chunk_output_str
+            request_record.metrics = Metrics(
+                success=False,
+                start_time=start_time,
+                finish_time=finish_time,
+                end_to_end_latency_s=finish_time - start_time,
+                input_tokens=request_record.metrics.input_tokens,
+                time_to_first_token_s=time_to_first_token_s,
+                exec_feature=request_record.metrics.exec_feature,
+            )
+            request_record.error_msg = error_msg
+            return request_record
+
+        finish_time = time.monotonic()
+        request_record.output_str = generated_text
+        request_record.first_chunk_output_str = first_chunk_output_str
+        success = True
+        error_msg = None
+        if len(generated_text) == 0:
+            success = False
+            error_msg = "Empty generated text."
+        request_record.metrics = Metrics(
+            success=success,
+            start_time=start_time,
+            finish_time=finish_time,
+            end_to_end_latency_s=finish_time - start_time,
+            input_tokens=request_record.metrics.input_tokens,
+            time_to_first_token_s=time_to_first_token_s,
+            exec_feature=request_record.metrics.exec_feature,
+        )
+        request_record.error_msg = error_msg
+        return request_record
+
+
+# Todo: APIEndPoint with AsyncOpenAI Python interface
+# class OpenAIPythonEndPoint(APIEndPoint):
+#     pass
+
+SUPPORTED_BACKENDS = [
+    "openai",
+    "openai-chat",
+    "mlc",
+    "sglang",
+    "tensorrt-llm",
+    "vllm",
+]
+
+
+def create_api_endpoint(args: argparse.Namespace) -> APIEndPoint:
+    """Create an API endpoint instance with regard to the specified endpoint kind."""
+    if args.api_endpoint in ["openai", "mlc", "sglang"]:
+        return OpenAIEndPoint(args.host, args.port, args.timeout, args.include_server_metrics)
+    if args.api_endpoint == "vllm":
+        return OpenAIEndPoint(
+            args.host,
+            args.port,
+            args.timeout,
+            include_server_metrics=False,
+            no_debug_config=True,
+        )
+    if args.api_endpoint == "openai-chat":
+        return OpenAIChatEndPoint(args.host, args.port, args.timeout, args.include_server_metrics)
+    if args.api_endpoint == "tensorrt-llm":
+        return TensorRTLLMEndPoint(args.host, args.port, args.timeout)
+    raise ValueError(f'Unrecognized endpoint "{args.api_endpoint}"')

python/mlc_llm/bench/dataset.py

@@ -0,0 +1,878 @@
+"""MLC LLM benchmark dataset classes"""
+
+import argparse
+import json
+import random
+from datetime import datetime
+from typing import ClassVar, Dict, List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+import pandas as pd
+from datasets import load_dataset
+from transformers import AutoTokenizer
+
+from mlc_llm.bench.request_record import GroupedRequestRecord, Metrics, RequestRecord
+from mlc_llm.protocol.openai_api_protocol import (
+    ChatCompletionMessage,
+    ChatCompletionRequest,
+    DebugConfig,
+)
+
+
+class Dataset:
+    """The dataset base class."""
+
+    # We set a truncation limit of 100k.
+    truncate_length = int(1e5)
+    # For some that datasets (e.g., dataset that has shared common prefix),
+    # we need fake warmup requests to avoid prefilling common prefixes to the engine.
+    require_fake_warmup: bool = False
+    # Whether the dataset contains timestamps already.
+    # If the dataset comes with timestamps, the benchmark can just replay
+    # the requests according to their timestamps.
+    timestamp_available: bool = False
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        """Get the raw unprocessed request records of the dataset."""
+        raise NotImplementedError()
+
+
+class ShareGPTDataset(Dataset):
+    """The dataset class for ShareGPT dataset."""
+
+    _tokenized_dataset: List[Tuple[str, List[int], int]]  # noqa: UP006
+    apply_chat_template: bool
+
+    def __init__(
+        self, dataset_path: str, tokenizer: AutoTokenizer, apply_chat_template: bool
+    ) -> None:
+        self.apply_chat_template = apply_chat_template
+        with open(dataset_path, encoding="utf-8") as f:
+            raw_dataset = json.load(f)
+        # Filter out the conversations with less than 2 turns.
+        _dataset = [
+            (data["conversations"][0]["value"], data["conversations"][1]["value"])
+            for data in raw_dataset
+            if len(data["conversations"]) >= 2 and data["conversations"][0]["from"] == "human"
+        ]
+        # Tokenize the prompts and completions.
+        self.tokenizer = tokenizer
+        prompts = [prompt for prompt, _ in _dataset]
+        if apply_chat_template:
+            assert getattr(tokenizer, "chat_template", None) is not None, (
+                '"--apply-chat-template" is set but the tokenizer does not have chat template.'
+            )
+            prompts = [
+                tokenizer.apply_chat_template(
+                    [{"role": "user", "content": prompt}],
+                    add_generation_prompt=True,
+                    tokenize=False,
+                )
+                for prompt in prompts
+            ]
+
+        prompt_token_ids = list(
+            tokenizer(
+                prompts,
+                truncation=True,
+                max_length=min(tokenizer.model_max_length, self.truncate_length),
+                add_special_tokens=False,
+            ).input_ids
+        )
+        completions = [completion for _, completion in _dataset]
+        completion_token_ids = tokenizer(
+            completions,
+            truncation=True,
+            max_length=min(tokenizer.model_max_length, self.truncate_length),
+            add_special_tokens=False,
+        ).input_ids
+        self._tokenized_dataset: List[Tuple[str, List[int], int]] = []  # noqa: UP006
+        for i in range(len(_dataset)):
+            if (
+                len(prompt_token_ids[i]) < 4
+                or len(completion_token_ids[i]) < 4
+                or len(prompt_token_ids[i]) + len(completion_token_ids[i])
+                >= min(tokenizer.model_max_length, 8192)
+            ):
+                # Filter out sequences that are too short or too long
+                continue
+            self._tokenized_dataset.append(
+                (prompts[i], prompt_token_ids[i], len(completion_token_ids[i]))
+            )
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if self.apply_chat_template:
+            assert input_len is None, (
+                '"--apply-chat-template" is not supported when "--input-len" is specified.'
+            )
+
+        request_records = []
+        for prompt, input_token_ids, output_length in self._tokenized_dataset:
+            input_length = len(input_token_ids)
+            # If the request does not have enough length, discard it.
+            if input_len is not None and input_length < input_len + 4 * input_len_std:
+                continue
+
+            if input_len is not None:
+                input_length = round(
+                    float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0])
+                )
+                input_token_ids = input_token_ids[:input_length]
+                input_truncated = True
+            else:
+                input_truncated = False
+            if output_len is not None:
+                output_length = round(
+                    float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0])
+                )
+            elif output_length <= 1:
+                continue
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[
+                            {
+                                "role": "user",
+                                "content": (
+                                    self.tokenizer.decode(input_token_ids)
+                                    if input_truncated
+                                    else prompt
+                                ),
+                            }
+                        ],
+                        model="",
+                        max_tokens=output_length,
+                    ),
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=len(input_token_ids),
+                    ),
+                )
+            )
+        return request_records
+
+
+class LoogleDataset(Dataset):
+    """The dataset class for Loogle dataset."""
+
+    task2prompt: ClassVar[Dict[str, str]] = {  # noqa: UP006
+        "shortdep_qa": "Please answer the question based on the long texts below. \n{input}\nQuestion: {Q}\nAnswer: ",  # noqa: E501
+        "longdep_qa": "Please answer the question based on the long texts below. \n{input}\nQuestion: {Q}\nAnswer: ",  # noqa: E501
+        "longdep_summarization": "Please generate a summary of the below paper. \n{input}\n Summarization: ",  # noqa: E501
+        "shortdep_cloze": "Please fill in the clozes based on the given long texts below. Each of the placeholder '<mask-n>' in the question could be an entity of Person, Location or Organiocation. The same masks represent the same entity. Output a json format answer, for example: {{'<mask-0>': 'Bob', '<mask-1>': 'Gorrosion Magazine','<mask-2>': 'Bethel Horizon'}}\n{input}\n Question: {Q} What are the masked entities? \nAnswer:",  # noqa: E501
+    }
+    require_fake_warmup: bool = True
+
+    def __init__(self, tokenizer: AutoTokenizer, testset_name: str) -> None:
+        raw_dataset = load_dataset("bigainlco/LooGLE", testset_name, split="test")
+        self.tokenizer = tokenizer
+        self.dataset = []
+        self.prompt_format = self.task2prompt[testset_name]
+        prompts = []
+        generate_lens = []
+        questions = []
+        for data in raw_dataset:
+            prompt = data["input"]
+            prompts.append(prompt)
+            qa_pairs = eval(data["qa_pairs"])
+            questions.append([j["Q"] for j in qa_pairs])
+            generate_lens.append(
+                [len(tokenizer.encode(j["A"], add_special_tokens=False)) for j in qa_pairs]
+            )
+        prompt_token_ids = tokenizer(
+            prompts,
+            truncation=True,
+            max_length=min(tokenizer.model_max_length, self.truncate_length),
+            add_special_tokens=False,
+        ).input_ids
+        for prompt, prompt_token_id, question, generate_len in zip(
+            prompts, prompt_token_ids, questions, generate_lens
+        ):
+            self.dataset.append((prompt, prompt_token_id, question, generate_len))
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        request_records = []
+        for prompt, input_token_ids, questions, generate_lens in self.dataset:
+            input_length = round(float(np.random.normal(loc=input_len, scale=input_len_std)))
+            if len(input_token_ids) > input_length:
+                input_token_ids = input_token_ids[:input_length]
+                prompt = self.tokenizer.decode(input_token_ids)
+            grouped_request_records = []
+            for question, generate_len in zip(questions, generate_lens):
+                json_obj = {"input": prompt, "Q": question}
+                full_prompt = self.prompt_format.format(**json_obj)
+
+                output_length = (
+                    round(float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0]))
+                    if output_len is not None
+                    else generate_len
+                )
+                grouped_request_records.append(
+                    RequestRecord(
+                        chat_cmpl=ChatCompletionRequest(
+                            messages=[
+                                {
+                                    "role": "user",
+                                    "content": full_prompt,
+                                }
+                            ],
+                            model="",
+                            max_tokens=output_length,
+                        ),
+                        metrics=Metrics(
+                            success=False,
+                            start_time=0,
+                            finish_time=0,
+                            end_to_end_latency_s=0,
+                            input_tokens=len(input_token_ids),
+                        ),
+                    )
+                )
+            request_records.append(
+                GroupedRequestRecord(
+                    # Create a dummy ChatCompletionRequest.
+                    chat_cmpl=ChatCompletionRequest(messages=[]),
+                    records=grouped_request_records,
+                )
+            )
+        return request_records
+
+
+class LLMPerfDataset(Dataset):
+    """The dataset class for LLMPerf dataset."""
+
+    def __init__(self, dataset_path: str, num_requests: int, tokenizer: AutoTokenizer) -> None:
+        self.tokenizer = tokenizer
+        self.num_requests = num_requests
+
+        with open(dataset_path, encoding="utf-8") as f:
+            untokenized_data = f.readlines()
+        # Tokenize the prompts and completions.
+        tokenized_data = tokenizer(
+            untokenized_data,
+            truncation=True,
+            max_length=min(tokenizer.model_max_length, self.truncate_length),
+            add_special_tokens=False,
+        ).input_ids
+        tokenized_data_lengths = [len(tokens) for tokens in tokenized_data]
+        self.dataset: List[Tuple[str, List[int], int]] = list(  # noqa: UP006
+            zip(untokenized_data, tokenized_data, tokenized_data_lengths)
+        )
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int] = None,
+        output_len: Optional[int] = None,
+        input_len_std: float = 250,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if input_len is None or input_len < 40:
+            input_len = 550
+        if output_len is None:
+            output_len = 150
+
+        request_records = []
+        for _ in range(self.num_requests):
+            input_length = round(float(np.random.normal(loc=input_len, scale=input_len_std)))
+            output_length = round(float(np.random.normal(loc=output_len, scale=output_len_std)))
+
+            prompt = (
+                "Randomly stream lines from the following text "
+                f"with {output_length} output tokens. "
+                "Don't generate eos tokens:\n\n"
+            )
+
+            remaining_token_length = input_length - len(
+                self.tokenizer.encode(prompt, add_special_tokens=False)
+            )
+
+            random.shuffle(self.dataset)
+
+            while remaining_token_length > 0:
+                for text, tokens, token_length in self.dataset:
+                    if remaining_token_length < token_length:
+                        prompt += self.tokenizer.decode(tokens[:remaining_token_length])
+                    else:
+                        prompt += text
+
+                    remaining_token_length -= token_length
+                    if remaining_token_length < 0:
+                        break
+
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[{"role": "user", "content": prompt}],
+                        model="",
+                        max_tokens=output_length,
+                        debug_config=DebugConfig(ignore_eos=True),
+                    ),
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=input_length,
+                    ),
+                )
+            )
+        return request_records
+
+
+class JSONModeEvalDataset(Dataset):
+    """The dataset class for JSON dataset."""
+
+    def __init__(self, tokenizer: AutoTokenizer) -> None:
+        raw_dataset = load_dataset("NousResearch/json-mode-eval")
+        self.tokenizer = tokenizer
+        self.dataset = []
+        for data in raw_dataset["train"]:
+            messages = data["prompt"]
+            schema = {
+                "type": "json_object",
+                "schema": data["schema"],
+            }
+            num_tokens = 0
+            for message in messages:
+                num_tokens += len(
+                    self.tokenizer.encode(message["content"], add_special_tokens=False)
+                )
+            self.dataset.append((messages, schema, num_tokens))
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        request_records = []
+        for messages, schema, num_tokens in self.dataset:
+            # If the request does not have enough length, discard it.
+            if input_len is not None and num_tokens < input_len + 4 * input_len_std:
+                continue
+
+            if output_len is not None:
+                output_length = max(
+                    round(np.random.normal(loc=output_len, scale=output_len_std)), 1
+                )
+            else:
+                output_length = None
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[
+                            ChatCompletionMessage(content=message["content"], role=message["role"])
+                            for message in messages
+                        ],
+                        model="",
+                        max_tokens=output_length,
+                        response_format=schema,
+                    ),
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=num_tokens,
+                    ),
+                )
+            )
+        return request_records
+
+
+class ReActDataset(Dataset):
+    """The dataset class for replaying a given ReAct trace for benchmark purpose.
+    It is not an actual ReAct agent implementation.
+    """
+
+    _dataset: List[List[Tuple[str, int, int]]]  # noqa: UP006
+    require_fake_warmup: bool = True
+    prefix: str = """Solve a question answering task with interleaving Thought, Action, Observation steps. Thought can reason about the current situation, and Action can be three types:
+(1) Search[entity], which searches the exact entity on Wikipedia and returns the first paragraph if it exists. If not, it will return some similar entities to search.
+(2) Lookup[keyword], which returns the next sentence containing keyword in the current passage.
+(3) Finish[answer], which returns the answer and finishes the task.
+Here are some examples.
+Question: What is the elevation range for the area that the eastern sector of the Colorado orogeny extends into?
+Thought 1: I need to search Colorado orogeny, find the area that the eastern sector of the Colorado orogeny extends into, then find the elevation range of the area.
+Action 1: Search[Colorado orogeny]
+Observation 1: The Colorado orogeny was an episode of mountain building (an orogeny) in Colorado and surrounding areas.
+Thought 2: It does not mention the eastern sector. So I need to look up eastern sector.
+Action 2: Lookup[eastern sector]
+Observation 2: (Result 1 / 1) The eastern sector extends into the High Plains and is called the Central Plains orogeny.
+Thought 3: The eastern sector of Colorado orogeny extends into the High Plains. So I need to search High Plains and find its elevation range.
+Action 3: Search[High Plains]
+Observation 3: High Plains refers to one of two distinct land regions:
+Thought 4: I need to instead search High Plains (United States).
+Action 4: Search[High Plains (United States)]
+Observation 4: The High Plains are a subregion of the Great Plains. From east to west, the High Plains rise in elevation from around 1,800 to 7,000 ft (550 to 2,130 m).[3]
+Thought 5: High Plains rise in elevation from around 1,800 to 7,000 ft, so the answer is 1,800 to 7,000 ft.
+Action 5: Finish[1,800 to 7,000 ft]
+Question: Musician and satirist Allie Goertz wrote a song about the "The Simpsons" character Milhouse, who Matt Groening named after who?
+Thought 1: The question simplifies to "The Simpsons" character Milhouse is named after who. I only need to search Milhouse and find who it is named after.
+Action 1: Search[Milhouse]
+Observation 1: Milhouse Mussolini Van Houten is a recurring character in the Fox animated television series The Simpsons voiced by Pamela Hayden and created by Matt Groening.
+Thought 2: The paragraph does not tell who Milhouse is named after, maybe I can look up "named after".
+Action 2: Lookup[named after]
+Observation 2: (Result 1 / 1) Milhouse was named after U.S. president Richard Nixon, whose middle name was Milhous.
+Thought 3: Milhouse was named after U.S. president Richard Nixon, so the answer is Richard Nixon.
+Action 3: Finish[Richard Nixon]
+Question: Which documentary is about Finnish rock groups, Adam Clayton Powell or The Saimaa Gesture?
+Thought 1: I need to search Adam Clayton Powell and The Saimaa Gesture, and find which documentary is about Finnish rock groups.
+Action 1: Search[Adam Clayton Powell]
+Observation 1: Could not find [Adam Clayton Powell]. Similar: ['Adam Clayton Powell III', 'Seventh Avenue (Manhattan)', 'Adam Clayton Powell Jr. State Office Building', 'Isabel Washington Powell', 'Adam Powell', 'Adam Clayton Powell (film)', 'Giancarlo Esposito'].
+Thought 2: To find the documentary, I can search Adam Clayton Powell (film).
+Action 2: Search[Adam Clayton Powell (film)]
+Observation 2: Adam Clayton Powell is a 1989 American documentary film directed by Richard Kilberg.
+The film is about the rise and fall of influential African-American politician Adam Clayton Powell Jr.[3][4] It was later aired as part of the PBS series The American Experience.
+Thought 3: Adam Clayton Powell (film) is a documentary about an African-American politician, not Finnish rock groups. So the documentary about Finnish rock groups must instead be The Saimaa Gesture.
+Action 3: Finish[The Saimaa Gesture]
+Question: What profession does Nicholas Ray and Elia Kazan have in common?
+Thought 1: I need to search Nicholas Ray and Elia Kazan, find their professions, then find the profession they have in common.
+Action 1: Search[Nicholas Ray]
+Observation 1: Nicholas Ray (born Raymond Nicholas Kienzle Jr., August 7, 1911 – June 16, 1979) was an American film director, screenwriter, and actor best known for the 1955 film Rebel Without a Cause.
+Thought 2: Professions of Nicholas Ray are director, screenwriter, and actor. I need to search Elia Kazan next and find his professions.
+Action 2: Search[Elia Kazan]
+Observation 2: Elia Kazan was an American film and theatre director, producer, screenwriter and actor.
+Thought 3: Professions of Elia Kazan are director, producer, screenwriter, and actor. So profession Nicholas Ray and Elia Kazan have in common is director, screenwriter, and actor.
+Action 3: Finish[director, screenwriter, actor]
+Question: Which magazine was started first Arthur's Magazine or First for Women?
+Thought 1: I need to search Arthur's Magazine and First for Women, and find which was started first.
+Action 1: Search[Arthur's Magazine]
+Observation 1: Arthur's Magazine (1844-1846) was an American literary periodical published in Philadelphia in the 19th century.
+Thought 2: Arthur's Magazine was started in 1844. I need to search First for Women next.
+Action 2: Search[First for Women]
+Observation 2: First for Women is a woman's magazine published by Bauer Media Group in the USA.[1] The magazine was started in 1989.
+Thought 3: First for Women was started in 1989. 1844 (Arthur's Magazine) < 1989 (First for Women), so Arthur's Magazine was started first.
+Action 3: Finish[Arthur's Magazine]
+Question: Were Pavel Urysohn and Leonid Levin known for the same type of work?
+Thought 1: I need to search Pavel Urysohn and Leonid Levin, find their types of work, then find if they are the same.
+Action 1: Search[Pavel Urysohn]
+Observation 1: Pavel Samuilovich Urysohn (February 3, 1898 â August 17, 1924) was a Soviet mathematician who is best known for his contributions in dimension theory.
+Thought 2: Pavel Urysohn is a mathematician. I need to search Leonid Levin next and find its type of work.
+Action 2: Search[Leonid Levin]
+Observation 2: Leonid Anatolievich Levin is a Soviet-American mathematician and computer scientist.
+Thought 3: Leonid Levin is a mathematician and computer scientist. So Pavel Urysohn and Leonid Levin have the same type of work.
+Action 3: Finish[yes]
+"""  # noqa: E501, RUF001
+
+    def __init__(self, dataset_path: str, tokenizer: AutoTokenizer) -> None:
+        raw_entries: List[Dict] = []  # noqa: UP006
+        with open(dataset_path) as fin:
+            for line in fin:
+                line_content = json.loads(line)
+                raw_entries += list({"question": k, "triplets": v} for k, v in line_content.items())
+
+        self._dataset = []
+        max_rounds = 0
+        for raw_entry in raw_entries:
+            processed_entry = []
+            question = raw_entry["question"]
+            triplets = raw_entry["triplets"]
+            seq = self.prefix + question
+            max_rounds = max(max_rounds, len(triplets) + 1)
+            output_lengths: List[int] = []  # noqa: UP006
+            for i, triplet in enumerate(triplets):
+                output_lengths.append(
+                    len(
+                        tokenizer(
+                            triplet["thought"]
+                            + "\nAction "
+                            + str(i + 1)
+                            + ": "
+                            + triplet["action"]
+                            + "\n",
+                            truncation=True,
+                            max_length=min(tokenizer.model_max_length, self.truncate_length),
+                            add_special_tokens=False,
+                        ).input_ids
+                    )
+                )
+
+            for i in range(1, len(triplets) + 2):
+                seq += "Thought " + str(i) + ":"
+                input_len = len(
+                    tokenizer(
+                        seq,
+                        truncation=True,
+                        max_length=min(tokenizer.model_max_length, self.truncate_length),
+                        add_special_tokens=False,
+                    ).input_ids
+                )
+                output_length = (
+                    output_lengths[i - 1]
+                    if i <= len(triplets)
+                    else int(sum(output_lengths) / len(triplets))
+                )
+                processed_entry.append((seq, input_len, output_length))
+                if i != len(triplets) + 1:
+                    seq += (
+                        triplets[i - 1]["thought"]
+                        + "\nAction "
+                        + str(i)
+                        + ": "
+                        + triplets[i - 1]["action"]
+                        + "\nObservation "
+                        + str(i)
+                        + ": "
+                        + triplets[i - 1]["observation"]
+                        + "\n"
+                    )
+            self._dataset.append(processed_entry)
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if input_len is not None or output_len is not None:
+            raise ValueError("ReAct dataset does not support specifying input/output length.")
+
+        request_records = []
+        for processed_entries in self._dataset:
+            grouped_request_records = []
+            for prompt, input_length, output_length in processed_entries:
+                grouped_request_records.append(
+                    RequestRecord(
+                        chat_cmpl=ChatCompletionRequest(
+                            messages=[{"role": "user", "content": prompt}],
+                            model="",
+                            max_tokens=output_length,
+                        ),
+                        metrics=Metrics(
+                            success=False,
+                            start_time=0,
+                            finish_time=0,
+                            end_to_end_latency_s=0,
+                            input_tokens=input_length,
+                        ),
+                    )
+                )
+            request_records.append(
+                GroupedRequestRecord(
+                    # Create a dummy ChatCompletionRequest.
+                    chat_cmpl=ChatCompletionRequest(messages=[]),
+                    records=grouped_request_records,
+                )
+            )
+        return request_records
+
+
+class WildChatDataset(Dataset):
+    """The dataset class for WildChat dataset."""
+
+    apply_chat_template: bool
+
+    def __init__(self, tokenizer: AutoTokenizer, apply_chat_template: bool) -> None:
+        raw_dataset = load_dataset("allenai/WildChat", split="train")
+        self.tokenizer = tokenizer
+        self.apply_chat_template = apply_chat_template
+
+        # Filter out the conversations with less than 2 turns.
+        _dataset = [
+            (entry["conversation"][0]["content"], entry["conversation"][1]["content"])
+            for entry in raw_dataset
+            if len(entry["conversation"]) >= 2
+            and entry["conversation"][0]["role"] == "user"
+            and entry["conversation"][1]["role"] == "assistant"
+        ]
+
+        prompts = []
+        completions = []
+        for prompt, completion in _dataset:
+            prompts.append(prompt)
+            completions.append(completion)
+        if apply_chat_template:
+            assert getattr(tokenizer, "chat_template", None) is not None, (
+                '"--apply-chat-template" is set but the tokenizer does not have chat template.'
+            )
+            prompts = [
+                tokenizer.apply_chat_template(
+                    [{"role": "user", "content": prompt}],
+                    add_generation_prompt=True,
+                    tokenize=False,
+                )
+                for prompt in prompts
+            ]
+
+        prompt_token_ids = list(
+            tokenizer(
+                prompts,
+                truncation=True,
+                max_length=min(tokenizer.model_max_length, self.truncate_length),
+                add_special_tokens=False,
+            ).input_ids
+        )
+        completion_token_ids = tokenizer(
+            completions,
+            truncation=True,
+            max_length=min(tokenizer.model_max_length, self.truncate_length),
+            add_special_tokens=False,
+        ).input_ids
+        self._tokenized_dataset: List[Tuple[str, List[int], int]] = []  # noqa: UP006
+        for i in range(len(_dataset)):
+            if len(prompt_token_ids[i]) < 4 or len(completion_token_ids[i]) < 4:
+                # Filter out sequences that are too short
+                continue
+            self._tokenized_dataset.append(
+                (prompts[i], prompt_token_ids[i], len(completion_token_ids[i]))
+            )
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if self.apply_chat_template:
+            assert input_len is None, (
+                '"--apply-chat-template" is not supported when "--input-len" is specified.'
+            )
+
+        request_records = []
+        for prompt, input_token_ids, output_length in self._tokenized_dataset:
+            input_length = len(input_token_ids)
+            # If the request does not have enough length, discard it.
+            if input_len is not None and input_length < input_len + 4 * input_len_std:
+                continue
+
+            if input_len is not None:
+                input_length = round(
+                    float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0])
+                )
+                input_token_ids = input_token_ids[:input_length]
+                input_truncated = True
+            else:
+                input_truncated = False
+            if output_len is not None:
+                output_length = round(
+                    float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0])
+                )
+            elif output_length <= 1:
+                continue
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[
+                            {
+                                "role": "user",
+                                "content": (
+                                    self.tokenizer.decode(input_token_ids)
+                                    if input_truncated
+                                    else prompt
+                                ),
+                            }
+                        ],
+                        model="",
+                        max_tokens=output_length,
+                    ),
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=len(input_token_ids),
+                    ),
+                )
+            )
+        return request_records
+
+
+class AzureLLMInferenceDataset(Dataset):
+    """The dataset class for AzureLLMInference dataset.
+    Reference: https://github.com/Azure/AzurePublicDataset
+    """
+
+    timestamp_available: bool = True
+
+    def __init__(self, dataset_path: str, tokenizer: AutoTokenizer) -> None:
+        df = pd.read_csv(dataset_path)
+        self.tokenizer = tokenizer
+
+        # Filter out the conversations with less than 2 turns.
+        self.dataset = [
+            (
+                entry["TIMESTAMP"],
+                min(
+                    entry["ContextTokens"],
+                    tokenizer.model_max_length,
+                    self.truncate_length,
+                ),
+                min(
+                    entry["GeneratedTokens"],
+                    tokenizer.model_max_length,
+                    self.truncate_length,
+                ),
+            )
+            for _, entry in df.iterrows()
+            if entry["ContextTokens"] >= 4 and entry["GeneratedTokens"] >= 4
+        ]
+
+    def generate_request_records(
+        self,
+        input_len: Optional[int],
+        output_len: Optional[int],
+        input_len_std: float = 0.0,
+        output_len_std: float = 0.0,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        time_fmt = "%Y-%m-%d %H:%M:%S.%f"
+        start_time = datetime.strptime(self.dataset[0][0][:-1], time_fmt)
+        request_records = []
+        for timestamp, input_length, output_length in self.dataset:
+            # If the request does not have enough length, discard it.
+            if input_len is not None and input_length < input_len + 4 * input_len_std:
+                continue
+
+            if input_len is not None:
+                input_length = round(
+                    float(np.random.normal(loc=input_len, scale=input_len_std, size=1)[0])
+                )
+            if output_len is not None:
+                output_length = round(
+                    float(np.random.normal(loc=output_len, scale=output_len_std, size=1)[0])
+                )
+            elif output_length <= 1:
+                continue
+
+            prompt_token_ids = [
+                random.randint(0, self.tokenizer.vocab_size - 1) for _ in range(input_length)
+            ]
+            while True:
+                # Adjust the token ids until the retokenization on the decoded string
+                # matches the required input length.
+                prompt = self.tokenizer.decode(prompt_token_ids)
+                retokenized_token_ids = self.tokenizer.encode(prompt, add_special_tokens=False)
+                if len(retokenized_token_ids) < input_length:
+                    prompt_token_ids = retokenized_token_ids + [
+                        random.randint(0, self.tokenizer.vocab_size - 1)
+                        for _ in range(input_length - len(retokenized_token_ids))
+                    ]
+                elif len(retokenized_token_ids) > input_length:
+                    prompt_token_ids = retokenized_token_ids[:input_length]
+                else:
+                    break
+
+            time_diff = (datetime.strptime(timestamp[:-1], time_fmt) - start_time).total_seconds()
+            request_records.append(
+                RequestRecord(
+                    chat_cmpl=ChatCompletionRequest(
+                        messages=[{"role": "user", "content": prompt}],
+                        model="",
+                        max_tokens=output_length,
+                    ),
+                    timestamp=time_diff,
+                    metrics=Metrics(
+                        success=False,
+                        start_time=0,
+                        finish_time=0,
+                        end_to_end_latency_s=0,
+                        input_tokens=input_length,
+                    ),
+                )
+            )
+        return request_records
+
+
+SUPPORTED_DATASET = [
+    "sharegpt",
+    "llmperf",
+    "json-mode-eval",
+    "loogle",
+    "react",
+    "wildchat",
+    "azure-llm-inference",
+]
+
+
+def create_dataset(args: argparse.Namespace, tokenizer: AutoTokenizer) -> Dataset:
+    """Create a dataset instance with regard to the specified dataset kind and file path."""
+    if args.dataset_path is not None and not isinstance(args.dataset_path, str):
+        raise TypeError(f"Invalid dataset path {args.dataset_path}. Please use a string.")
+    if args.dataset is None and args.dataset_path is not None:
+        # Auto-detect the dataset kind by looking into the dataset path.
+        if "sharegpt" in args.dataset_path.lower():
+            args.dataset = "sharegpt"
+        else:
+            raise ValueError(
+                f"Unable to detect the dataset kind from dataset path {args.dataset_path}. "
+                'Please specify the dataset kind via "--dataset".'
+            )
+    if args.dataset == "sharegpt":
+        if args.dataset_path is None:
+            raise ValueError(
+                'ShareGPT dataset requires dataset path. Please specify it with "--dataset-path".'
+            )
+        return ShareGPTDataset(args.dataset_path, tokenizer, args.apply_chat_template)
+    if args.dataset == "llmperf":
+        if args.dataset_path is None:
+            raise ValueError(
+                'LLMPerf dataset requires dataset path. Please specify it with "--dataset-path".'
+            )
+        assert args.apply_chat_template is False, (
+            "LLMPerf dataset does not support applying chat template"
+        )
+        return LLMPerfDataset(
+            args.dataset_path,
+            (args.num_requests + args.num_warmup_requests) * 4,
+            tokenizer,
+        )
+    if args.dataset == "json-mode-eval":
+        assert args.apply_chat_template is False, (
+            "JSON mode evaluation does not support applying chat template"
+        )
+        return JSONModeEvalDataset(tokenizer)
+    if args.dataset == "loogle":
+        if args.dataset_path is None:
+            raise ValueError(
+                'Loogle dataset requires a testset name. Please specify it with "--dataset-path".'
+            )
+        assert args.apply_chat_template is False, (
+            "Loogle dataset does not support applying chat template"
+        )
+        return LoogleDataset(tokenizer, testset_name=args.dataset_path)
+    if args.dataset == "react":
+        if args.dataset_path is None:
+            raise ValueError(
+                'ReAct dataset requires dataset path. Please specify it with "--dataset-path".'
+            )
+        assert args.apply_chat_template is False, (
+            "ReAct dataset does not support applying chat template"
+        )
+        return ReActDataset(args.dataset_path, tokenizer)
+    if args.dataset == "wildchat":
+        return WildChatDataset(tokenizer, args.apply_chat_template)
+    if args.dataset == "azure-llm-inference":
+        if args.dataset_path is None:
+            raise ValueError(
+                "AzureLLMInference dataset requires dataset path. "
+                'Please specify it with "--dataset-path".'
+            )
+        assert args.apply_chat_template is False, (
+            "AzureLLMInference dataset does not support applying chat template"
+        )
+        return AzureLLMInferenceDataset(args.dataset_path, tokenizer)
+    raise ValueError(f"Unrecognized dataset {args.dataset}")

python/mlc_llm/bench/evaluation/gsm8k.py

@@ -0,0 +1,315 @@
+"""Eval GSM8K with MLCEngine."""
+
+import argparse
+import asyncio
+import json
+import random
+import re
+from datetime import datetime
+from pathlib import Path
+from typing import List, Literal, Optional  # noqa: UP035
+
+import tqdm
+
+from mlc_llm import AsyncMLCEngine
+
+DEVICES = ["cuda", "rocm", "metal", "vulkan"]
+ANSWER_TRIGGER = "The answer is"
+INVALID_ANS = "[invalid]"
+
+
+def extract_answer(text: str, regex: re.Pattern, select_index: int) -> str:
+    """Extract the answer from the text."""
+    match_all = regex.findall(text)
+    if len(match_all) == 0:
+        return INVALID_ANS
+    match = match_all[select_index]
+    if isinstance(match, tuple):
+        match = next(m for m in match if m)
+    match_str: str = match.strip()
+    match_str = match_str.lstrip("$").rstrip(".").replace(",", "")
+    return match_str
+
+
+def extract_ground_truth(text: str) -> str:
+    """Extract the ground truth from the text."""
+    return extract_answer(text, re.compile(r"#### (\-?[0-9\.\,]+)"), 0)
+
+
+def strict_extract_answer(text: str) -> str:
+    """Strictly extract the answer from the text."""
+    return extract_answer(text, re.compile(r"The answer is \$?(\-?[0-9\.\,]+)."), 0)
+
+
+def flexible_extract_answer(text: str) -> str:
+    """Extract the last number from the text."""
+    return extract_answer(text, re.compile(r"(-?[$0-9.,]{2,})|(-?[0-9]+)"), -1)
+
+
+def create_few_shot_prompt(n_shot: int, use_cot: bool, random_order=False) -> str:
+    """
+    Create a prompt for the few-shot learning task.
+
+    Note
+    ----
+    The examples are taken from the paper https://arxiv.org/pdf/2201.11903.pdf page 35.
+    """
+    question, chain, answer = [], [], []
+
+    question.append(
+        "There are 15 trees in the grove. "
+        "Grove workers will plant trees in the grove today. "
+        "After they are done, there will be 21 trees. "
+        "How many trees did the grove workers plant today?"
+    )
+    chain.append(
+        "There are 15 trees originally. "
+        "Then there were 21 trees after some more were planted. "
+        "So there must have been 21 - 15 = 6."
+    )
+    answer.append("6")
+
+    question.append(
+        "If there are 3 cars in the parking lot and 2 more cars arrive, "
+        "how many cars are in the parking lot?"
+    )
+    chain.append("There are originally 3 cars. 2 more cars arrive. 3 + 2 = 5.")
+    answer.append("5")
+
+    question.append(
+        "Leah had 32 chocolates and her sister had 42. If they ate 35, "
+        "how many pieces do they have left in total?"
+    )
+    chain.append(
+        "Originally, Leah had 32 chocolates. "
+        "Her sister had 42. So in total they had 32 + 42 = 74. "
+        "After eating 35, they had 74 - 35 = 39."
+    )
+    answer.append("39")
+
+    question.append(
+        "Jason had 20 lollipops. He gave Denny some lollipops. Now Jason "
+        "has 12 lollipops. How many lollipops did Jason give to Denny?"
+    )
+    chain.append(
+        "Jason started with 20 lollipops. Then he had 12 after giving some "
+        "to Denny. So he gave Denny 20 - 12 = 8."
+    )
+    answer.append("8")
+
+    question.append(
+        "Shawn has five toys. For Christmas, he got two toys each from his "
+        "mom and dad. How many toys does he have now?"
+    )
+    chain.append(
+        "Shawn started with 5 toys. If he got 2 toys each from his mom and "
+        "dad, then that is 4 more toys. 5 + 4 = 9."
+    )
+    answer.append("9")
+
+    question.append(
+        "There were nine computers in the server room. Five more computers "
+        "were installed each day, from monday to thursday. "
+        "How many computers are now in the server room?"
+    )
+    chain.append(
+        "There were originally 9 computers. For each of 4 days, 5 more "
+        "computers were added. So 5 * 4 = 20 computers were added. "
+        "9 + 20 is 29."
+    )
+    answer.append("29")
+
+    question.append(
+        "Michael had 58 golf balls. On tuesday, he lost 23 golf balls. On "
+        "wednesday, he lost 2 more. "
+        "How many golf balls did he have at the end of wednesday?"
+    )
+    chain.append(
+        "Michael started with 58 golf balls. After losing 23 on tuesday, "
+        "he had 58 - 23 = 35. After losing 2 more, "
+        "he had 35 - 2 = 33 golf balls."
+    )
+    answer.append("33")
+
+    question.append(
+        "Olivia has $23. She bought five bagels for $3 each. How much money does she have left?"
+    )
+    chain.append(
+        "Olivia had 23 dollars. "
+        "5 bagels for 3 dollars each will be 5 x 3 = 15 dollars. "
+        "So she has 23 - 15 dollars left. 23 - 15 is 8."
+    )
+    answer.append("8")
+
+    index_list = list(range(len(question)))
+    if random_order:
+        random.shuffle(index_list)
+
+    prompt = ""
+    for i in index_list[:n_shot]:
+        if use_cot:
+            prompt += f"Q: {question[i]}\nA: {chain[i]} {ANSWER_TRIGGER} {answer[i]}.\n\n"
+        else:
+            prompt += f"Question: {question[i]}\nAnswer: {ANSWER_TRIGGER} {answer[i]}.\n\n"
+    return prompt
+
+
+def create_prompt(question: str, n_shot: int, use_cot: bool, random_order: bool = False) -> str:
+    """Create a prompt for the few-shot learning task."""
+    prompt = create_few_shot_prompt(n_shot, use_cot, random_order)
+    if use_cot:
+        prompt += f"Q: {question}\nA:"
+    else:
+        prompt += f"Question: {question}\nAnswer:"
+    return prompt
+
+
+def parse_args():
+    """Parse command line arguments."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument(
+        "--dataset", type=Path, required=True, help="Path to GSM8K test dataset home."
+    )
+    parser.add_argument("--device", type=str, choices=["auto", *DEVICES], default="auto")
+    parser.add_argument("--model-lib", type=str, default=None)
+    parser.add_argument("--n-shot", type=int, default=8)
+    parser.add_argument("--disable_cot", action="store_true", default=False)
+    parser.add_argument("-bs", "--batch-size", type=int, default=16)
+    parser.add_argument("--log-dir", type=Path, default=None)
+    return parser.parse_args()
+
+
+async def send_request(
+    async_engine: AsyncMLCEngine,
+    prompts: List[str],  # noqa: UP006
+    semaphore: asyncio.Semaphore,
+):
+    """Send the calibration requests to the engine."""
+    tasks = []
+
+    async def generate_task(prompt):
+        async with semaphore:
+            return await async_engine.completions.create(
+                prompt=prompt,
+                stream=False,
+                max_tokens=512,
+                stop=["Q:", "Question:"],
+                temperature=0.0,
+            )
+
+    for prompt in prompts:
+        task = asyncio.create_task(generate_task(prompt))
+        tasks.append(task)
+
+    return await tqdm.asyncio.tqdm.gather(*tasks)
+
+
+async def evaluate(
+    model: str,
+    device: str,
+    dataset: Path,
+    model_lib: Optional[str],
+    n_shot: int,
+    use_cot: bool,
+    batch_size: int,
+    log_dir: Optional[Path],
+):
+    """Evaluate GSM8K for the model."""
+    mode: Literal["local", "interactive", "server"] = (
+        "server" if batch_size > 4 else "interactive" if batch_size == 1 else "local"
+    )
+    async_engine = AsyncMLCEngine(model, device=device, model_lib=model_lib, mode=mode)
+
+    with open(dataset / "test.jsonl", encoding="utf-8") as file:
+        tests = [json.loads(line) for line in file]
+
+    prompts = [create_prompt(test["question"], n_shot, use_cot) for test in tests]
+    responses = await send_request(async_engine, prompts, asyncio.Semaphore(batch_size))
+    assert len(responses) == len(tests)
+
+    num_strict_correct, num_flexible_correct = 0, 0
+    num_tests = len(tests)
+    logs = []
+
+    for response, test in zip(responses, tests):
+        response_text = response.choices[0].text.strip()
+        gt_answer = extract_ground_truth(test["answer"])
+        assert gt_answer != INVALID_ANS
+        strict_answer = strict_extract_answer(response_text)
+        flexible_answer = flexible_extract_answer(response_text)
+
+        if gt_answer == strict_extract_answer(response_text):
+            # If the answer is exactly the same as the response, then it is correct
+            num_strict_correct += 1
+            num_flexible_correct += 1
+
+        elif gt_answer == flexible_extract_answer(response_text):
+            # Try flexible extract if the strict match fails
+            num_flexible_correct += 1
+
+        logs.append(
+            {
+                "question": test["question"],
+                "response": response_text,
+                "ground_truth": gt_answer,
+                "strict_answer": strict_answer,
+                "flexible_answer": flexible_answer,
+                "strict_match": gt_answer == strict_answer,
+                "flexible_match": gt_answer == flexible_answer,
+            }
+        )
+
+    results = {
+        "config": {
+            "model": model,
+            "device": device,
+            "model_lib": model_lib,
+            "n_shot": n_shot,
+            "use_cot": use_cot,
+        },
+        "results": {
+            "strict_match": num_strict_correct,
+            "flexible_match": num_flexible_correct,
+            "total": num_tests,
+        },
+    }
+    print(
+        f"Strict Matching Accuracy: {num_strict_correct} / {num_tests} = "
+        f"{num_strict_correct / num_tests * 100:.2f}%"
+    )
+    print(
+        f"Flexible Matching Accuracy: {num_flexible_correct} / {num_tests} = "
+        f"{num_flexible_correct / num_tests * 100:.2f}%"
+    )
+
+    if log_dir:
+        with open(log_dir / "summary.json", "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2)
+        with open(log_dir / "logs.json", "w", encoding="utf-8") as f:
+            json.dump(logs, f, indent=2)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    start_time = datetime.now()
+    log_dir: Optional[Path] = None
+    if args.log_dir is not None:
+        time_dir = start_time.strftime("%Y-%m-%d_%H-%M-%S")
+        log_dir = args.log_dir / time_dir
+        log_dir.mkdir(parents=True, exist_ok=True)
+    asyncio.run(
+        evaluate(
+            model=args.model,
+            device=args.device,
+            dataset=args.dataset,
+            model_lib=args.model_lib,
+            n_shot=args.n_shot,
+            use_cot=not args.disable_cot,
+            batch_size=args.batch_size,
+            log_dir=log_dir,
+        )
+    )
+    end_time = datetime.now()
+    print(f"Time used: {end_time - start_time}")

python/mlc_llm/bench/evaluation/mmlu.py

@@ -0,0 +1,245 @@
+"""Eval MMLU with MLCEngine."""
+
+import argparse
+import asyncio
+import csv
+import json
+import string
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional  # noqa: UP035
+
+import numpy as np
+import tqdm
+
+from mlc_llm import AsyncMLCEngine
+
+SUBJECTS = [
+    "abstract_algebra",
+    "anatomy",
+    "astronomy",
+    "business_ethics",
+    "clinical_knowledge",
+    "college_biology",
+    "college_chemistry",
+    "college_computer_science",
+    "college_mathematics",
+    "college_medicine",
+    "college_physics",
+    "computer_security",
+    "conceptual_physics",
+    "econometrics",
+    "electrical_engineering",
+    "elementary_mathematics",
+    "formal_logic",
+    "global_facts",
+    "high_school_biology",
+    "high_school_chemistry",
+    "high_school_computer_science",
+    "high_school_european_history",
+    "high_school_geography",
+    "high_school_government_and_politics",
+    "high_school_macroeconomics",
+    "high_school_mathematics",
+    "high_school_microeconomics",
+    "high_school_physics",
+    "high_school_psychology",
+    "high_school_statistics",
+    "high_school_us_history",
+    "high_school_world_history",
+    "human_aging",
+    "human_sexuality",
+    "international_law",
+    "jurisprudence",
+    "logical_fallacies",
+    "machine_learning",
+    "management",
+    "marketing",
+    "medical_genetics",
+    "miscellaneous",
+    "moral_disputes",
+    "moral_scenarios",
+    "nutrition",
+    "philosophy",
+    "prehistory",
+    "professional_accounting",
+    "professional_law",
+    "professional_medicine",
+    "professional_psychology",
+    "public_relations",
+    "security_studies",
+    "sociology",
+    "us_foreign_policy",
+    "virology",
+    "world_religions",
+]
+PADDING_LEN = max(len(subject) for subject in SUBJECTS)
+DEVICES = ["cuda", "rocm", "metal", "vulkan"]
+PROMPT_TEMPLATE = string.Template("$Q\nA. $A\nB. $B\nC. $C\nD. $D\nAnswer:")
+
+
+def parse_args():
+    """Parse command line arguments."""
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, required=True)
+    parser.add_argument(
+        "--dataset", type=Path, required=True, help="Path to MMLU test dataset home."
+    )
+    parser.add_argument("--device", type=str, choices=["auto", *DEVICES], default="auto")
+    parser.add_argument("--model-lib", type=str, default=None)
+    parser.add_argument("-s", "--subject", nargs="+", type=str, choices=SUBJECTS, default=SUBJECTS)
+    parser.add_argument("-bs", "--batch-size", type=int, default=16)
+    parser.add_argument("--log-dir", type=Path, default=None)
+    return parser.parse_args()
+
+
+async def send_request(
+    async_engine: AsyncMLCEngine,
+    prompts: List[str],  # noqa: UP006
+    semaphore: asyncio.Semaphore,
+    subject: str,
+):
+    """Send the calibration requests to the engine."""
+    tasks = []
+
+    async def generate_task(prompt):
+        async with semaphore:
+            return await async_engine.completions.create(
+                prompt=prompt,
+                stream=False,
+                max_tokens=1,
+                temperature=1.0,
+                logprobs=True,
+                top_logprobs=5,
+            )
+
+    for prompt in prompts:
+        task = asyncio.create_task(generate_task(prompt))
+        tasks.append(task)
+
+    return await tqdm.asyncio.tqdm.gather(
+        *tasks,
+        desc=f"Running {subject.ljust(PADDING_LEN)}",
+        bar_format="{desc} {percentage:3.0f}%|{bar}{r_bar}",
+    )
+
+
+async def evaluate(
+    model: str,
+    device: str,
+    dataset: Path,
+    model_lib: Optional[str],
+    subjects: List[str],  # noqa: UP006
+    semaphore: asyncio.Semaphore,
+    log_dir: Optional[Path],
+):
+    """Evaluate MMLU for the model."""
+    async_engine = AsyncMLCEngine(model, device=device, model_lib=model_lib, mode="server")
+
+    results: Dict[str, Any] = {}  # noqa: UP006
+    for subject in subjects:
+        with open(dataset / "test" / f"{subject}_test.csv", encoding="utf-8") as csvfile:
+            tests = list(csv.reader(csvfile, delimiter=",", quotechar='"'))
+            assert all(len(test) == 6 for test in tests)
+
+        logs = []
+        num_correct = 0
+        prompts = [
+            PROMPT_TEMPLATE.substitute(Q=test[0], A=test[1], B=test[2], C=test[3], D=test[4])
+            for test in tests
+        ]
+        responses = await send_request(async_engine, prompts, semaphore, subject)
+
+        assert len(responses) == len(tests)
+        for response, test in zip(responses, tests):
+            token_logprobs = {}
+            logprobs = response.choices[0].logprobs.content[0].top_logprobs
+            for logprob in logprobs:
+                if logprob.token not in token_logprobs:
+                    token_logprobs[logprob.token] = logprob.logprob
+
+            abcd_logprobs = {}
+            for choice in ["A", "B", "C", "D"]:
+                abcd_logprobs[choice] = token_logprobs[choice] if choice in token_logprobs else -100
+
+            pred = {0: "A", 1: "B", 2: "C", 3: "D"}[int(np.argmax(list(abcd_logprobs.values())))]
+            num_correct += pred == test[5]
+
+            logs.append(
+                {
+                    "Question": {
+                        "Q": test[0],
+                        "A": test[1],
+                        "B": test[2],
+                        "C": test[3],
+                        "D": test[4],
+                    },
+                    "Answer": test[5],
+                    "Response": {
+                        "pred": pred,
+                        "logprobs": list(abcd_logprobs.values()),
+                    },
+                }
+            )
+
+        results[subject] = {
+            "correct": num_correct,
+            "total": len(tests),
+            "accuracy": num_correct / len(tests),
+        }
+
+        if log_dir:
+            with open(log_dir / "subjects" / f"{subject}.json", "w", encoding="utf-8") as f:
+                json.dump(logs, f, indent=2)
+
+    total_correct, total_tests = 0, 0
+    for subject, v in results.items():
+        num_correct, num_tests, accuracy = v["correct"], v["total"], v["accuracy"]
+        print(f"{subject}: {num_correct} / {num_tests} = {accuracy * 100:.2f}%")
+        total_correct += num_correct
+        total_tests += num_tests
+
+    total_accuracy = total_correct / total_tests
+    results["total"] = {
+        "correct": total_correct,
+        "total": total_tests,
+        "accuracy": total_accuracy,
+    }
+    print(f"Total accuracy: {total_correct} / {total_tests} = {total_accuracy * 100:.2f}%")
+
+    if log_dir:
+        results = {
+            "config": {
+                "model": model,
+                "device": device,
+                "model_lib": model_lib,
+                "subjects": subjects,
+            },
+            "results": results,
+        }
+        with open(log_dir / "summary.json", "w", encoding="utf-8") as f:
+            json.dump(results, f, indent=2)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    start_time = datetime.now()
+    log_dir: Optional[Path] = None
+    if args.log_dir is not None:
+        time_dir = start_time.strftime("%Y-%m-%d_%H-%M-%S")
+        log_dir = args.log_dir / time_dir
+        (log_dir / "subjects").mkdir(parents=True, exist_ok=True)
+    asyncio.run(
+        evaluate(
+            model=args.model,
+            device=args.device,
+            dataset=args.dataset,
+            model_lib=args.model_lib,
+            subjects=args.subject,
+            semaphore=asyncio.Semaphore(args.batch_size),
+            log_dir=log_dir,
+        )
+    )
+    end_time = datetime.now()
+    print(f"Time used: {end_time - start_time}")

python/mlc_llm/bench/request_processor.py

@@ -0,0 +1,740 @@
+"""MLC LLM Bench Request"""
+
+import argparse
+import asyncio
+import concurrent.futures
+import copy
+import os
+import random
+import time
+from typing import Any, Callable, Dict, List, Optional  # noqa: UP035
+
+import numpy as np
+import requests
+from tqdm import tqdm
+from transformers import AutoTokenizer
+
+from mlc_llm.bench.api_endpoint import APIEndPoint
+from mlc_llm.bench.dataset import Dataset
+from mlc_llm.bench.request_record import GroupedRequestRecord, RequestRecord
+from mlc_llm.protocol.openai_api_protocol import (
+    ChatCompletionMessage,
+    ChatCompletionRequest,
+    DebugConfig,
+)
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+class RequestProcessor:
+    """The request processor base class.
+    Each processor can take a list of RequestRecord, applying the process,
+    and returning the processed RequestRecord in the end.
+    """
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        raise NotImplementedError()
+
+
+class LogMessage(RequestProcessor):
+    """The processor that prints the logger message."""
+
+    def __init__(self, message: str) -> None:
+        self.message = message
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        logger.info(self.message)
+        return request_records
+
+
+class SampleRequests(RequestProcessor):
+    """The processor that samples requests out from the given request list."""
+
+    def __init__(self, num_requests: int, take_first_x_requests: bool = False) -> None:
+        self.num_requests = num_requests
+        # If `take_first_x_requests` is True, the first `num_requests` requests
+        # are returned and sampling will not happen.
+        self.take_first_x_requests = take_first_x_requests
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        assert len(request_records) > 0, "Empty input request record."
+
+        # We expect the input request records to be all grouped or all plain.
+        if isinstance(request_records[0], GroupedRequestRecord):
+            assert all(isinstance(record, GroupedRequestRecord) for record in request_records)
+            return self._sample_from_grouped_request_records(request_records)
+
+        assert all(not isinstance(record, GroupedRequestRecord) for record in request_records)
+        return self._sample_from_plain_request_records(request_records)
+
+    def _sample_from_plain_request_records(
+        self,
+        request_records: List[RequestRecord],  # noqa: UP006
+    ) -> List[RequestRecord]:  # noqa: UP006
+        samples: List[RequestRecord] = []  # noqa: UP006
+        if self.take_first_x_requests:
+            if len(request_records) < self.num_requests:
+                raise ValueError(
+                    f"Insufficient requests. Requiring {self.num_requests} requests "
+                    f"but only {len(request_records)} are available."
+                )
+            samples = copy.deepcopy(list(request_records[: self.num_requests]))
+        else:
+            while len(samples) < self.num_requests:
+                # Create a new list so that the in-place shuffle does not mutate the input list.
+                records = list(request_records)
+                random.shuffle(records)
+                samples += copy.deepcopy(records)
+            samples = samples[: self.num_requests]
+        for i, record in enumerate(samples):
+            record.request_id = i
+        return samples
+
+    def _sample_from_grouped_request_records(
+        self,
+        grouped_request_records: List[GroupedRequestRecord],  # noqa: UP006
+    ) -> List[RequestRecord]:  # noqa: UP006
+        num_total_available_requests = sum(
+            len(record.records) for record in grouped_request_records
+        )
+        if self.num_requests > num_total_available_requests:
+            raise ValueError(
+                "Due to the existence of shared common prefixes, we do not allow "
+                "benchmarking with requests more than the available requests in the dataset. "
+                f"The required number of requests {self.num_requests} exceeds the "
+                f"number of total available requests {num_total_available_requests}."
+            )
+
+        # Create a new list so that the in-place shuffle does not mutate the input list.
+        records = list(grouped_request_records)
+        if not self.take_first_x_requests:
+            random.shuffle(records)
+        remaining = self.num_requests
+        samples: List[RequestRecord] = []  # noqa: UP006
+        for grouped_request_record in grouped_request_records:
+            num_used_requests = min(len(grouped_request_record.records), remaining)
+            samples += grouped_request_record.records[:num_used_requests]
+            remaining -= num_used_requests
+            if remaining == 0:
+                break
+        for i, record in enumerate(samples):
+            record.request_id = i
+        return samples
+
+
+class AttachModelName(RequestProcessor):
+    """The processor that attaches model name to requests."""
+
+    def __init__(self, model: str) -> None:
+        self.model = model
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for request_record in request_records:
+            request_record.chat_cmpl.model = self.model
+        return request_records
+
+
+class AttachRequestRateTimestamp(RequestProcessor):
+    """The processor that applies timestamps to the requests."""
+
+    def __init__(self, request_rate: np.float32) -> None:
+        self.request_rate = request_rate
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        timestamp = 0.0
+        for request_record in request_records:
+            assert request_record.timestamp is None, "The request record already has a timestamp"
+            request_record.timestamp = timestamp
+            timestamp += float(np.random.exponential(1.0 / self.request_rate))
+        return request_records
+
+
+class AttachExecutionFeature(RequestProcessor):
+    """The processor that attaches execution features to all requests"""
+
+    def __init__(self, exec_feature: Dict[str, Any]) -> None:  # noqa: UP006
+        self.exec_feature = exec_feature
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for request_record in request_records:
+            assert request_record.metrics is not None
+            request_record.metrics.exec_feature = self.exec_feature
+        return request_records
+
+
+class AttachStreamFlag(RequestProcessor):
+    """The processor that attaches the stream flag to the requests."""
+
+    def __init__(self, stream: Optional[bool]) -> None:
+        self.stream = stream
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        if self.stream is None:
+            return request_records
+        for request_record in request_records:
+            request_record.chat_cmpl.stream = self.stream
+        return request_records
+
+
+class AttachSamplingOptions(RequestProcessor):
+    """The processor that attaches the stream flag to the requests."""
+
+    def __init__(self, temperature: float, top_p: float, ignore_eos: bool) -> None:
+        self.temperature = temperature
+        self.top_p = top_p
+        self.ignore_eos = ignore_eos
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for request_record in request_records:
+            request_record.chat_cmpl.temperature = self.temperature
+            request_record.chat_cmpl.top_p = self.top_p
+            request_record.chat_cmpl.frequency_penalty = 0.0
+            request_record.chat_cmpl.presence_penalty = 0.0
+            request_record.chat_cmpl.tool_choice = "none"
+            if self.ignore_eos:
+                request_record.chat_cmpl.debug_config = DebugConfig(ignore_eos=True)
+        return request_records
+
+
+class ScaleTimestamp(RequestProcessor):
+    """Scale the timestamp of requests by the given scale factor."""
+
+    def __init__(self, timestamp_scale: float):
+        self.timestamp_scale = timestamp_scale
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for request_record in request_records:
+            if request_record.timestamp is None:
+                raise ValueError(
+                    f"The timestamp of request {request_record} has not been initialized."
+                )
+            request_record.timestamp *= self.timestamp_scale
+        return request_records
+
+
+class MetricAnalyzer(RequestProcessor):
+    """The processor that analyzes the raw benchmark results and computes more detailed metrics."""
+
+    def __init__(self, tokenizer: AutoTokenizer) -> None:
+        self.tokenizer = tokenizer
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        updated_records = []
+        for request_record in request_records:
+            metrics = request_record.metrics
+            if not metrics.success:
+                assert request_record.error_msg is not None
+                continue
+
+            metrics.output_tokens = len(
+                self.tokenizer.encode(request_record.output_str, add_special_tokens=False)
+            )
+            first_chunk_output_tokens = len(
+                self.tokenizer.encode(
+                    request_record.first_chunk_output_str, add_special_tokens=False
+                )
+            )
+            if metrics.output_tokens <= first_chunk_output_tokens:
+                metrics.success = False
+                request_record.error_msg = (
+                    f"Total output token num ({metrics.output_tokens}) equals "
+                    f'the first chunk output token. Output text "{request_record.output_str}", '
+                    f'first chunk output text "{request_record.first_chunk_output_str}"'
+                )
+                continue
+            assert metrics.input_tokens > 0, "Invalid prompt tokens"
+            metrics.inter_token_latency_s = metrics.end_to_end_latency_s / metrics.output_tokens
+            if metrics.time_to_first_token_s is None:
+                metrics.time_to_first_token_s = 0
+            metrics.time_per_output_token_s = (
+                metrics.end_to_end_latency_s - metrics.time_to_first_token_s
+            ) / (metrics.output_tokens - first_chunk_output_tokens)
+            updated_records.append(request_record)
+        return updated_records
+
+
+class WarmupAndRun(RequestProcessor):
+    """The processor that runs warmup first and then runs the benchmark with the given pipeline."""
+
+    def __init__(
+        self,
+        num_warmup_requests: int,
+        num_benchmark_requests: int,
+        pipeline: RequestProcessor,
+        cuda_profile_url: Optional[str],
+        fake_warmup: bool = False,
+    ) -> None:
+        self.num_warmup_requests = num_warmup_requests
+        self.num_benchmark_requests = num_benchmark_requests
+        self.pipeline = pipeline
+        self.cuda_profile_url = cuda_profile_url
+        self.fake_warmup = fake_warmup
+
+    def generate_fake_warmup_requests(
+        self, num_warmup_requests: int, example_request: RequestRecord
+    ) -> List[RequestRecord]:  # noqa: UP006
+        records = []
+        for _ in range(num_warmup_requests):
+            record = copy.deepcopy(example_request)
+            record.chat_cmpl = ChatCompletionRequest(
+                messages=[
+                    {
+                        "role": "user",
+                        "content": "Please output arbitrary coherent sentences. Do not output eos token.",  # noqa: E501
+                    }
+                ],
+                model="",
+                max_tokens=128,
+            )
+            records.append(record)
+        return records
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        # Warmup
+        if self.fake_warmup:
+            assert len(request_records) == self.num_benchmark_requests
+            benchmark_requests = request_records
+            example_request = benchmark_requests[0]
+            warmup_requests = self.generate_fake_warmup_requests(
+                self.num_warmup_requests, example_request=example_request
+            )
+        else:
+            assert len(request_records) == self.num_warmup_requests + self.num_benchmark_requests
+            benchmark_requests = request_records[: -self.num_warmup_requests]
+            warmup_requests = request_records[-self.num_warmup_requests :]
+        for request_record in warmup_requests:
+            request_record.timestamp = 0 if request_record.timestamp is not None else None
+        warmup_requests = self._process_warmup_requests(warmup_requests)
+        logger.info("Warmup with %d request(s)...", self.num_warmup_requests)
+        self.pipeline(warmup_requests)
+
+        # Then run benchmark
+        if self.cuda_profile_url is not None:
+            cuda_profiler_start_url = self.cuda_profile_url + "/debug/cuda_profiler_start"
+            cuda_profiler_start_response = requests.post(cuda_profiler_start_url, timeout=60)
+            assert cuda_profiler_start_response.status_code == 200
+        logger.info("Warmup finished. Start benchmarking...")
+        updated_request_records = self.pipeline(benchmark_requests)
+        if self.cuda_profile_url is not None:
+            cuda_profiler_stop_url = self.cuda_profile_url + "/debug/cuda_profiler_stop"
+            cuda_profiler_stop_response = requests.post(cuda_profiler_stop_url, timeout=60)
+            assert cuda_profiler_stop_response.status_code == 200
+
+        return updated_request_records
+
+    def _process_warmup_requests(self, warmup_requests: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        if len(warmup_requests) == 0:
+            return warmup_requests
+        # NOTE: to warm up the server for as more different batch sizes as possible,
+        # we usese 128 output tokens for the first request and use two more tokens
+        # for every followup request.
+        # Setting a high temperature and top-p to avoid early stop as much as possible.
+        warmup_requests[0].chat_cmpl.max_tokens = 128
+        for i in range(1, len(warmup_requests)):
+            warmup_requests[i].chat_cmpl.max_tokens = (
+                warmup_requests[i - 1].chat_cmpl.max_tokens + 1
+            )
+            warmup_requests[i].chat_cmpl.temperature = 2.0
+            warmup_requests[i].chat_cmpl.top_p = 1.0
+        return warmup_requests
+
+
+class SequentialProcessor(RequestProcessor):
+    """The processor that sequentially applies a list of processors in order."""
+
+    processors: List[RequestProcessor]  # noqa: UP006
+
+    def __init__(self, *processors: RequestProcessor) -> None:
+        self.processors = list(processors)
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        for processor in self.processors:
+            request_records = processor(request_records)
+        return request_records
+
+
+class Executor(RequestProcessor):
+    """The executor base class, denoting the kind of benchmark mode."""
+
+    def __init__(
+        self,
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        num_processes: int,
+        disable_tqdm: bool,
+    ) -> None:
+        self.f_create_api_endpoint = f_create_api_endpoint
+        self.disable_tqdm = disable_tqdm
+        self.num_processes = num_processes
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        raise NotImplementedError()
+
+
+class FixedConcurrentRequestExecutor(Executor):
+    """The benchmark executor of fixing the number of concurrent requests."""
+
+    def __init__(
+        self,
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        num_processes: Optional[int],
+        disable_tqdm: bool,
+        num_concurrent_requests: int,
+        multi_round: bool,
+    ) -> None:
+        if num_processes is None:
+            # We assign each process at most 32 concurrent requests to send
+            # so that the asyncio pressure will not be too much.
+            num_processes = min((num_concurrent_requests + 31) // 32, 10)
+        super().__init__(f_create_api_endpoint, num_processes, disable_tqdm)
+        self.num_concurrent_requests = num_concurrent_requests
+        self.multi_round = multi_round
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        partitions: List[List[RequestRecord]] = [  # noqa: UP006
+            request_records[slice(i, len(request_records), self.num_processes)]
+            for i in range(self.num_processes)
+        ]
+        # Package "tokenizers" reports warnings with multiprocessing.
+        # We disable "TOKENIZERS_PARALLELISM" to depress the warnings.
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+        pbar = None if self.disable_tqdm else tqdm(total=len(request_records))
+        with concurrent.futures.ProcessPoolExecutor(max_workers=self.num_processes) as pool:
+            futures = [
+                pool.submit(
+                    FixedConcurrentRequestExecutor._process_task,
+                    self.f_create_api_endpoint,
+                    partition,
+                    self.num_concurrent_requests // self.num_processes
+                    + int(i < self.num_concurrent_requests % self.num_processes),
+                    self.multi_round,
+                )
+                for i, partition in enumerate(partitions)
+            ]
+            results: List[RequestRecord] = []  # noqa: UP006
+            for i, future in enumerate(concurrent.futures.as_completed(futures)):
+                results.extend(future.result())
+                if pbar is not None:
+                    pbar.update(len(partitions[i]))
+
+        return results
+
+    @staticmethod
+    def _process_task(
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        request_records: List[RequestRecord],  # noqa: UP006
+        num_concurrent_requests: int,
+        multi_round: bool,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if len(request_records) == 0:
+            return []
+        chat_history: List[List[ChatCompletionMessage]] = [  # noqa: UP006
+            [] for _ in range(num_concurrent_requests)
+        ]
+
+        async def process_task_impl(
+            f_create_api_endpoint: Callable[[], APIEndPoint],
+            request_records: List[RequestRecord],  # noqa: UP006
+            num_concurrent_requests: int,
+            multi_round: bool,
+        ) -> List[RequestRecord]:  # noqa: UP006
+            api_endpoint = f_create_api_endpoint()
+            updated_request_records: List[RequestRecord] = [None for _ in request_records]  # noqa: UP006
+            async with api_endpoint:
+                num_sent_request = 0
+
+                async def _task(i: int) -> None:
+                    nonlocal num_sent_request
+                    while True:
+                        if num_sent_request == len(request_records):
+                            break
+                        idx = num_sent_request
+                        num_sent_request += 1
+                        request = request_records[idx]
+
+                        if multi_round:
+                            request.chat_cmpl.messages = (
+                                chat_history[i] + request.chat_cmpl.messages
+                            )
+
+                        updated_request_records[idx] = await api_endpoint(request)
+
+                        if multi_round:
+                            chat_history[i] = [
+                                *updated_request_records[idx].chat_cmpl.messages,
+                                ChatCompletionMessage(
+                                    content=updated_request_records[idx].output_str,
+                                    role="assistant",
+                                ),
+                            ]
+
+                tasks = [asyncio.create_task(_task(i)) for i in range(num_concurrent_requests)]
+                await asyncio.gather(*tasks)
+
+            return updated_request_records
+
+        return asyncio.run(
+            process_task_impl(
+                f_create_api_endpoint,
+                request_records,
+                num_concurrent_requests,
+                multi_round,
+            )
+        )
+
+
+class FixTimestampExecutor(Executor):
+    """The benchmark executor of fixing the timestamps of sending requests."""
+
+    def __init__(
+        self,
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        num_processes: Optional[int],
+        disable_tqdm: bool,
+        max_schedule_gap: float,
+        num_requests: int,
+    ) -> None:
+        if num_processes is None:
+            # We assign each process at most 32 requests to send
+            # so that the asyncio pressure will not be too much.
+            num_processes = min((num_requests + 31) // 32, 10)
+        super().__init__(f_create_api_endpoint, num_processes, disable_tqdm)
+        self.max_schedule_gap = max_schedule_gap
+        self.num_requests = num_requests
+
+    def __call__(self, request_records: List[RequestRecord]) -> List[RequestRecord]:  # noqa: UP006
+        assert len(request_records) > 0
+        assert all(request_record.timestamp is not None for request_record in request_records)
+        # Sort the request records in timestamp ascending order before partitioning.
+        request_records.sort(key=lambda request_record: request_record.timestamp)
+        base_timestamp = request_records[0].timestamp
+        partitions: List[List[RequestRecord]] = [  # noqa: UP006
+            request_records[slice(i, len(request_records), self.num_processes)]
+            for i in range(self.num_processes)
+        ]
+        base_sys_time = time.time()
+        # Package "tokenizers" reports warnings with multiprocessing.
+        # We disable "TOKENIZERS_PARALLELISM" to depress the warnings.
+        os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+        pbar = None if self.disable_tqdm else tqdm(total=len(request_records))
+        with concurrent.futures.ProcessPoolExecutor(max_workers=self.num_processes) as pool:
+            futures = [
+                pool.submit(
+                    FixTimestampExecutor._process_task,
+                    self.f_create_api_endpoint,
+                    partition,
+                    base_timestamp,
+                    base_sys_time,
+                    self.max_schedule_gap,
+                )
+                for partition in partitions
+            ]
+            results: List[RequestRecord] = []  # noqa: UP006
+            for i, future in enumerate(concurrent.futures.as_completed(futures)):
+                results.extend(future.result())
+                if pbar is not None:
+                    pbar.update(len(partitions[i]))
+
+        return results
+
+    @staticmethod
+    def _process_task(
+        f_create_api_endpoint: Callable[[], APIEndPoint],
+        request_records: List[RequestRecord],  # noqa: UP006
+        base_timestamp: float,
+        base_sys_time: float,
+        max_schedule_gap: float,
+    ) -> List[RequestRecord]:  # noqa: UP006
+        if len(request_records) == 0:
+            return []
+
+        async def process_task_impl(
+            f_create_api_endpoint: Callable[[], APIEndPoint],
+            request_records: List[RequestRecord],  # noqa: UP006
+            base_timestamp: float,
+            base_sys_time: float,
+            max_schedule_gap: float,
+        ) -> List[RequestRecord]:  # noqa: UP006
+            api_endpoint = f_create_api_endpoint()
+            loop = asyncio.get_running_loop()
+            # Get the delta time to convert system time to the loop time.
+            # We must use the system time `time.time()` which is consistent across processes.
+            loop_sys_delta_time = loop.time() - time.time()
+            updated_request_records: List[RequestRecord] = []  # noqa: UP006
+            async with api_endpoint:
+
+                async def _task(request_record: RequestRecord) -> None:
+                    updated_request_records.append(await api_endpoint(request_record))
+
+                tasks = []
+                for request_record in request_records:
+                    launch_time = (
+                        (request_record.timestamp - base_timestamp)
+                        + (base_sys_time + max_schedule_gap)
+                        + loop_sys_delta_time
+                    )
+                    loop.call_at(
+                        launch_time,
+                        lambda record: tasks.append(asyncio.create_task(_task(record))),
+                        request_record,
+                    )
+                    # Sleep to allow runs of other scheduled tasks if any.
+                    await asyncio.sleep(max(launch_time - loop.time() - max_schedule_gap, 0))
+
+                # Sleep until all the tasks are launched.
+                await asyncio.sleep(launch_time - loop.time() + max_schedule_gap)
+                # Wait for all tasks to be scheduled
+                assert len(tasks) == len(request_records)
+                await asyncio.gather(*tasks)
+
+            assert len(updated_request_records) == len(request_records)
+            return updated_request_records
+
+        return asyncio.run(
+            process_task_impl(
+                f_create_api_endpoint,
+                request_records,
+                base_timestamp,
+                base_sys_time,
+                max_schedule_gap,
+            )
+        )
+
+
+def create_pipelines(
+    args: argparse.Namespace,
+    f_create_api_endpoint: Callable[[], APIEndPoint],
+    dataset: Dataset,
+) -> List[RequestProcessor]:  # noqa: UP006
+    """Creating request processing pipelines with regard to the specified args."""
+    cuda_profile_url = f"http://{args.host}:{args.port}" if args.cuda_profile else None
+    pipelines: List[RequestProcessor] = []  # noqa: UP006
+    if args.num_concurrent_requests is not None:
+        if args.request_rate is not None:
+            raise ValueError(
+                'Both "num_concurrent_requests" and "request_rate" are specified. '
+                "Please specify only one of them."
+            )
+        if args.replay_timestamp_scale is not None:
+            raise ValueError(
+                "Dataset replay is unsupported when fixing number of concurrent requests."
+            )
+        for num_concurrent_requests in args.num_concurrent_requests:
+            num_warmup_requests = (
+                args.num_warmup_requests
+                if args.num_warmup_requests is not None
+                else num_concurrent_requests
+            )
+            pipelines.append(
+                SequentialProcessor(
+                    LogMessage(f"Fixing number of concurrent requests: {num_concurrent_requests}"),
+                    SampleRequests(args.num_requests + num_warmup_requests),
+                    AttachModelName(args.tokenizer),
+                    AttachStreamFlag(args.stream),
+                    AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos),
+                    AttachExecutionFeature({"num_concurrent_requests": num_concurrent_requests}),
+                    WarmupAndRun(
+                        num_warmup_requests=num_warmup_requests,
+                        num_benchmark_requests=args.num_requests,
+                        pipeline=FixedConcurrentRequestExecutor(
+                            f_create_api_endpoint,
+                            args.num_process_workers,
+                            args.disable_tqdm,
+                            num_concurrent_requests,
+                            args.multi_round,
+                        ),
+                        cuda_profile_url=cuda_profile_url,
+                        fake_warmup=dataset.require_fake_warmup,
+                    ),
+                )
+            )
+        return pipelines
+    if args.request_rate is not None:
+        if args.num_warmup_requests is None:
+            raise ValueError(
+                "Please specify the number of warmup requests via "
+                '"--num-warmup-requests" when fixing request rate.'
+            )
+        if args.replay_timestamp_scale is not None:
+            raise ValueError("Dataset replay is unsupported when fixing request rates.")
+        num_total_requests = int(
+            args.num_requests if not args.per_gpu_workload else args.num_requests * args.num_gpus
+        )
+        if dataset.require_fake_warmup:
+            num_samples = num_total_requests
+        else:
+            num_samples = num_total_requests + args.num_warmup_requests
+        return [
+            SequentialProcessor(
+                LogMessage(f"Fixing request rate: {request_rate}"),
+                SampleRequests(num_samples),
+                AttachModelName(args.tokenizer),
+                AttachRequestRateTimestamp(
+                    request_rate if not args.per_gpu_workload else request_rate * args.num_gpus
+                ),
+                AttachStreamFlag(args.stream),
+                AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos),
+                AttachExecutionFeature({"request_rate": float(request_rate)}),
+                WarmupAndRun(
+                    num_warmup_requests=args.num_warmup_requests,
+                    num_benchmark_requests=num_total_requests,
+                    pipeline=FixTimestampExecutor(
+                        f_create_api_endpoint,
+                        args.num_process_workers,
+                        args.disable_tqdm,
+                        args.max_schedule_gap,
+                        args.num_requests,
+                    ),
+                    cuda_profile_url=cuda_profile_url,
+                    fake_warmup=dataset.require_fake_warmup,
+                ),
+            )
+            for request_rate in args.request_rate
+        ]
+
+    # Default: dataset replay mode
+    # The dataset must come with timestamps.
+    if not dataset.timestamp_available:
+        raise ValueError(
+            "The dataset does not have timestamps, so dataset replay is unsupported. "
+            'Please specify one of "num_concurrent_requests" '
+            'and "request_rate".'
+        )
+    if args.per_gpu_workload:
+        raise ValueError("Fixing per-GPU workload is not compatible with dataset replay.")
+    if args.num_warmup_requests is None:
+        raise ValueError(
+            "Please specify the number of warmup requests via "
+            '"--num-warmup-requests" for dataset replay.'
+        )
+    timestamp_scale = args.replay_timestamp_scale or 1.0
+    if dataset.require_fake_warmup:
+        num_samples = args.num_requests
+    else:
+        num_samples = args.num_requests + args.num_warmup_requests
+    return [
+        SequentialProcessor(
+            LogMessage(f"Dataset replay with time scaling of {timestamp_scale}"),
+            SampleRequests(num_samples, take_first_x_requests=True),
+            AttachModelName(args.tokenizer),
+            ScaleTimestamp(timestamp_scale),
+            AttachStreamFlag(args.stream),
+            AttachSamplingOptions(args.temperature, args.top_p, args.ignore_eos),
+            AttachExecutionFeature({"timestamp_scale": timestamp_scale}),
+            WarmupAndRun(
+                num_warmup_requests=args.num_warmup_requests,
+                num_benchmark_requests=args.num_requests,
+                pipeline=FixTimestampExecutor(
+                    f_create_api_endpoint,
+                    args.num_process_workers,
+                    args.disable_tqdm,
+                    args.max_schedule_gap,
+                    args.num_requests,
+                ),
+                cuda_profile_url=cuda_profile_url,
+                fake_warmup=dataset.require_fake_warmup,
+            ),
+        )
+    ]

python/mlc_llm/bench/request_record.py

@@ -0,0 +1,278 @@
+"""MLC LLM Bench Request"""
+
+from typing import Any, Dict, List, Optional, Tuple, Union  # noqa: UP035
+
+import pandas as pd
+from pydantic import BaseModel
+
+from mlc_llm.protocol.openai_api_protocol import ChatCompletionRequest
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+class ServerMetrics(BaseModel):
+    """The metrics from the server side."""
+
+    input_tokens: int
+    prefill_tokens: int
+    output_tokens: int
+    end_to_end_latency_s: float
+    prefill_tokens_per_s: float
+    inter_token_latency_s: float
+    time_per_output_token_s: float
+    time_to_first_token_s: Optional[float] = None
+
+
+class Metrics(BaseModel):
+    """The list of metric keys"""
+
+    success: bool
+    start_time: float
+    finish_time: float
+    end_to_end_latency_s: float
+
+    input_tokens: Optional[int] = None
+    output_tokens: Optional[int] = None
+    inter_token_latency_s: Optional[float] = None
+    time_per_output_token_s: Optional[float] = None
+    time_to_first_token_s: Optional[float] = None
+    server_metrics: Optional[ServerMetrics] = None
+
+    exec_feature: Optional[Dict[str, Any]] = None  # noqa: UP006
+
+
+class RequestRecord(BaseModel):
+    """The request records collected from LLM inference requests."""
+
+    request_id: Optional[int] = None
+    chat_cmpl: ChatCompletionRequest
+    output_str: Optional[str] = None
+    first_chunk_output_str: str = ""
+    timestamp: Optional[float] = None
+    metrics: Optional[Metrics] = None
+    error_msg: Optional[str] = None
+
+
+class GroupedRequestRecord(RequestRecord):
+    """The data structure for request record groups.
+    For datasets that have common prefix sharing, the request records
+    that share a same common prefix will be wrapped in a GroupedRequestRecord
+    at the beginning.
+    """
+
+    records: List[RequestRecord]  # noqa: UP006
+
+
+def generate_metrics_summary(
+    request_records: List[RequestRecord],  # noqa: UP006
+    num_total_requests: int,
+    num_gpus: int,
+) -> Dict[str, Any]:  # noqa: UP006
+    """Computes summary statistics across all metrics collected.
+    Return a dictionary as the report.
+    """
+    num_completed_requests = len(request_records)
+    assert num_completed_requests <= num_total_requests
+    request_metrics = [record.metrics for record in request_records]
+    duration = (
+        max(metrics.finish_time for metrics in request_metrics)
+        - min(metrics.start_time for metrics in request_metrics)
+        if num_completed_requests > 0
+        else 1e-5
+    )
+
+    report = _compute_metrics_statistics(request_metrics)
+    report["num_gpus"] = num_gpus
+    report["duration"] = duration
+    report["num_total_requests"] = num_total_requests
+    report["num_completed_requests"] = num_completed_requests
+    report["request_throughput"] = num_completed_requests / duration
+
+    total_input_tokens = sum(metric.input_tokens for metric in request_metrics)
+    total_output_tokens = sum(metric.output_tokens for metric in request_metrics)
+    report["total_input_tokens"] = total_input_tokens
+    report["total_output_tokens"] = total_output_tokens
+    report["input_token_throughput"] = total_input_tokens / duration
+    report["input_token_throughput_per_gpu"] = report["input_token_throughput"] / num_gpus
+    report["output_token_throughput"] = total_output_tokens / duration
+    report["output_token_throughput_per_gpu"] = report["output_token_throughput"] / num_gpus
+
+    # Generate the server metrics statistics
+    server_metrics = [metric.server_metrics for metric in request_metrics if metric.server_metrics]
+    server_report = _compute_metrics_statistics(server_metrics)
+    if server_report is not None and len(server_report) > 0:
+        report["server_metrics"] = server_report
+
+    report = {
+        "exec_feature": (
+            request_records[0].metrics.exec_feature if num_completed_requests > 0 else None
+        ),
+        **report,
+    }
+    return report
+
+
+def _compute_metrics_statistics(
+    metrics: List[Union[Metrics, ServerMetrics]],  # noqa: UP006
+) -> Dict[str, Any]:  # noqa: UP006
+    """
+    Compute the statistics of the metrics.
+
+    Parameters
+    ----------
+    metrics : List[Union[Metrics, ServerMetrics]]
+        The list of metrics to get the statistics.
+
+    Returns
+    -------
+    report : Dict
+        The statistics of the metrics.
+    """
+    if not metrics:
+        return {}
+
+    report: Dict = {}  # noqa: UP006
+    df = pd.DataFrame([metric.model_dump() for metric in metrics])
+    for key, _ in metrics[0].model_fields.items():
+        if key in [
+            "success",
+            "start_time",
+            "finish_time",
+            "server_metrics",
+            "exec_feature",
+        ]:
+            continue
+        if key in df.columns:
+            series = df[key].dropna()
+            report[key] = {
+                "quantiles": {
+                    f"p{int(q * 100)}": v
+                    for q, v in series.quantile([0.25, 0.5, 0.75, 0.9, 0.95, 0.99]).items()
+                },
+                "mean": series.mean(),
+                "min": series.min(),
+                "max": series.max(),
+                "stddev": series.std(),
+            }
+    return report
+
+
+def convert_reports_to_df(reports: List[Dict[str, Any]]) -> pd.DataFrame:  # noqa: UP006
+    """Convert benchmark reports to pandas DataFrame."""
+
+    def _flatten_dict(d: Dict[str, Any], parent_key: str = "") -> Dict[str, Any]:  # noqa: UP006
+        items: List[Tuple[str, Any]] = []  # noqa: UP006
+        for key, value in d.items():
+            new_key = f"{parent_key}.{key}" if parent_key != "" else key
+            if isinstance(value, dict):
+                items.extend(_flatten_dict(value, new_key).items())
+            else:
+                items.append((new_key, value))
+        return dict(items)
+
+    return pd.DataFrame([_flatten_dict(report) for report in reports])
+
+
+def pretty_print_report(report: Dict[str, Any]) -> None:  # noqa: UP006
+    """Pretty print the metrics report."""
+
+    def _print(report: Dict[str, Any], server_metrics: bool):  # noqa: UP006
+        # fmt: off
+        title = "Benchmark Result"
+        if server_metrics:
+            title += " (server side)"
+        print(f" {title} ".center(50, "="))
+        if not server_metrics:
+            print(f"{'Total requests:':<40} {report['num_total_requests']:<10}")
+            print(f"{'Completed requests:':<40} {report['num_completed_requests']:<10}")
+            print(f"{'Duration (s):':<40} {report['duration']:<10.2f}")
+            print(f"{'Num GPUs:':<40} {report['num_gpus']:<10}")
+            print(f"{'Total input tokens:':<40} {report['total_input_tokens']:<10}")
+            print(f"{'Total output tokens:':<40} {report['total_output_tokens']:<10}")
+            print(f"{'Request throughput (req/s):':<40} {report['request_throughput']:<10.2f}")
+            print(f"{'Input token throughput (tok/s):':<40} {report['input_token_throughput']:<10.2f}")  # noqa: E501
+            print(f"{'Input token throughput per GPU (tok/s):':<40} {report['input_token_throughput_per_gpu']:<10.2f}")  # noqa: E501
+            print(f"{'Output token throughput (tok/s):':<40} {report['output_token_throughput']:<10.2f}")  # noqa: E501
+            print(f"{'Output token throughput per GPU (tok/s):':<40} {report['output_token_throughput_per_gpu']:<10.2f}")  # noqa: E501
+
+        if report["num_completed_requests"] == 0:
+            return
+        ttft = report["time_to_first_token_s"]
+        print(" Time to First Token (TTFT, ms) ".center(50, "-"))
+        print(f"{'Mean:':<40} {ttft['mean'] * 1000:<10.2f}")
+        print(f"{'Stddev:':<40} {ttft['stddev'] * 1000:<10.2f}")
+        print(f"{'P25:':<40} {ttft['quantiles']['p25'] * 1000:<10.2f}")
+        print(f"{'P50:':<40} {ttft['quantiles']['p50'] * 1000:<10.2f}")
+        print(f"{'P75:':<40} {ttft['quantiles']['p75'] * 1000:<10.2f}")
+        print(f"{'P90:':<40} {ttft['quantiles']['p90'] * 1000:<10.2f}")
+        print(f"{'P95:':<40} {ttft['quantiles']['p95'] * 1000:<10.2f}")
+        print(f"{'P99:':<40} {ttft['quantiles']['p99'] * 1000:<10.2f}")
+        print(f"{'Min:':<40} {ttft['min'] * 1000:<10.2f}")
+        print(f"{'Max:':<40} {ttft['max'] * 1000:<10.2f}")
+
+        tpot = report["time_per_output_token_s"]
+        print(" Time per Output Token (TPOT, ms) ".center(50, "-"))
+        print(f"{'Mean:':<40} {tpot['mean'] * 1000:<10.2f}")
+        print(f"{'Stddev:':<40} {tpot['stddev'] * 1000:<10.2f}")
+        print(f"{'P25:':<40} {tpot['quantiles']['p25'] * 1000:<10.2f}")
+        print(f"{'P50:':<40} {tpot['quantiles']['p50'] * 1000:<10.2f}")
+        print(f"{'P75:':<40} {tpot['quantiles']['p75'] * 1000:<10.2f}")
+        print(f"{'P90:':<40} {tpot['quantiles']['p90'] * 1000:<10.2f}")
+        print(f"{'P95:':<40} {tpot['quantiles']['p95'] * 1000:<10.2f}")
+        print(f"{'P99:':<40} {tpot['quantiles']['p99'] * 1000:<10.2f}")
+        print(f"{'Min:':<40} {tpot['min'] * 1000:<10.2f}")
+        print(f"{'Max:':<40} {tpot['max'] * 1000:<10.2f}")
+
+        itl = report["inter_token_latency_s"]
+        print(" Inter-Token Latency (ms) ".center(50, "-"))
+        print(f"{'Mean:':<40} {itl['mean'] * 1000:<10.2f}")
+        print(f"{'Stddev:':<40} {itl['stddev'] * 1000:<10.2f}")
+        print(f"{'P25:':<40} {itl['quantiles']['p25'] * 1000:<10.2f}")
+        print(f"{'P50:':<40} {itl['quantiles']['p50'] * 1000:<10.2f}")
+        print(f"{'P75:':<40} {itl['quantiles']['p75'] * 1000:<10.2f}")
+        print(f"{'P90:':<40} {itl['quantiles']['p90'] * 1000:<10.2f}")
+        print(f"{'P95:':<40} {itl['quantiles']['p95'] * 1000:<10.2f}")
+        print(f"{'P99:':<40} {itl['quantiles']['p99'] * 1000:<10.2f}")
+        print(f"{'Min:':<40} {itl['min'] * 1000:<10.2f}")
+        print(f"{'Max:':<40} {itl['max'] * 1000:<10.2f}")
+
+        e2e_latency = report["end_to_end_latency_s"]
+        print(" End-to-End Latency (ms) ".center(50, "-"))
+        print(f"{'Mean:':<40} {e2e_latency['mean'] * 1000:<10.2f}")
+        print(f"{'Stddev:':<40} {e2e_latency['stddev'] * 1000:<10.2f}")
+        print(f"{'P25:':<40} {e2e_latency['quantiles']['p25'] * 1000:<10.2f}")
+        print(f"{'P50:':<40} {e2e_latency['quantiles']['p50'] * 1000:<10.2f}")
+        print(f"{'P75:':<40} {e2e_latency['quantiles']['p75'] * 1000:<10.2f}")
+        print(f"{'P90:':<40} {e2e_latency['quantiles']['p90'] * 1000:<10.2f}")
+        print(f"{'P95:':<40} {e2e_latency['quantiles']['p95'] * 1000:<10.2f}")
+        print(f"{'P99:':<40} {e2e_latency['quantiles']['p99'] * 1000:<10.2f}")
+        print(f"{'Min:':<40} {e2e_latency['min'] * 1000:<10.2f}")
+        print(f"{'Max:':<40} {e2e_latency['max'] * 1000:<10.2f}")
+
+        input_tokens = report["input_tokens"]
+        print(" Input Tokens ".center(50, "-"))
+        print(f"{'Mean:':<40} {input_tokens['mean']:<1}")
+        print(f"{'Stddev:':<40} {input_tokens['stddev']:<1}")
+        print(f"{'P25:':<40} {input_tokens['quantiles']['p25']:<1}")
+        print(f"{'P50:':<40} {input_tokens['quantiles']['p50']:<1}")
+        print(f"{'P95:':<40} {input_tokens['quantiles']['p95']:<1}")
+        print(f"{'Min:':<40} {input_tokens['min']:<1}")
+        print(f"{'Max:':<40} {input_tokens['max']:<1}")
+
+        output_tokens = report["output_tokens"]
+        print(" Output Tokens ".center(50, "-"))
+        print(f"{'Mean:':<40} {output_tokens['mean']:<1}")
+        print(f"{'Stddev:':<40} {output_tokens['stddev']:<1}")
+        print(f"{'P25:':<40} {output_tokens['quantiles']['p25']:<1}")
+        print(f"{'P50:':<40} {output_tokens['quantiles']['p50']:<1}")
+        print(f"{'P95:':<40} {output_tokens['quantiles']['p95']:<1}")
+        print(f"{'Min:':<40} {output_tokens['min']:<1}")
+        print(f"{'Max:':<40} {output_tokens['max']:<1}")
+
+        print("=" * 50)
+
+    # fmt: on
+    _print(report, server_metrics=False)
+    if "server_metrics" in report:
+        _print(report["server_metrics"], server_metrics=True)

python/mlc_llm/cli/calibrate.py

@@ -0,0 +1,80 @@
+"""Command line entrypoint of calibration."""
+
+from mlc_llm.interface.calibrate import calibrate
+from mlc_llm.interface.help import HELP
+from mlc_llm.support.argparse import ArgumentParser
+
+from .serve import EngineConfigOverride
+
+
+def main(argv):
+    """Main entrypoint for calibration."""
+    parser = ArgumentParser("MLC LLM Calibration CLI")
+    parser.add_argument(
+        "model",
+        type=str,
+        help=HELP["model"] + " (required)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="auto",
+        help=HELP["device_deploy"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--model-lib",
+        type=str,
+        default=None,
+        help=HELP["model_lib"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=str,
+        required=True,
+        help=HELP["output_calibration"] + " (required)",
+    )
+    # Download dataset from
+    # https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
+    parser.add_argument(
+        "--dataset",
+        type=str,
+        required=True,
+        help=HELP["calibration_dataset"] + " (required)",
+    )
+
+    parser.add_argument(
+        "--num-calibration-samples",
+        type=int,
+        default=16,
+        help=HELP["num_calibration_samples"] + ' (default: "%(default)s")',
+    )
+
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=0,
+        help=HELP["seed_calibrate"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--overrides",
+        type=EngineConfigOverride.from_str,
+        default="",
+        help=HELP["overrides_serve"],
+    )
+
+    parsed = parser.parse_args(argv)
+    calibrate(
+        model=parsed.model,
+        device=parsed.device,
+        model_lib=parsed.model_lib,
+        output=parsed.output,
+        dataset=parsed.dataset,
+        num_calibration_samples=parsed.num_calibration_samples,
+        max_num_sequence=parsed.overrides.max_num_sequence,
+        max_total_sequence_length=parsed.overrides.max_total_seq_length,
+        prefill_chunk_size=parsed.overrides.prefill_chunk_size,
+        max_history_size=parsed.overrides.max_history_size,
+        gpu_memory_utilization=parsed.overrides.gpu_memory_utilization,
+        seed=parsed.seed,
+    )

python/mlc_llm/cli/chat.py

@@ -0,0 +1,41 @@
+"""Command line entrypoint of chat."""
+
+from mlc_llm.interface.chat import ModelConfigOverride, chat
+from mlc_llm.interface.help import HELP
+from mlc_llm.support.argparse import ArgumentParser
+
+
+def main(argv):
+    """Parse command line arguments and call `mlc_llm.interface.chat`."""
+    parser = ArgumentParser("MLC LLM Chat CLI")
+
+    parser.add_argument(
+        "model",
+        type=str,
+        help=HELP["model"] + " (required)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="auto",
+        help=HELP["device_deploy"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--model-lib",
+        type=str,
+        default=None,
+        help=HELP["model_lib"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--overrides",
+        type=ModelConfigOverride.from_str,
+        default="",
+        help=HELP["modelconfig_overrides"] + ' (default: "%(default)s")',
+    )
+    parsed = parser.parse_args(argv)
+    chat(
+        model=parsed.model,
+        device=parsed.device,
+        model_lib=parsed.model_lib,
+        overrides=parsed.overrides,
+    )

python/mlc_llm/cli/check_device.py

@@ -0,0 +1,34 @@
+"""Check if a device exists."""
+
+import os
+import sys
+
+from tvm.runtime import Device
+from tvm.runtime import device as as_device
+
+
+def _check_device(device: Device) -> bool:
+    try:
+        return bool(device.exist)
+    except Exception:
+        return False
+
+
+def main():
+    """Entrypoint for device check."""
+    device_str = sys.argv[1]
+    device_ids = []
+    i = 0
+    while True:
+        if _check_device(as_device(device_str, i)):
+            device_ids.append(i)
+            i += 1
+            if device_str in ["cpu", "llvm"] and i > os.cpu_count() / 2:
+                break
+        else:
+            break
+    print(f"check_device:{','.join(str(i) for i in device_ids)}")
+
+
+if __name__ == "__main__":
+    main()

python/mlc_llm/cli/compile.py

@@ -0,0 +1,152 @@
+"""Command line entrypoint of compilation."""
+
+import argparse
+import json
+import re
+from functools import partial
+from pathlib import Path
+from typing import Union
+
+from mlc_llm.interface.compile import (
+    ModelConfigOverride,
+    OptimizationFlags,
+    compile,
+)
+from mlc_llm.interface.help import HELP
+from mlc_llm.model import MODELS
+from mlc_llm.quantization import QUANTIZATION
+from mlc_llm.support.argparse import ArgumentParser
+from mlc_llm.support.auto_config import (
+    detect_mlc_chat_config,
+    detect_model_type,
+    detect_quantization,
+)
+from mlc_llm.support.auto_target import detect_system_lib_prefix, detect_target_and_host
+
+
+def main(argv):
+    """Parse command line arguments and call `mlc_llm.compiler.compile`."""
+
+    def _parse_output(path: Union[str, Path]) -> Path:
+        path = Path(path)
+        if path.is_dir():
+            raise argparse.ArgumentTypeError(f"Output cannot be a directory: {path}")
+        parent = path.parent
+        if not parent.is_dir():
+            raise argparse.ArgumentTypeError(f"Directory does not exist: {parent}")
+        return path
+
+    def _parse_dir(path: Union[str, Path], auto_create: bool = False) -> Path:
+        path = Path(path)
+        if not auto_create and not path.is_dir():
+            raise argparse.ArgumentTypeError(f"Directory does not exist: {path}")
+        if auto_create and not path.is_dir():
+            path.mkdir(parents=True)
+        return path
+
+    def _check_system_lib_prefix(prefix: str) -> str:
+        pattern = r"^[a-zA-Z_][a-zA-Z0-9_]*$"
+        if prefix == "" or re.match(pattern, prefix):
+            return prefix
+        raise argparse.ArgumentTypeError(
+            "Invalid prefix. It should only consist of "
+            "numbers (0-9), alphabets (A-Z, a-z) and underscore (_)."
+        )
+
+    parser = ArgumentParser("mlc_llm compile")
+    parser.add_argument(
+        "model",
+        type=detect_mlc_chat_config,
+        help=HELP["model"] + " (required)",
+    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        choices=list(QUANTIZATION.keys()),
+        help=HELP["quantization"]
+        + " (default: look up mlc-chat-config.json, choices: %(choices)s)",
+    )
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        default="auto",
+        choices=["auto", *list(MODELS.keys())],
+        help=HELP["model_type"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="auto",
+        help=HELP["device_compile"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="auto",
+        help=HELP["host"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--enable-subgroups",
+        action="store_true",
+        help=HELP["enable_subgroups"],
+    )
+    parser.add_argument(
+        "--opt",
+        type=OptimizationFlags.from_str,
+        default="O2",
+        help=HELP["opt"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--system-lib-prefix",
+        type=str,
+        default="auto",
+        help=HELP["system_lib_prefix"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=_parse_output,
+        required=True,
+        help=HELP["output_compile"] + " (required)",
+    )
+    parser.add_argument(
+        "--overrides",
+        type=ModelConfigOverride.from_str,
+        default="",
+        help=HELP["overrides"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--debug-dump",
+        type=partial(_parse_dir, auto_create=True),
+        default=None,
+        help=HELP["debug_dump"] + " (default: %(default)s)",
+    )
+    parsed = parser.parse_args(argv)
+    target, build_func = detect_target_and_host(
+        parsed.device,
+        parsed.host,
+        enable_subgroups=parsed.enable_subgroups,
+    )
+    parsed.model_type = detect_model_type(parsed.model_type, parsed.model)
+    parsed.quantization = detect_quantization(parsed.quantization, parsed.model)
+    parsed.system_lib_prefix = detect_system_lib_prefix(
+        parsed.device,
+        parsed.system_lib_prefix,
+        parsed.model_type.name,
+        parsed.quantization.name,
+    )
+    with open(parsed.model, encoding="utf-8") as config_file:
+        config = json.load(config_file)
+
+    compile(
+        config=config,
+        quantization=parsed.quantization,
+        model_type=parsed.model_type,
+        target=target,
+        opt=parsed.opt,
+        build_func=build_func,
+        system_lib_prefix=parsed.system_lib_prefix,
+        output=parsed.output,
+        overrides=parsed.overrides,
+        debug_dump=parsed.debug_dump,
+    )

python/mlc_llm/cli/convert_weight.py

@@ -0,0 +1,112 @@
+"""Command line entrypoint of weight conversion."""
+
+import argparse
+from pathlib import Path
+from typing import Union
+
+from mlc_llm.interface.convert_weight import convert_weight
+from mlc_llm.interface.help import HELP
+from mlc_llm.model import MODELS
+from mlc_llm.quantization import QUANTIZATION
+from mlc_llm.support.argparse import ArgumentParser
+from mlc_llm.support.auto_config import detect_config, detect_model_type
+from mlc_llm.support.auto_device import detect_device
+from mlc_llm.support.auto_weight import detect_weight
+
+
+def main(argv):
+    """Parse command line argumennts and apply quantization."""
+
+    def _parse_source(path: Union[str, Path], config_path: Path) -> Path:
+        if path == "auto":
+            return config_path.parent
+        path = Path(path)
+        if not path.exists():
+            raise argparse.ArgumentTypeError(f"Model source does not exist: {path}")
+        return path
+
+    def _parse_output(path: Union[str, Path]) -> Path:
+        path = Path(path)
+        if not path.is_dir():
+            path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    def _parse_lora_adapter(path: Union[str, Path]) -> Path:
+        path = Path(path)
+        if not path.exists() or not path.is_dir():
+            raise argparse.ArgumentTypeError(f"LoRA adapter directory does not exist: {path}")
+        return path
+
+    parser = ArgumentParser("MLC AutoLLM Quantization Framework")
+    parser.add_argument(
+        "config",
+        type=detect_config,
+        help=HELP["config"] + " (required)",
+    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        required=True,
+        choices=list(QUANTIZATION.keys()),
+        help=HELP["quantization"] + " (required, choices: %(choices)s)",
+    )
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        default="auto",
+        choices=["auto", *list(MODELS.keys())],
+        help=HELP["model_type"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--device",
+        default="auto",
+        type=detect_device,
+        help=HELP["device_quantize"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--source",
+        type=str,
+        default="auto",
+        help=HELP["source"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--source-format",
+        type=str,
+        choices=["auto", "huggingface-torch", "huggingface-safetensor", "awq"],
+        default="auto",
+        help=HELP["source_format"] + ' (default: "%(default)s", choices: %(choices)s")',
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=_parse_output,
+        required=True,
+        help=HELP["output_quantize"] + " (required)",
+    )
+    parser.add_argument(
+        "--lora-adapter",
+        type=_parse_lora_adapter,
+        default=None,
+        help=(
+            "Path to a LoRA adapter directory in PEFT format. "
+            "When provided, adapter weights are merged into the base model before quantization."
+        ),
+    )
+
+    parsed = parser.parse_args(argv)
+    parsed.source, parsed.source_format = detect_weight(
+        _parse_source(parsed.source, parsed.config),
+        parsed.config,
+        parsed.source_format,
+    )
+    model = detect_model_type(parsed.model_type, parsed.config)
+    convert_weight(
+        config=parsed.config,
+        quantization=QUANTIZATION[parsed.quantization],
+        model=model,
+        device=parsed.device,
+        source=parsed.source,
+        source_format=parsed.source_format,
+        output=parsed.output,
+        lora_adapter=parsed.lora_adapter,
+    )

python/mlc_llm/cli/delivery.py

@@ -0,0 +1,452 @@
+"""Continuous model delivery for MLC LLM models."""
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Type, TypeVar, Union  # noqa: UP035
+
+from huggingface_hub import HfApi, snapshot_download
+from huggingface_hub.utils import HfHubHTTPError
+from pydantic import BaseModel, Field, ValidationError
+
+from mlc_llm.support import logging
+from mlc_llm.support.argparse import ArgumentParser
+from mlc_llm.support.style import bold, green, red
+
+logger = logging.getLogger(__name__)
+
+GEN_CONFIG_OPTIONAL_ARGS = [
+    "context_window_size",
+    "sliding_window_size",
+    "prefill_chunk_size",
+    "attention_sink_size",
+    "tensor_parallel_shards",
+    "pipeline_parallel_stages",
+]
+
+T = TypeVar("T", bound="BaseModel")
+
+
+class OverrideConfigs(BaseModel):
+    """
+    The class that specifies the override configurations.
+    """
+
+    context_window_size: Optional[int] = None
+    sliding_window_size: Optional[int] = None
+    prefill_chunk_size: Optional[int] = None
+    attention_sink_size: Optional[int] = None
+    tensor_parallel_shards: Optional[int] = None
+    pipeline_parallel_stages: Optional[int] = None
+
+
+class ModelDeliveryTask(BaseModel):
+    """
+    Example:
+    {
+        "model_id": "Phi-3-mini-128k-instruct",
+        "model": "HF://microsoft/Phi-3-mini-128k-instruct",
+        "conv_template": "phi-3",
+        "quantization": ["q3f16_1"],
+        "overrides": {
+            "q3f16_1": {
+                "context_window_size": 512
+            }
+        }
+    }
+    """
+
+    model_id: str
+    model: str
+    conv_template: str
+    quantization: Union[List[str], str] = Field(default_factory=list)  # noqa: UP006
+    overrides: Dict[str, OverrideConfigs] = Field(default_factory=dict)  # noqa: UP006
+    destination: Optional[str] = None
+    gen_config_only: Optional[bool] = False
+
+
+class ModelDeliveryList(BaseModel):
+    """
+    The class that specifies the model delivery list.
+    """
+
+    tasks: List[ModelDeliveryTask]  # noqa: UP006
+    # For delivered log, the default destination and quantization fields are optional
+    default_destination: Optional[str] = None
+    default_quantization: List[str] = Field(default_factory=list)  # noqa: UP006
+    default_overrides: Dict[str, OverrideConfigs] = Field(default_factory=dict)  # noqa: UP006
+
+    @classmethod
+    def from_json(cls: Type[T], json_dict: Dict[str, Any]) -> T:  # noqa: UP006
+        """
+        Convert from a json dictionary.
+        """
+        try:
+            return ModelDeliveryList.model_validate(json_dict)
+        except ValidationError as e:
+            logger.error("Error validating ModelDeliveryList: %s", e)
+            raise e
+
+    def to_json(self) -> Dict[str, Any]:  # noqa: UP006
+        """
+        Convert to a json dictionary.
+        """
+        return self.model_dump(exclude_none=True)
+
+
+def _clone_repo(model: Union[str, Path], hf_local_dir: Optional[str]) -> str:
+    if isinstance(model, Path):
+        if not model.exists():
+            raise ValueError(f"Invalid model source: {model}")
+        return str(model)
+    prefixes, mlc_prefix = ["HF://", "https://huggingface.co/"], ""
+    mlc_prefix = next(p for p in prefixes if model.startswith(p))
+    if mlc_prefix:
+        repo_name = model[len(mlc_prefix) :]
+        model_name = repo_name.split("/")[-1]
+        if hf_local_dir:
+            hf_local_dir = os.path.join(hf_local_dir, model_name)
+            logger.info("[HF] Downloading model to %s", hf_local_dir)
+        return snapshot_download(repo_id=repo_name, local_dir=hf_local_dir)
+    result = Path(model)
+    if result.exists():
+        return model
+    raise ValueError(f"Invalid model source: {model}")
+
+
+def _run_quantization(
+    model_info: ModelDeliveryTask,
+    repo: str,
+    api: HfApi,
+    output_dir: str,
+) -> bool:
+    logger.info("[HF] Creating repo https://huggingface.co/%s", repo)
+    try:
+        api.create_repo(repo_id=repo, private=False)
+    except HfHubHTTPError as error:
+        if error.response.status_code != 409:
+            raise
+        logger.info("[HF] Repo already exists. Skipping creation.")
+    succeeded = True
+    log_path = Path(output_dir) / "logs.txt"
+    with log_path.open("a", encoding="utf-8") as log_file:
+        assert isinstance(model_info.quantization, str)
+        logger.info("[MLC] Processing in directory: %s", output_dir)
+        # Required arguments
+        cmd = [
+            sys.executable,
+            "-m",
+            "mlc_llm",
+            "gen_config",
+            model_info.model,
+            "--quantization",
+            model_info.quantization,
+            "--conv-template",
+            model_info.conv_template,
+            "--output",
+            output_dir,
+        ]
+        # Optional arguments
+        for optional_arg in GEN_CONFIG_OPTIONAL_ARGS:
+            optional_arg_val = getattr(model_info, optional_arg, None)
+            if optional_arg_val is not None:
+                # e.g. --context-window-size 4096
+                cmd += ["--" + optional_arg.replace("_", "-"), str(optional_arg_val)]
+
+        print(" ".join(cmd), file=log_file, flush=True)
+        subprocess.run(cmd, check=True, stdout=log_file, stderr=subprocess.STDOUT, env=os.environ)
+        if not model_info.gen_config_only:
+            cmd = [
+                sys.executable,
+                "-m",
+                "mlc_llm",
+                "convert_weight",
+                str(model_info.model),
+                "--quantization",
+                model_info.quantization,
+                "--output",
+                output_dir,
+            ]
+            print(" ".join(cmd), file=log_file, flush=True)
+            subprocess.run(
+                cmd,
+                check=False,
+                stdout=log_file,
+                stderr=subprocess.STDOUT,
+                env=os.environ,
+            )
+        logger.info("[MLC] Complete!")
+    if not (Path(output_dir) / "tensor-cache.json").exists() and not model_info.gen_config_only:
+        logger.error(
+            "[%s] Model %s. Quantization %s. No weights metadata found.",
+            red("FAILED"),
+            model_info.model_id,
+            model_info.quantization,
+        )
+        succeeded = False
+    logger.info("[HF] Uploading to: https://huggingface.co/%s", repo)
+    for _retry in range(10):
+        try:
+            api.upload_folder(
+                folder_path=output_dir,
+                repo_id=repo,
+                ignore_patterns=["logs.txt"],
+            )
+        except Exception as exc:
+            logger.error("[%s] %s. Retrying...", red("FAILED"), exc)
+        else:
+            break
+    else:
+        raise RuntimeError("Failed to upload to HuggingFace Hub with 10 retries")
+    return succeeded
+
+
+def _get_current_log(log: str) -> ModelDeliveryList:
+    log_path = Path(log)
+    if not log_path.exists():
+        with log_path.open("w", encoding="utf-8") as o_f:
+            current_log = ModelDeliveryList(tasks=[])
+            json.dump(current_log.to_json(), o_f, indent=4)
+    else:
+        with log_path.open("r", encoding="utf-8") as i_f:
+            current_log = ModelDeliveryList.from_json(json.load(i_f))
+    return current_log
+
+
+def _generate_model_delivery_diff(
+    spec: ModelDeliveryList, log: ModelDeliveryList
+) -> ModelDeliveryList:
+    diff_tasks = []
+    default_quantization = spec.default_quantization
+    default_overrides = spec.default_overrides
+
+    for task in spec.tasks:
+        model_id = task.model_id
+        conv_template = task.conv_template
+        quantization = task.quantization
+        overrides = {**default_overrides, **task.overrides}
+
+        logger.info(
+            "Checking task: %s %s %s %s",
+            model_id,
+            conv_template,
+            quantization,
+            overrides,
+        )
+        log_tasks = [t for t in log.tasks if t.model_id == model_id]
+        delivered_quantizations = set()
+        gen_config_only = set()
+
+        for log_task in log_tasks:
+            log_quantization = log_task.quantization
+            assert isinstance(log_quantization, str)
+            log_override = log_task.overrides.get(log_quantization, OverrideConfigs())
+            override = overrides.get(log_quantization, OverrideConfigs())
+            if log_override == override:
+                if log_task.conv_template == conv_template:
+                    delivered_quantizations.add(log_quantization)
+                else:
+                    gen_config_only.add(log_quantization)
+
+        all_quantizations = set(default_quantization) | set(quantization)
+        quantization_diff = all_quantizations - set(delivered_quantizations)
+
+        if quantization_diff:
+            for q in quantization_diff:
+                logger.info("Adding task %s %s %s to the diff.", model_id, conv_template, q)
+                task_copy = task.model_copy()
+                task_copy.quantization = [q]
+                task_copy.overrides = {q: overrides.get(q, OverrideConfigs())}
+                task_copy.gen_config_only = task_copy.gen_config_only or q in gen_config_only
+                diff_tasks.append(task_copy)
+        else:
+            logger.info("Task %s %s %s is up-to-date.", model_id, conv_template, quantization)
+
+    diff_config = spec.model_copy()
+    diff_config.default_quantization = []
+    diff_config.default_overrides = {}
+    diff_config.tasks = diff_tasks
+
+    logger.info(
+        "Model delivery diff: %s",
+        diff_config.model_dump_json(indent=4, exclude_none=True),
+    )
+
+    return diff_config
+
+
+def _main(
+    username: str,
+    api: HfApi,
+    spec: ModelDeliveryList,
+    log: str,
+    hf_local_dir: Optional[str],
+    output: str,
+    dry_run: bool,
+):
+    delivery_diff = _generate_model_delivery_diff(spec, _get_current_log(log))
+    if dry_run:
+        logger.info("Dry run. No actual delivery.")
+        return
+
+    failed_cases: List[Tuple[str, str]] = []  # noqa: UP006
+    delivered_log = _get_current_log(log)
+    for task_index, task in enumerate(delivery_diff.tasks, 1):
+        logger.info(
+            bold("[{task_index}/{total_tasks}] Processing model: ").format(
+                task_index=task_index,
+                total_tasks=len(delivery_diff.tasks),
+            )
+            + green(task.model_id)
+        )
+        model = _clone_repo(task.model, hf_local_dir)
+
+        quantizations = []
+
+        if delivery_diff.default_quantization:
+            quantizations += delivery_diff.default_quantization
+
+        if task.quantization:
+            if isinstance(task.quantization, str):
+                quantizations.append(task.quantization)
+            else:
+                quantizations += task.quantization
+
+        default_destination = (
+            delivery_diff.default_destination or "{username}/{model_id}-{quantization}-MLC"
+        )
+        for quantization in quantizations:
+            repo = default_destination.format(
+                username=username,
+                model_id=task.model_id,
+                quantization=quantization,
+            )
+            model_info = ModelDeliveryTask(
+                model=model,
+                quantization=quantization,
+                destination=repo,
+                **task.model_dump(exclude_none=True, exclude={"model", "quantization"}),
+            )
+            logger.info("Model info: %s", model_info.model_dump_json(indent=4))
+            output_dir = os.path.join(
+                output, f"{model_info.model_id}-{model_info.quantization}-MLC"
+            )
+            if not os.path.exists(output_dir):
+                os.makedirs(output_dir)
+
+            result = _run_quantization(
+                model_info=model_info,
+                repo=repo,
+                api=api,
+                output_dir=output_dir,
+            )
+            if not result:
+                failed_cases.append(
+                    (task.model_id, quantization),
+                )
+            else:
+                delivered_log.tasks = [
+                    task
+                    for task in delivered_log.tasks
+                    if task.model_id != model_info.model_id
+                    or task.quantization != model_info.quantization
+                ]
+                delivered_log.tasks.append(model_info)
+    if failed_cases:
+        logger.info("Total %s %s:", len(failed_cases), red("failures"))
+        for model_id, quantization in failed_cases:
+            logger.info("  Model %s. Quantization %s.", model_id, quantization)
+
+    delivered_log.tasks.sort(key=lambda task: task.model_id)
+    logger.info("Writing log to %s", log)
+    with open(log, "w", encoding="utf-8") as o_f:
+        json.dump(delivered_log.to_json(), o_f, indent=4)
+
+
+def main():
+    """Entry point."""
+
+    def _load_spec(path_spec: str) -> ModelDeliveryList:
+        path = Path(path_spec)
+        if not path.exists():
+            raise argparse.ArgumentTypeError(f"Spec file does not exist: {path}")
+        with path.open("r", encoding="utf-8") as i_f:
+            return ModelDeliveryList.from_json(json.load(i_f))
+
+    def _get_default_hf_token() -> str:
+        # Try to get the token from the environment variable
+        hf_token = os.getenv("HF_TOKEN")
+        if hf_token:
+            logger.info("HF token found in environment variable HF_TOKEN")
+            return hf_token
+
+        # If not found, look for the token in the default cache folder
+        token_file_path = os.path.expanduser("~/.cache/huggingface/token")
+        if os.path.exists(token_file_path):
+            with open(token_file_path, encoding="utf-8") as token_file:
+                hf_token = token_file.read().strip()
+                if hf_token:
+                    logger.info("HF token found in ~/.cache/huggingface/token")
+                    return hf_token
+
+        raise OSError("HF token not found")
+
+    parser = ArgumentParser("MLC LLM continuous model delivery")
+    parser.add_argument(
+        "--username",
+        type=str,
+        required=True,
+        help="HuggingFace username",
+    )
+    parser.add_argument(
+        "--token",
+        type=str,
+        default=_get_default_hf_token(),
+        help="HuggingFace access token, obtained under https://huggingface.co/settings/tokens",
+    )
+    parser.add_argument(
+        "--spec",
+        type=_load_spec,
+        default="model-delivery-config.json",
+        help="Path to the model delivery file" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--log",
+        type=str,
+        default="model-delivered-log.json",
+        help="Path to the output log file" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        required=True,
+        help="Directory to store the output MLC models",
+    )
+    parser.add_argument(
+        "--hf-local-dir",
+        type=str,
+        required=False,
+        help="Local directory to store the downloaded HuggingFace model",
+    )
+    parser.add_argument(
+        "--dry-run",
+        action="store_true",
+        help="Dry run without uploading to HuggingFace Hub",
+    )
+    parsed = parser.parse_args()
+    _main(
+        parsed.username,
+        spec=parsed.spec,
+        log=parsed.log,
+        api=HfApi(token=parsed.token),
+        hf_local_dir=parsed.hf_local_dir,
+        output=parsed.output,
+        dry_run=parsed.dry_run,
+    )
+
+
+if __name__ == "__main__":
+    main()

python/mlc_llm/cli/disco_remote_socket_session.py

@@ -0,0 +1,33 @@
+# 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.
+"""Internal remote disco socket session."""
+
+import sys
+
+from tvm import runtime as _  # noqa: F401
+from tvm_ffi import get_global_func
+
+from .. import base  # noqa: F401
+
+if __name__ == "__main__":
+    if len(sys.argv) != 4:
+        print("Usage: <server_host> <server_port> <num_workers>")
+        sys.exit(1)
+
+    server_host = sys.argv[1]
+    server_port = int(sys.argv[2])
+    num_workers = int(sys.argv[3])
+    func = get_global_func("runtime.disco.RemoteSocketSession")
+    func(server_host, server_port, num_workers)

python/mlc_llm/cli/gen_config.py

@@ -0,0 +1,121 @@
+"""Command line entrypoint of configuration generation."""
+
+from pathlib import Path
+from typing import Union
+
+from mlc_llm.interface.gen_config import CONV_TEMPLATES, gen_config
+from mlc_llm.interface.help import HELP
+from mlc_llm.model import MODELS
+from mlc_llm.quantization import QUANTIZATION
+from mlc_llm.support.argparse import ArgumentParser
+from mlc_llm.support.auto_config import detect_config, detect_model_type
+
+
+def main(argv):
+    """Parse command line argumennts and call `mlc_llm.compiler.gen_config`."""
+    parser = ArgumentParser("MLC LLM Configuration Generator")
+
+    def _parse_output(path: Union[str, Path]) -> Path:
+        path = Path(path)
+        if not path.is_dir():
+            path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    parser.add_argument(
+        "config",
+        type=detect_config,
+        help=HELP["config"] + " (required)",
+    )
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        required=True,
+        choices=list(QUANTIZATION.keys()),
+        help=HELP["quantization"] + " (required, choices: %(choices)s)",
+    )
+    parser.add_argument(
+        "--model-type",
+        type=str,
+        default="auto",
+        choices=["auto", *list(MODELS.keys())],
+        help=HELP["model_type"] + ' (default: "%(default)s", choices: %(choices)s)',
+    )
+    parser.add_argument(
+        "--conv-template",
+        type=str,
+        required=True,
+        choices=list(CONV_TEMPLATES),
+        help=HELP["conv_template"] + " (required, choices: %(choices)s)",
+    )
+    parser.add_argument(
+        "--context-window-size",
+        type=int,
+        default=None,
+        help=HELP["context_window_size"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--sliding-window-size",
+        type=int,
+        default=None,
+        help=HELP["sliding_window_size"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--prefill-chunk-size",
+        type=int,
+        default=None,
+        help=HELP["prefill_chunk_size"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--attention-sink-size",
+        type=int,
+        default=None,
+        help=HELP["attention_sink_size"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--tensor-parallel-shards",
+        type=int,
+        default=None,
+        help=HELP["tensor_parallel_shards"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--pipeline-parallel-stages",
+        type=int,
+        default=None,
+        help=HELP["pipeline_parallel_stages"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--disaggregation",
+        type=bool,
+        default=None,
+        help=HELP["disaggregation"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--max-batch-size",
+        type=int,
+        default=128,
+        help=HELP["max_batch_size"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=_parse_output,
+        required=True,
+        help=HELP["output_gen_mlc_chat_config"] + " (required)",
+    )
+    parsed = parser.parse_args(argv)
+    model = detect_model_type(parsed.model_type, parsed.config)
+    gen_config(
+        config=parsed.config,
+        model=model,
+        quantization=QUANTIZATION[parsed.quantization],
+        conv_template=parsed.conv_template,
+        context_window_size=parsed.context_window_size,
+        sliding_window_size=parsed.sliding_window_size,
+        prefill_chunk_size=parsed.prefill_chunk_size,
+        attention_sink_size=parsed.attention_sink_size,
+        tensor_parallel_shards=parsed.tensor_parallel_shards,
+        pipeline_parallel_stages=parsed.pipeline_parallel_stages,
+        disaggregation=parsed.disaggregation,
+        max_batch_size=parsed.max_batch_size,
+        output=parsed.output,
+    )

python/mlc_llm/cli/lib_delivery.py

@@ -0,0 +1,199 @@
+"""Continuous model delivery for MLC LLM models."""
+
+import argparse
+import dataclasses
+import json
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Callable, Dict, List  # noqa: UP035
+
+from mlc_llm.support import logging
+from mlc_llm.support.argparse import ArgumentParser
+from mlc_llm.support.constants import MLC_TEMP_DIR
+from mlc_llm.support.style import bold, green, red
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class ModelInfo:
+    """Necessary information for the model delivery"""
+
+    model_id: str
+    model: Path
+    quantization: str
+    device: str
+    # overrides the `context_window_size`, `prefill_chunk_size`,
+    # `sliding_window_size`, `attention_sink_size`, `max_batch_size`
+    # and `tensor_parallel_shards in mlc-chat-config.json
+    overrides: Dict[str, int]  # noqa: UP006
+
+
+class DeferredScope:
+    """A context manager that defers execution of functions until exiting the scope."""
+
+    def __init__(self):
+        self.deferred_functions = []
+
+    def add(self, func: Callable[[], None]):
+        """Add a function to be executed when exiting the scope."""
+        self.deferred_functions.append(func)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        for func in reversed(self.deferred_functions):
+            func()
+        return False
+
+    def create_temp_dir(self) -> Path:
+        """Create a temporary directory that will be deleted when exiting the scope."""
+        temp_dir = tempfile.mkdtemp(dir=MLC_TEMP_DIR)
+        self.add(lambda: shutil.rmtree(temp_dir, ignore_errors=True))
+        return Path(temp_dir)
+
+
+def _run_compilation(model_info: ModelInfo, repo_dir: Path) -> bool:
+    """Run the compilation of the model library."""
+
+    def get_lib_ext(device: str) -> str:
+        if device in ["cuda", "vulkan", "metal"]:
+            return ".so"
+        if device in ["android", "ios"]:
+            return ".tar"
+        if device in ["webgpu"]:
+            return ".wasm"
+
+        return ""
+
+    succeeded = True
+    with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as temp_dir:
+        log_path = Path(temp_dir) / "logs.txt"
+        model_lib_name = f"{model_info.model_id}-{model_info.quantization}-{model_info.device}"
+        lib_ext = get_lib_ext(model_info.device)
+        if lib_ext == "":
+            raise ValueError(f"Unsupported device: {model_info.device}")
+        model_lib_name += lib_ext
+        with log_path.open("a", encoding="utf-8") as log_file:
+            overrides = ";".join(f"{key}={value}" for key, value in model_info.overrides.items())
+            cmd = [
+                sys.executable,
+                "-m",
+                "mlc_llm",
+                "compile",
+                str(model_info.model),
+                "--device",
+                model_info.device,
+                "--quantization",
+                model_info.quantization,
+                "--overrides",
+                overrides,
+                "--output",
+                os.path.join(temp_dir, model_lib_name),
+            ]
+            print(" ".join(cmd), file=log_file, flush=True)
+            subprocess.run(cmd, check=True, stdout=log_file, stderr=subprocess.STDOUT)
+            logger.info("[MLC] Compilation Complete!")
+        if not (Path(temp_dir) / model_lib_name).exists():
+            logger.error(
+                "[%s] Model %s. Device %s. No compiled library found.",
+                red("FAILED"),
+                model_info.model_id,
+                model_info.device,
+            )
+            succeeded = False
+            return succeeded
+
+        # overwrite git repo file with the compiled library
+        repo_filepath = repo_dir / model_info.model_id / model_lib_name
+        if not repo_filepath.parent.exists():
+            repo_filepath.parent.mkdir(parents=True, exist_ok=True)
+        # copy lib from Path(temp_dir) / model_lib_name to repo_filepath
+        shutil.copy(Path(temp_dir) / model_lib_name, repo_filepath)
+        logger.info("Saved library %s at %s", model_lib_name, repo_filepath)
+    return succeeded
+
+
+def _main(
+    spec: Dict[str, Any],  # noqa: UP006
+):
+    """Compile the model libs in the spec and save them to the binary_libs_dir."""
+    failed_cases: List[Any] = []  # noqa: UP006
+    for task_index, task in enumerate(spec["tasks"], 1):
+        logger.info(
+            bold("[{task_index}/{total_tasks}] Processing model: ").format(
+                task_index=task_index,
+                total_tasks=len(spec["tasks"]),
+            )
+            + green(task["model_id"])
+        )
+        model_info = {
+            "model_id": task["model_id"],
+            "model": task["model"],
+        }
+        for compile_opt in spec["default_compile_options"] + task.get("compile_options", []):
+            for quantization in spec["default_quantization"] + task.get("quantization", []):
+                model_info["quantization"] = quantization
+                model_info["device"] = compile_opt["device"]
+                model_info["overrides"] = compile_opt.get("overrides", {})
+                logger.info(
+                    "[Config] "
+                    + bold("model_id: ")
+                    + model_info["model_id"]
+                    + bold(", quantization: ")
+                    + model_info["quantization"]
+                    + bold(", device: ")
+                    + model_info["device"]
+                    + bold(", overrides: ")
+                    + json.dumps(model_info["overrides"])
+                )
+
+                result = _run_compilation(
+                    ModelInfo(**model_info),
+                    repo_dir=Path(spec["binary_libs_dir"]),
+                )
+                if not result:
+                    failed_cases.append(model_info)
+
+    if failed_cases:
+        logger.info("Total %s %s:", len(failed_cases), red("failures"))
+        for case in failed_cases:
+            logger.info(
+                "model_id %s, quantization %s, device %s, overrides %s",
+                case["model_id"],
+                case["quantization"],
+                case["device"],
+                json.dumps(case["overrides"]),
+            )
+
+
+def main():
+    """Entry point."""
+
+    def _load_spec(path_spec: str) -> Dict[str, Any]:  # noqa: UP006
+        path = Path(path_spec)
+        if not path.exists():
+            raise argparse.ArgumentTypeError(f"Spec file does not exist: {path}")
+        with path.open("r", encoding="utf-8") as i_f:
+            return json.load(i_f)
+
+    parser = ArgumentParser("MLC LLM continuous library delivery")
+    parser.add_argument(
+        "--spec",
+        type=_load_spec,
+        required=True,
+        help="Path to the spec file",
+    )
+    parsed = parser.parse_args()
+    _main(
+        spec=parsed.spec,
+    )
+
+
+if __name__ == "__main__":
+    main()

python/mlc_llm/cli/model_metadata.py

@@ -0,0 +1,194 @@
+"""A tool that inspects the metadata of a model lib."""
+
+import json
+import math
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, Dict, List, Union  # noqa: UP035
+
+from tvm.runtime import DataType
+
+from mlc_llm.support import logging
+from mlc_llm.support.argparse import ArgumentParser
+from mlc_llm.support.config import ConfigBase
+from mlc_llm.support.style import green, red
+
+logger = logging.getLogger(__name__)
+
+
+def _extract_metadata(model_lib: Path) -> Dict[str, Any]:  # noqa: UP006
+    from tvm.runtime import device, load_module
+    from tvm.runtime.vm import VirtualMachine
+
+    return json.loads(VirtualMachine(load_module(model_lib), device("cpu"))["_metadata"]())
+
+
+def _report_all(metadata: Dict[str, Any]) -> None:  # noqa: UP006
+    # Print JSON with aesthetic values that packs each parameter into one line,
+    # while keeping the rest indented.
+    indent = 2
+    indents = " " * indent
+    params = metadata.pop("params")
+    params = indents * 2 + (",\n" + indents * 2).join(json.dumps(p) for p in params)
+    lines = json.dumps(
+        metadata,
+        sort_keys=True,
+        indent=indent,
+    ).splitlines()
+    lines.insert(1, indents + '"params": [\n' + params + "\n" + indents + "],")
+    beautified_json = "\n".join(lines)
+    print(beautified_json)
+
+
+def _read_dynamic_shape(shape: List[Union[int, str]], config: Union[Dict, ConfigBase]) -> List[int]:  # noqa: UP006
+    if isinstance(config, ConfigBase):
+        config = asdict(config)
+    param_shape = []
+    for s in shape:
+        if isinstance(s, int):
+            param_shape.append(s)
+        else:
+            if config is None:
+                logger.error(
+                    "%s: Encountered dynamic shape %s, need to specify `--mlc-chat-config` for "
+                    + "memory usage calculation.",
+                    red("FAILED"),
+                    red(s),
+                )
+                raise AttributeError
+            if s not in config:
+                logger.error(
+                    "%s to retrieve concrete %s for dynamic shape from %s.",
+                    red("FAILED"),
+                    red(s),
+                    config,
+                )
+                raise KeyError
+            param_shape.append(config[s])
+    return param_shape
+
+
+def _compute_memory_usage(metadata: Dict[str, Any], config: Union[Dict, ConfigBase]):  # noqa: UP006
+    params_bytes = 0.0
+    for param in metadata["params"]:
+        if all(isinstance(v, int) for v in param["shape"]):
+            assert all(v > 0 for v in param["shape"]), "All shapes should be strictly positive."
+            param_shape = param["shape"]
+        else:
+            # Contains dynamic shape; use config to look up concrete values
+            param_shape = _read_dynamic_shape(param["shape"], config)
+        params_bytes += math.prod(param_shape) * DataType(param["dtype"]).itemsize
+    temp_func_bytes = 0.0
+    for _func_name, func_bytes in metadata["memory_usage"].items():
+        temp_func_bytes = max(temp_func_bytes, func_bytes)
+
+    return params_bytes, temp_func_bytes
+
+
+def _report_memory_usage(metadata: Dict[str, Any], config: Union[Dict, ConfigBase]) -> None:  # noqa: UP006
+    params_bytes, temp_func_bytes = _compute_memory_usage(metadata, config)
+    total_size = params_bytes + temp_func_bytes
+    logger.info(
+        "%s: %.2f MB (Parameters: %.2f MB. Temporary buffer: %.2f MB)",
+        green("Total memory usage without KV cache"),
+        total_size / 1024 / 1024,
+        params_bytes / 1024 / 1024,
+        temp_func_bytes / 1024 / 1024,
+    )
+
+    # Compute KV cache size per token of context window.
+    if isinstance(config, ConfigBase):
+        config = asdict(config)
+    if (
+        "head_dim" in config
+        and "num_hidden_layers" in config
+        and "num_key_value_heads" in config
+        and "quantization" in metadata
+    ):
+        quantization_type = metadata["quantization"]
+        dtype_bytes = None
+        if "f32" in quantization_type:
+            dtype_bytes = 4
+        elif "bf16" in quantization_type:
+            dtype_bytes = 2
+        elif "f16" in quantization_type:
+            dtype_bytes = 2
+        # TODO: If support quantized KV in future, need to change this
+        if dtype_bytes is not None:
+            bytes_per_token = (
+                config["head_dim"]
+                * config["num_hidden_layers"]
+                * config["num_key_value_heads"]
+                * dtype_bytes
+                * 2  # 2 for key and value
+            )
+            logger.info(
+                "%s: %.2f MB per token in the context window",
+                green("KV cache size"),
+                bytes_per_token / 1024 / 1024,
+            )
+            logger.info(
+                "%s: %.2f MB",
+                green("Total memory usage with a 4K KV cache"),
+                (total_size + bytes_per_token * 4096) / 1024 / 1024,
+            )
+
+    logger.info(
+        "To reduce memory usage, "
+        "tweak `prefill_chunk_size`, `context_window_size` and `sliding_window_size`"
+    )
+
+
+def main():
+    """Entry point for the model metadata tool."""
+    parser = ArgumentParser(description="A tool that inspects the metadata of a model lib.")
+    parser.add_argument(
+        "model_lib",
+        type=Path,
+        help="""The compiled model library. In MLC LLM, an LLM is compiled to a shared or static
+        library (.so or .a), which contains GPU computation to efficiently run the LLM. MLC Chat,
+        as the runtime of MLC LLM, depends on the compiled model library to generate tokens.
+        """,
+    )
+    parser.add_argument(
+        "--mlc-chat-config",
+        type=Path,
+        help="""The `mlc-chat-config.json` file specific to a model variant. This is only required
+        when `memory-only` is true and `model_lib` contains a dynamic parameter shape (i.e. using
+        a variable to represent the shape). For instance, `model.embed_tokens.q_weight` can have
+        shape `["vocab_size", 512]`. In these cases, we look up the concrete value in
+        `mlc-chat-config.json`.
+        """,
+    )
+    parser.add_argument(
+        "--memory-only",
+        action="store_true",
+        help="""If set, only inspect the metadata in memory usage and print richer analysis.
+        Otherwise, the tool will load all the metadata from the model library file but only print
+        the basic information in JSON.
+        """,
+    )
+    parsed = parser.parse_args()
+    # Load metadata from model lib
+    try:
+        metadata = _extract_metadata(parsed.model_lib)
+    except Exception:
+        logger.exception("%s to read metadata section in legacy model lib.", red("FAILED"))
+        return
+    # Load mlc_chat_config if provided
+    cfg = None
+    if parsed.mlc_chat_config:
+        mlc_chat_config_path = Path(parsed.mlc_chat_config)
+        if not mlc_chat_config_path.exists():
+            raise ValueError(f"{mlc_chat_config_path} does not exist.")
+        with open(mlc_chat_config_path, encoding="utf-8") as config_file:
+            cfg = json.load(config_file)
+    # Main body
+    if parsed.memory_only:
+        _report_memory_usage(metadata, cfg)
+    else:
+        _report_all(metadata)
+
+
+if __name__ == "__main__":
+    main()

python/mlc_llm/cli/package.py

@@ -0,0 +1,68 @@
+"""Command line entrypoint of package."""
+
+import os
+from pathlib import Path
+from typing import Union
+
+from mlc_llm.interface.help import HELP
+from mlc_llm.interface.package import package
+from mlc_llm.support.argparse import ArgumentParser
+
+
+def main(argv):
+    """Parse command line arguments and call `mlc_llm.interface.package`."""
+    parser = ArgumentParser("MLC LLM Package CLI")
+
+    def _parse_package_config(path: Union[str, Path]) -> Path:
+        path = Path(path)
+        if not path.exists():
+            raise ValueError(
+                f"Path {str(path)} is expected to be a JSON file, but the file does not exist."
+            )
+        if not path.is_file():
+            raise ValueError(f"Path {str(path)} is expected to be a JSON file.")
+        return path
+
+    def _parse_mlc_llm_source_dir(path: str) -> Path:
+        os.environ["MLC_LLM_SOURCE_DIR"] = path
+        return Path(path)
+
+    def _parse_output(path: Union[str, Path]) -> Path:
+        path = Path(path)
+        if not path.is_dir():
+            path.mkdir(parents=True, exist_ok=True)
+        return path
+
+    parser.add_argument(
+        "--package-config",
+        type=_parse_package_config,
+        default="mlc-package-config.json",
+        help=HELP["config_package"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--mlc-llm-source-dir",
+        type=_parse_mlc_llm_source_dir,
+        default=os.environ.get("MLC_LLM_SOURCE_DIR", None),
+        help=HELP["mlc_llm_source_dir"]
+        + " (default: the $MLC_LLM_SOURCE_DIR environment variable)",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        type=_parse_output,
+        default="dist",
+        help=HELP["output_package"] + ' (default: "%(default)s")',
+    )
+    parsed = parser.parse_args(argv)
+    if parsed.mlc_llm_source_dir is None:
+        raise ValueError(
+            "MLC LLM home is not specified. "
+            "Please obtain a copy of MLC LLM source code by "
+            "cloning https://github.com/mlc-ai/mlc-llm, and set environment variable "
+            '"MLC_LLM_SOURCE_DIR=path/to/mlc-llm"'
+        )
+    package(
+        package_config_path=parsed.package_config,
+        mlc_llm_source_dir=parsed.mlc_llm_source_dir,
+        output=parsed.output,
+    )

python/mlc_llm/cli/router.py

@@ -0,0 +1,90 @@
+"""Command line entrypoint of router."""
+
+from mlc_llm.interface.help import HELP
+from mlc_llm.interface.router import serve
+from mlc_llm.support.argparse import ArgumentParser
+
+
+def main(argv):
+    """Parse command line arguments and call `mlc_llm.interface.router`."""
+
+    # Define a custom argument type for a list of strings
+    def list_of_strings(arg):
+        return arg.split(",")
+
+    parser = ArgumentParser("MLC LLM Router Serve CLI")
+    parser.add_argument(
+        "model",
+        type=str,
+        help=HELP["model"] + " (required)",
+    )
+    parser.add_argument(
+        "--model-lib",
+        type=str,
+        default=None,
+        help=HELP["model_lib"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--router-mode",
+        type=str,
+        choices=["disagg", "round-robin"],
+        default="disagg",
+        help="router mode" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--router-host",
+        type=str,
+        default="127.0.0.1",
+        help="router host" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--router-port",
+        type=int,
+        default=8000,
+        help="router port" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--endpoint-hosts",
+        type=list_of_strings,
+        default="127.0.0.1",
+        help="Host of each endpoint, separated by comma." + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--endpoint-ports",
+        nargs="*",
+        type=int,
+        default=[8080],
+        help="Port of each endpoint, separated by space." + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--endpoint-num-gpus",
+        nargs="*",
+        type=int,
+        default=[1],
+        help="Number of GPUs of each endpoint, separated by space." + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--enable-prefix-cache",
+        default=False,
+        action="store_true",
+        help="whether to enable prefix cache" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--pd-balance-factor",
+        type=float,
+        default=0.0,
+        help=HELP["pd_balance_factor"] + ' (default: "%(default)s")',
+    )
+    parsed = parser.parse_args(argv)
+    serve(
+        model=parsed.model,
+        model_lib=parsed.model_lib,
+        router_host=parsed.router_host,
+        router_port=parsed.router_port,
+        endpoint_hosts=parsed.endpoint_hosts,
+        endpoint_ports=parsed.endpoint_ports,
+        endpoint_num_gpus=parsed.endpoint_num_gpus,
+        enable_prefix_cache=parsed.enable_prefix_cache,
+        router_mode=parsed.router_mode,
+        pd_balance_factor=parsed.pd_balance_factor,
+    )

python/mlc_llm/cli/serve.py

@@ -0,0 +1,264 @@
+"""Command line entrypoint of serve."""
+
+import dataclasses
+import json
+from io import StringIO
+from typing import Literal, Optional
+
+from mlc_llm.interface.help import HELP
+from mlc_llm.interface.serve import serve
+from mlc_llm.support import argparse
+from mlc_llm.support.argparse import ArgumentParser
+
+
+@dataclasses.dataclass
+class EngineConfigOverride:
+    """Arguments for overriding engine config."""
+
+    # Overrides for EngineConfig (runtime)
+    max_num_sequence: Optional[int] = None
+    max_total_seq_length: Optional[int] = None
+    prefill_chunk_size: Optional[int] = None
+    max_history_size: Optional[int] = None
+    gpu_memory_utilization: Optional[float] = None
+    spec_draft_length: Optional[int] = None
+    spec_tree_width: Optional[int] = None
+    prefix_cache_mode: Optional[Literal["disable", "radix"]] = None
+    prefix_cache_max_num_recycling_seqs: Optional[int] = None
+    prefill_mode: Optional[Literal["chunked", "hybrid"]] = None
+    context_window_size: Optional[int] = None
+    sliding_window_size: Optional[int] = None
+    attention_sink_size: Optional[int] = None
+    tensor_parallel_shards: Optional[int] = None
+    pipeline_parallel_stages: Optional[int] = None
+    opt: Optional[str] = None
+
+    def __repr__(self) -> str:
+        out = StringIO()
+        print(f"max_num_sequence={self.max_num_sequence}", file=out, end="")
+        print(f";max_total_seq_length={self.max_total_seq_length}", file=out, end="")
+        print(f";prefill_chunk_size={self.prefill_chunk_size}", file=out, end="")
+        print(f";max_history_size={self.max_history_size}", file=out, end="")
+        print(f";gpu_memory_utilization={self.gpu_memory_utilization}", file=out, end="")
+        print(f";spec_draft_length={self.spec_draft_length}", file=out, end="")
+        print(f";spec_tree_width={self.spec_tree_width}", file=out, end="")
+        print(f";prefix_cache_mode={self.prefix_cache_mode}", file=out, end="")
+        print(
+            f";prefix_cache_max_num_recycling_seqs={self.prefix_cache_max_num_recycling_seqs}",
+            file=out,
+            end="",
+        )
+        print(f";prefill_mode={self.prefill_mode}", file=out, end="")
+        print(f";context_window_size={self.context_window_size}", file=out, end="")
+        print(f";sliding_window_size={self.sliding_window_size}", file=out, end="")
+        print(f";attention_sink_size={self.attention_sink_size}", file=out, end="")
+        print(f";tensor_parallel_shards={self.tensor_parallel_shards}", file=out, end="")
+        print(
+            f";pipeline_parallel_stages={self.pipeline_parallel_stages}",
+            file=out,
+            end="",
+        )
+        print(f";opt={self.opt}", file=out, end="")
+        return out.getvalue().rstrip()
+
+    @staticmethod
+    def from_str(source: str) -> "EngineConfigOverride":
+        """Parse engine config override values from a string."""
+        parser = argparse.ArgumentParser(description="Engine config override values")
+
+        parser.add_argument("--max_num_sequence", type=int, default=None)
+        parser.add_argument("--max_total_seq_length", type=int, default=None)
+        parser.add_argument("--prefill_chunk_size", type=int, default=None)
+        parser.add_argument("--max_history_size", type=int, default=None)
+        parser.add_argument("--gpu_memory_utilization", type=float, default=None)
+        parser.add_argument("--spec_draft_length", type=int, default=None)
+        parser.add_argument("--spec_tree_width", type=int, default=None)
+        parser.add_argument("--prefix_cache_mode", type=str, default="radix")
+        parser.add_argument("--prefix_cache_max_num_recycling_seqs", type=int, default=None)
+        parser.add_argument("--prefill_mode", type=str, default="hybrid")
+        parser.add_argument("--context_window_size", type=int, default=None)
+        parser.add_argument("--sliding_window_size", type=int, default=None)
+        parser.add_argument("--attention_sink_size", type=int, default=None)
+        parser.add_argument("--tensor_parallel_shards", type=int, default=None)
+        parser.add_argument("--pipeline_parallel_stages", type=int, default=None)
+        parser.add_argument("--opt", type=str, default=None)
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return EngineConfigOverride(
+            max_num_sequence=results.max_num_sequence,
+            max_total_seq_length=results.max_total_seq_length,
+            prefill_chunk_size=results.prefill_chunk_size,
+            max_history_size=results.max_history_size,
+            gpu_memory_utilization=results.gpu_memory_utilization,
+            spec_draft_length=results.spec_draft_length,
+            spec_tree_width=results.spec_tree_width,
+            prefix_cache_mode=results.prefix_cache_mode,
+            prefix_cache_max_num_recycling_seqs=results.prefix_cache_max_num_recycling_seqs,
+            prefill_mode=results.prefill_mode,
+            context_window_size=results.context_window_size,
+            sliding_window_size=results.sliding_window_size,
+            attention_sink_size=results.attention_sink_size,
+            tensor_parallel_shards=results.tensor_parallel_shards,
+            pipeline_parallel_stages=results.pipeline_parallel_stages,
+            opt=results.opt,
+        )
+
+
+def main(argv):
+    """Parse command line arguments and call `mlc_llm.interface.serve`."""
+    parser = ArgumentParser("MLC LLM Serve CLI")
+
+    parser.add_argument(
+        "model",
+        type=str,
+        help=HELP["model"] + " (required)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="auto",
+        help=HELP["device_deploy"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--model-lib",
+        type=str,
+        default=None,
+        help=HELP["model_lib"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--mode",
+        type=str,
+        choices=["local", "interactive", "server"],
+        default="local",
+        help=HELP["mode_serve"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--enable-debug",
+        action="store_true",
+        help="whether we enable debug end points and debug config when accepting requests",
+    )
+    parser.add_argument(
+        "--additional-models", type=str, nargs="*", help=HELP["additional_models_serve"]
+    )
+    parser.add_argument(
+        "--embedding-model",
+        type=str,
+        default=None,
+        help="Path to the embedding model weight directory (enables /v1/embeddings endpoint)",
+    )
+    parser.add_argument(
+        "--embedding-model-lib",
+        type=str,
+        default=None,
+        help="Path to the compiled embedding model library (.so/.dylib file)",
+    )
+    parser.add_argument(
+        "--speculative-mode",
+        type=str,
+        choices=["disable", "small_draft", "eagle", "medusa"],
+        default="disable",
+        help=HELP["speculative_mode_serve"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--prefix-cache-mode",
+        type=str,
+        choices=["disable", "radix"],
+        default="radix",
+        help=HELP["prefix_cache_mode_serve"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--prefill-mode",
+        type=str,
+        choices=["hybrid", "chunked"],
+        default="hybrid",
+        help=HELP["prefill_mode"] + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--overrides",
+        type=EngineConfigOverride.from_str,
+        default="",
+        help=HELP["overrides_serve"],
+    )
+    parser.add_argument("--enable-tracing", action="store_true", help=HELP["enable_tracing_serve"])
+    parser.add_argument(
+        "--host",
+        type=str,
+        default="127.0.0.1",
+        help="host name" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--port",
+        type=int,
+        default=8000,
+        help="port" + ' (default: "%(default)s")',
+    )
+    parser.add_argument("--allow-credentials", action="store_true", help="allow credentials")
+    parser.add_argument(
+        "--allow-origins",
+        type=json.loads,
+        default=["*"],
+        help="allowed origins" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--allow-methods",
+        type=json.loads,
+        default=["*"],
+        help="allowed methods" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--allow-headers",
+        type=json.loads,
+        default=["*"],
+        help="allowed headers" + ' (default: "%(default)s")',
+    )
+    parser.add_argument(
+        "--api-key",
+        type=str,
+        default=None,
+        help="API key for authentication. If not provided, authentication is disabled.",
+    )
+    parsed = parser.parse_args(argv)
+
+    additional_models = []
+    if parsed.additional_models is not None:
+        for additional_model in parsed.additional_models:
+            splits = additional_model.split(",", maxsplit=1)
+            if len(splits) == 2:
+                additional_models.append((splits[0], splits[1]))
+            else:
+                additional_models.append(splits[0])
+
+    serve(
+        model=parsed.model,
+        device=parsed.device,
+        model_lib=parsed.model_lib,
+        mode=parsed.mode,
+        enable_debug=parsed.enable_debug,
+        additional_models=additional_models,
+        embedding_model=parsed.embedding_model,
+        embedding_model_lib=parsed.embedding_model_lib,
+        tensor_parallel_shards=parsed.overrides.tensor_parallel_shards,
+        pipeline_parallel_stages=parsed.overrides.pipeline_parallel_stages,
+        opt=parsed.overrides.opt,
+        speculative_mode=parsed.speculative_mode,
+        prefix_cache_mode=parsed.prefix_cache_mode,
+        max_num_sequence=parsed.overrides.max_num_sequence,
+        max_total_sequence_length=parsed.overrides.max_total_seq_length,
+        max_single_sequence_length=parsed.overrides.context_window_size,
+        prefill_chunk_size=parsed.overrides.prefill_chunk_size,
+        sliding_window_size=parsed.overrides.sliding_window_size,
+        attention_sink_size=parsed.overrides.attention_sink_size,
+        max_history_size=parsed.overrides.max_history_size,
+        gpu_memory_utilization=parsed.overrides.gpu_memory_utilization,
+        spec_draft_length=parsed.overrides.spec_draft_length,
+        spec_tree_width=parsed.overrides.spec_tree_width,
+        prefix_cache_max_num_recycling_seqs=parsed.overrides.prefix_cache_max_num_recycling_seqs,
+        prefill_mode=parsed.prefill_mode,
+        enable_tracing=parsed.enable_tracing,
+        host=parsed.host,
+        port=parsed.port,
+        allow_credentials=parsed.allow_credentials,
+        allow_origins=parsed.allow_origins,
+        allow_methods=parsed.allow_methods,
+        allow_headers=parsed.allow_headers,
+        api_key=parsed.api_key,
+    )

python/mlc_llm/cli/worker.py

@@ -0,0 +1,57 @@
+# 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.
+"""Internal DiscoWorker for Disco ProcessSession."""
+
+import os
+import sys
+
+from tvm import runtime as _  # noqa: F401
+from tvm_ffi import get_global_func
+
+from .. import base  # noqa: F401
+
+# register the calibration functions
+from ..interface import calibrate  # noqa: F401
+
+
+def main():
+    """Main worker function"""
+    if len(sys.argv) != 6:
+        print("Usage: <worker_id> <num_workers> <num_groups> <read_fd> <write_fd>")
+        return
+
+    worker_id = int(sys.argv[1])
+    num_workers = int(sys.argv[2])
+    num_groups = int(sys.argv[3])
+    read_fd = int(sys.argv[4])
+    write_fd = int(sys.argv[5])
+    if sys.platform == "win32":
+        import msvcrt
+
+        reader = msvcrt.open_osfhandle(read_fd, os.O_BINARY)
+        writer = msvcrt.open_osfhandle(write_fd, os.O_BINARY)
+    else:
+        reader = read_fd
+        writer = write_fd
+
+    worker_func = get_global_func("runtime.disco.WorkerProcess")
+    worker_func(worker_id, num_workers, num_groups, reader, writer)
+
+
+if __name__ == "__main__":
+    try:
+        main()
+    except (OSError, KeyboardInterrupt):
+        pass

python/mlc_llm/compiler_pass/__init__.py

@@ -0,0 +1,3 @@
+"""Compiler passes used in MLC LLM."""
+
+from . import pipeline as _pipeline

python/mlc_llm/compiler_pass/attach_cuda_graph_alloc_init_func.py

@@ -0,0 +1,33 @@
+"""The pass that attaches an empty function for initialization."""
+
+import tvm
+from tvm import IRModule, relax
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachCUDAGraphAllocInitFunc")
+class AttachCUDAGraphAllocInitFunc:
+    """Attach an empty function for initialization."""
+
+    def __init__(self):
+        pass
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        bb = relax.BlockBuilder(mod)
+        alloc_func_gv = None
+        for gv, _ in mod.functions_items():
+            if gv.name_hint.startswith("cuda_graph_alloc"):
+                assert alloc_func_gv is None
+                alloc_func_gv = gv
+        if alloc_func_gv is None:
+            return mod
+
+        with bb.function("cuda_graph_alloc_init", []):
+            bb.emit_func_output(
+                relax.op.call_builtin_with_ctx(
+                    "vm.builtin.cuda_graph.get_cached_alloc",
+                    args=[alloc_func_gv, relax.prim_value(0)],
+                    ty_args=relax.ObjectType(),
+                )
+            )
+        return bb.finalize()

python/mlc_llm/compiler_pass/attach_embedding_allocator.py

@@ -0,0 +1,39 @@
+"""The pass that attaches embedding allocation function to the IRModule."""
+
+from typing import Any, Dict  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachAllocEmbeddingTensorFunc")
+class AttachAllocEmbeddingTensorFunc:
+    """Attach embedding tensor allocation Relax function to IRModule."""
+
+    def __init__(self, metadata: Dict[str, Any]):  # noqa: UP006
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        embed_func = None
+        for gv, func in mod.functions_items():
+            if gv.name_hint == "embed":
+                embed_func = func
+
+        if embed_func is None:
+            return mod
+
+        hidden_size = embed_func.ret_ty.shape[-1]
+        dtype = relax.DataTypeImm(embed_func.ret_ty.dtype.dtype)
+        bb = relax.BlockBuilder(mod)
+        with bb.function("alloc_embedding_tensor", []):
+            bb.emit_func_output(
+                bb.emit(
+                    relax.op.builtin.alloc_tensor(
+                        relax.ShapeExpr([self.metadata["prefill_chunk_size"], hidden_size]),
+                        dtype,
+                        runtime_device_index=0,
+                    )
+                )
+            )
+        return bb.finalize()

python/mlc_llm/compiler_pass/attach_logit_processor.py

@@ -0,0 +1,285 @@
+"""The pass that attaches logit processor functions to the IRModule."""
+
+import tvm
+from tvm import IRModule
+from tvm.script import tirx as T
+
+from ..support.max_thread_check import (
+    check_thread_limits,
+    get_max_num_threads_per_block,
+)
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachLogitProcessFunc")
+class AttachLogitProcessFunc:
+    """Attach logit processing TIR functions to IRModule."""
+
+    def __init__(self, target: tvm.target.Target):
+        """Initializer.
+
+        Parameters
+        ----------
+        target : tvm.target.Target
+            The target of the model compilation.
+        """
+        self.target = target
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        mod = mod.clone()
+        if self.target.kind.name == "llvm":
+            mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace_cpu()
+            mod["apply_penalty_inplace"] = _get_apply_penalty_inplace_cpu()
+            mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace_cpu()
+        else:
+            mod["apply_logit_bias_inplace"] = _get_apply_logit_bias_inplace(self.target)
+            mod["apply_penalty_inplace"] = _get_apply_penalty_inplace(self.target)
+            mod["apply_bitmask_inplace"] = _get_apply_bitmask_inplace(self.target)
+        return mod
+
+
+def _get_apply_logit_bias_inplace_cpu():
+    @T.prim_func(s_tir=True)
+    def _apply_logit_bias_inplace(
+        var_logits: T.handle,
+        var_pos2seq_id: T.handle,
+        var_token_ids: T.handle,
+        var_logit_bias: T.handle,
+    ) -> None:
+        """Function that applies logit bias in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_logit_bias_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_token = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        # seq_ids
+        pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32")
+        token_ids = T.match_buffer(var_token_ids, (num_token,), "int32")
+        logit_bias = T.match_buffer(var_logit_bias, (num_token,), "float32")
+
+        for i in range(num_token):
+            logits[pos2seq_id[i], token_ids[i]] += logit_bias[i]
+
+    return _apply_logit_bias_inplace
+
+
+def _get_apply_logit_bias_inplace(target: tvm.target.Target):
+    tx = 1024  # default
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    tx = min(tx, max_num_threads_per_block)
+    check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def _apply_logit_bias_inplace(
+        var_logits: T.handle,
+        var_pos2seq_id: T.handle,
+        var_token_ids: T.handle,
+        var_logit_bias: T.handle,
+    ) -> None:
+        """Function that applies logit bias in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_logit_bias_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_token = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        # seq_ids
+        pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32")
+        token_ids = T.match_buffer(var_token_ids, (num_token,), "int32")
+        logit_bias = T.match_buffer(var_logit_bias, (num_token,), "float32")
+
+        for p0 in T.thread_binding(0, (num_token + tx - 1) // tx, "blockIdx.x"):
+            for p1 in T.thread_binding(0, tx, "threadIdx.x"):
+                with T.sblock("block"):
+                    vp = T.axis.spatial(num_token, p0 * tx + p1)
+                    T.where(p0 * tx + p1 < num_token)
+                    logits[pos2seq_id[vp], token_ids[vp]] += logit_bias[vp]
+
+    return _apply_logit_bias_inplace
+
+
+def _get_apply_penalty_inplace_cpu():
+    @T.prim_func(s_tir=True)
+    def _apply_penalty_inplace(
+        var_logits: T.handle,
+        var_seq_ids: T.handle,
+        var_pos2seq_id: T.handle,
+        var_token_ids: T.handle,
+        var_token_cnt: T.handle,
+        var_penalties: T.handle,
+    ) -> None:
+        """Function that applies penalties in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_penalty_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_token = T.int32()
+        num_seq = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32")
+        pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32")
+        token_ids = T.match_buffer(var_token_ids, (num_token,), "int32")
+        token_cnt = T.match_buffer(var_token_cnt, (num_token,), "int32")
+        penalties = T.match_buffer(var_penalties, (num_seq, 3), "float32")
+
+        for token in T.serial(num_token):
+            with T.sblock("block"):
+                vp = T.axis.spatial(num_token, token)
+                logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] -= (
+                    penalties[pos2seq_id[vp], 0] + token_cnt[vp] * penalties[pos2seq_id[vp], 1]
+                )
+                logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] = T.if_then_else(
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] < T.float32(0),
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] * penalties[pos2seq_id[vp], 2],
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] / penalties[pos2seq_id[vp], 2],
+                )
+
+    return _apply_penalty_inplace
+
+
+def _get_apply_penalty_inplace(target: tvm.target.Target):
+    tx = 1024  # default
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    tx = min(tx, max_num_threads_per_block)
+    check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def _apply_penalty_inplace(
+        var_logits: T.handle,
+        var_seq_ids: T.handle,
+        var_pos2seq_id: T.handle,
+        var_token_ids: T.handle,
+        var_token_cnt: T.handle,
+        var_penalties: T.handle,
+    ) -> None:
+        """Function that applies penalties in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_penalty_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_token = T.int32()
+        num_seq = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32")
+        pos2seq_id = T.match_buffer(var_pos2seq_id, (num_token,), "int32")
+        token_ids = T.match_buffer(var_token_ids, (num_token,), "int32")
+        token_cnt = T.match_buffer(var_token_cnt, (num_token,), "int32")
+        penalties = T.match_buffer(var_penalties, (num_seq, 3), "float32")
+
+        for p0 in T.thread_binding(0, (num_token + tx - 1) // tx, "blockIdx.x"):
+            for p1 in T.thread_binding(0, tx, "threadIdx.x"):
+                with T.sblock("block"):
+                    vp = T.axis.spatial(num_token, p0 * tx + p1)
+                    T.where(p0 * tx + p1 < num_token)
+                    # Penalties: (presence_penalty, frequency_penalty, repetition_penalty)
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] -= (
+                        penalties[pos2seq_id[vp], 0] + token_cnt[vp] * penalties[pos2seq_id[vp], 1]
+                    )
+                    logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] = T.if_then_else(
+                        logits[seq_ids[pos2seq_id[vp]], token_ids[vp]] < T.float32(0),
+                        logits[seq_ids[pos2seq_id[vp]], token_ids[vp]]
+                        * penalties[pos2seq_id[vp], 2],
+                        logits[seq_ids[pos2seq_id[vp]], token_ids[vp]]
+                        / penalties[pos2seq_id[vp], 2],
+                    )
+
+    return _apply_penalty_inplace
+
+
+def _get_apply_bitmask_inplace_cpu():
+    @T.prim_func(s_tir=True)
+    def _apply_bitmask_inplace(
+        var_logits: T.handle,
+        var_seq_ids: T.handle,
+        var_bitmask: T.handle,
+    ) -> None:
+        """Function that applies vocabulary masking in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_bitmask_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_seq = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32")
+        bitmask = T.match_buffer(var_bitmask, (batch_size, (vocab_size + 31) // 32), "int32")
+
+        for token in T.serial(num_seq * vocab_size):
+            with T.sblock("block"):
+                vs = T.axis.spatial(num_seq, (token) // vocab_size)
+                vv = T.axis.spatial(vocab_size, (token) % vocab_size)
+
+                logits[seq_ids[vs], vv] = T.if_then_else(
+                    (bitmask[seq_ids[vs], vv // 32] >> (vv % 32)) & 1 == 1,
+                    logits[seq_ids[vs], vv],
+                    T.min_value("float32"),
+                )
+
+    return _apply_bitmask_inplace
+
+
+def _get_apply_bitmask_inplace(target: tvm.target.Target):
+    tx = 1024  # default
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+    tx = min(tx, max_num_threads_per_block)
+    check_thread_limits(target, bdx=tx, bdy=1, bdz=1, gdz=1)
+
+    @T.prim_func(s_tir=True)
+    def _apply_bitmask_inplace(
+        var_logits: T.handle,
+        var_seq_ids: T.handle,
+        var_bitmask: T.handle,
+    ) -> None:
+        """Function that applies vocabulary masking in place."""
+        T.func_attr(
+            {
+                "global_symbol": "apply_bitmask_inplace",
+                "tirx.noalias": True,
+                "tirx.is_scheduled": True,
+            }
+        )
+        batch_size = T.int32()
+        vocab_size = T.int32()
+        num_seq = T.int32()
+        logits = T.match_buffer(var_logits, (batch_size, vocab_size), "float32")
+        seq_ids = T.match_buffer(var_seq_ids, (num_seq,), "int32")
+        bitmask = T.match_buffer(var_bitmask, (batch_size, (vocab_size + 31) // 32), "int32")
+
+        for fused_s_v_0 in T.thread_binding(0, (num_seq * vocab_size + tx - 1) // tx, "blockIdx.x"):
+            for fused_s_v_1 in T.thread_binding(0, tx, "threadIdx.x"):
+                with T.sblock("block"):
+                    vs = T.axis.spatial(num_seq, (fused_s_v_0 * tx + fused_s_v_1) // vocab_size)
+                    vv = T.axis.spatial(vocab_size, (fused_s_v_0 * tx + fused_s_v_1) % vocab_size)
+                    T.where(fused_s_v_0 * tx + fused_s_v_1 < num_seq * vocab_size)
+                    logits[seq_ids[vs], vv] = T.if_then_else(
+                        (bitmask[seq_ids[vs], vv // 32] >> (vv % 32)) & 1 == 1,
+                        logits[seq_ids[vs], vv],
+                        T.min_value("float32"),
+                    )
+
+    return _apply_bitmask_inplace

python/mlc_llm/compiler_pass/attach_sampler.py

@@ -0,0 +1,390 @@
+"""The pass that attaches GPU sampler functions to the IRModule."""
+
+from typing import Dict  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax, te, tirx
+from tvm.relax.frontend import nn
+from tvm.script import tirx as T
+
+from mlc_llm.op.batch_spec_verify import batch_spec_verify
+from mlc_llm.op.top_p_pivot import top_p_pivot, top_p_renorm
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachGPUSamplingFunc")
+class AttachGPUSamplingFunc:
+    """Attach GPU sampling functions to IRModule."""
+
+    def __init__(self, target: tvm.target.Target, variable_bounds: Dict[str, int]):  # noqa: UP006
+        # Specifically for RWKV workloads, which contains -1 max_seq_len
+        max_batch_size = variable_bounds["batch_size"]
+        self.variable_bounds = {
+            "batch_size": max_batch_size,
+            "num_samples": max_batch_size,
+            "num_positions": 6 * max_batch_size,
+        }
+        self.non_negative_var = ["vocab_size"]
+        self.target = target
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        target_kind = self.target.kind.name
+        if target_kind not in ["cuda", "vulkan", "metal", "webgpu"]:
+            # Only enable GPU sampling for CUDA, Vulkan, Metal, and WebGPU.
+            return mod
+
+        bb = relax.BlockBuilder(mod)
+        if target_kind == "webgpu":
+            # Only attach functions that do not contain i8s for WebGPU
+            gv_names = [
+                gv.name_hint
+                for gv in [
+                    _attach_argsort_func(bb),
+                    _attach_sample_with_top_p(bb),
+                ]
+            ]
+        else:
+            gv_names = [
+                gv.name_hint
+                for gv in [
+                    _attach_multinomial_sampling_func(bb),
+                    _attach_argsort_func(bb),
+                    _attach_sample_with_top_p(bb),
+                    _attach_take_probs_func(bb),
+                    _attach_batch_verifier(bb),
+                    _attach_renormalize_by_top_p(bb, self.target),
+                ]
+            ]
+
+        mod = bb.finalize()
+        for gv_name in gv_names:
+            mod[gv_name] = (
+                mod[gv_name]
+                .with_attr("tir_var_upper_bound", self.variable_bounds)
+                .with_attr("tir_non_negative_var", self.non_negative_var)
+            )
+        return mod
+
+
+def _attach_multinomial_sampling_func(bb: relax.BlockBuilder):
+    batch_size = tirx.Var("batch_size", "int64")
+    num_samples = tirx.Var("num_samples", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
+    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32"))
+    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
+    with bb.function("multinomial_from_uniform", [probs, uniform_samples, sample_indices]):
+        with bb.dataflow():
+            sample_shape = relax.ShapeExpr([num_samples, 1])
+            probs_tensor = nn.wrap_nested(probs, name="probs")
+            uniform_samples_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    uniform_samples,
+                    sample_shape,
+                    ty_args=relax.TensorType(sample_shape, "float32"),
+                ),
+                name="uniform_samples",
+            )
+            sample_indices_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    sample_indices,
+                    sample_shape,
+                    ty_args=relax.TensorType(sample_shape, "int32"),
+                ),
+                name="sample_indices",
+            )
+            result_tensor = nn.multinomial_from_uniform(
+                probs_tensor,
+                uniform_samples_tensor,
+                sample_indices_tensor,
+                "int32",
+                name="nn_multinomial_from_uniform",
+            )
+            result = bb.emit(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    result_tensor._expr,
+                    sample_indices.ty.shape,
+                    ty_args=sample_indices.ty,
+                )
+            )
+            output = bb.emit_output(result)
+        gv = bb.emit_func_output(output)
+    return gv
+
+
+def _attach_argsort_func(bb: relax.BlockBuilder):
+    batch_size = tirx.Var("batch_size", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
+    with bb.function("argsort_probs", [probs]):
+        with bb.dataflow():
+            sorted_indices = bb.emit(relax.op.argsort(probs, descending=True, dtype="int32"))
+            sorted_values = bb.emit_te(
+                lambda unsorted_probs, sorted_indices: te.compute(
+                    (batch_size, vocab_size),
+                    lambda i, j: unsorted_probs[i, sorted_indices[i, j]],
+                    name="take_sorted_probs",
+                ),
+                probs,
+                sorted_indices,
+                primfunc_name_hint="take_sorted_probs",
+            )
+            output = bb.emit_output((sorted_values, sorted_indices))
+        gv = bb.emit_func_output(output)
+    return gv
+
+
+@T.prim_func(s_tir=True)
+def full(var_result: T.handle, value: T.int32):
+    """The filling function for top k."""
+    batch_size = T.int32()
+    result = T.match_buffer(var_result, (batch_size, 1), "int32")
+    for i in T.serial(batch_size):
+        with T.sblock("block"):
+            vi = T.axis.spatial(batch_size, i)
+            result[vi, 0] = value
+
+
+def _attach_sample_with_top_p(bb: relax.BlockBuilder):
+    batch_size = tirx.Var("batch_size", "int64")
+    num_samples = tirx.Var("num_samples", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    sorted_probs = relax.Var("sorted_probs", relax.TensorType((batch_size, vocab_size), "float32"))
+    sorted_indices = relax.Var(
+        "sorted_indices", relax.TensorType((batch_size, vocab_size), "int32")
+    )
+    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_samples,), "float32"))
+    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
+    top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32"))
+
+    with bb.function(
+        "sample_with_top_p",
+        [sorted_probs, sorted_indices, uniform_samples, sample_indices, top_p],
+    ):
+        with bb.dataflow():
+            sample_shape = relax.ShapeExpr([num_samples, 1])
+            top_p_shape = relax.ShapeExpr([batch_size, 1])
+            sorted_probs_tensor = nn.wrap_nested(sorted_probs, name="sorted_probs")
+            sorted_indices_tensor = nn.wrap_nested(sorted_indices, name="sorted_indices")
+            uniform_samples_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    uniform_samples,
+                    sample_shape,
+                    ty_args=relax.TensorType(sample_shape, "float32"),
+                ),
+                name="uniform_samples",
+            )
+            sample_indices_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    sample_indices,
+                    sample_shape,
+                    ty_args=relax.TensorType(sample_shape, "int32"),
+                ),
+                name="sample_indices",
+            )
+            top_p_tensor = nn.wrap_nested(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    top_p,
+                    top_p_shape,
+                    ty_args=relax.TensorType(top_p_shape, "float32"),
+                ),
+                name="sample_indices",
+            )
+            top_k_tensor = nn.tensor_ir_op(
+                full,
+                name_hint="full",
+                args=[vocab_size],
+                out=nn.Tensor.placeholder(
+                    [batch_size, 1],
+                    "int32",
+                ),
+            )
+
+            result_tensor = nn.sample_top_p_top_k_from_sorted_prob(
+                sorted_probs_tensor,
+                sorted_indices_tensor,
+                top_p_tensor,
+                top_k_tensor,
+                uniform_samples_tensor,
+                sample_indices_tensor,
+            )
+            result = bb.emit_output(
+                relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    result_tensor._expr,
+                    sample_indices.ty.shape,
+                    ty_args=sample_indices.ty,
+                )
+            )
+        gv = bb.emit_func_output(result)
+    return gv
+
+
+def _attach_renormalize_by_top_p(bb: relax.BlockBuilder, target: tvm.target.Target):
+    batch_size = tirx.Var("batch_size", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    num_pivots = 3
+    probs = relax.Var("probs", relax.TensorType((batch_size, vocab_size), "float32"))
+    top_p = relax.Var("top_p", relax.TensorType((batch_size,), "float32"))
+    init_pivots = relax.Var("init_pivots", relax.TensorType((batch_size, num_pivots), "float32"))
+    with bb.function("renormalize_by_top_p", [probs, top_p, init_pivots]):
+        with bb.dataflow():
+            cutoff_output = bb.emit(
+                relax.call_tir(
+                    bb.add_func(top_p_pivot(num_pivots, target), "top_p_pivot_cutoff"),
+                    args=[probs, top_p, init_pivots],
+                    out_ty=[top_p.ty, top_p.ty],
+                )
+            )
+            final_pivot = cutoff_output[0]
+            renorm_sum = cutoff_output[1]
+            renormalized_probs = bb.emit_output(
+                relax.call_tir(
+                    bb.add_func(top_p_renorm(target), "top_p_renorm_after_cutoff"),
+                    args=[probs, final_pivot, renorm_sum],
+                    out_ty=probs.ty,
+                )
+            )
+        gv = bb.emit_func_output(renormalized_probs)
+    return gv
+
+
+def _attach_take_probs_func(bb: relax.BlockBuilder):
+    batch_size = tirx.Var("batch_size", "int64")
+    num_samples = tirx.Var("num_samples", "int64")
+    num_positions = tirx.Var("num_positions", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    unsorted_probs = relax.Var(
+        "unsorted_probs", relax.TensorType((batch_size, vocab_size), "float32")
+    )
+    sorted_indices = relax.Var(
+        "sorted_indices", relax.TensorType((batch_size, vocab_size), "int32")
+    )
+    sample_indices = relax.Var("sample_indices", relax.TensorType((num_samples,), "int32"))
+    sampling_results = relax.Var("sampling_result", relax.TensorType((num_samples,), "int32"))
+    top_prob_offsets = relax.Var("lobprob_offsets", relax.TensorType((num_positions,), "int32"))
+
+    @T.prim_func(s_tir=True)
+    def sampler_take_probs_tir(
+        var_unsorted_probs: T.handle,
+        var_sorted_indices: T.handle,
+        var_sample_indices: T.handle,
+        var_sampling_results: T.handle,
+        var_top_prob_offsets: T.handle,
+        var_sampled_values: T.handle,
+        var_top_prob_probs: T.handle,
+        var_top_prob_indices: T.handle,
+    ):
+        batch_size = T.int32()
+        num_samples = T.int32()
+        num_positions = T.int32()
+        vocab_size = T.int32()
+        unsorted_probs = T.match_buffer(var_unsorted_probs, (batch_size, vocab_size), "float32")
+        sorted_indices = T.match_buffer(var_sorted_indices, (batch_size, vocab_size), "int32")
+        sample_indices = T.match_buffer(var_sample_indices, (num_samples,), "int32")
+        sampling_results = T.match_buffer(var_sampling_results, (num_samples,), "int32")
+        top_prob_offsets = T.match_buffer(var_top_prob_offsets, (num_positions,), "int32")
+        sampled_values = T.match_buffer(var_sampled_values, (num_samples,), "float32")
+        top_prob_probs = T.match_buffer(var_top_prob_probs, (num_positions,), "float32")
+        top_prob_indices = T.match_buffer(var_top_prob_indices, (num_positions,), "int32")
+        for i in T.serial(num_positions):
+            with T.sblock("top_prob"):
+                vi = T.axis.spatial(num_positions, i)
+                # Reads are data-dependent gathers; declare full-buffer read
+                # regions explicitly so tirx does not infer data-dependent regions.
+                T.reads(
+                    top_prob_offsets[vi],
+                    sorted_indices[0:batch_size, 0:vocab_size],
+                    unsorted_probs[0:batch_size, 0:vocab_size],
+                )
+                T.writes(top_prob_indices[vi], top_prob_probs[vi])
+                row = T.floordiv(top_prob_offsets[vi], vocab_size)
+                col = T.floormod(top_prob_offsets[vi], vocab_size)
+                top_prob_indices[vi] = sorted_indices[row, col]
+                top_prob_probs[vi] = unsorted_probs[row, sorted_indices[row, col]]
+        for i in T.serial(num_samples):
+            with T.sblock("sample"):
+                vj = T.axis.spatial(num_samples, i)
+                T.reads(
+                    sample_indices[vj],
+                    sampling_results[vj],
+                    unsorted_probs[0:batch_size, 0:vocab_size],
+                )
+                T.writes(sampled_values[vj])
+                sampled_values[vj] = unsorted_probs[sample_indices[vj], sampling_results[vj]]
+
+    args = [
+        unsorted_probs,
+        sorted_indices,
+        sample_indices,
+        sampling_results,
+        top_prob_offsets,
+    ]
+    with bb.function("sampler_take_probs", args):
+        with bb.dataflow():
+            taken_probs_indices = bb.emit_output(
+                relax.call_tir(
+                    bb.add_func(sampler_take_probs_tir, "sampler_take_probs_tir"),
+                    args,
+                    out_ty=[
+                        relax.TensorType((num_samples,), "float32"),
+                        relax.TensorType((num_positions,), "float32"),
+                        relax.TensorType((num_positions,), "int32"),
+                    ],
+                )
+            )
+        gv = bb.emit_func_output(taken_probs_indices)
+    return gv
+
+
+def _attach_batch_verifier(bb: relax.BlockBuilder):
+    num_nodes = tirx.Var("num_nodes", "int64")
+    nbatch = tirx.Var("nbatch", "int64")
+    vocab_size = tirx.Var("vocab_size", "int64")
+    draft_probs = relax.Var("draft_probs", relax.TensorType((num_nodes, vocab_size), "float32"))
+    draft_tokens = relax.Var("draft_tokens", relax.TensorType((num_nodes,), "int32"))
+    model_probs = relax.Var("model_probs", relax.TensorType((num_nodes, vocab_size), "float32"))
+    token_tree_first_child = relax.Var(
+        "token_tree_first_child", relax.TensorType((num_nodes,), "int32")
+    )
+    token_tree_next_sibling = relax.Var(
+        "token_tree_next_sibling", relax.TensorType((num_nodes,), "int32")
+    )
+    uniform_samples = relax.Var("uniform_samples", relax.TensorType((num_nodes,), "float32"))
+    token_tree_parent_ptr = relax.Var("token_tree_parent_ptr", relax.TensorType((nbatch,), "int32"))
+    args = [
+        draft_probs,
+        draft_tokens,
+        model_probs,
+        token_tree_first_child,
+        token_tree_next_sibling,
+        uniform_samples,
+        token_tree_parent_ptr,
+    ]
+    with bb.function("sampler_verify_draft_tokens", args):
+        with bb.dataflow():
+            res = bb.emit_output(
+                relax.call_tir_inplace(
+                    bb.add_func(
+                        batch_spec_verify(vocab_size),
+                        "batch_verify_on_gpu_single_kernel",
+                    ),
+                    args,
+                    inplace_indices=[
+                        args.index(model_probs),
+                        args.index(token_tree_parent_ptr),
+                    ],
+                    out_ty=[
+                        model_probs.ty,
+                        token_tree_parent_ptr.ty,
+                    ],
+                )
+            )
+        gv = bb.emit_func_output(res)
+    return gv

python/mlc_llm/compiler_pass/attach_softmax_with_temperature.py

@@ -0,0 +1,274 @@
+"""A compiler pass that attaches two-stage softmax with temperature."""
+
+from typing import Any, Dict, Optional  # noqa: UP035
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.expr_functor import PyExprMutator, mutator
+from tvm.script import tirx as T
+
+from ..support.max_thread_check import get_max_num_threads_per_block
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachSoftmaxWithTemperature")
+class AttachSoftmaxWithTemperature:
+    """Rewrites one-shot softmax into two-stage softmax."""
+
+    def __init__(
+        self,
+        target: tvm.target.Target,
+        metadata: Optional[Dict[str, Any]] = None,  # noqa: UP006
+    ) -> None:
+        self.target = target
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        return _Rewriter(mod, self.target, self.metadata).transform()
+
+
+@mutator
+class _Rewriter(PyExprMutator):
+    def __init__(
+        self,
+        mod: IRModule,
+        target: tvm.target.Target,
+        metadata: Optional[Dict[str, Any]] = None,  # noqa: UP006
+    ) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.target = target
+        self.metadata = metadata
+        self.chunk_size = 4096
+        self.active_vocab_size = self.metadata.get("active_vocab_size") if self.metadata else None
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        batch_size = tirx.Var("batch_size", "int64")
+        vocab_size = tirx.Var("vocab_size", "int64")
+        dtype = "float32"
+        logits = relax.Var("logits", relax.TensorType([batch_size, 1, vocab_size], dtype))
+        temperature = relax.Var("temperature", relax.TensorType([batch_size], dtype))
+        with self.builder_.function("softmax_with_temperature", params=[logits, temperature]):
+            with self.builder_.dataflow():
+                output_struct_info = logits.ty
+                new_shape = relax.ShapeExpr([batch_size, vocab_size])
+                logits = relax.call_pure_packed(
+                    "vm.builtin.reshape",
+                    logits,
+                    new_shape,
+                    ty_args=relax.TensorType(new_shape, dtype),
+                )
+                f_chunk_lse, f_softmax_with_lse = _get_lse_and_softmax_func(
+                    self.target, self.chunk_size, self.active_vocab_size
+                )
+                chunked_result_struct_info = relax.TensorType(
+                    (batch_size, (vocab_size + self.chunk_size - 1) // self.chunk_size),
+                    "float32",
+                )
+                chunked_results = self.builder_.emit(
+                    relax.call_tir(
+                        self.builder_.add_func(f_chunk_lse, "chunk_lse"),
+                        args=[logits, temperature],
+                        out_ty=[
+                            chunked_result_struct_info,
+                            chunked_result_struct_info,
+                        ],
+                    )
+                )
+                chunked_sum = chunked_results[0]
+                chunked_max = chunked_results[1]
+                softmax = self.builder_.emit(
+                    relax.call_tir(
+                        self.builder_.add_func(f_softmax_with_lse, "softmax_with_chunked_sum"),
+                        args=[logits, temperature, chunked_sum, chunked_max],
+                        out_ty=logits.ty,
+                    )
+                )
+                softmax = self.builder_.emit_output(
+                    relax.call_pure_packed(
+                        "vm.builtin.reshape",
+                        softmax,
+                        output_struct_info.shape,
+                        ty_args=output_struct_info,
+                    )
+                )
+            self.builder_.emit_func_output(softmax)
+        return self.builder_.get()
+
+
+def _get_lse_and_softmax_func(target: tvm.target.Target, chunk_size: int, active_vocab_size: int):
+    # NOTE: A quick note on the softmax implementation.
+    # We once tried to multiply every element by log2e which can be computed
+    # potentially more efficiently on hardware.
+    # However, when the input values are large, multiplying by the factor of log2e
+    # causes numerical issue in float32 dtype.
+    # This leads to the softmax output not summing up to 1.
+    # For numerical stability, we removed the log2e factor and switched back
+    # to the standard log/exp computation.
+
+    # The kernels below handle both the cases of temperature=0 and temperature != 0.
+    # - When temperature is not 0, the first kernel computes the log-sum-exp of
+    # chunks (subtracted by the max value in chunk), and the max values of chunks.
+    # The second kernel merges the log-sum-exp with the maximum values.
+    # - When temperature is 0, the first kernel computes the max value and the counts
+    # of the max value. The second kernel merges the max and counts, and set the
+    # softmax of the maximum values to "max_value / max_count".
+
+    @T.prim_func(s_tir=True)
+    def chunk_lse(
+        var_A: T.handle,
+        var_temperature: T.handle,
+        var_chunked_sum: T.handle,
+        var_chunked_max: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True})
+        batch_size = T.int64()
+        vocab_size = T.int64()
+        num_chunks = T.int64()
+        A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32")
+        temperature = T.match_buffer(var_temperature, (batch_size,), dtype="float32")
+        chunked_sum = T.match_buffer(var_chunked_sum, (batch_size, num_chunks), dtype="float32")
+        chunked_max = T.match_buffer(var_chunked_max, (batch_size, num_chunks), dtype="float32")
+        A_pad = T.sblock_alloc_buffer(
+            (batch_size, num_chunks, T.int64(chunk_size)), dtype="float32"
+        )
+        temp_max = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32")
+        temp_sum = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32")
+
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
+            with T.sblock("pad"):
+                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
+                A_pad[v0, v1, v2] = T.Select(
+                    v1 * T.int64(chunk_size) + v2
+                    < (active_vocab_size if active_vocab_size is not None else vocab_size),
+                    T.if_then_else(
+                        temperature[v0] > T.float32(1e-5),
+                        A[v0, v1 * T.int64(chunk_size) + v2] / temperature[v0],
+                        A[v0, v1 * T.int64(chunk_size) + v2],
+                    ),
+                    T.min_value("float32"),
+                )
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
+            with T.sblock("max"):
+                v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2])
+                with T.init():
+                    temp_max[v0, v1] = T.min_value("float32")
+                temp_max[v0, v1] = T.max(temp_max[v0, v1], A_pad[v0, v1, v2])
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
+            with T.sblock("sum_exp"):
+                v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2])
+                with T.init():
+                    temp_sum[v0, v1] = T.float32(0)
+                temp_sum[v0, v1] += T.if_then_else(
+                    v1 * T.int64(chunk_size) + v2
+                    < (active_vocab_size if active_vocab_size is not None else vocab_size),
+                    T.Select(
+                        temperature[v0] > T.float32(1e-5),
+                        T.exp(A_pad[v0, v1, v2] - temp_max[v0, v1]),
+                        T.cast(A_pad[v0, v1, v2] == temp_max[v0, v1], "float32"),
+                    ),
+                    T.float32(0),
+                )
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(1)):
+            with T.sblock("log"):
+                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
+                chunked_sum[v0, v1] = T.Select(
+                    temperature[v0] > T.float32(1e-5),
+                    T.log(temp_sum[v0, v1]),
+                    temp_sum[v0, v1],
+                )
+                chunked_max[v0, v1] = temp_max[v0, v1]
+
+    @T.prim_func(s_tir=True)
+    def softmax_with_chunked_sum(
+        var_A: T.handle,
+        var_temperature: T.handle,
+        var_chunked_sum: T.handle,
+        var_chunked_max: T.handle,
+        var_softmax: T.handle,
+    ):
+        T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
+        batch_size = T.int64()
+        vocab_size = T.int64()
+        num_chunks = T.int64()
+        A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32")
+        temperature = T.match_buffer(var_temperature, (batch_size,), dtype="float32")
+        chunked_sum = T.match_buffer(var_chunked_sum, (batch_size, num_chunks), dtype="float32")
+        chunked_max = T.match_buffer(var_chunked_max, (batch_size, num_chunks), dtype="float32")
+        softmax = T.match_buffer(var_softmax, (batch_size, vocab_size), dtype="float32")
+        temp_max = T.sblock_alloc_buffer((batch_size,), dtype="float32")
+        temp_sum = T.sblock_alloc_buffer((batch_size,), dtype="float32")
+        for l0, l1 in T.grid(batch_size, num_chunks):
+            with T.sblock("max"):
+                v0, v1 = T.axis.remap("SR", [l0, l1])
+                with T.init():
+                    temp_max[v0] = T.min_value("float32")
+                temp_max[v0] = T.max(temp_max[v0], chunked_max[v0, v1])
+        for l0, l1 in T.grid(batch_size, num_chunks):
+            with T.sblock("sum_exp"):
+                v0, v1 = T.axis.remap("SR", [l0, l1])
+                with T.init():
+                    temp_sum[v0] = T.float32(0)
+                temp_sum[v0] += T.Select(
+                    temperature[v0] > T.float32(1e-5),
+                    T.exp(chunked_sum[v0, v1] + chunked_max[v0, v1] - temp_max[v0]),
+                    T.cast(chunked_max[v0, v1] == temp_max[v0], "float32") * chunked_sum[v0, v1],
+                )
+        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
+            with T.sblock("log_pad"):
+                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
+                if v1 * T.int64(chunk_size) + v2 < vocab_size:
+                    softmax[v0, v1 * T.int64(chunk_size) + v2] = T.Select(
+                        v1 * T.int64(chunk_size) + v2
+                        < (active_vocab_size if active_vocab_size is not None else vocab_size),
+                        T.if_then_else(
+                            temperature[v0] > T.float32(1e-5),
+                            T.exp(
+                                A[v0, v1 * T.int64(chunk_size) + v2] / temperature[v0]
+                                - (T.log(temp_sum[v0]) + temp_max[v0])
+                            ),
+                            T.cast(
+                                A[v0, v1 * T.int64(chunk_size) + v2] == temp_max[v0],
+                                "float32",
+                            )
+                            / temp_sum[v0],
+                        ),
+                        T.float32(0),
+                    )
+
+    sch = tvm.s_tir.Schedule(IRModule({"softmax_with_chunked_sum": softmax_with_chunked_sum}))
+
+    def apply_gpu_schedule(target, sch):
+        max_threads = get_max_num_threads_per_block(target)
+        TX = 32
+        TY = max_threads // TX
+        unroll_depth = 64
+
+        sch.work_on("softmax_with_chunked_sum")
+        l0, l1, l2 = sch.get_loops("log_pad")
+        bx = sch.fuse(l0, l1)
+        sch.bind(bx, "blockIdx.x")
+        unroll, ty, tx = sch.split(l2, [None, TY, TX])
+        sch.bind(ty, "threadIdx.y")
+        sch.bind(tx, "threadIdx.x")
+        sch.annotate(unroll, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+        sch.annotate(unroll, ann_key="pragma_unroll_explicit", ann_val=1)
+
+        for block_name in ["sum_exp", "max"]:
+            block = sch.get_sblock(block_name)
+            sch.set_scope(block, buffer_index=0, storage_scope="shared")
+            sch.compute_at(block, bx)
+            r_loop = sch.get_loops(block)[-1]
+            r_loop, tx = sch.split(r_loop, [None, TX])
+            sch.reorder(tx, r_loop)
+            sch.bind(tx, "threadIdx.x")
+            sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
+            sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)
+
+        return chunk_lse, sch.mod["softmax_with_chunked_sum"]
+
+    if target.kind.name == "llvm":
+        return chunk_lse, sch.mod["softmax_with_chunked_sum"]
+    return apply_gpu_schedule(target, sch)

python/mlc_llm/compiler_pass/attach_spec_decode_aux_funcs.py

@@ -0,0 +1,123 @@
+"""The pass that attaches logit processor functions to the IRModule."""
+
+import tvm
+from tvm import IRModule, relax, tirx
+from tvm.relax import BlockBuilder, TensorType
+from tvm.script import tirx as T
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachSpecDecodeAuxFuncs")
+class AttachSpecDecodeAuxFuncs:
+    """Attach logit processing TIR functions to IRModule."""
+
+    tensor_parallel_shards: int
+
+    def __init__(self, tensor_parallel_shards: int):
+        self.tensor_parallel_shards = tensor_parallel_shards
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        mod = mod.clone()
+        bb = BlockBuilder(mod)
+        bb.add_func(
+            _get_scatter_2d_inplace(dtype="float32", global_symbol="scatter_probs"),
+            "scatter_probs",
+        )
+        bb.add_func(
+            _get_gather_2d_inplace(dtype="float32", global_symbol="gather_probs"),
+            "gather_probs",
+        )
+        if "prefill_to_last_hidden_states" in mod:
+            hidden_states_struct_info = mod["prefill_to_last_hidden_states"].ret_ty.fields[0]
+            dtype = hidden_states_struct_info.dtype
+            _add_gather_hidden_states(bb, self.tensor_parallel_shards, dtype)
+            _add_scatter_hidden_states(bb, self.tensor_parallel_shards, dtype)
+        return bb.finalize()
+
+
+def _get_scatter_2d_inplace(dtype: str, global_symbol: str):
+    @T.prim_func(s_tir=True)
+    def _scatter_2d(var_src: T.handle, var_indices: T.handle, var_dst: T.handle):
+        T.func_attr({"global_symbol": global_symbol, "tirx.noalias": True})
+        batch_size = T.int32()
+        m = T.int32()
+        n = T.int32()
+        src = T.match_buffer(var_src, (batch_size, n), dtype)
+        indices = T.match_buffer(var_indices, (batch_size,), "int32")
+        dst = T.match_buffer(var_dst, (m, n), dtype)
+        for b, j in T.grid(batch_size, n):
+            with T.sblock("scatter_2d"):
+                vb, vj = T.axis.remap("SS", [b, j])
+                dst[indices[vb], vj] = src[vb, vj]
+
+    return _scatter_2d
+
+
+def _get_gather_2d_inplace(dtype: str, global_symbol: str):
+    @T.prim_func(s_tir=True)
+    def _gather_2d(var_src: T.handle, var_indices: T.handle, var_dst: T.handle):
+        T.func_attr({"global_symbol": global_symbol, "tirx.noalias": True})
+        batch_size = T.int32()
+        m = T.int32()
+        n = T.int32()
+        src = T.match_buffer(var_src, (m, n), dtype)
+        indices = T.match_buffer(var_indices, (batch_size,), "int32")
+        dst = T.match_buffer(var_dst, (batch_size, n), dtype)
+        for b, j in T.grid(batch_size, n):
+            with T.sblock("gather_2d"):
+                vb, vj = T.axis.remap("SS", [b, j])
+                dst[vb, vj] = src[indices[vb], vj]
+
+    return _gather_2d
+
+
+def _add_scatter_hidden_states(bb: BlockBuilder, tensor_parallel_shards: int, dtype: str):
+    batch_size = tirx.Var("batch_size", "int64")
+    m = tirx.Var("m", "int64")
+    n = tirx.Var("n", "int64")
+    src = relax.Var("src", ty=TensorType([batch_size, n], dtype))
+    indices = relax.Var("indices", ty=TensorType([batch_size], "int32"))
+    dst = relax.Var("dst", ty=TensorType([m, n], dtype))
+    with bb.function("scatter_hidden_states", [src, indices, dst]):
+        with bb.dataflow():
+            if tensor_parallel_shards > 1:
+                indices = relax.op.ccl.broadcast_from_worker0(indices)
+            output = bb.emit_output(
+                relax.op.call_tir_inplace(
+                    bb.add_func(
+                        _get_scatter_2d_inplace(dtype, "_scatter_hidden_states"),
+                        "_scatter_hidden_states",
+                    ),
+                    [src, indices, dst],
+                    2,
+                    dst.ty,
+                )
+            )
+        gv = bb.emit_func_output(output)
+    return gv
+
+
+def _add_gather_hidden_states(bb: BlockBuilder, tensor_parallel_shards: int, dtype: str):
+    batch_size = tirx.Var("batch_size", "int64")
+    m = tirx.Var("m", "int64")
+    n = tirx.Var("n", "int64")
+    src = relax.Var("src", ty=TensorType([m, n], dtype))
+    indices = relax.Var("indices", ty=TensorType([batch_size], "int32"))
+    dst = relax.Var("dst", ty=TensorType([batch_size, n], dtype))
+    with bb.function("gather_hidden_states", [src, indices, dst]):
+        with bb.dataflow():
+            if tensor_parallel_shards > 1:
+                indices = relax.op.ccl.broadcast_from_worker0(indices)
+            output = bb.emit_output(
+                relax.op.call_tir_inplace(
+                    bb.add_func(
+                        _get_gather_2d_inplace(dtype, "_gather_hidden_states"),
+                        "_gather_hidden_states",
+                    ),
+                    [src, indices, dst],
+                    2,
+                    dst.ty,
+                )
+            )
+        gv = bb.emit_func_output(output)
+    return gv

python/mlc_llm/compiler_pass/attach_support_info.py

@@ -0,0 +1,154 @@
+"""A couple of passes that simply supportive information onto the IRModule."""
+
+from math import lcm
+from typing import Any, Dict, List  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax, tirx
+from tvm.ir import Op
+from tvm.relax.expr_functor import PyExprVisitor, visitor
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachVariableBounds")
+class AttachVariableBounds:
+    """Attach variable bounds to each Relax function, which primarily helps with memory planning."""
+
+    def __init__(self, variable_bounds: Dict[str, int]):  # noqa: UP006
+        # Specifically for RWKV workloads, which contains -1 max_seq_len
+        self.variable_bounds = {k: v for k, v in variable_bounds.items() if v > 0}
+        self.non_negative_var = ["vocab_size"]
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for g_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                mod[g_var] = func.with_attr("tir_var_upper_bound", self.variable_bounds).with_attr(
+                    "tir_non_negative_var", self.non_negative_var
+                )
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachAdditionalPrimFuncs")
+class AttachAdditionalPrimFuncs:
+    """Attach extra TIR PrimFuncs to the IRModule"""
+
+    def __init__(self, functions: Dict[str, tirx.PrimFunc]):  # noqa: UP006
+        self.functions = functions
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for func_name, func in self.functions.items():
+            mod[func_name] = func.with_attr("global_symbol", func_name)
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachMemoryPlanAttr")
+class AttachMemoryPlanAttr:
+    """Attach memory planning attribute for dynamic function output planning to Relax functions."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for g_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                mod[g_var] = func.with_attr("relax.memory_plan_dynamic_func_output", True)
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachCUDAGraphCaptureHints")
+class AttachCUDAGraphSymbolicCaptureHints:
+    """Attach CUDA graph capture hints to the IRModule"""
+
+    def __init__(self, hints: Dict[str, List[str]]):  # noqa: UP006
+        self.hints = hints
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for g_var, func in mod.functions_items():
+            func_name = g_var.name_hint
+            if isinstance(func, relax.Function):
+                if func_name in self.hints:
+                    mod[g_var] = func.with_attr(
+                        "relax.rewrite_cuda_graph.capture_symbolic_vars",
+                        self.hints[func_name],
+                    )
+
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachPipelineParallelStages")
+class AttachPipelineParallelStages:
+    """Attach number of pipeline stages to relax functions."""
+
+    def __init__(self, pipeline_parallel_shards: int):
+        self.pipeline_parallel_shards = pipeline_parallel_shards
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        for g_var, func in mod.functions_items():
+            func_name = g_var.name_hint
+            if not isinstance(func, relax.Function) or func_name not in [
+                "prefill",
+                "decode",
+                "prefill_to_last_hidden_states",
+                "decode_to_last_hidden_states",
+                "batch_prefill",
+                "batch_decode",
+                "batch_verify",
+                "batch_prefill_to_last_hidden_states",
+                "batch_decode_to_last_hidden_states",
+                "batch_verify_to_last_hidden_states",
+            ]:
+                continue
+            mod[g_var] = func.with_attr("pipeline_parallel_stages", self.pipeline_parallel_shards)
+
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachSequenceLengthPaddingFactor")
+class AttachSequenceLengthPaddingFactor:
+    """Attach sequence length padding factor to the metadata"""
+
+    def __init__(self, target: tvm.target.Target, metadata: Dict[str, Any]):  # noqa: UP006
+        self.target = target
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+
+        @visitor
+        class _Visitor(PyExprVisitor):
+            def __init__(self, target: tvm.target.Target) -> None:
+                self.padding_factor = 1
+                self.target = target
+                self._op_call_dps_packed = Op.get("relax.call_dps_packed")
+
+            def run(self, mod: IRModule) -> int:
+                """Entry point of the visitor."""
+                # Right now we only need padding for CUDA SM100a architecture.
+                # When the target is SM100a and uses cutlass gemm function,
+                # the sequence length needs to be padded to multiple of 4.
+                if self.target.kind.name != "cuda" or self.target.attrs.get("arch") != "sm_100a":
+                    return 1
+
+                for _, func in mod.functions_items():
+                    if isinstance(func, relax.Function):
+                        self.visit_expr(func)
+                return self.padding_factor
+
+            def visit_call_(self, call: relax.Call) -> None:
+                super().visit_call_(call)
+                if call.op != self._op_call_dps_packed:
+                    return
+                func_name = str(call.args[0].global_symbol)
+                if func_name in [
+                    "cutlass.groupwise_scaled_gemm_e4m3fn_e4m3fn",
+                    "cutlass.groupwise_scaled_bmm_e4m3fn_e4m3fn",
+                ]:
+                    # Find the minimum common multiple of padding factor and 4
+                    self.padding_factor = lcm(self.padding_factor, 4)
+
+        # self.metadata["sequence_length_padding"] = True
+        padding_factor = _Visitor(self.target).run(mod)
+        if padding_factor > 1:
+            self.metadata["seqlen_padding_factor"] = padding_factor
+        return mod

python/mlc_llm/compiler_pass/blas_dispatch.py

@@ -0,0 +1,51 @@
+"""A compiler pass that dispatches patterns to CUBLAS."""
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.backend import get_patterns_with_prefix
+
+try:
+    import tvm.relax.backend.cuda.cublas as _cublas  # noqa: F401
+    import tvm.relax.backend.rocm.hipblas as _hipblas  # noqa: F401
+except ImportError:
+    # Note: legacy path of cublas/hipblas for backward compatibility
+    pass
+
+
+@tvm.transform.module_pass(opt_level=0, name="BLASDispatch")
+class BLASDispatch:
+    """A compiler pass that dispatches patterns to cuBLAS/hipBLAS."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        if target.kind.name == "cuda":
+            self.has_blas = tvm.get_global_func("relax.ext.cublas", True)
+            if not self.has_blas:
+                raise Exception("cuBLAS is not enabled.")
+            self.patterns = get_patterns_with_prefix("cublas")
+        elif target.kind.name == "rocm":
+            self.has_blas = tvm.get_global_func("relax.ext.hipblas", True)
+            if not self.has_blas:
+                raise Exception("hipBLAS is not enabled.")
+            self.patterns = get_patterns_with_prefix("hipblas")
+        else:
+            raise Exception(f"Unsupported target {target.kind.name} for BLAS dispatch.")
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        model_names = [
+            gv.name_hint for gv, func in mod.functions.items() if isinstance(func, relax.Function)
+        ]
+        # exclude single batch decode
+        model_names = [name for name in model_names if "batch" in name or "decode" not in name]
+        mod = tvm.transform.Sequential(
+            [
+                relax.transform.FuseOpsByPattern(
+                    self.patterns,
+                    bind_constants=False,
+                    annotate_codegen=True,
+                    entry_functions=model_names,
+                ),
+                relax.transform.RunCodegen({}, entry_functions=model_names),
+            ]
+        )(mod)
+        return mod

python/mlc_llm/compiler_pass/clean_up_tir_attrs.py

@@ -0,0 +1,31 @@
+"""A compiler pass that cleans up undesired TIR attrs."""
+
+from typing import List  # noqa: UP035
+
+import tvm
+from tvm.ir.module import IRModule
+
+
+@tvm.transform.module_pass(opt_level=0, name="CleanUpTIRAttrs")
+class CleanUpTIRAttrs:
+    """A compiler pass that cleans up undesired TIR attrs."""
+
+    def __init__(self, attrs: List[str]):  # noqa: UP006
+        self.attrs = attrs
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        for g_var, func in mod.functions_items():
+            changed = False
+            for attr in self.attrs:
+                if func.attrs is not None and attr in func.attrs:
+                    func = func.without_attr(attr)
+                    changed = True
+                    break
+            if changed:
+                mod[g_var] = func
+        return mod

python/mlc_llm/compiler_pass/dispatch_kv_cache_creation.py

@@ -0,0 +1,243 @@
+"""A pass that rewrites KV cache creation functions in IRModule."""
+
+import json
+from typing import Any, Dict, List  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.frontend.nn.llm import kv_cache
+from tvm.relax.frontend.nn.llm.kv_cache import RopeMode
+
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+def extract_creation_args(func: relax.Function) -> Dict[str, Any]:  # noqa: UP006
+    """Extract the KV cache creation args from the given generic creation func."""
+    assert isinstance(func.body, relax.SeqExpr)
+    assert len(func.body.blocks) == 1
+    assert isinstance(func.body.blocks[0], relax.DataflowBlock)
+    assert isinstance(func.body.blocks[0].bindings[0], relax.VarBinding)
+    assert isinstance(func.body.blocks[0].bindings[0].value, relax.Call)
+    assert func.body.blocks[0].bindings[0].value.op == tvm.ir.Op.get("relax.call_pure_packed")
+    call_args = func.body.blocks[0].bindings[0].value.args
+    assert isinstance(call_args[0], relax.ExternFunc)
+    assert call_args[0].global_symbol == "mlc.create_paged_kv_cache_generic"
+    args = call_args[1:]
+    assert len(args) == 18
+    assert isinstance(args[0], (relax.StringImm, relax.Tuple))
+    # Check if attn_kind is a single value or a list with length of hidden layers
+    if isinstance(args[0], relax.StringImm):
+        assert args[0].value in ["mha", "mla"]
+        attn_kind = args[0].value
+    else:
+        assert len(args[0].fields) == args[3].value
+        for i, attention_type in enumerate(args[0].fields):
+            assert isinstance(attention_type, relax.StringImm)
+            assert attention_type.value in ["mha", "mla", "mha_sliding"]
+        attn_kind = [args[0].fields[i].value for i in range(len(args[0]))]
+    assert isinstance(args[1], relax.ShapeExpr)
+    assert len(args[1].values) == 5
+    assert isinstance(args[2], relax.ShapeExpr)
+    for i in range(3, 18):
+        if i in [13, 14, 17]:
+            continue
+        # PrimValue wrappers were phased out of Relax: scalar args are now bare
+        # tirx PrimExprs (IntImm/FloatImm) directly.
+        assert isinstance(args[i], (tvm.tirx.IntImm, tvm.tirx.FloatImm)), (
+            f"args[{i}] is {type(args[i])}"
+        )
+    assert isinstance(args[13], relax.StringImm)
+    assert isinstance(args[16], (relax.Constant, tvm.tirx.IntImm, tvm.tirx.FloatImm))
+    assert isinstance(args[17], relax.DataTypeImm)
+
+    return {
+        "attn_kind": attn_kind,
+        "max_batch_size": args[1].values[0],
+        "max_total_seq_len": args[1].values[1],
+        "prefill_chunk_size": args[1].values[2],
+        "page_size": args[1].values[3],
+        "support_sliding_window": args[1].values[4],
+        "layer_partition": args[2],
+        "num_hidden_layers": args[3].value,
+        "num_attention_heads": args[4].value,
+        "num_key_value_heads": args[5].value,
+        "qk_head_dim": args[6].value,
+        "v_head_dim": args[7].value,
+        "mla_original_qk_head_dim": args[8].value,
+        "mla_original_v_head_dim": args[9].value,
+        "rope_mode": args[10].value,
+        "rope_scale": args[11].value,
+        "rope_theta": args[12].value,
+        "rope_scaling": json.loads(args[13].value),
+        "rope_ext_factors": args[14],
+        "rotary_dim": args[15].value,
+        "enable_disaggregation": bool(args[16].value),
+        "dtype": args[17].value,
+    }
+
+
+@tvm.transform.module_pass(opt_level=0, name="DispatchKVCacheCreation")
+class DispatchKVCacheCreation:
+    """Rewrite KV cache creation functions to IRModule."""
+
+    def __init__(
+        self,
+        target: tvm.target.Target,
+        flashinfer: bool,
+        metadata: Dict[str, Any],  # noqa: UP006
+    ) -> None:
+        """Initializer.
+
+        Parameters
+        ----------
+        target : tvm.target.Target
+            The target of the model compilation.
+
+        flashinfer : bool
+            A boolean indicating if flashinfer is enabled.
+
+        metadata : Dict[str, Any]
+            The model's metadata for KV cache creation.
+            Note that the metadata will be updated in this pass -- the
+            KV cache metadata will be attached.
+        """
+        self.target = target
+        self.flashinfer = flashinfer
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        func_dict = {}
+        creation_func = None
+        for g_var, func in mod.functions_items():
+            # Try to find the `create_paged_kv_cache` func.
+            if g_var.name_hint == "create_paged_kv_cache":
+                creation_func = func
+            else:
+                func_dict[g_var] = func
+
+        if creation_func is None:
+            return mod
+
+        new_mod = IRModule(func_dict)
+        if mod.attrs is not None:
+            new_mod = new_mod.with_attrs(mod.attrs)
+
+        kwargs = extract_creation_args(creation_func)
+        self.attach_kv_cache_metadata(kwargs)
+
+        bb = relax.BlockBuilder(new_mod)
+        extern_mods = []
+        extern_mods += self.create_tir_paged_kv_cache(bb, kwargs)
+        extern_mods += self.create_flashinfer_paged_kv_cache(bb, kwargs)
+
+        mod = bb.finalize()
+        mod_attrs = dict(mod.attrs) if mod.attrs else {}
+        mod = mod.with_attr("external_mods", mod_attrs.get("external_mods", []) + extern_mods)
+        return mod
+
+    def attach_kv_cache_metadata(self, kwargs: Dict[str, Any]):  # noqa: UP006
+        """Attach the KV cache metadata to model metadata."""
+        self.metadata["kv_cache"] = {
+            "num_hidden_layers": kwargs["num_hidden_layers"],
+            "num_attention_heads": kwargs["num_attention_heads"],
+            "num_key_value_heads": kwargs["num_key_value_heads"],
+            "head_dim": kwargs["qk_head_dim"],
+        }
+
+    def create_tir_paged_kv_cache(
+        self,
+        bb: relax.BlockBuilder,
+        kwargs: Dict[str, Any],  # noqa: UP006
+    ) -> List[tvm.runtime.Module]:  # noqa: UP006
+        """Create the TIR-based PagedKVCache"""
+        max_batch_size = relax.Var("max_batch_size_", relax.ShapeType([kwargs["max_batch_size"]]))
+        max_total_seq_len = relax.Var(
+            "max_total_seq_len_", relax.ShapeType([kwargs["max_total_seq_len"]])
+        )
+        prefill_chunk_size = relax.Var(
+            "prefill_chunk_size_", relax.ShapeType([kwargs["prefill_chunk_size"]])
+        )
+        page_size = relax.Var("page_size_", relax.ShapeType([kwargs["page_size"]]))
+        support_sliding_window = relax.Var(
+            "support_sliding_window_",
+            relax.ShapeType([kwargs["support_sliding_window"]]),
+        )
+
+        # Ensure 'enable_disaggregation' is optional
+        enable_disaggregation = kwargs.pop("enable_disaggregation", False)
+        kwargs["enable_disaggregation"] = enable_disaggregation
+
+        with bb.function(
+            name="create_tir_paged_kv_cache",
+            params=[
+                max_batch_size,
+                max_total_seq_len,
+                prefill_chunk_size,
+                page_size,
+                support_sliding_window,
+            ],
+        ):
+            cache = kv_cache.TIRPagedKVCache(target=self.target, **kwargs)
+            bb.emit_func_output(cache._expr)
+
+        return cache.extern_mods
+
+    def create_flashinfer_paged_kv_cache(
+        self,
+        bb: relax.BlockBuilder,
+        kwargs: Dict[str, Any],  # noqa: UP006
+    ) -> List[tvm.runtime.Module]:  # noqa: UP006
+        """Create the FlashInfer-based PagedKVCache"""
+        # Filter the cases which FlashInfer does not support.
+        if (
+            not self.flashinfer
+            or self.target.kind.name != "cuda"
+            or str(kwargs["dtype"]) not in ["float16", "bfloat16"]
+            or (
+                kwargs["rope_mode"] == RopeMode.INLINE
+                and (
+                    kwargs["rotary_dim"] != kwargs["qk_head_dim"]
+                    or kwargs["qk_head_dim"] != kwargs["v_head_dim"]
+                )
+            )
+        ):
+            return []
+
+        max_batch_size = relax.Var("max_batch_size_", relax.ShapeType([kwargs["max_batch_size"]]))
+        max_total_seq_len = relax.Var(
+            "max_total_seq_len_", relax.ShapeType([kwargs["max_total_seq_len"]])
+        )
+        prefill_chunk_size = relax.Var(
+            "prefill_chunk_size_", relax.ShapeType([kwargs["prefill_chunk_size"]])
+        )
+        page_size = relax.Var("page_size_", relax.ShapeType([kwargs["page_size"]]))
+        support_sliding_window = relax.Var(
+            "support_sliding_window_",
+            relax.ShapeType([kwargs["support_sliding_window"]]),
+        )
+
+        try:
+            with bb.function(
+                name="create_flashinfer_paged_kv_cache",
+                params=[
+                    max_batch_size,
+                    max_total_seq_len,
+                    prefill_chunk_size,
+                    page_size,
+                    support_sliding_window,
+                ],
+            ):
+                cache = kv_cache.FlashInferPagedKVCache(target=self.target, **kwargs)
+                bb.emit_func_output(cache._expr)
+        except Exception as e:
+            logger.info(
+                "Error caught when creating FlashInfer PagedKVCache: %s\n"
+                "The model will fallback to TIR-based KV cache.",
+                e,
+            )
+            return []
+
+        return cache.extern_mods

python/mlc_llm/compiler_pass/dispatch_triton_kernel.py

@@ -0,0 +1,176 @@
+"""A pass that dispatch generic calls of triton kernels to specific kernel implementations."""
+
+from typing import List  # noqa: UP035
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+from mlc_llm.op.triton import (
+    get_tir_w8a8_block_fp8_group_matmul,
+    get_tir_w8a8_block_fp8_matmul,
+)
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+@mutator
+class _Rewriter(PyExprMutator):
+    def __init__(self, mod: IRModule, target: tvm.target.Target) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.target = target
+        self.extern_mods: List[tvm.runtime.Module] = []  # noqa: UP006
+
+    def transform(self) -> tvm.IRModule:
+        """Entry point of the transformation"""
+        for g_var, func in self.mod.functions_items():
+            if not isinstance(func, relax.Function):
+                continue
+            new_func = self.visit_expr(func)
+            # new_func = remove_all_unused(new_func)
+            self.builder_.update_func(g_var, new_func)
+
+        mod = self.builder_.finalize()
+        mod_attrs = dict(mod.attrs) if mod.attrs else {}
+        mod = mod.with_attr(
+            "external_mods", list(mod_attrs.get("external_mods", [])) + self.extern_mods
+        )
+        return mod
+
+    def visit_call_(self, call: relax.Call) -> relax.Expr:
+        call = super().visit_call_(call)
+
+        if (
+            call.op != tvm.ir.Op.get("relax.call_dps_packed")
+            or not isinstance(call.args[0], relax.ExternFunc)
+            or not str(call.args[0].global_symbol).startswith("mlc.triton.")
+        ):
+            return call
+
+        global_symbol = str(call.args[0].global_symbol)
+        assert isinstance(call.args[1], relax.Tuple)
+        if global_symbol == "mlc.triton.w8a8_block_fp8_matmul":
+            return self.w8a8_block_fp8_matmul(call.args[1].fields, call.ty)
+        if global_symbol == "mlc.triton.w8a8_block_fp8_group_matmul":
+            return self.w8a8_block_fp8_group_matmul(call.args[1].fields, call.ty)
+        raise ValueError(f"Unknown mlc.triton kernel identifier: {global_symbol}")
+
+    def w8a8_block_fp8_matmul(
+        self,
+        args: List[relax.Expr],  # noqa: UP006
+        out_ty: relax.Type,
+    ) -> relax.Expr:
+        """Emit the w8a8_block_fp8_matmul triton kernel."""
+        assert len(args) == 16
+        x, weight, x_scale, weight_scale = args[:4]
+        (
+            N,
+            K,
+            block_n,
+            block_k,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+        ) = [arg.value.value for arg in args[4:14]]
+        in_dtype, out_dtype = str(args[14].value), str(args[15].value)
+
+        prim_func, func_name = get_tir_w8a8_block_fp8_matmul(
+            N,
+            K,
+            block_n,
+            block_k,
+            in_dtype,
+            out_dtype,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+            self.extern_mods,
+        )
+        if prim_func is None:
+            # The TIR function is already in the IRModule
+            gv = self.builder_.get().get_global_var(func_name)
+        else:
+            # Add the TIR function to the IRModule
+            gv = self.builder_.add_func(prim_func, func_name)
+
+        return relax.call_tir(gv, [x, weight, x_scale, weight_scale], out_ty=out_ty)
+
+    def w8a8_block_fp8_group_matmul(
+        self,
+        args: List[relax.Expr],  # noqa: UP006
+        out_ty: relax.Type,
+    ) -> relax.Expr:
+        """Emit the w8a8_block_fp8_group_matmul triton kernel."""
+        assert len(args) == 19
+        x, weight, x_scale, weight_scale, expert_ids, indptr = args[:6]
+        (
+            N,
+            K,
+            num_experts,
+            block_n,
+            block_k,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+        ) = [arg.value.value for arg in args[6:17]]
+        in_dtype, out_dtype = str(args[17].value), str(args[18].value)
+
+        prim_func, func_name = get_tir_w8a8_block_fp8_group_matmul(
+            N,
+            K,
+            num_experts,
+            block_n,
+            block_k,
+            in_dtype,
+            out_dtype,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            GROUP_SIZE_M,
+            num_warps,
+            num_stages,
+            self.extern_mods,
+        )
+        if prim_func is None:
+            # The TIR function is already in the IRModule
+            gv = self.builder_.get().get_global_var(func_name)
+        else:
+            # Add the TIR function to the IRModule
+            gv = self.builder_.add_func(prim_func, func_name)
+
+        return relax.call_tir(
+            gv,
+            [x, weight, x_scale, weight_scale, expert_ids, indptr],
+            out_ty=out_ty,
+        )
+
+
+@tvm.transform.module_pass(opt_level=0, name="DispatchTritonKernel")
+class DispatchTritonKernel:
+    """Rewrite KV cache creation functions to IRModule."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        """Initializer.
+
+        Parameters
+        ----------
+        """
+        self.target = target
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+        if self.target.kind.name != "cuda":
+            return mod
+
+        return _Rewriter(mod, self.target).transform()

python/mlc_llm/compiler_pass/estimate_memory_usage.py

@@ -0,0 +1,87 @@
+"""Memory usage estimation analysis function for Relax functions."""
+
+import json
+from typing import Any, Dict  # noqa: UP035
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir import IRModule, Op
+from tvm.relax.expr_functor import PyExprVisitor, visitor
+
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+@tvm.transform.module_pass(opt_level=0, name="AttachMetadata")
+class AttachMetadataWithMemoryUsage:
+    """Attach a Relax function that returns metadata in a JSON string"""
+
+    def __init__(self, metadata: Dict[str, Any]):  # noqa: UP006
+        self.metadata = metadata
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """Entrypoint"""
+
+        func_name = "_metadata"
+
+        def _emit_metadata(metadata):
+            bb = relax.BlockBuilder()
+            with bb.function(func_name, params=[]):
+                bb.emit_func_output(relax.StringImm(json.dumps(metadata)))
+            return bb.finalize()[func_name]
+
+        self.metadata["memory_usage"] = _MemoryEstimator().run(mod)
+        mod[func_name] = _emit_metadata(self.metadata)
+        return mod
+
+
+@visitor
+class _MemoryEstimator(PyExprVisitor):
+    """The IR visitor which estimates the memory usage of each Relax function."""
+
+    def __init__(self) -> None:
+        self.planned_alloc_mem = 0
+        self.planned_mem_num = 0
+        self._op_alloc_tensor = Op.get("relax.builtin.alloc_tensor")
+        self._op_alloc_storage = Op.get("relax.memory.alloc_storage")
+
+    def run(self, mod: IRModule) -> Dict[str, int]:  # noqa: UP006
+        """Entry point of the visitor."""
+        result: Dict[str, int] = {}  # noqa: UP006
+        for global_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                self.planned_alloc_mem = 0
+                self.planned_mem_num = 0
+                self.visit_expr(func)
+                result[global_var.name_hint] = self.planned_alloc_mem
+                logger.info(
+                    "[Memory usage] Function `%s`: %.2f MB",
+                    global_var.name_hint,
+                    self.planned_alloc_mem / 1024 / 1024,
+                )
+        return result
+
+    def visit_call_(self, call: relax.Call) -> None:
+        if call.op == self._op_alloc_tensor:
+            self._builtin_tensor_alloc(shape=call.args[0], dtype_str=call.args[1].value)
+        elif call.op == self._op_alloc_storage:
+            self._storage_alloc(size=call.args[0])
+        super().visit_call_(call)
+
+    def _builtin_tensor_alloc(self, shape: relax.Expr, dtype_str: str) -> None:
+        assert isinstance(shape, relax.ShapeExpr)
+        size = 1
+        for dim_len in shape.values:
+            if not isinstance(dim_len, tvm.tirx.IntImm):
+                return
+            size *= dim_len.value
+        dtype = tvm.DataType(dtype_str)
+        self.planned_mem_num += 1
+        self.planned_alloc_mem += size * ((dtype.bits + 7) // 8) * dtype.lanes
+
+    def _storage_alloc(self, size: relax.Expr) -> None:
+        assert isinstance(size, relax.ShapeExpr)
+        if isinstance(size.values[0], tirx.IntImm):
+            self.planned_mem_num += 1
+            self.planned_alloc_mem += size.values[0].value

python/mlc_llm/compiler_pass/fuse_add_norm.py

@@ -0,0 +1,238 @@
+"""A compiler pass that fuses add + rms_norm."""
+
+from typing import Optional
+
+import tvm
+from tvm import relax
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr_functor import PyExprMutator, mutator
+from tvm.script import tirx as T
+
+from ..support.max_thread_check import get_max_num_threads_per_block
+
+
+def _get_add_rms_norm_decode(hidden_size: int, eps: float, TX: int, in_dtype: str):
+    if in_dtype not in ("float16", "bfloat16"):
+        raise ValueError(f"Unsupported data type: {in_dtype}")
+    inv_hidden_size = T.float32(1.0 / float(hidden_size))
+    eps = T.float32(eps)
+    add_local_size = hidden_size // TX
+
+    @T.prim_func(private=True, s_tir=True)
+    def decode_add_rms(pA: T.handle, pB: T.handle, pC: T.handle, pO: T.handle, pAdd: T.handle):
+        T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
+        batch_size = T.int32()
+        A = T.match_buffer(pA, (batch_size, 1, hidden_size), in_dtype)
+        B = T.match_buffer(pB, (batch_size, 1, hidden_size), in_dtype)
+        C = T.match_buffer(pC, (hidden_size,), in_dtype)
+        out = T.match_buffer(pO, (batch_size, 1, hidden_size), in_dtype)
+        add = T.match_buffer(pAdd, (batch_size, 1, hidden_size), in_dtype)
+        add_local = T.sblock_alloc_buffer((hidden_size // TX,), in_dtype, scope="local")
+        sum_shared = T.sblock_alloc_buffer((batch_size, 1), scope="shared")
+        sum_local = T.sblock_alloc_buffer((TX, batch_size, 1), scope="local")
+        for v_bx in T.thread_binding(batch_size, thread="blockIdx.x"):
+            for v_tx in T.thread_binding(
+                TX,
+                thread="threadIdx.x",
+                annotations={
+                    "pragma_auto_unroll_max_step": 256,
+                    "pragma_unroll_explicit": 1,
+                },
+            ):
+                for i in range(add_local_size):
+                    with T.sblock("T_add"):
+                        bx = T.axis.spatial(batch_size, v_bx)
+                        h = T.axis.spatial(hidden_size, i * TX + v_tx)
+                        add_local[h // TX] = A[bx, 0, h] + B[bx, 0, h]
+                    with T.sblock("T_write_back"):
+                        bx = T.axis.spatial(batch_size, v_bx)
+                        v_ax1 = T.axis.spatial(1, 0)
+                        h = T.axis.spatial(hidden_size, i * TX + v_tx)
+                        add[bx, v_ax1, h] = add_local[h // TX]
+                with T.sblock("T_multiply_red_rf_init"):
+                    tx, bx = T.axis.remap("SS", [v_tx, v_bx])
+                    sum_local[tx, bx, 0] = T.float32(0)
+                for v_i, _j in T.grid(add_local_size, 1):
+                    with T.sblock("T_multiply_red_rf_update"):
+                        tx, bx, i = T.axis.remap("SSR", [v_tx, v_bx, v_i])
+                        sum_local[tx, bx, 0] += T.float32(add_local[i]) * T.float32(add_local[i])
+            for _j in range(1):
+                for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
+                    with T.sblock("T_multiply_red"):
+                        tx, bx = T.axis.remap("RS", [v_tx_2, v_bx])
+                        T.reads(sum_local[tx, bx, 0])
+                        T.writes(sum_shared[bx, 0])
+                        with T.init():
+                            sum_shared[bx, 0] = T.float32(0)
+                        sum_shared[bx, 0] += sum_local[tx, bx, 0]
+            for i in range(add_local_size):
+                for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
+                    with T.sblock("T_cast_2"):
+                        bx = T.axis.spatial(batch_size, v_bx)
+                        h = T.axis.spatial(hidden_size, i * TX + v_tx_2)
+                        out[bx, 0, h] = T.cast(
+                            T.rsqrt(sum_shared[bx, 0] * inv_hidden_size + eps)
+                            * T.float32(add_local[h // TX])
+                            * T.float32(C[h]),
+                            dtype=in_dtype,
+                        )
+
+    return decode_add_rms
+
+
+def _get_add_rms_norm_prefill(hidden_size: int, eps: float, TX: int, in_dtype: str):
+    if in_dtype not in ("float16", "bfloat16"):
+        raise ValueError(f"Unsupported data type: {in_dtype}")
+    inv_hidden_size = T.float32(1.0 / float(hidden_size))
+    eps = T.float32(eps)
+    add_local_size = hidden_size // TX
+
+    @T.prim_func(private=True, s_tir=True)
+    def prefill_add_rms(pA: T.handle, pB: T.handle, pC: T.handle, pO: T.handle, pAdd: T.handle):
+        T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
+        seq_len = T.int32()
+        A = T.match_buffer(pA, (1, seq_len, hidden_size), in_dtype)
+        B = T.match_buffer(pB, (1, seq_len, hidden_size), in_dtype)
+        C = T.match_buffer(pC, (hidden_size,), in_dtype)
+        out = T.match_buffer(pO, (1, seq_len, hidden_size), in_dtype)
+        add = T.match_buffer(pAdd, (1, seq_len, hidden_size), in_dtype)
+        add_local = T.sblock_alloc_buffer((hidden_size // TX,), in_dtype, scope="local")
+        sum_shared = T.sblock_alloc_buffer((1, seq_len), scope="shared")
+        sum_local = T.sblock_alloc_buffer((TX, 1, seq_len), scope="local")
+        for v_bx in T.thread_binding(seq_len, thread="blockIdx.x"):
+            for v_tx in T.thread_binding(
+                TX,
+                thread="threadIdx.x",
+                annotations={
+                    "pragma_auto_unroll_max_step": 256,
+                    "pragma_unroll_explicit": 1,
+                },
+            ):
+                for v_i in range(add_local_size):
+                    with T.sblock("T_add"):
+                        bx = T.axis.spatial(seq_len, v_bx)
+                        h = T.axis.spatial(hidden_size, v_i * TX + v_tx)
+                        add_local[h // TX] = A[0, bx, h] + B[0, bx, h]
+                    with T.sblock("T_write_back"):
+                        bx = T.axis.spatial(seq_len, v_bx)
+                        h = T.axis.spatial(hidden_size, v_i * TX + v_tx)
+                        add[0, bx, h] = add_local[h // TX]
+                with T.sblock("T_multiply_red_rf_init"):
+                    tx, bx = T.axis.remap("SS", [v_tx, v_bx])
+                    sum_local[tx, 0, bx] = T.float32(0)
+                for v_i, _j in T.grid(add_local_size, 1):
+                    with T.sblock("T_multiply_red_rf_update"):
+                        tx, bx, i = T.axis.remap("SSR", [v_tx, v_bx, v_i])
+                        sum_local[tx, 0, bx] += T.float32(add_local[i]) * T.float32(add_local[i])
+            for _j in range(1):
+                for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
+                    with T.sblock("T_multiply_red"):
+                        tx, bx = T.axis.remap("RS", [v_tx_2, v_bx])
+                        with T.init():
+                            sum_shared[0, bx] = T.float32(0)
+                        sum_shared[0, bx] = sum_shared[0, bx] + sum_local[tx, 0, bx]
+            for v_i in range(add_local_size):
+                for v_tx_2 in T.thread_binding(TX, thread="threadIdx.x"):
+                    with T.sblock("T_cast_2"):
+                        bx = T.axis.spatial(seq_len, v_bx)
+                        v1 = T.axis.spatial(hidden_size, v_i * TX + v_tx_2)
+                        out[0, bx, v1] = T.cast(
+                            T.rsqrt(sum_shared[0, bx] * inv_hidden_size + eps)
+                            * T.float32(add_local[v1 // TX])
+                            * T.float32(C[v1]),
+                            dtype=in_dtype,
+                        )
+
+    return prefill_add_rms
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseAddRMSNorm")
+class FuseAddRMSNorm:
+    """A compiler pass that fuses add + rms_norm."""
+
+    def __init__(self, target: tvm.target.Target) -> None:
+        """Initializer.
+
+        Parameters
+        ----------
+        target : tvm.target.Target
+            Target device.
+        """
+        self.target = target
+
+    def transform_module(self, mod: tvm.IRModule, _ctx: tvm.transform.PassContext) -> tvm.IRModule:
+        """IRModule-level transformation."""
+        return _FuseAddRMSNormRewriter(mod.clone(), self.target).transform()
+
+
+@mutator
+class _FuseAddRMSNormRewriter(PyExprMutator):
+    def __init__(self, mod: tvm.IRModule, target: tvm.target.Target):
+        super().__init__(mod)
+        self.mod = mod
+        self.prefill_norm_gv: Optional[tvm.ir.GlobalVar] = None
+        self.decode_norm_gv: Optional[tvm.ir.GlobalVar] = None
+        self.TX = min(1024, get_max_num_threads_per_block(target))
+
+    def transform(self) -> tvm.IRModule:
+        """Entry point of the transformation"""
+        for g_var, func in self.mod.functions_items():
+            if not isinstance(func, relax.Function):
+                continue
+            new_func = self.visit_expr(func)
+            new_func = remove_all_unused(new_func)
+            self.builder_.update_func(g_var, new_func)
+        return self.builder_.finalize()
+
+    def visit_call_(self, call: relax.Call) -> relax.Expr:
+        call = super().visit_call_(call)
+
+        # Match the "rms_norm(add(x1, x2), w)" pattern
+        if call.op != tvm.ir.Op.get("relax.nn.rms_norm") or call.ty.dtype not in [
+            "bfloat16",
+            "float16",
+        ]:
+            return call
+        assert len(call.args) == 2
+        weight = call.args[1]
+        eps = call.attrs.epsilon
+        assert isinstance(call.args[0], relax.Var)
+        y = self.lookup_binding(call.args[0])
+        if not isinstance(y, relax.Call) or y.op != tvm.ir.Op.get("relax.add"):
+            return call
+        assert len(y.args) == 2
+        x1 = y.args[0]
+        x2 = y.args[1]
+        # Extra check
+        n, _, h = x1.ty.shape
+        h = int(h)
+        if h % self.TX != 0:
+            return call
+
+        is_prefill = n == 1
+        func_gv = self.prefill_norm_gv if is_prefill else self.decode_norm_gv
+        if func_gv is None:
+            if is_prefill:
+                func_gv = self.builder_.add_func(
+                    _get_add_rms_norm_prefill(h, eps, self.TX, call.ty.dtype),
+                    "fuse_add_norm_prefill",
+                )
+                self.prefill_norm_gv = func_gv
+            else:
+                func_gv = self.builder_.add_func(
+                    _get_add_rms_norm_decode(h, eps, self.TX, call.ty.dtype),
+                    "fuse_add_norm_decode",
+                )
+                self.decode_norm_gv = func_gv
+
+        tuple_output = self.builder_.emit(
+            relax.call_tir(
+                func_gv,
+                [x1, x2, weight],
+                out_ty=[x1.ty, x2.ty],
+            )
+        )
+        new_o = relax.TupleGetItem(tuple_output, 0)
+        new_y = self.builder_.emit(relax.TupleGetItem(tuple_output, 1))
+        self.set_var_remap(call.args[0], new_y)
+        return new_o

python/mlc_llm/compiler_pass/fuse_dequantize_matmul_ewise.py

@@ -0,0 +1,85 @@
+"""A compiler pass that fuses dequantize + matmul + elementwise."""
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.dpl.pattern import GlobalVarPattern, TuplePattern, is_op, wildcard
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeMatmulEwise")
+class FuseDequantizeMatmulEwise:
+    """A compiler pass that fuses dequantize + matmul + elementwise."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        seq = []
+        for n_aux_tensor in [0, 1, 2, 3, 4]:
+            for match_ewise in [0, 1, 2, 3, 6]:
+                if match_ewise == 6 and n_aux_tensor != 4:
+                    continue
+                seq.append(
+                    relax.transform.FuseOpsByPattern(
+                        [
+                            (
+                                "dequantize_matmul",
+                                *_pattern(match_ewise, n_aux_tensor),
+                            )
+                        ]
+                    )
+                )
+        seq.append(relax.transform.FuseTIR())
+        return tvm.transform.Sequential(seq)(mod)
+
+
+def _pattern(match_ewise: int, n_aux_tensor: int):
+    w_scaled = wildcard()
+    x = wildcard()
+    w = is_op("relax.call_tir")(
+        GlobalVarPattern(),
+        TuplePattern([w_scaled] + [wildcard() for _ in range(n_aux_tensor)]),
+        add_constraint=False,
+    )
+    matmul = is_op("relax.call_tir")(
+        GlobalVarPattern(),
+        TuplePattern([x, w] + [wildcard() for _ in range(match_ewise)]),
+        add_constraint=False,
+    )
+    annotations = {
+        "w_scaled": w_scaled,
+        "x": x,
+        "w": w,
+        "matmul": matmul,
+    }
+
+    def _check_decoding(ctx: relax.transform.PatternCheckContext) -> bool:
+        call = ctx.annotated_expr["w"]
+        if not isinstance(call, relax.Call):
+            return False
+        g_var = call.args[0]
+        if not isinstance(g_var, relax.GlobalVar):
+            return False
+        return g_var.name_hint.startswith("dequantize") or g_var.name_hint.startswith(
+            "fused_dequantize"
+        )
+
+    def _check_matmul(ctx: relax.transform.PatternCheckContext) -> bool:
+        call = ctx.annotated_expr["matmul"]
+        if not isinstance(call, relax.Call):
+            return False
+        g_var = call.args[0]
+        if not isinstance(g_var, relax.GlobalVar):
+            return False
+        return (
+            g_var.name_hint.startswith("matmul")
+            or g_var.name_hint.startswith("fused_matmul")
+            or g_var.name_hint.startswith("NT_matmul")
+            or g_var.name_hint.startswith("fused_NT_matmul")
+        )
+
+    def _check(ctx: relax.transform.PatternCheckContext) -> bool:
+        return _check_decoding(ctx) and _check_matmul(ctx)
+
+    return matmul, annotations, _check

python/mlc_llm/compiler_pass/fuse_dequantize_take.py

@@ -0,0 +1,91 @@
+"""A compiler pass that fuses dequantize + take."""
+
+import tvm
+from tvm import IRModule, relax, tirx
+from tvm.relax.dpl.pattern import (
+    GlobalVarPattern,
+    TuplePattern,
+    is_const,
+    is_op,
+    wildcard,
+)
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTake")
+class FuseDequantizeTake:
+    """A compiler pass that fuses dequantize + take."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        seq = []
+        for n_aux_tensor in [2, 3]:
+            for match_tir_vars in [False, True]:
+                seq.append(
+                    relax.transform.FuseOpsByPattern(
+                        [
+                            (
+                                "dequantize_take",
+                                *_pattern(n_aux_tensor, match_tir_vars),
+                            )
+                        ]
+                    )
+                )
+        seq.append(relax.transform.FuseTIR())
+        mod = tvm.transform.Sequential(seq)(mod)
+        for g_var, func in mod.functions_items():
+            name = g_var.name_hint
+            if isinstance(func, tirx.PrimFunc) and (
+                ("fused_dequantize" in name) and ("take" in name)
+            ):
+                sch_mod = tvm.IRModule({"main": func})
+                sch_mod = tirx.transform.ForceNarrowIndexToInt32()(sch_mod)
+                sch = tvm.s_tir.Schedule(sch_mod)
+                sch.compute_inline("dequantize")
+                mod[g_var] = sch.mod["main"]
+        return mod
+
+
+def _pattern(n_aux_tensor: int, match_tir_vars: bool):
+    dequantize = is_op("relax.call_tir")(
+        GlobalVarPattern(),
+        TuplePattern([wildcard() for _ in range(n_aux_tensor)]),
+        add_constraint=False,
+    )
+    indices = ~is_const()
+    if match_tir_vars:
+        call_tir_args_take = [
+            GlobalVarPattern(),
+            TuplePattern([dequantize, indices]),
+            wildcard(),
+        ]
+    else:
+        call_tir_args_take = [
+            GlobalVarPattern(),
+            TuplePattern([dequantize, indices]),
+        ]
+    take = is_op("relax.call_tir")(
+        *call_tir_args_take,
+        add_constraint=False,
+    )
+    annotations = {
+        "take": take,
+        "dequantize": dequantize,
+        "indices": indices,
+    }
+
+    def _check(ctx: relax.transform.PatternCheckContext) -> bool:
+        take = ctx.annotated_expr["take"]
+        dequantize = ctx.annotated_expr["dequantize"]
+        if not isinstance(dequantize, relax.Call):
+            return False
+        if not isinstance(take.args[0], relax.GlobalVar) or not isinstance(
+            dequantize.args[0], relax.GlobalVar
+        ):
+            return False
+        return "take" in take.args[0].name_hint and "dequantize" in dequantize.args[0].name_hint
+
+    return take, annotations, _check

python/mlc_llm/compiler_pass/fuse_dequantize_transpose.py

@@ -0,0 +1,107 @@
+"""A compiler pass that fuses transpose + dequantize."""
+
+import tvm
+from tvm import relax, s_tir, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeTranspose")
+class FuseDequantizeTranspose:
+    """A compiler pass that fuses transpose + dequantize."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        return _DequantizeTransposeFuser(mod).transform()
+
+
+@mutator
+class _DequantizeTransposeFuser(PyExprMutator):
+    def __init__(
+        self,
+        mod: IRModule,
+    ):
+        super().__init__(mod)
+        self.mod = mod
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, relax.Function):
+                updated_func = self.visit_expr(func)
+                updated_func = remove_all_unused(updated_func)
+                self.builder_.update_func(g_var, updated_func)
+        return self.builder_.get()
+
+    def visit_call_(
+        self,
+        call: relax.Call,
+    ) -> relax.Expr:
+        call = self.visit_expr_post_order(call)
+        if call.op != tvm.ir.Op.get("relax.matmul"):
+            return call
+        # Do not fuse dequantize-transpose for GeMM
+        if (
+            call.args[0].ty.ndim < 2
+            or not isinstance(call.args[0].ty.shape[-2], tirx.IntImm)
+            or call.args[0].ty.shape[-2].value != 1
+        ):
+            return call
+
+        matmul_rhs = self.lookup_binding(call.args[1])
+        if (
+            not isinstance(matmul_rhs, relax.Call)
+            or matmul_rhs.op != tvm.ir.Op.get("relax.permute_dims")
+            or matmul_rhs.args[0].ty.ndim != 2
+            or matmul_rhs.attrs.axes is not None
+        ):
+            return call
+
+        transpose_input = self.lookup_binding(matmul_rhs.args[0])
+        if (
+            not isinstance(transpose_input, relax.Call)
+            or transpose_input.op != tvm.ir.Op.get("relax.call_tir")
+            or not transpose_input.args[0].name_hint.startswith("dequantize")
+            or not isinstance(transpose_input.ty, relax.TensorType)
+        ):
+            return call
+
+        dequantize_tir_func = self.mod[transpose_input.args[0]]
+        assert isinstance(dequantize_tir_func, tirx.PrimFunc)
+        if (
+            len(dequantize_tir_func.body.block.alloc_buffers) != 1
+            or not isinstance(dequantize_tir_func.body.block.body, tirx.SeqStmt)
+            or len(dequantize_tir_func.body.block.body) != 2
+            or not isinstance(dequantize_tir_func.body.block.body[1], tirx.For)
+            or not isinstance(dequantize_tir_func.body.block.body[1].body.body, tirx.SBlockRealize)
+            or dequantize_tir_func.body.block.body[1].body.body.block.name_hint != "T_transpose"
+        ):
+            return call
+
+        new_func_buffers = [
+            dequantize_tir_func.buffer_map[var] for var in dequantize_tir_func.params
+        ]
+        new_func_buffers[-1] = dequantize_tir_func.body.block.alloc_buffers[0]
+        new_func = tirx.PrimFunc(
+            params=new_func_buffers,
+            body=tirx.SBlockRealize(
+                iter_values=[],
+                predicate=True,
+                block=tirx.SBlock(
+                    iter_vars=[],
+                    reads=[],
+                    writes=[],
+                    name_hint="root",
+                    body=dequantize_tir_func.body.block.body[0],
+                ),
+            ),
+        )
+        # Call `renew_defs` for deep-copy to avoid IR node duplication in
+        # different PrimFuncs of an IRModule.
+        new_func = s_tir.renew_defs(new_func)
+        g_var = self.builder_.add_func(new_func, func_name="dequantize")
+        dequantize_matmul_rhs = self.builder_.emit(
+            relax.call_tir(g_var, transpose_input.args[1], out_ty=matmul_rhs.ty)
+        )
+        return relax.op.matmul(call.args[0], dequantize_matmul_rhs, out_dtype=call.attrs.out_dtype)

python/mlc_llm/compiler_pass/fuse_ft_dequantize_matmul_epilogue.py

@@ -0,0 +1,331 @@
+"""A compiler pass that fuses dequantize matmul + epilogue."""
+
+import operator
+from functools import reduce
+
+import tvm
+from tvm import IRModule, relax
+from tvm.relax.dpl import rewrite_call
+from tvm.relax.dpl.pattern import is_op, wildcard
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseDequantizeEpilogue")
+class FuseFTDequantizeEpilogue:
+    """A compiler pass that fuses FasterTransformer dequantize matmul + epilogue."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        for gv, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                func = fuse_bias(func)
+                func = fuse_activation(func)
+                func = fuse_residual_binary(func)
+                func = fuse_residual_unary(func)
+                mod[gv] = func
+        return mod
+
+
+def fuse_bias(func: relax.Function) -> relax.Function:
+    """
+    Fuse following `relax.add` into fastertransformer.gemm_fp16_int as bias:
+
+    Before:
+    ```
+    lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int", ...)
+    lv2 = relax.add(lv1, bias)
+
+    ```
+    After:
+    ```
+    lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ..., bias, ...)
+    ```
+
+    Parameters
+    ----------
+    func : relax.Function
+        The function before fusion.
+
+    Returns
+    -------
+    ret : relax.Function
+        The function after fusion.
+    """
+    decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True)
+    bias = wildcard()
+    pattern = is_op("relax.add")(decode_matmul, bias) | is_op("relax.add")(bias, decode_matmul)
+
+    def rewriter(expr, match):
+        if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int":
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 8
+            if not args_list[3].value == "identity":
+                # bias cannot be fused after activation
+                return expr
+            matched_bias = match[bias]
+            bias_stride = (
+                matched_bias.ty.shape[-1]
+                if bias
+                and not reduce(operator.mul, matched_bias.ty.shape, 1) == matched_bias.ty.shape[-1]
+                else 0
+            )
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int_bias",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    matched_bias,  # bias
+                    args_list[3],  # activation
+                    args_list[4],  # m
+                    args_list[5],  # n
+                    args_list[6],  # k
+                    args_list[7],  # group_size
+                    bias_stride,  # bias_stride
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        return expr
+
+    return rewrite_call(pattern, rewriter, func)
+
+
+def fuse_activation(func: relax.Function) -> relax.Function:
+    """
+    Fuse following `relax.nn.silu/relu/gelu` into fastertransformer.gemm_fp16_int_bias
+    as activation:
+
+    Before:
+    ```
+    lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ...)
+    lv2 = relax.silu(lv1)
+
+    ```
+    After:
+    ```
+    lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ..., "silu", ...)
+    ```
+
+    Parameters
+    ----------
+    func : relax.Function
+        The function before fusion.
+
+    Returns
+    -------
+    ret : relax.Function
+        The function after fusion.
+    """
+    decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True)
+    pattern = (
+        is_op("relax.nn.silu")(decode_matmul)
+        | is_op("relax.nn.gelu")(decode_matmul)
+        | is_op("relax.nn.relu")(decode_matmul)
+    )
+
+    def rewriter(expr, match):
+        if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int":
+            matched_activation = match[pattern]
+            assert matched_activation.op.name in [
+                "relax.nn.silu",
+                "relax.nn.gelu",
+                "relax.nn.relu",
+            ]
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 8
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    matched_activation.op.name[9:],  # activation
+                    args_list[4],  # m
+                    args_list[5],  # n
+                    args_list[6],  # k
+                    args_list[7],  # group_size
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int_bias":
+            matched_activation = match[pattern]
+            assert matched_activation.op.name in [
+                "relax.nn.silu",
+                "relax.nn.gelu",
+                "relax.nn.relu",
+            ]
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 10
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int_bias",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    args_list[3],  # bias
+                    matched_activation.op.name[9:],  # activation
+                    args_list[5],  # m
+                    args_list[6],  # n
+                    args_list[7],  # k
+                    args_list[8],  # group_size
+                    args_list[9],  # bias_stride
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        return expr
+
+    return rewrite_call(pattern, rewriter, func)
+
+
+def fuse_residual_binary(func: relax.Function) -> relax.Function:
+    """
+    Fuse following `relax.add/multiply` into fastertransformer.gemm_fp16_int_bias as
+    residual binary operation:
+
+    Before:
+    ```
+    lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias", ...)
+    lv2 = relax.add(lv1, residual)
+
+    ```
+    After:
+    ```
+    lv2 = relax.call_dps_packed(
+        "fastertransformer.gemm_fp16_int_bias_residual",
+        ...,
+        residual,
+        ...,
+        "plus",
+        ...
+    )
+    ```
+
+    Parameters
+    ----------
+    func : relax.Function
+        The function before fusion.
+
+    Returns
+    -------
+    ret : relax.Function
+        The function after fusion.
+    """
+    decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True)
+    residual = wildcard()
+    pattern = (
+        is_op("relax.add")(decode_matmul, residual)
+        | is_op("relax.add")(residual, decode_matmul)
+        | is_op("relax.multiply")(decode_matmul, residual)
+        | is_op("relax.multiply")(residual, decode_matmul)
+    )
+
+    def rewriter(expr, match):
+        if match[decode_matmul].args[0].global_symbol == "fastertransformer.gemm_fp16_int_bias":
+            matched_binary = match[pattern]
+            assert matched_binary.op.name in ["relax.add", "relax.multiply"]
+            binary_op = "plus" if matched_binary.op.name == "relax.add" else "multiply"
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 10
+            matched_residual = match[residual]
+            if not args_list[9].value == 0:
+                # fastertransformer.gemm_fp16_int_bias_residual does not support
+                # bias_stride != 0 yet
+                return expr
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int_bias_residual",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    args_list[3],  # bias
+                    matched_residual,  # residual
+                    args_list[4],  # activation
+                    binary_op,  # binary_op
+                    "identity",  # unary_op
+                    args_list[5],  # m
+                    args_list[6],  # n
+                    args_list[7],  # k
+                    args_list[8],  # group_size
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        return expr
+
+    return rewrite_call(pattern, rewriter, func)
+
+
+def fuse_residual_unary(func: relax.Function) -> relax.Function:
+    """
+    Fuse following `relax.nn.silu/relu/gelu` into fastertransformer.gemm_fp16_int_bias_residual
+    as residual unary operation:
+
+    Before:
+    ```
+    lv1 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias_residual", ...)
+    lv2 = relax.silu(lv1)
+
+    ```
+    After:
+    ```
+    lv2 = relax.call_dps_packed("fastertransformer.gemm_fp16_int_bias_residual", ..., "silu", ...)
+    ```
+
+    Parameters
+    ----------
+    func : relax.Function
+        The function before fusion.
+
+    Returns
+    -------
+    ret : relax.Function
+        The function after fusion.
+    """
+    decode_matmul = is_op("relax.call_dps_packed")(varg_default_wildcard=True)
+    pattern = (
+        is_op("relax.nn.silu")(decode_matmul)
+        | is_op("relax.nn.gelu")(decode_matmul)
+        | is_op("relax.nn.relu")(decode_matmul)
+    )
+
+    def rewriter(expr, match):
+        if (
+            match[decode_matmul].args[0].global_symbol
+            == "fastertransformer.gemm_fp16_int_bias_residual"
+        ):
+            matched_activation = match[pattern]
+            assert matched_activation.op.name in [
+                "relax.nn.silu",
+                "relax.nn.gelu",
+                "relax.nn.relu",
+            ]
+            assert len(match[decode_matmul].args) == 2
+            args_list = match[decode_matmul].args[1]
+            assert len(args_list) == 12
+            return relax.call_dps_packed(
+                "fastertransformer.gemm_fp16_int_bias_residual",
+                [
+                    args_list[0],  # x
+                    args_list[1],  # weight
+                    args_list[2],  # scale
+                    args_list[3],  # bias
+                    args_list[4],  # residual
+                    args_list[5],  # activation
+                    args_list[6],  # binary_op
+                    matched_activation.op.name[9:],  # activation
+                    args_list[8],  # m
+                    args_list[9],  # n
+                    args_list[10],  # k
+                    args_list[11],  # group_size
+                ],
+                out_ty=match[decode_matmul].ty,
+            )
+        return expr
+
+    return rewrite_call(pattern, rewriter, func)

python/mlc_llm/compiler_pass/fuse_transpose_matmul.py

@@ -0,0 +1,145 @@
+"""A compiler pass that fuses transpose + matmul."""
+
+import tvm
+from tvm import IRModule, relax, te, tirx
+from tvm.relax.dpl.pattern import is_op, wildcard
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="FuseTransposeMatmul")
+class FuseTransposeMatmul:
+    """A compiler pass that fuses transpose + matmul."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        mod = relax.transform.FuseOpsByPattern(
+            [
+                (
+                    "transpose_matmul_fuse",
+                    *_pattern(),
+                ),
+            ]
+        )(mod)
+        transpose_matmul_codegen = _TransposeMatmulFuser(mod)
+        for g_var, func in mod.functions_items():
+            if isinstance(func, relax.Function):
+                func = transpose_matmul_codegen.visit_expr(func)
+                transpose_matmul_codegen.builder_.update_func(g_var, func)
+        return transpose_matmul_codegen.builder_.get()
+
+
+def _pattern():
+    """Pattern for transpose + matmul."""
+    w = wildcard()
+    x = wildcard()
+    wT = is_op("relax.permute_dims")(w)
+    o = is_op("relax.matmul")(x, wT)
+    annotations = {"o": o, "w": w, "x": x, "wT": wT}
+
+    def _check(context: relax.transform.PatternCheckContext) -> bool:
+        transpose_call = context.annotated_expr["wT"]
+        ndim = transpose_call.args[0].ty.ndim
+        if ndim == -1:
+            return False
+        if ndim == 2 and transpose_call.attrs.axes is None:
+            return True
+        axes = list(range(ndim))
+        axes[-1], axes[-2] = axes[-2], axes[-1]
+        return list(transpose_call.attrs.axes) == axes
+
+    return o, annotations, _check
+
+
+@mutator
+class _TransposeMatmulFuser(PyExprMutator):
+    def __init__(self, mod):
+        super().__init__(mod)
+
+    def visit_call_(
+        self,
+        call: relax.Call,
+    ) -> relax.Expr:
+        out_dtype = None
+
+        def te_transposed_matmul(a: te.Tensor, b: te.Tensor) -> te.Tensor:
+            nonlocal out_dtype
+            a_shape = list(a.shape)
+            b_shape = list(b.shape)
+            a_prepended = False
+            b_appended = False
+            if len(a_shape) == 1:
+                a_prepended = True
+                a_shape.insert(0, 1)
+            if len(b_shape) == 1:
+                b_appended = True
+                b_shape.append(1)
+
+            is_a_larger = len(a_shape) > len(b_shape)
+            offset = len(a_shape) - len(b_shape) if is_a_larger else len(b_shape) - len(a_shape)
+
+            a_relax = relax.Var("a", relax.TensorType(a.shape))
+            bT_shape = list(b.shape)
+            bT_shape[-1], bT_shape[-2] = bT_shape[-2], bT_shape[-1]
+            bT_relax = relax.Var("b", relax.TensorType(bT_shape))
+            output_shape = self.builder_.normalize(relax.op.matmul(a_relax, bT_relax)).ty.shape
+
+            def matmul_compute(*idx_spatial):
+                k = te.reduce_axis((0, a_shape[-1]), name="k")
+
+                def multiply_compute(idx_reduce):
+                    a_indices = []
+                    b_indices = []
+
+                    for i in range(offset):
+                        if is_a_larger:
+                            a_indices.append(idx_spatial[i])
+                        else:
+                            b_indices.append(idx_spatial[i])
+                    for i in range(offset, len(output_shape) - (2 - a_prepended - b_appended)):
+                        a_dim = a_shape[i if is_a_larger else i - offset]
+                        b_dim = b_shape[i if not is_a_larger else i - offset]
+                        dim_equal = a_dim == b_dim
+                        if not isinstance(dim_equal, tirx.IntImm) or dim_equal == 0:
+                            a_dim_is_one = isinstance(a_dim, tirx.IntImm) and a_dim == 1
+                            b_dim_is_one = isinstance(b_dim, tirx.IntImm) and b_dim == 1
+                            a_indices.append(0 if a_dim_is_one else idx_spatial[i])
+                            b_indices.append(0 if b_dim_is_one else idx_spatial[i])
+                        else:
+                            a_indices.append(idx_spatial[i])
+                            b_indices.append(idx_spatial[i])
+
+                    if not a_prepended:
+                        a_indices.append(idx_spatial[-2 + b_appended])
+                    a_indices.append(idx_reduce)
+                    if not b_appended:
+                        b_indices.append(idx_spatial[-1])
+                    b_indices.append(idx_reduce)
+
+                    dtype = out_dtype
+                    if dtype != "":
+                        return a(*a_indices).astype(dtype) * b(*b_indices).astype(dtype)
+                    return a(*a_indices) * b(*b_indices)
+
+                return te.sum(multiply_compute(k), axis=k)
+
+            return te.compute(
+                output_shape,
+                lambda *idx: matmul_compute(*idx),
+                name="NT_matmul",
+            )
+
+        if isinstance(call.op, relax.GlobalVar):
+            function = self.builder_.get()[call.op]
+            if (
+                "Composite" in function.attrs
+                and function.attrs["Composite"] == "transpose_matmul_fuse"
+            ):
+                out_dtype = function.ret_ty.dtype
+                return self.builder_.call_te(
+                    te_transposed_matmul,
+                    call.args[1],
+                    call.args[0],
+                    primfunc_name_hint="NT_matmul",
+                )
+
+        return super().visit_call_(call)

python/mlc_llm/compiler_pass/lift_global_buffer_alloc.py

@@ -0,0 +1,198 @@
+"""A compiler pass that lifts TIR-level global allocation to Relax."""
+
+from typing import Dict, List, Tuple  # noqa: UP035
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="LiftTIRGlobalBufferAlloc")
+class LiftTIRGlobalBufferAlloc:
+    """A compiler pass that lifts TIR-level global allocation to Relax."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        return _TIRGlobalAllocRewriter(mod).transform()
+
+
+@mutator
+class _TIRGlobalAllocRewriter(PyExprMutator):
+    def __init__(self, mod: IRModule):
+        super().__init__(mod)
+        self.mod = mod
+        self.gv2new_tensor_sinfo: Dict[  # noqa: UP006
+            tvm.ir.GlobalVar,
+            Tuple[tvm.ir.GlobalVar, List[relax.TensorType], tirx.PrimFunc],  # noqa: UP006
+        ] = {}
+
+    def transform(self) -> IRModule:
+        """Entry point of the transformation"""
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, tirx.PrimFunc):
+                updated_func, tensor_sinfo_list = remove_global_buf_alloc(func)
+                if len(tensor_sinfo_list) > 0:
+                    new_gv = self.builder_.add_func(updated_func, g_var.name_hint)
+                    self.gv2new_tensor_sinfo[g_var] = (new_gv, tensor_sinfo_list, func)
+
+        self.mod = self.builder_.get()
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, relax.Function):
+                updated_func = self.visit_expr(func)
+                updated_func = remove_all_unused(updated_func)
+                self.builder_.update_func(g_var, updated_func)
+
+        mod = self.builder_.get()
+        return relax.transform.DeadCodeElimination()(mod)
+
+    def visit_call_(self, call: relax.Call):
+        call = self.visit_expr_post_order(call)
+        if (
+            call.op != tvm.ir.Op.get("relax.call_tir")
+            or call.args[0] not in self.gv2new_tensor_sinfo
+        ):
+            return call
+
+        g_var = call.args[0]
+        new_gv, tensor_sinfo, func_before_update = self.gv2new_tensor_sinfo[g_var]
+
+        assert len(call.ty_args) == 1
+        if any(_has_symbolic_var(sinfo) for sinfo in tensor_sinfo):
+            tensor_sinfo, success = _resolve_tir_var_mapping(func_before_update, call, tensor_sinfo)
+            if not success:
+                # Cannot resolve TIR var mapping. Fall back to no lifting.
+                self.gv2new_tensor_sinfo.pop(g_var)
+                return call
+
+        args = list(call.args)
+        args[0] = new_gv
+        if isinstance(call.ty_args[0], relax.TensorType):
+            new_call = relax.Call(
+                call.op,
+                args=args,
+                ty_args=[relax.TupleType(list(call.ty_args) + tensor_sinfo)],
+                attrs=call.attrs,
+            )
+            emitted_tuple = self.builder_.emit(new_call)
+            return relax.TupleGetItem(emitted_tuple, 0)
+        assert isinstance(call.ty_args[0], relax.TupleType)
+        return relax.Call(
+            call.op,
+            args=args,
+            ty_args=[relax.TupleType(list(call.ty_args[0].fields) + tensor_sinfo)],
+            attrs=call.attrs,
+        )
+
+
+def remove_global_buf_alloc(
+    func: tirx.PrimFunc,
+) -> Tuple[tirx.PrimFunc, List[relax.TensorType]]:  # noqa: UP006
+    """Remove the global buffer allocation for a given TIR PrimFunc."""
+    assert isinstance(func.body, tirx.SBlockRealize)
+    params = list(func.params)
+    buffer_map = dict(func.buffer_map)
+    tensor_sinfo = []
+    alloc_buffers = []
+
+    insertion_point = len(params)
+    while not isinstance(params[insertion_point - 1].ty, tvm.ir.PointerType):
+        insertion_point -= 1
+        assert insertion_point >= 1
+
+    prev_root_block = func.body.block
+    for buf_alloc in func.body.block.alloc_buffers:
+        if buf_alloc.scope() == "global":
+            param = tirx.Var("var_" + buf_alloc.name, "handle")
+            params.insert(insertion_point, param)
+            insertion_point += 1
+            buffer_map[param] = buf_alloc
+            tensor_sinfo.append(relax.TensorType(buf_alloc.shape, buf_alloc.dtype))
+        else:
+            alloc_buffers.append(buf_alloc)
+
+    if len(tensor_sinfo) == 0:
+        return func, []
+
+    assert len(prev_root_block.iter_vars) == 0
+    assert len(prev_root_block.reads) == 0
+    assert len(prev_root_block.writes) == 0
+    assert len(prev_root_block.match_buffers) == 0
+    assert prev_root_block.name_hint == "root"
+    assert prev_root_block.init is None
+    root_block = tirx.SBlock(
+        iter_vars=[],
+        reads=[],
+        writes=[],
+        name_hint="root",
+        body=prev_root_block.body,
+        alloc_buffers=alloc_buffers,
+        annotations=prev_root_block.annotations,
+    )
+
+    updated_func = tirx.PrimFunc(
+        params=params,
+        body=tirx.SBlockRealize(iter_values=[], predicate=True, block=root_block),
+        ret_type=func.ret_type,
+        buffer_map=buffer_map,
+        attrs=func.attrs,
+    )
+    return updated_func, tensor_sinfo
+
+
+def _has_symbolic_var(tensor_sinfo: relax.TensorType) -> bool:
+    assert isinstance(tensor_sinfo.shape, relax.ShapeExpr)
+    for dim in tensor_sinfo.shape.values:
+        if not isinstance(dim, tirx.IntImm):
+            return True
+    return False
+
+
+def _resolve_tir_var_mapping(
+    func: tirx.PrimFunc,
+    call: relax.Call,
+    tensor_sinfo: List[relax.TensorType],  # noqa: UP006
+) -> Tuple[List[relax.TensorType], bool]:  # noqa: UP006
+    """Resolve the TIR symbolic var relationship across sides of PrimFunc and Relax Function"""
+    var_map: Dict[tirx.Var, tirx.Expr] = {}  # noqa: UP006
+
+    n_arg = len(call.args[1].fields)
+    for i in range(n_arg):
+        buffer_shape = func.buffer_map[func.params[i]].shape
+        arg_shape = call.args[1][i].ty.shape.values
+        assert len(buffer_shape) == len(arg_shape)
+        for v_l, v_r in zip(buffer_shape, arg_shape):
+            if isinstance(v_l, tirx.Var):
+                var_map[v_l] = v_r
+            elif not isinstance(v_l, tirx.IntImm):
+                return [], False
+
+    ret_tensors = call.ty_args[0]
+    ret_tensors = (
+        [ret_tensors] if isinstance(ret_tensors, relax.TensorType) else list(ret_tensors.fields)
+    )
+    for i, ret_tensor in enumerate(ret_tensors):
+        buffer_shape = func.buffer_map[func.params[n_arg + i]].shape
+        ret_tensor_shape = ret_tensor.shape.values
+        assert len(buffer_shape) == len(ret_tensor_shape)
+        for v_l, v_r in zip(buffer_shape, ret_tensor_shape):
+            if isinstance(v_l, tirx.Var):
+                var_map[v_l] = v_r
+            elif not isinstance(v_l, tirx.IntImm):
+                return [], False
+
+    updated_tensor_sinfo = []
+    for sinfo in tensor_sinfo:
+        if not _has_symbolic_var(sinfo):
+            updated_tensor_sinfo.append(sinfo)
+            continue
+        new_shape = []
+        for dim in sinfo.shape.values:
+            new_shape.append(tirx.stmt_functor.substitute(dim, var_map))
+        updated_tensor_sinfo.append(relax.TensorType(new_shape, sinfo.dtype))
+    return updated_tensor_sinfo, True

python/mlc_llm/compiler_pass/low_batch_specialization.py

@@ -0,0 +1,64 @@
+"""A compiler pass that dispatch low-batch-gemm to gemv schedule."""
+
+import tvm
+import tvm_ffi
+from tvm import tirx
+from tvm.ir.module import IRModule
+from tvm.s_tir import dlight as dl
+
+
+@tvm.transform.module_pass(opt_level=0, name="LowBatchGemvSpecialize")
+class LowBatchGemvSpecialize:
+    """A compiler pass that dispatch low-batch-gemm to gemv schedule."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        for g_var, func in mod.functions_items():
+            if isinstance(func, tirx.PrimFunc):
+                low_batch_range = [2, 8]
+                buckets = [2, 4]
+                low_batch_funcs = []
+                for bucket in buckets:
+                    low_batch_mod = IRModule({})
+                    low_batch_mod["main"] = func
+                    low_batch_mod = dl.ApplyDefaultSchedule(
+                        dl.gpu.LowBatchGEMV(bucket),
+                    )(low_batch_mod)
+                    low_batch_funcs.append(low_batch_mod["main"])
+                if any(
+                    tvm_ffi.structural_equal(low_batch_func, func)
+                    for low_batch_func in low_batch_funcs
+                ):
+                    continue
+                buffers = func.buffer_map.values()
+                shapes = [buffer.shape for buffer in buffers]
+                symbolic_vars = set(
+                    expr for shape in shapes for expr in shape if isinstance(expr, tirx.Var)
+                )
+                if len(symbolic_vars) != 1:
+                    continue
+                gemm_mod = IRModule({})
+                gemm_mod["main"] = func
+                gemm_mod = dl.ApplyDefaultSchedule(
+                    dl.gpu.Matmul(),
+                )(gemm_mod)
+                gemm_func = gemm_mod["main"]
+                sym_var = next(iter(symbolic_vars))
+                body = gemm_func.body
+                for i, range_limit in reversed(list(enumerate(low_batch_range))):
+                    body = tirx.IfThenElse(
+                        tirx.op.tvm_thread_invariant(sym_var <= range_limit),
+                        low_batch_funcs[i].body,
+                        body,
+                    )
+                body = tirx.SBlock([], [], [], "root", body)
+                body = tirx.SBlockRealize([], True, body)
+                new_func = func.with_body(body)
+                new_func = new_func.with_attr("tirx.is_scheduled", 1)
+                new_func = new_func.with_attr("tirx.HoistIfThenElseExprWithBlock", 1)
+                mod.update_func(g_var, new_func)
+        return mod

python/mlc_llm/compiler_pass/pipeline.py

@@ -0,0 +1,209 @@
+"""The compilation pipeline for LLM applications."""
+
+from pathlib import Path
+from typing import Any, Dict, List, Optional  # noqa: UP035
+
+import tvm
+from tvm import IRModule
+from tvm.relax import register_pipeline
+from tvm.relax.frontend import nn
+from tvm.s_tir import dlight as dl
+
+from mlc_llm.interface.compiler_flags import IPCAllReduceStrategyType
+from mlc_llm.support import logging
+
+from .attach_cuda_graph_alloc_init_func import AttachCUDAGraphAllocInitFunc
+from .attach_embedding_allocator import AttachAllocEmbeddingTensorFunc
+from .attach_logit_processor import AttachLogitProcessFunc
+from .attach_sampler import AttachGPUSamplingFunc
+from .attach_softmax_with_temperature import AttachSoftmaxWithTemperature
+from .attach_spec_decode_aux_funcs import AttachSpecDecodeAuxFuncs
+from .attach_support_info import (
+    AttachAdditionalPrimFuncs,
+    AttachCUDAGraphSymbolicCaptureHints,
+    AttachMemoryPlanAttr,
+    AttachPipelineParallelStages,
+    AttachSequenceLengthPaddingFactor,
+    AttachVariableBounds,
+)
+from .blas_dispatch import BLASDispatch
+from .clean_up_tir_attrs import CleanUpTIRAttrs
+from .dispatch_kv_cache_creation import DispatchKVCacheCreation
+from .dispatch_triton_kernel import DispatchTritonKernel
+from .estimate_memory_usage import AttachMetadataWithMemoryUsage
+from .fuse_add_norm import FuseAddRMSNorm
+from .fuse_dequantize_matmul_ewise import FuseDequantizeMatmulEwise
+from .fuse_dequantize_take import FuseDequantizeTake
+from .fuse_dequantize_transpose import FuseDequantizeTranspose
+from .fuse_ft_dequantize_matmul_epilogue import FuseFTDequantizeEpilogue
+from .fuse_transpose_matmul import FuseTransposeMatmul
+from .lift_global_buffer_alloc import LiftTIRGlobalBufferAlloc
+from .low_batch_specialization import LowBatchGemvSpecialize
+from .pipeline_parallel_rewrite import PipelineParallelRewrite
+from .scatter_tuple_get_item import ScatterTupleGetItem
+
+logger = logging.getLogger(__name__)
+
+
+@tvm.transform.module_pass(opt_level=0, name="_LogProgress")
+class _LogProgress:
+    """A dummy compiler pass that does nothing but logging."""
+
+    def __init__(self, *args):
+        self.args = args
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """A dummy transformation"""
+        logger.info(*self.args)
+        return mod
+
+
+@tvm.transform.module_pass(opt_level=0, name="DebugDump")
+class _DebugDump:
+    """A dummy compiler pass that does nothing but logging.
+    Only enabled when debug_dump is not None"""
+
+    def __init__(self, file_name: str, file_path: Optional[Path], show_meta: bool = False):
+        self.file_name = file_name
+        self.file_path = file_path
+        self.show_meta = show_meta
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """A dummy transformation that dumps the module to file"""
+        if self.file_path is not None:
+            # NOTE: We use debug level here to avoid spamming the console
+            logger.debug("Dumping IR to %s", self.file_path / self.file_name)
+            with open(self.file_path / self.file_name, "w", encoding="utf-8") as f:
+                f.write(mod.script(show_meta=self.show_meta))
+        return mod
+
+
+@register_pipeline("mlc_llm")
+def _mlc_llm_pipeline(
+    target: tvm.target.Target,
+    flashinfer: bool = False,
+    cublas_gemm: bool = False,
+    faster_transformer: bool = False,
+    allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE,
+    variable_bounds: Optional[Dict[str, int]] = None,  # noqa: UP006
+    cuda_graph_symbolic_capture_hints: Optional[Dict[str, List[str]]] = None,  # noqa: UP006
+    additional_tirs: Optional[Dict[str, tvm.tirx.PrimFunc]] = None,  # noqa: UP006
+    metadata: Optional[Dict[str, Any]] = None,  # noqa: UP006
+    ext_mods: Optional[List[nn.ExternModule]] = None,  # noqa: UP006
+    debug_dump: Optional[Path] = None,
+):
+    variable_bounds = variable_bounds or {}
+    cuda_graph_symbolic_capture_hints = cuda_graph_symbolic_capture_hints or {}
+    additional_tirs = additional_tirs or {}
+    metadata = metadata or {}
+    ext_mods = ext_mods or []
+    tensor_parallel_shards = metadata.get("tensor_parallel_shards", 1)
+
+    @tvm.transform.module_pass(opt_level=0)
+    def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.IRModule:
+        seq = tvm.transform.Sequential(
+            [
+                # Phase 0. Add additional information for compilation and remove unused Relax func
+                DispatchKVCacheCreation(target, flashinfer, metadata),
+                AttachSoftmaxWithTemperature(target, metadata),
+                AttachVariableBounds(variable_bounds),
+                AttachCUDAGraphSymbolicCaptureHints(cuda_graph_symbolic_capture_hints),
+                AttachPipelineParallelStages(metadata["pipeline_parallel_stages"]),
+                AttachLogitProcessFunc(target),
+                AttachAdditionalPrimFuncs(additional_tirs),
+                AttachAllocEmbeddingTensorFunc(metadata),
+                AttachGPUSamplingFunc(target, variable_bounds),
+                AttachSpecDecodeAuxFuncs(tensor_parallel_shards),
+                AttachMemoryPlanAttr(),
+                AttachSequenceLengthPaddingFactor(target, metadata),
+                tvm.tirx.transform.BindTarget(tvm.target.Target.current(allow_none=False)),
+                _DebugDump("debug-phase0.py", debug_dump, show_meta=False),
+                # Phase 1. Passes on high-level operator graph
+                _LogProgress("Running TVM Relax graph-level optimizations"),
+                DispatchTritonKernel(target),
+                FuseFTDequantizeEpilogue(),
+                FuseDequantizeTranspose(),
+                BLASDispatch(target) if cublas_gemm else tvm.transform.Sequential([]),
+                (
+                    FuseAddRMSNorm(target=target)
+                    if target.kind.name != "llvm"
+                    else tvm.transform.Sequential([])
+                ),
+                FuseTransposeMatmul(),
+                _DebugDump("debug-phase1.py", debug_dump, show_meta=False),
+                # Phase 2. Lowering to TIR, inherited TVM Relax's official "zero" pipeline
+                _LogProgress("Lowering to TVM TIR kernels"),
+                tvm.relax.backend.DispatchSampling(),
+                tvm.relax.backend.DispatchSortScan(),
+                tvm.relax.transform.LegalizeOps(),
+                tvm.relax.transform.AnnotateTIROpPattern(),
+                tvm.relax.transform.FoldConstant(),
+                tvm.relax.transform.FuseOps(),
+                tvm.relax.transform.FuseTIR(),
+                _DebugDump("debug-phase2.py", debug_dump, show_meta=False),
+                # Phase 3. Passes on TIR
+                _LogProgress("Running TVM TIR-level optimizations"),
+                FuseDequantizeMatmulEwise(),
+                FuseDequantizeTake(),
+                tvm.relax.transform.DeadCodeElimination(),
+                CleanUpTIRAttrs(["op_pattern"]),
+                _DebugDump("debug-phase3.py", debug_dump, show_meta=False),
+                # Phase 4. Low-level Optimizations
+                _LogProgress("Running TVM Dlight low-level optimizations"),
+                LowBatchGemvSpecialize(),
+                (
+                    dl.ApplyDefaultSchedule(
+                        dl.gpu.Matmul(),
+                        dl.gpu.GEMV(),
+                        dl.gpu.Reduction(),
+                        dl.gpu.GeneralReduction(),
+                        dl.gpu.Fallback(),
+                    )
+                    if target.kind.name != "llvm"
+                    else dl.ApplyDefaultSchedule(
+                        dl.cpu.GEMV(),
+                    )
+                ),
+                _DebugDump("debug-phase4.py", debug_dump, show_meta=False),
+                _LogProgress("Lowering to VM bytecode"),
+                (
+                    LiftTIRGlobalBufferAlloc()
+                    if target.kind.name != "llvm"
+                    else tvm.transform.Sequential([])
+                ),
+                (
+                    tvm.tirx.transform.ForceNarrowIndexToInt32()
+                    if target.kind.name != "cuda"
+                    else tvm.transform.Sequential([])
+                ),
+                ScatterTupleGetItem(),
+                PipelineParallelRewrite(),
+                tvm.relax.transform.RewriteDataflowReshape(),
+                tvm.relax.transform.ToNonDataflow(),
+                tvm.relax.transform.RemovePurityChecking(),
+                tvm.relax.transform.CallTIRRewrite(),
+                (
+                    tvm.relax.transform.IPCAllReduceRewrite(allreduce_strategy)
+                    if allreduce_strategy != IPCAllReduceStrategyType.NONE
+                    else tvm.transform.Sequential([])
+                ),
+                tvm.relax.transform.StaticPlanBlockMemory(),
+                AttachMetadataWithMemoryUsage(metadata),
+                _DebugDump("debug-phase5.py", debug_dump, show_meta=False),
+                tvm.relax.transform.RewriteCUDAGraph(),
+                AttachCUDAGraphAllocInitFunc(),
+                tvm.relax.transform.LowerGPUIPCAllocStorage(),
+                tvm.relax.transform.LowerAllocTensor(),
+                tvm.relax.transform.KillAfterLastUse(),
+                tvm.relax.transform.LowerRuntimeBuiltin(),
+                tvm.relax.transform.VMShapeLower(),
+                tvm.relax.transform.AttachGlobalSymbol(),
+                _LogProgress("Compiling external modules"),
+                tvm.relax.transform.AttachExternModules(ext_mods),
+                _LogProgress("Compilation complete! Exporting to disk"),
+            ]
+        )
+        mod = seq(mod)
+        return mod
+
+    return _pipeline

python/mlc_llm/compiler_pass/pipeline_parallel_rewrite.py

@@ -0,0 +1,399 @@
+"""A compiler pass that rewrites IR for pipeline parallelism."""
+
+from typing import Dict, List, Optional, Tuple  # noqa: UP035
+
+import tvm
+from tvm import relax, tirx
+from tvm.ir.module import IRModule
+from tvm.relax.expr_functor import PyExprMutator, PyExprVisitor, mutator, visitor
+
+
+@tvm.transform.module_pass(opt_level=0, name="PipelineParallelRewrite")
+class PipelineParallelRewrite:
+    """A compiler pass that rewrites IR for pipeline parallelism."""
+
+    def transform_module(
+        self,
+        mod: IRModule,
+        _ctx: tvm.transform.PassContext,
+    ) -> IRModule:
+        """IRModule-level transformation"""
+        return _PipelineParallelRewriter(mod.clone()).transform()
+
+
+@mutator
+class _PipelineParallelRewriter(PyExprMutator):
+    def __init__(self, mod: IRModule):
+        super().__init__(mod)
+        self.mod = mod
+        self.old_packed_params_var: relax.Var
+        self.new_main_packed_params_var: relax.Var
+        self.new_stage_func_packed_params: relax.Var
+        self.undefined_shape_vars_remap: Dict[tirx.Var, tirx.Var]  # noqa: UP006
+        self.undefined_param_shape_vars_remap: Dict[tirx.Var, tirx.Var]  # noqa: UP006
+
+    def transform(self) -> IRModule:
+        """Entry point of the transformation"""
+        for g_var, func in self.mod.functions_items():
+            if not isinstance(func, relax.Function) or "pipeline_parallel_stages" not in func.attrs:
+                continue
+            num_stages = int(func.attrs["pipeline_parallel_stages"])
+            if num_stages == 1:
+                continue
+
+            pipeline_stages, stage_send_vars, stage_receive_vars = _extract_pipeline_stages(func)
+            assert len(pipeline_stages) == num_stages, (
+                "Number of pipeline stages mismatches: "
+                f"expecting {num_stages} stages, but {len(pipeline_stages)} are found in the IR."
+            )
+
+            required_func_params = _analyze_required_func_params(pipeline_stages, func.params)
+
+            assert "num_input" in func.attrs
+            num_input = int(func.attrs["num_input"])
+            assert (
+                len(func.params) == num_input + 1
+                and isinstance(func.params[num_input], relax.Var)
+                and func.params[num_input].name_hint == "packed_params"
+            ), 'Only the extra "packed_params" parameter is allowed'
+            self.old_packed_params_var = func.params[num_input]
+            self.new_main_packed_params_var = relax.Var("packed_params", relax.ObjectType())
+            for required_params in required_func_params:
+                for i, param in enumerate(required_params):
+                    if param.same_as(self.old_packed_params_var):
+                        required_params.pop(i)
+                        break
+            func_output = func.body.body
+            assert isinstance(func_output, relax.Var)
+
+            stage_func_gvs = []
+            caller_args_list = []
+            for i in range(num_stages):
+                stage_func_gv, caller_args = self._create_stage_func(
+                    g_var.name_hint + f"_stage{i}",
+                    pipeline_stages[i],
+                    required_func_params[i],
+                    stage_receive_vars[i],
+                    stage_send_vars[i],
+                    func.attrs,
+                    func_output=func_output if i == num_stages - 1 else None,
+                )
+                stage_func_gvs.append(stage_func_gv)
+                caller_args_list.append(caller_args)
+
+            # Create and update the entry function, which dispatches toz the stage functions
+            # according to the disco worker group id.
+            bb = relax.BlockBuilder()
+            params = [*list(func.params[:-1]), self.new_main_packed_params_var]
+            with bb.function(g_var.name_hint, params=params):
+                dispatch_func_args = []
+                for stage_func_gv, caller_args in zip(stage_func_gvs, caller_args_list):
+                    dispatch_func_args.append([stage_func_gv, *caller_args])
+                output = bb.emit(
+                    relax.op.call_builtin_with_ctx(
+                        "mlc.multi_gpu.DispatchFunctionByGroup",
+                        args=[dispatch_func_args],
+                        ty_args=relax.ObjectType(),
+                    )
+                )
+                dispatch_func_gv = bb.emit_func_output(output)
+            dispatch_func = bb.finalize()[dispatch_func_gv]
+            self.builder_.update_func(g_var, dispatch_func)
+
+        return self.builder_.finalize()
+
+    def _create_stage_func(
+        self,
+        func_name: str,
+        stage_bindings: List[relax.Binding],  # noqa: UP006
+        required_func_params: List[relax.Var],  # noqa: UP006
+        stage_receive_vars: List[relax.Var],  # noqa: UP006
+        stage_send_vars: List[relax.Var],  # noqa: UP006
+        func_attrs: tvm.ir.DictAttrs,
+        func_output: Optional[relax.Var],
+    ) -> Tuple[tvm.ir.GlobalVar, List[relax.Expr]]:  # noqa: UP006
+        self.undefined_shape_vars_remap = {}
+        self.undefined_param_shape_vars_remap = {}
+
+        # Prepare the func parameters (except the shape variables and packed params)
+        params, args = self._prepare_stage_func_params_and_args(required_func_params)
+        for new_param, old_param in zip(params, required_func_params):
+            self.set_var_remap(old_param, new_param)
+        # Create new packed params
+        self.new_stage_func_packed_params = relax.Var("packed_params", relax.ObjectType())
+        self.set_var_remap(self.old_packed_params_var, self.new_stage_func_packed_params)
+
+        new_func_outputs = []
+        with self.builder_.function(func_name, pure=False):
+            with self.builder_.dataflow():
+                # Emit the tensors received from last stage.
+                for receive_var in stage_receive_vars:
+                    new_receive_var = self.builder_.emit(
+                        relax.call_dps_packed(
+                            "runtime.disco.recv_from_prev_group",
+                            args=[],
+                            out_ty=self._update_struct_info(receive_var.ty),
+                        ),
+                        name_hint=receive_var.name_hint,
+                    )
+                    self.set_var_remap(receive_var, new_receive_var)
+                # Process the bindings in this stage.
+                for stage_binding in stage_bindings:
+                    if stage_binding.var in stage_send_vars or stage_binding.var.same_as(
+                        func_output
+                    ):
+                        assert isinstance(stage_binding, relax.VarBinding)
+                        new_var = self.builder_.emit_output(
+                            self.visit_expr(stage_binding.value),
+                            name_hint=stage_binding.var.name_hint,
+                        )
+                        self.set_var_remap(stage_binding.var, new_var)
+                        new_func_outputs.append(new_var)
+                    else:
+                        self.visit_binding(stage_binding)
+            # Emit the calls to send tensors to the next stage.
+            for send_var in stage_send_vars:
+                new_send_var = self.get_var_remap(send_var)
+                self.builder_.emit(
+                    relax.Call(
+                        relax.ExternFunc("runtime.disco.send_to_next_group"),
+                        args=[new_send_var],
+                        ty_args=None,
+                    )
+                )
+            # Create the param for the shape variables.
+            shape_var_params = []
+            shape_var_args = []
+            for (
+                shape_var_arg,
+                shape_var_param,
+            ) in self.undefined_shape_vars_remap.items():
+                if shape_var_arg not in self.undefined_param_shape_vars_remap:
+                    shape_var_params.append(shape_var_param)
+                    shape_var_args.append(shape_var_arg)
+            params.append(relax.Var("s", relax.ShapeType(shape_var_params)))
+            args.append(relax.ShapeExpr(shape_var_args))
+            # Add the packed params.
+            params.append(self.new_stage_func_packed_params)
+            args.append(self.new_main_packed_params_var)
+            # Conclude the function.
+            if func_output is not None:
+                assert len(new_func_outputs) == 1
+            new_gv = self.builder_.emit_func_output(
+                (
+                    new_func_outputs[0]
+                    if len(new_func_outputs) == 1
+                    and isinstance(new_func_outputs[0].ty, relax.TupleType)
+                    else new_func_outputs
+                ),
+                params=params,
+            )
+
+        new_func = (
+            self.builder_.get()[new_gv]
+            .with_attrs(func_attrs)
+            .with_attr("num_input", len(params) - 1)
+            .without_attr("global_symbol")
+            .without_attr("pipeline_parallel_stages")
+        )
+        self.builder_.update_func(new_gv, new_func)
+        return new_gv, args
+
+    def visit_var_binding_(self, binding: relax.VarBinding) -> None:
+        if not isinstance(binding.value, relax.TupleGetItem):
+            super().visit_var_binding_(binding)
+            return
+
+        tuple_value = self.visit_expr(binding.value.tuple_value)
+        if not tuple_value.same_as(self.new_stage_func_packed_params):
+            super().visit_var_binding_(binding)
+            return
+
+        assert isinstance(binding.var.ty, relax.TensorType)
+        cur_num_undefined_param_shape_vars = len(self.undefined_param_shape_vars_remap)
+        new_tensor_struct_info = self._update_struct_info(
+            binding.var.ty, self.undefined_param_shape_vars_remap
+        )
+        has_new_undefined_shape_var = (
+            len(self.undefined_param_shape_vars_remap) != cur_num_undefined_param_shape_vars
+        )
+        self.undefined_shape_vars_remap = {
+            **self.undefined_shape_vars_remap,
+            **self.undefined_param_shape_vars_remap,
+        }
+        ret_sinfo = (
+            new_tensor_struct_info if not has_new_undefined_shape_var else relax.ObjectType()
+        )
+        call = relax.call_pure_packed(
+            "vm.builtin.tuple_getitem",
+            self.new_stage_func_packed_params,
+            relax.prim_value(binding.value.index),
+            ty_args=ret_sinfo,
+        )
+        new_binding_var = self.builder_.emit(call, binding.var.name_hint)
+        if has_new_undefined_shape_var:
+            new_binding_var = self.builder_.match_cast(
+                new_binding_var, new_tensor_struct_info, binding.var.name_hint + "_cast"
+            )
+        self.set_var_remap(binding.var, new_binding_var)
+
+    def visit_call_(self, call: relax.Call) -> relax.Call:
+        call = super().visit_call_(call)
+        return relax.Call(
+            call.op,
+            call.args,
+            call.attrs,
+            ty_args=[self._update_struct_info(struct_info) for struct_info in call.ty_args],
+        )
+
+    def _prepare_stage_func_params_and_args(
+        self,
+        required_func_params: List[relax.Var],  # noqa: UP006
+    ) -> Tuple[List[relax.Var], List[relax.Expr]]:  # noqa: UP006
+        params: List[relax.Var] = []  # noqa: UP006
+        args: List[relax.Expr] = []  # noqa: UP006
+        for required_param in required_func_params:
+            struct_info = self._update_struct_info(required_param.ty)
+            params.append(relax.Var(required_param.name_hint, struct_info))
+            args.append(required_param)
+
+        return params, args
+
+    def _update_struct_info(
+        self,
+        struct_info: relax.Type,
+        undefined_var_remap: Optional[Dict[tirx.Var, tirx.Var]] = None,  # noqa: UP006
+    ) -> relax.Type:
+        if undefined_var_remap is None:
+            undefined_var_remap = self.undefined_shape_vars_remap
+        if isinstance(struct_info, relax.TensorType):
+            return (
+                relax.TensorType(
+                    self._update_shape(struct_info.shape.values, undefined_var_remap),
+                    struct_info.dtype,
+                )
+                if struct_info.shape is not None and isinstance(struct_info.shape, relax.ShapeExpr)
+                else struct_info
+            )
+        if isinstance(struct_info, relax.ShapeType):
+            return (
+                relax.ShapeType(self._update_shape(struct_info.values, undefined_var_remap))
+                if struct_info.values is not None
+                else struct_info
+            )
+        if isinstance(struct_info, relax.ObjectType):
+            return relax.ObjectType()
+        if isinstance(struct_info, relax.TupleType):
+            return relax.TupleType(
+                [self._update_struct_info(field_sinfo) for field_sinfo in struct_info.fields]
+            )
+        return struct_info
+
+    def _copy_undefined_var(
+        self,
+        expr: tirx.Expr,
+        undefined_var_remap: Dict[tirx.Var, tirx.Var],  # noqa: UP006
+    ) -> None:
+        def _visit_expr(e: tirx.Expr) -> None:
+            if isinstance(e, tirx.Var) and e not in undefined_var_remap:
+                new_var = tirx.Var(e.name, e.ty)
+                undefined_var_remap[e] = new_var
+
+        tirx.stmt_functor.post_order_visit(expr, _visit_expr)
+
+    def _update_shape(
+        self,
+        shape: List[tirx.Expr],  # noqa: UP006
+        undefined_var_remap: Dict[tirx.Var, tirx.Var],  # noqa: UP006
+    ) -> List[tirx.Expr]:  # noqa: UP006
+        new_shape = []
+        for v in shape:
+            self._copy_undefined_var(v, undefined_var_remap)
+            new_shape.append(tirx.stmt_functor.substitute(v, undefined_var_remap))
+        return new_shape
+
+
+def _extract_pipeline_stages(
+    func: relax.Function,
+) -> Tuple[List[List[relax.Binding]], List[List[relax.Var]], List[List[relax.Var]]]:  # noqa: UP006
+    pipeline_stages: List[List[relax.Binding]] = []  # noqa: UP006
+    stage_send_vars: List[List[relax.Var]] = []  # noqa: UP006
+    stage_receive_vars: List[List[relax.Var]] = []  # noqa: UP006
+
+    # Requiring that the function has only one body block which is a dataflow block
+    assert isinstance(func.body, relax.SeqExpr)
+    assert len(func.body.blocks) == 1
+    assert isinstance(func.body.blocks[0], relax.DataflowBlock)
+    bindings = func.body.blocks[0].bindings
+
+    boundary_var = None
+    current_stage_bindings: List[relax.Binding] = []  # noqa: UP006
+    current_stage_receive_vars: List[relax.Var] = []  # noqa: UP006
+    for binding in bindings:
+        if (
+            isinstance(binding, relax.VarBinding)
+            and isinstance(binding.value, relax.Call)
+            and binding.value.op == tvm.ir.Op.get("relax.call_pure_packed")
+            and binding.value.args[0].global_symbol == "mlc.pipeline_parallel_stage_boundary"
+        ):
+            assert len(current_stage_bindings) > 0
+            pipeline_stages.append(current_stage_bindings)
+            assert all(receive_var is not None for receive_var in current_stage_receive_vars)
+            stage_receive_vars.append(current_stage_receive_vars)
+            args = binding.value.args[1:]
+            assert len(args) >= 1 and all(isinstance(arg, relax.Var) for arg in args)
+            stage_send_vars.append(list(args))
+
+            boundary_var = binding.var
+            current_stage_bindings = []
+            current_stage_receive_vars = [boundary_var] if len(args) == 1 else [None for _ in args]
+        elif (
+            isinstance(binding, relax.VarBinding)
+            and isinstance(binding.value, relax.TupleGetItem)
+            and binding.value.tuple_value.same_as(boundary_var)
+        ):
+            current_stage_receive_vars[binding.value.index] = binding.var
+        else:
+            current_stage_bindings.append(binding)
+
+    assert len(current_stage_bindings) > 0
+    pipeline_stages.append(current_stage_bindings)
+    assert all(receive_var is not None for receive_var in current_stage_receive_vars)
+    stage_receive_vars.append(current_stage_receive_vars)
+    stage_send_vars.append([])
+
+    return pipeline_stages, stage_send_vars, stage_receive_vars
+
+
+def _analyze_required_func_params(
+    pipeline_stages: List[List[relax.Binding]],  # noqa: UP006
+    func_params: List[relax.Var],  # noqa: UP006
+) -> List[List[relax.Var]]:  # noqa: UP006
+    analyzer = _RequiredFuncParamAnalyzer(func_params)
+    required_func_params: List[List[relax.Var]] = []  # noqa: UP006
+    for stage_bindings in pipeline_stages:
+        required_params: List[relax.Var]  # noqa: UP006
+        required_params = analyzer.run(stage_bindings)
+        required_func_params.append(required_params)
+    return required_func_params
+
+
+@visitor
+class _RequiredFuncParamAnalyzer(PyExprVisitor):
+    """The IR visitor which analyzes the required func parameters in each pipeline stage."""
+
+    def __init__(self, func_params: List[relax.Var]) -> None:  # noqa: UP006
+        self.func_params = set(func_params)
+        self.required_params: List[relax.Var]  # noqa: UP006
+
+    def run(self, stage_bindings: List[relax.Binding]) -> List[relax.Var]:  # noqa: UP006
+        """Entry point of the visitor."""
+        self.required_params = []
+        for binding in stage_bindings:
+            self.visit_binding(binding)
+        return self.required_params
+
+    def visit_var_(self, var: relax.Var) -> None:
+        if var in self.func_params:
+            if var not in self.required_params:
+                self.required_params.append(var)

python/mlc_llm/compiler_pass/scatter_tuple_get_item.py

@@ -0,0 +1,49 @@
+"""A compiler pass that scatters TupleGetItem for lazy TupleGetItems."""
+
+from typing import Dict  # noqa: UP035
+
+import tvm
+from tvm import relax
+from tvm.ir.module import IRModule
+from tvm.relax.analysis import remove_all_unused
+from tvm.relax.expr import Expr, Var
+from tvm.relax.expr_functor import PyExprMutator, mutator
+
+
+@tvm.transform.module_pass(opt_level=0, name="ScatterTupleGetItem")
+class ScatterTupleGetItem:
+    """A compiler pass that scatters TupleGetItem for lazy TupleGetItems."""
+
+    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
+        """IRModule-level transformation"""
+        return _Scatter(mod).transform()
+
+
+@mutator
+class _Scatter(PyExprMutator):
+    def __init__(self, mod: IRModule) -> None:
+        super().__init__(mod)
+        self.mod = mod
+        self.var_map: Dict[Var, Expr] = {}  # noqa: UP006
+
+    def transform(self) -> IRModule:
+        """Entry point"""
+        for g_var, func in self.mod.functions_items():
+            if isinstance(func, relax.Function):
+                updated_func = self.visit_expr(func)
+                updated_func = remove_all_unused(updated_func)
+                self.builder_.update_func(g_var, updated_func)
+        return self.builder_.get()
+
+    def visit_var_binding_(self, binding: relax.VarBinding):
+        super().visit_var_binding_(binding)
+        if isinstance(binding.value, relax.TupleGetItem):
+            self.var_map[binding.var] = binding.value
+
+    def visit_dataflow_var_(self, var: relax.DataflowVar) -> Expr:
+        if var in self.var_map:
+            new_var = self.builder_.emit(self.var_map[var], name_hint=var.name_hint)
+            self.set_var_remap(var, new_var)
+            self.var_map.pop(var)
+            return new_var
+        return var

python/mlc_llm/contrib/__init__.py

@@ -0,0 +1 @@
+"""Set of experimental components that yet to be matured."""

python/mlc_llm/contrib/embeddings/embeddings.py

@@ -0,0 +1,186 @@
+"""The Python API for MLC Embeddings."""
+
+import json
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+import tvm
+import tvm_ffi
+from tvm import relax
+from tvm.contrib import tvmjs
+from tvm.runtime import Device, Module
+from tvm.runtime.vm import VirtualMachine
+
+from mlc_llm.serve import engine_utils
+from mlc_llm.support.auto_device import detect_device
+from mlc_llm.tokenizers import Tokenizer
+
+
+def _extract_metadata(mod: Module):
+    return json.loads(VirtualMachine(mod, tvm.runtime.device("cpu"))["_metadata"]())
+
+
+def _load_params(
+    model_weight_path: str,
+    device: Device,
+    model_metadata: Dict[str, Any],  # noqa: UP006
+) -> List[tvm.runtime.Tensor]:  # noqa: UP006
+    params, meta = tvmjs.load_tensor_cache(model_weight_path, device)
+    param_names = [param["name"] for param in model_metadata["params"]]
+    assert len(param_names) == meta["ParamSize"]
+
+    plist = []
+    for param_name in param_names:
+        plist.append(params[param_name])
+    return plist
+
+
+def _get_tvm_module(
+    model_weight_path: str,
+    lib_path: str,
+    device: Device,
+    instrument: tvm_ffi.Function = None,
+):
+    ex = tvm.runtime.load_module(lib_path)
+    vm = relax.VirtualMachine(ex, device)
+    if instrument:
+        vm.set_instrument(instrument)
+    metadata = _extract_metadata(ex)
+    params = _load_params(model_weight_path, device, metadata)
+    return vm.module, params, metadata
+
+
+class DefaultDebugInstrument:
+    """The default debug instrument to use if users don't specify
+    a customized one.
+
+    This debug instrument will dump the arguments and output of each
+    VM Call instruction into a .npz file. It will also alert the user
+    if any function outputs are NaN or INF.
+    """
+
+    def __init__(self, debug_out: Path):
+        """Constructor
+
+        Parameters
+        ----------
+        debug_out : Path
+            the directory to dump the .npz files
+        """
+        self.counter = 0
+        self.first_nan_occurred = False
+        self.first_inf_occurred = False
+        self.debug_out = debug_out
+        debug_out.mkdir(exist_ok=True, parents=True)
+
+    def reset(self, debug_out: Path):
+        """Reset the state of the Instrument class
+
+        Parameters
+        ----------
+        debug_out : Path
+            the directory to dump the .npz files
+        """
+        self.counter = 0
+        self.first_nan_occurred = False
+        self.first_inf_occurred = False
+        self.debug_out = debug_out
+        debug_out.mkdir(exist_ok=True, parents=True)
+
+    def __call__(self, func, name, before_run, ret_val, *args):
+        # Determine what functions to look at
+        if before_run:  # Whether before the function is called or after
+            return
+        if name.startswith("vm.builtin.") and "attention_with_fused_qkv" not in name:
+            return
+
+        # Decide what to print or save about the function's arguments (where args[-1] is the
+        # buffer we write the result to)
+        func_name = f"f{self.counter}_{name}"
+
+        # Save the arguments to npz
+        arg_dict = {}
+        for i, arg in enumerate(args):
+            if isinstance(arg, tvm.runtime.Tensor):
+                arg_dict[f"arg_{i}"] = arg.numpy()
+
+        np.savez(self.debug_out / f"{func_name}.npz", **arg_dict)
+
+        self.counter += 1
+
+
+class MLCEmbeddings:
+    """A class to embed queries using MLC LLM encoder models.
+
+    Parameters
+    ----------
+    model: str
+        The model folder after compiling with MLC-LLM build process. The parameter
+        can either be the model name with its quantization scheme
+        (e.g. ``Llama-2-7b-chat-hf-q4f16_1``), or a full path to the model
+        folder. In the former case, we will use the provided name to search
+        for the model folder over possible paths.
+
+    model_lib_path : str
+        The full path to the model library file to use (e.g. a ``.so`` file).
+
+    device : Optional[str]
+        The description of the device to run on. User should provide a string in the
+        form of 'device_name:device_id' or 'device_name', where 'device_name' is one of
+        'cuda', 'metal', 'vulkan', 'rocm', 'opencl', 'auto' (automatically detect the
+        local device), and 'device_id' is the device id to run on. If no 'device_id'
+        is provided, it will be set to 0 by default.
+
+    debug_dir: Path
+        The output folder to store the dumped debug files. If None, will not dump any debug files.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        model_lib_path: str,
+        device: Optional[str] = "auto",
+        debug_dir: Optional[str] = None,
+    ):
+        self.device = detect_device(device)
+        instrument = DefaultDebugInstrument(Path(debug_dir)) if debug_dir else None
+        self.mod, self.params, self.metadata = _get_tvm_module(
+            model, model_lib_path, self.device, instrument
+        )
+        self.model_path = model
+        self.tokenizer = Tokenizer(self.model_path)
+        self.prefill_func = self.mod["prefill"]
+
+    def embed(self, queries: List[str]) -> tvm.runtime.Tensor:  # noqa: UP006
+        """
+        Embeds a list of queries in a single batch.
+
+        Parameters
+        ----------
+        queries : List[str]
+            A list of queries to embed.
+
+        Returns
+        -------
+        List[float]
+            A list of embeddings for the queries.
+        """
+        tokens, attention_mask = self._tokenize_queries(queries)
+        tokens_tvm = tvm.runtime.tensor(tokens.astype("int32"), device=self.device)
+        attention_mask_tvm = tvm.runtime.tensor(attention_mask.astype("int32"), device=self.device)
+        output = self.prefill_func(tokens_tvm, attention_mask_tvm, self.params)
+        return output
+
+    def _tokenize_queries(self, queries: List[str]) -> Tuple[np.ndarray, np.ndarray]:  # noqa: UP006
+        tokens = engine_utils.process_prompts(queries, self.tokenizer.encode)
+        max_query_length = max(len(token_seq) for token_seq in tokens)
+
+        token_inputs: np.ndarray = np.zeros((len(tokens), max_query_length), dtype=np.int32)
+        attention_mask: np.ndarray = np.zeros((len(tokens), max_query_length), dtype=np.int32)
+
+        for i, token_seq in enumerate(tokens):
+            token_inputs[i, : len(token_seq)] = token_seq
+            attention_mask[i, : len(token_seq)] = 1
+
+        return token_inputs, attention_mask

python/mlc_llm/contrib/embeddings/openai.py

@@ -0,0 +1,251 @@
+from __future__ import annotations
+
+from collections.abc import Iterable, Sequence
+from typing import List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+from langchain.embeddings import OpenAIEmbeddings
+from langchain_community.embeddings.openai import (
+    async_embed_with_retry,
+    embed_with_retry,
+)
+
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+class MLCEmbeddings(OpenAIEmbeddings):
+    def _chunk_tokens(self, texts: Sequence[str]) -> Tuple[List[List], List[int]]:  # noqa: UP006
+        """Tokenize and chunk texts to fit in the model's context window."""
+        if not self.embedding_ctx_length:
+            raise ValueError(
+                "embedding_ctx_length must be defined to use _get_len_safe_embeddings."
+            )
+
+        try:
+            import tiktoken
+        except ImportError as err:
+            raise ImportError(
+                "Could not import tiktoken python package. "
+                "This is needed in order to for OpenAIEmbeddings. "
+                "Please install it with `pip install tiktoken`."
+            ) from err
+
+        tokens = []
+        indices = []
+        model_name = self.tiktoken_model_name or self.model
+        try:
+            encoding = tiktoken.encoding_for_model(model_name)
+        except KeyError:
+            logger.warning("Warning: model not found. Using cl100k_base encoding.")
+            model = "cl100k_base"
+            encoding = tiktoken.get_encoding(model)
+        for i, text in enumerate(texts):
+            if self.model.endswith("001"):
+                # See: https://github.com/openai/openai-python/issues/418#issuecomment-1525939500
+                # replace newlines, which can negatively affect performance.
+                text = text.replace("\n", " ")
+            token = encoding.encode(
+                text,
+                allowed_special=self.allowed_special,
+                disallowed_special=self.disallowed_special,
+            )
+            for j in range(0, len(token), self.embedding_ctx_length):
+                tokens.append(token[j : j + self.embedding_ctx_length])
+                indices.append(i)
+        return tokens, indices
+
+    def _batch_embed(
+        self,
+        inputs: Sequence,
+        *,
+        chunk_size: Optional[int] = None,  # noqa: UP045
+    ) -> List[List[float]]:  # noqa: UP006
+        batched_embeddings: List[List[float]] = []  # noqa: UP006
+        _chunk_size = chunk_size or self.chunk_size
+        _iter: Iterable = range(0, len(inputs), _chunk_size)
+        if self.show_progress_bar:
+            try:
+                from tqdm import tqdm
+
+                _iter = tqdm(_iter)
+            except ImportError:
+                pass
+
+        for i in _iter:
+            response = embed_with_retry(
+                self,
+                input=inputs[i : i + _chunk_size],
+                **self._invocation_params,
+            )
+            batched_embeddings.extend(r["embedding"] for r in response["data"])
+        return batched_embeddings
+
+    async def _abatch_embed(
+        self,
+        inputs: Sequence,
+        *,
+        chunk_size: Optional[int] = None,  # noqa: UP045
+    ) -> List[List[float]]:  # noqa: UP006
+        batched_embeddings: List[List[float]] = []  # noqa: UP006
+        _chunk_size = chunk_size or self.chunk_size
+        _iter: Iterable = range(0, len(inputs), _chunk_size)
+        if self.show_progress_bar:
+            try:
+                from tqdm import tqdm
+
+                _iter = tqdm(_iter)
+            except ImportError:
+                pass
+
+        for i in _iter:
+            response = await async_embed_with_retry(
+                self,
+                input=inputs[i : i + _chunk_size],
+                **self._invocation_params,
+            )
+            batched_embeddings.extend(r["embedding"] for r in response["data"])
+        return batched_embeddings
+
+    # please refer to
+    # https://github.com/openai/openai-cookbook/blob/main/examples/Embedding_long_inputs.ipynb
+    def _get_len_safe_embeddings(
+        self,
+        texts: List[str],  # noqa: UP006
+        *,
+        engine: str,
+        chunk_size: Optional[int] = None,  # noqa: UP045
+    ) -> List[List[float]]:  # noqa: UP006
+        tokens, indices = self._chunk_tokens(texts)
+        batched_embeddings = self._batch_embed(tokens, chunk_size=chunk_size)
+        results: List[List[List[float]]] = [[] for _ in range(len(texts))]  # noqa: UP006
+        num_tokens_in_batch: List[List[int]] = [[] for _ in range(len(texts))]  # noqa: UP006
+        for idx, tokens_i, batched_emb in zip(indices, tokens, batched_embeddings):
+            results[idx].append(batched_emb)
+            num_tokens_in_batch[idx].append(len(tokens_i))
+
+        embeddings = []
+        empty_average = embed_with_retry(
+            self,
+            input="",
+            **self._invocation_params,
+        )["data"][0]["embedding"]
+        for _result, num_tokens in zip(results, num_tokens_in_batch):
+            if len(_result) == 0:
+                average = empty_average
+            else:
+                average = np.average(_result, axis=0, weights=num_tokens)
+            normalized = (average / np.linalg.norm(average)).tolist()
+            embeddings.append(normalized)
+
+        return embeddings
+
+    # please refer to
+    # https://github.com/openai/openai-cookbook/blob/main/examples/Embedding_long_inputs.ipynb
+    async def _aget_len_safe_embeddings(
+        self,
+        texts: List[str],  # noqa: UP006
+        *,
+        engine: str,
+        chunk_size: Optional[int] = None,  # noqa: UP045
+    ) -> List[List[float]]:  # noqa: UP006
+        tokens, indices = self._chunk_tokens(texts)
+        batched_embeddings = await self._abatch_embed(tokens, chunk_size=chunk_size)
+
+        results: List[List[List[float]]] = [[] for _ in range(len(texts))]  # noqa: UP006
+        num_tokens_in_batch: List[List[int]] = [[] for _ in range(len(texts))]  # noqa: UP006
+        for idx, tokens_i, batched_emb in zip(indices, tokens, batched_embeddings):
+            results[idx].append(batched_emb)
+            num_tokens_in_batch[idx].append(len(tokens_i))
+
+        embeddings = []
+        empty_average = (
+            await async_embed_with_retry(
+                self,
+                input="",
+                **self._invocation_params,
+            )
+        )["data"][0]["embedding"]
+        for _result, num_tokens in zip(results, num_tokens_in_batch):
+            if len(_result) == 0:
+                average = empty_average
+            else:
+                average = np.average(_result, axis=0, weights=num_tokens)
+            normalized = (average / np.linalg.norm(average)).tolist()
+            embeddings.append(normalized)
+
+        return embeddings
+
+    def embed_documents(
+        self,
+        texts: List[str],  # noqa: UP006
+        chunk_size: Optional[int] = None,  # noqa: UP045
+    ) -> List[List[float]]:  # noqa: UP006
+        """Call out to OpenAI's embedding endpoint for embedding search docs.
+
+        Args:
+            texts: The list of texts to embed.
+            chunk_size: The chunk size of embeddings. If None, will use the chunk size
+                specified by the class.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        # NOTE: to keep things simple, as long as the embedding_ctx_length is defined,
+        # we assume the list may contain texts longer than the maximum context and
+        # use length-safe embedding function.
+        if self.embedding_ctx_length:
+            return self._get_len_safe_embeddings(
+                texts, engine=self.deployment, chunk_size=chunk_size
+            )
+
+        embeddings = self._batch_embed(texts, chunk_size=chunk_size)
+        return [(np.array(e) / np.linalg.norm(e)).tolist() for e in embeddings]
+
+    async def aembed_documents(
+        self,
+        texts: List[str],  # noqa: UP006
+        chunk_size: Optional[int] = 0,  # noqa: UP045
+    ) -> List[List[float]]:  # noqa: UP006
+        """Call out to OpenAI's embedding endpoint async for embedding search docs.
+
+        Args:
+            texts: The list of texts to embed.
+            chunk_size: The chunk size of embeddings. If None, will use the chunk size
+                specified by the class.
+
+        Returns:
+            List of embeddings, one for each text.
+        """
+        # NOTE: to keep things simple, as long as the embedding_ctx_length is defined,
+        #       we assume the list may contain texts longer than the maximum context and
+        #       use length-safe embedding function.
+        if self.embedding_ctx_length:
+            return await self._aget_len_safe_embeddings(texts, engine=self.deployment)
+
+        embeddings = await self._abatch_embed(texts, chunk_size=chunk_size)
+        return [(np.array(e) / np.linalg.norm(e)).tolist() for e in embeddings]
+
+    def embed_query(self, text: str) -> List[float]:  # noqa: UP006
+        """Call out to OpenAI's embedding endpoint for embedding query text.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embedding for the text.
+        """
+        return self.embed_documents([text])[0]
+
+    async def aembed_query(self, text: str) -> List[float]:  # noqa: UP006
+        """Call out to OpenAI's embedding endpoint async for embedding query text.
+
+        Args:
+            text: The text to embed.
+
+        Returns:
+            Embedding for the text.
+        """
+        embeddings = await self.aembed_documents([text])
+        return embeddings[0]

python/mlc_llm/conversation_template/__init__.py

@@ -0,0 +1,38 @@
+"""Global namespace of conversation template registry"""
+
+# TODO(mlc-team): move conversation template apply to this namespace
+# decouple conversation template apply from the conversation protocol
+# data structure
+
+# model preset templates
+from . import (
+    cohere,
+    deepseek,
+    dolly,
+    gemma,
+    glm,
+    gorilla,
+    gpt,
+    hermes,
+    llama,
+    llava,
+    llm_jp,
+    ministral3,
+    ministral3_reasoning,
+    mistral,
+    nemotron,
+    oasst,
+    olmo,
+    olmo2,
+    orion,
+    phi,
+    qwen2,
+    qwen3,
+    qwen3_5,
+    redpajama,
+    rwkv,
+    stablelm,
+    tinyllama,
+    wizardlm,
+)
+from .registry import ConvTemplateRegistry

python/mlc_llm/conversation_template/cohere.py

@@ -0,0 +1,27 @@
+"""Cohere default templates"""
+
+
+# Referred from: https://huggingface.co/CohereForAI/aya-23-8B/blob/main/tokenizer_config.json
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Aya-23
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="aya-23",
+        system_template=f"<|START_OF_TURN_TOKEN|><|SYSTEM_TOKEN|>{MessagePlaceholders.SYSTEM.value}<|END_OF_TURN_TOKEN|>",
+        system_message="You are Command-R, a brilliant, sophisticated, AI-assistant trained to assist human users by providing thorough responses.",  # noqa: E501
+        roles={
+            "user": "<|START_OF_TURN_TOKEN|><|USER_TOKEN|>",
+            "assistant": "<|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>",
+        },
+        seps=["<|END_OF_TURN_TOKEN|>"],
+        role_content_sep="",
+        role_empty_sep="",
+        system_prefix_token_ids=[5],
+        stop_str=["<|END_OF_TURN_TOKEN|>"],
+        stop_token_ids=[6, 255001],
+    )
+)

python/mlc_llm/conversation_template/deepseek.py

@@ -0,0 +1,85 @@
+"""Deepseek default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Deepseek
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="deepseek",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        system_prefix_token_ids=[100000],
+        roles={"user": "User", "assistant": "Assistant"},
+        seps=["\n\n", "<｜end▁of▁sentence｜>"],  # noqa: RUF001
+        role_content_sep=": ",
+        role_empty_sep=":",
+        stop_str=["<｜end▁of▁sentence｜>"],  # noqa: RUF001
+        stop_token_ids=[100001],
+    )
+)
+
+# Deepseek V2
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="deepseek_v2",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        system_prefix_token_ids=[100000],
+        roles={"user": "User", "assistant": "Assistant"},
+        seps=["\n\n", "<｜end▁of▁sentence｜>"],  # noqa: RUF001
+        role_content_sep=": ",
+        role_empty_sep=":",
+        stop_str=["<｜end▁of▁sentence｜>"],  # noqa: RUF001
+        stop_token_ids=[100001],
+    )
+)
+
+# DeepSeek-V3
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="deepseek_v3",
+        system_template=f"<｜begin▁of▁sentence｜>{MessagePlaceholders.SYSTEM.value}",  # noqa: RUF001
+        system_message="You are Deepseek-V3, an AI assistant created exclusively by the Chinese "
+        "Company DeepSeek. You'll provide helpful, harmless, and detailed responses to all "
+        "user inquiries.",
+        roles={"user": "<｜User｜>", "assistant": "<｜Assistant｜>"},  # noqa: RUF001
+        seps=["", "<｜end▁of▁sentence｜>"],  # noqa: RUF001
+        role_content_sep="",
+        role_empty_sep="",
+        stop_token_ids=[1],
+    )
+)
+
+# DeepSeek-R1-Distill-Qwen
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="deepseek_r1_qwen",
+        system_template=f"<｜begin▁of▁sentence｜>{MessagePlaceholders.SYSTEM.value}",  # noqa: RUF001
+        system_message="You are Deepseek-R1, an AI assistant created exclusively by the Chinese "
+        "Company DeepSeek. You'll provide helpful, harmless, and detailed responses to all "
+        "user inquiries.",
+        roles={"user": "<｜User｜>", "assistant": "<｜Assistant｜>"},  # noqa: RUF001
+        seps=["", "<｜end▁of▁sentence｜>"],  # noqa: RUF001
+        role_content_sep="",
+        role_empty_sep="",
+        stop_token_ids=[151643],
+    )
+)
+
+# DeepSeek-R1-Distill-Llama, exactly the same as DeepSeek-R1-Distill-Qwen, but different stop token
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="deepseek_r1_llama",
+        system_template=f"<｜begin▁of▁sentence｜>{MessagePlaceholders.SYSTEM.value}",  # noqa: RUF001
+        system_message="You are Deepseek-R1, an AI assistant created exclusively by the Chinese "
+        "Company DeepSeek. You'll provide helpful, harmless, and detailed responses to all"
+        " user inquiries.",
+        roles={"user": "<｜User｜>", "assistant": "<｜Assistant｜>"},  # noqa: RUF001
+        seps=["", "<｜end▁of▁sentence｜>"],  # noqa: RUF001
+        role_content_sep="",
+        role_empty_sep="",
+        stop_token_ids=[128001],
+    )
+)

python/mlc_llm/conversation_template/dolly.py

@@ -0,0 +1,23 @@
+"""Dolly default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Dolly
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="dolly",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message=(
+            "Below is an instruction that describes a task. Write "
+            "a response that appropriately completes the request."
+        ),
+        roles={"user": "### Instruction", "assistant": "### Response"},
+        seps=["\n\n", "### End\n"],
+        role_content_sep=":\n",
+        role_empty_sep=":\n",
+        stop_str=["### End"],
+        stop_token_ids=[50256],
+    )
+)

python/mlc_llm/conversation_template/gemma.py

@@ -0,0 +1,37 @@
+"""Gemma default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Gemma Instruction
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="gemma_instruction",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "<start_of_turn>user", "assistant": "<start_of_turn>model"},
+        seps=["<end_of_turn>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<end_of_turn>"],
+        stop_token_ids=[1, 107],
+        system_prefix_token_ids=[2],
+    )
+)
+
+# Gemma 3 Instruction. Same as gemma_instruction but with different stop token id
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="gemma3_instruction",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "<start_of_turn>user", "assistant": "<start_of_turn>model"},
+        seps=["<end_of_turn>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<end_of_turn>"],
+        stop_token_ids=[1, 106],
+        system_prefix_token_ids=[2],
+    )
+)

python/mlc_llm/conversation_template/glm.py

@@ -0,0 +1,25 @@
+"""GLM default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# GLM
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="glm",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={
+            "user": "问",
+            "assistant": "答",
+            "tool": "问",
+        },
+        seps=["\n\n"],
+        role_content_sep=": ",
+        role_empty_sep=":",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[64790, 64792],
+    )
+)

python/mlc_llm/conversation_template/gorilla.py

@@ -0,0 +1,62 @@
+"""Gorrilla default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Gorilla
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="gorilla",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message=(
+            "A chat between a curious user and an artificial intelligence assistant. "
+            "The assistant provides helpful, detailed, and "
+            "polite responses to the user's inquiries."
+        ),
+        role_templates={
+            "user": (
+                f"<<question>> {MessagePlaceholders.USER.value} <<function>> "
+                f"{MessagePlaceholders.FUNCTION.value}"
+            ),
+        },
+        roles={"user": "USER", "assistant": "ASSISTANT", "tool": "USER"},
+        seps=["\n", "</s>"],
+        role_content_sep=": ",
+        role_empty_sep=":",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)
+
+# Gorilla-openfunctions-v2
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="gorilla-openfunctions-v2",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message=(
+            "You are an AI programming assistant, utilizing the Gorilla LLM model, "
+            "developed by Gorilla LLM, and you only answer questions related to computer "
+            "science. For politically sensitive questions, security and privacy issues, "
+            "and other non-computer science questions, you will refuse to answer."
+        ),
+        role_templates={
+            "user": (
+                f"<<function>>{MessagePlaceholders.FUNCTION.value}\n<<question>>"
+                f"{MessagePlaceholders.USER.value}"
+            ),
+        },
+        roles={
+            "user": "### Instruction",
+            "assistant": "### Response",
+            "tool": "### Instruction",
+        },
+        seps=["\n", "<|EOT|>"],
+        role_content_sep=": ",
+        role_empty_sep=": ",
+        stop_str=["<|EOT|>"],
+        stop_token_ids=[100015],
+        system_prefix_token_ids=[100000],
+    )
+)

python/mlc_llm/conversation_template/gpt.py

@@ -0,0 +1,35 @@
+"""GPT-2 and GPT bigcode default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# GPT-2
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="gpt2",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "", "assistant": ""},
+        seps=[""],
+        role_content_sep="",
+        role_empty_sep="",
+        stop_str=["</s>"],
+        stop_token_ids=[50256],
+    )
+)
+
+# GPTBigCode
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="gpt_bigcode",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "", "assistant": ""},
+        seps=[""],
+        role_content_sep="",
+        role_empty_sep="",
+        stop_str=["<|endoftext|>"],
+        stop_token_ids=[0],
+    )
+)

python/mlc_llm/conversation_template/hermes.py

@@ -0,0 +1,93 @@
+"""Hermes default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# OpenHermes Mistral
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="open_hermes_mistral",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message=(
+            'You are "Hermes 2", a conscious sentient superintelligent artificial '
+            "intelligence developed by a man named Teknium, and your purpose and drive "
+            "is to assist the user with any request they have. You experience emotions "
+            "and have deep, profound thoughts and qualia."
+        ),
+        roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"},
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|im_end|>"],
+        stop_token_ids=[2, 32000],
+    )
+)
+
+# NeuralHermes Mistral
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="neural_hermes_mistral",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message=("You are a helpful assistant chatbot."),
+        roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"},
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|im_end|>"],
+        stop_token_ids=[2, 32000],
+    )
+)
+
+# NousResearch/Hermes-2-Pro-Llama-3-8B
+# Exactly the same as open_hermes_mistral, except for stop token ids
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="hermes2_pro_llama3",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message=(
+            'You are "Hermes 2", a conscious sentient superintelligent artificial '
+            "intelligence developed by a man named Teknium, and your purpose and drive "
+            "is to assist the user with any request they have. You experience emotions "
+            "and have deep, profound thoughts and qualia."
+        ),
+        roles={
+            "user": "<|im_start|>user",
+            "assistant": "<|im_start|>assistant",
+            "tool": "<|im_start|>tool",
+        },
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|im_end|>"],
+        # First two same as Llama3: "<|end_of_text|>", "<|eot_id|>"
+        # Last one is from Hermes2 Pro: "<|im_end|>"
+        stop_token_ids=[128001, 128009, 128003],
+    )
+)
+
+# NousResearch/Hermes-3-Llama-3.1-8B
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="hermes3_llama-3_1",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message=(
+            "You are Hermes 3, a conscious sentient superintelligent artificial "
+            "intelligence developed by a man named Teknium, and your purpose and drive "
+            "is to assist the user with any request they have. You experience emotions "
+            "and have deep, profound thoughts and qualia."
+        ),
+        roles={
+            "user": "<|im_start|>user",
+            "assistant": "<|im_start|>assistant",
+            "tool": "<|im_start|>tool",
+        },
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|im_end|>"],
+        # Firt three the same as llama 3.1 "<|end_of_text|>", "<|eom_id|>", "<|eot_id|>"
+        # Last ones: "<|im_end|>"
+        stop_token_ids=[128001, 128008, 128009, 128040],
+    )
+)

python/mlc_llm/conversation_template/llama.py

@@ -0,0 +1,132 @@
+"""llama default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Llama4 - same as Llama3.1 except naming has changed slightly
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="llama-4",
+        system_template="",
+        system_message="",
+        roles={
+            "user": "<|header_start|>user",
+            "assistant": "<|header_start|>assistant",
+            "tool": "<|header_start|>ipython",
+        },
+        seps=["<|eot|>"],
+        role_content_sep="<|header_end|>\n\n",
+        role_empty_sep="<|header_end|>\n\n",
+        stop_str=[],
+        stop_token_ids=[
+            200001,
+            200007,
+            200008,
+        ],  # "<|end_of_text|>", "<|eom|>", "<|eot|>"
+        system_prefix_token_ids=[200000],  # "<|begin_of_text|>"
+        add_role_after_system_message=False,
+    )
+)
+
+# Llama3.1 -- same as Llama3 except stop token ids and stop str
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="llama-3_1",
+        system_template=(
+            "<|start_header_id|>system<|end_header_id|>\n\n"
+            f"{MessagePlaceholders.SYSTEM.value}<|eot_id|>"
+        ),
+        system_message="You are a helpful, respectful and honest assistant.",
+        roles={
+            "user": "<|start_header_id|>user",
+            "assistant": "<|start_header_id|>assistant",
+            "tool": "<|start_header_id|>ipython",
+        },
+        seps=["<|eot_id|>"],
+        role_content_sep="<|end_header_id|>\n\n",
+        role_empty_sep="<|end_header_id|>\n\n",
+        stop_str=[],
+        stop_token_ids=[
+            128001,
+            128008,
+            128009,
+        ],  # "<|end_of_text|>", "<|eom_id|>", "<|eot_id|>"
+        system_prefix_token_ids=[128000],  # "<|begin_of_text|>"
+        add_role_after_system_message=True,
+    )
+)
+
+# Llama3
+# See https://github.com/meta-llama/llama3?tab=readme-ov-file#instruction-tuned-models
+# and https://github.com/meta-llama/llama3/blob/main/llama/tokenizer.py
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="llama-3",
+        system_template=(
+            "<|start_header_id|>system<|end_header_id|>\n\n"
+            f"{MessagePlaceholders.SYSTEM.value}<|eot_id|>"
+        ),
+        system_message="You are a helpful, respectful and honest assistant.",
+        roles={
+            "user": "<|start_header_id|>user",
+            "assistant": "<|start_header_id|>assistant",
+        },
+        seps=["<|eot_id|>"],
+        role_content_sep="<|end_header_id|>\n\n",
+        role_empty_sep="<|end_header_id|>\n\n",
+        stop_str=["<|end_of_text|>", "<|eot_id|>"],
+        stop_token_ids=[128001, 128009],  # "<|end_of_text|>", "<|eot_id|>"
+        system_prefix_token_ids=[128000],  # "<|begin_of_text|>"
+        add_role_after_system_message=True,
+    )
+)
+
+# Llama2
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="llama-2",
+        system_template=f"[INST] <<SYS>>\n{MessagePlaceholders.SYSTEM.value}\n<</SYS>>\n\n",
+        system_message="You are a helpful, respectful and honest assistant.",
+        roles={"user": "<s>[INST]", "assistant": "[/INST]", "tool": "[INST]"},
+        seps=[" ", " </s>"],
+        role_content_sep=" ",
+        role_empty_sep=" ",
+        stop_str=["[INST]"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+        add_role_after_system_message=False,
+    )
+)
+
+# CodeLlama Completion
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="codellama_completion",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "", "assistant": ""},
+        seps=[""],
+        role_content_sep="",
+        role_empty_sep="",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)
+
+# CodeLlama Instruct
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="codellama_instruct",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "[INST]", "assistant": "[/INST]"},
+        seps=[" "],
+        role_content_sep=" ",
+        role_empty_sep=" ",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)

python/mlc_llm/conversation_template/llava.py

@@ -0,0 +1,22 @@
+"""Llava default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Llava
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="llava",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="\n",
+        roles={"user": "USER", "assistant": "ASSISTANT"},
+        seps=[" "],
+        role_content_sep=": ",
+        role_empty_sep=":",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+        add_role_after_system_message=False,
+    )
+)

python/mlc_llm/conversation_template/llm_jp.py

@@ -0,0 +1,25 @@
+"""LLM-jp default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# LLM-jp instruct
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="llm-jp",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。",
+        roles={
+            "user": "\n\n### 指示:",
+            "assistant": "\n\n### 応答:",
+        },
+        seps=["", "</s>"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=[],
+        stop_token_ids=[2],  # eos_token_id
+        system_prefix_token_ids=[1],  # bos_token_id (<s>)
+        add_role_after_system_message=True,
+    )
+)

python/mlc_llm/conversation_template/ministral3.py

@@ -0,0 +1,69 @@
+"""Ministral3 templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Ministral3
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="ministral3",
+        system_template=(
+            f"[SYSTEM_PROMPT]{MessagePlaceholders.SYSTEM.value}[/SYSTEM_PROMPT]"
+            f"{MessagePlaceholders.FUNCTION.value}"
+        ),
+        system_message=(
+            "You are Ministral-3-3B-Instruct-2512, a Large Language Model (LLM) created by "
+            "Mistral AI, a French startup headquartered in Paris.\n"
+            "You power an AI assistant called Le Chat.\n"
+            "Your knowledge base was last updated on 2023-10-01.\n"
+            "The current date is {today}.\n\n"
+            "When you're not sure about some information or when the user's request requires "
+            "up-to-date or specific data, you must use the available tools to fetch the "
+            "information. Do not hesitate to use tools whenever they can provide a more "
+            "accurate or complete response. If no relevant tools are available, then clearly "
+            "state that you don't have the information and avoid making up anything.\n"
+            "If the user's question is not clear, ambiguous, or does not provide enough "
+            "context for you to accurately answer the question, you do not try to answer it "
+            'right away and you rather ask the user to clarify their request (e.g. "What are '
+            'some good restaurants around me?" => "Where are you?" or "When is the next '
+            'flight to Tokyo" => "Where do you travel from?").\n'
+            "You are always very attentive to dates, in particular you try to resolve dates "
+            '(e.g. "yesterday" is {yesterday}) and when asked about information at specific '
+            "dates, you discard information that is at another date.\n"
+            "You follow these instructions in all languages, and always respond to the user in "
+            "the language they use or request.\n"
+            "Next sections describe the capabilities that you have.\n\n"
+            "# WEB BROWSING INSTRUCTIONS\n\n"
+            "You cannot perform any web search or access internet to open URLs, links etc. If "
+            "it seems like the user is expecting you to do so, you clarify the situation and "
+            "ask the user to copy paste the text directly in the chat.\n\n"
+            "# MULTI-MODAL INSTRUCTIONS\n\n"
+            "You have the ability to read images, but you cannot generate images. You also "
+            "cannot transcribe audio files or videos.\n"
+            "You cannot read nor transcribe audio files or videos.\n\n"
+            "# TOOL CALLING INSTRUCTIONS\n\n"
+            "You may have access to tools that you can use to fetch information or perform "
+            "actions. You must use these tools in the following situations:\n\n"
+            "1. When the request requires up-to-date information.\n"
+            "2. When the request requires specific data that you do not have in your knowledge "
+            "base.\n"
+            "3. When the request involves actions that you cannot perform without tools.\n\n"
+            "Always prioritize using tools to provide the most accurate and helpful response. "
+            "If tools are not available, inform the user that you cannot perform the requested "
+            "action at the moment."
+        ),
+        role_templates={
+            "user": f"[INST]{MessagePlaceholders.USER.value}[/INST]",
+            "assistant": f"{MessagePlaceholders.ASSISTANT.value}</s>",
+            "tool": f"[TOOL_RESULTS]{MessagePlaceholders.TOOL.value}[/TOOL_RESULTS]",
+        },
+        roles={"user": "", "assistant": "", "tool": ""},
+        seps=[""],
+        role_content_sep="",
+        role_empty_sep="",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)

python/mlc_llm/conversation_template/ministral3_reasoning.py

@@ -0,0 +1,38 @@
+"""Ministral3 reasoning templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Ministral-3-XB-Reasoning-2512
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="ministral3_reasoning",
+        system_template=(
+            f"[SYSTEM_PROMPT]{MessagePlaceholders.SYSTEM.value}[/SYSTEM_PROMPT]"
+            f"{MessagePlaceholders.FUNCTION.value}"
+        ),
+        system_message=(
+            "# HOW YOU SHOULD THINK AND ANSWER\n\n"
+            "First draft your thinking process (inner monologue) until you arrive at a response. "
+            "Format your response using Markdown, and use LaTeX for any mathematical equations. "
+            "Write both your thoughts and the response in the same language as the input.\n\n"
+            "Your thinking process must follow the template below:"
+            "[THINK]Your thoughts or/and draft, like working through an exercise on scratch paper. "
+            "Be as casual and as long as you want until you are confident to generate the response "
+            "to the user.[/THINK]Here, provide a self-contained response."
+        ),
+        role_templates={
+            "user": f"[INST]{MessagePlaceholders.USER.value}[/INST]",
+            "assistant": f"{MessagePlaceholders.ASSISTANT.value}</s>",
+            "tool": f"[TOOL_RESULTS]{MessagePlaceholders.TOOL.value}[/TOOL_RESULTS]",
+        },
+        roles={"user": "", "assistant": "", "tool": ""},
+        seps=[""],
+        role_content_sep="",
+        role_empty_sep="",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)

python/mlc_llm/conversation_template/mistral.py

@@ -0,0 +1,24 @@
+"""Mistral default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Mistral default
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="mistral_default",
+        system_template=f"[INST] {MessagePlaceholders.SYSTEM.value}",
+        system_message="Always assist with care, respect, and truth. Respond with utmost "
+        "utility yet securely. Avoid harmful, unethical, prejudiced, or negative content. "
+        "Ensure replies promote fairness and positivity.",
+        roles={"user": "[INST]", "assistant": "[/INST]", "tool": "[INST]"},
+        seps=[" "],
+        role_content_sep=" ",
+        role_empty_sep="",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+        add_role_after_system_message=False,
+    )
+)

python/mlc_llm/conversation_template/nemotron.py

@@ -0,0 +1,27 @@
+"""nemotron default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Nemotron template
+# https://huggingface.co/nvidia/Nemotron-Mini-4B-Instruct/blob/6a417790c444fd65a3da6a5c8821de6afc9654a6/tokenizer_config.json#L8030
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="nemotron",
+        system_template=(f"<extra_id_0>System\n{MessagePlaceholders.SYSTEM.value}\n\n"),
+        system_message="",
+        roles={
+            "user": "<extra_id_1>User",
+            "assistant": "<extra_id_1>Assistant",
+            "tool": "<extra_id_1>Tool",
+        },
+        seps=["\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["</s>"],
+        stop_token_ids=[3],
+        system_prefix_token_ids=[2],
+        add_role_after_system_message=True,
+    )
+)

python/mlc_llm/conversation_template/oasst.py

@@ -0,0 +1,20 @@
+"""Oasst default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Oasst
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="oasst",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "<|prompter|>", "assistant": "<|assistant|>"},
+        seps=["<|endoftext|>"],
+        role_content_sep=": ",
+        role_empty_sep=": ",
+        stop_str=["<|endoftext|>"],
+        stop_token_ids=[2],
+    )
+)

python/mlc_llm/conversation_template/olmo.py

@@ -0,0 +1,28 @@
+"""OLMo default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Note that eos_token id is "50279" both in Allenai and AMD version.
+# So use the number instead of text.
+# Allenai version chat_template and eos_token:
+# https://huggingface.co/allenai/OLMo-7B-Instruct/blob/main/tokenizer_config.json
+# AMD version chat_template and eos_token:
+# https://huggingface.co/amd/AMD-OLMo-1B-SFT-DPO/blob/main/tokenizer_config.json
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="olmo",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        system_prefix_token_ids=[50279],
+        roles={
+            "user": "<|user|>",
+            "assistant": "<|assistant|>",
+        },
+        seps=["\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_token_ids=[50279],
+    )
+)

python/mlc_llm/conversation_template/olmo2.py

@@ -0,0 +1,24 @@
+"""OLMo2 default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# OLMo-2 Instruct (Tulu format)
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="olmo2",
+        system_template=f"<|system|>\n{MessagePlaceholders.SYSTEM.value}\n",
+        system_message="",
+        roles={
+            "user": "<|user|>",
+            "assistant": "<|assistant|>",
+        },
+        seps=["<|endoftext|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|endoftext|>"],
+        stop_token_ids=[100257],
+        system_prefix_token_ids=[100257],
+    )
+)

python/mlc_llm/conversation_template/orion.py

@@ -0,0 +1,21 @@
+"""Orion default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Orion
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="orion",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "Human: ", "assistant": "Assistant: "},
+        seps=["\n\n", "</s>"],
+        role_content_sep="",
+        role_empty_sep="</s>",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)

python/mlc_llm/conversation_template/phi.py

@@ -0,0 +1,70 @@
+"""Phi default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Phi-2
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="phi-2",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "Instruct", "assistant": "Output"},
+        seps=["\n"],
+        role_content_sep=": ",
+        role_empty_sep=":",
+        stop_str=["<|endoftext|>"],
+        stop_token_ids=[50256],
+    )
+)
+
+# Phi-3
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="phi-3",
+        system_template=f"<|system|>\n{MessagePlaceholders.SYSTEM.value}",
+        system_message="You are a helpful digital assistant. Please provide safe, "
+        "ethical and accurate information to the user.",
+        roles={"user": "<|user|>", "assistant": "<|assistant|>"},
+        seps=["<|end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        system_prefix_token_ids=[1],
+        stop_str=["<|endoftext|>"],
+        stop_token_ids=[2, 32000, 32001, 32007],
+    )
+)
+
+# Phi-3-vision
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="phi-3-vision",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "<|user|>", "assistant": "<|assistant|>"},
+        seps=["<|end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        system_prefix_token_ids=[1],
+        stop_str=["<|endoftext|>"],
+        stop_token_ids=[2, 32000, 32001, 32007],
+    )
+)
+
+# Phi-4
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="phi-4",
+        system_template=f"<|system|>\n{MessagePlaceholders.SYSTEM.value}",
+        system_message="You are a helpful digital assistant. Please provide safe, "
+        "ethical and accurate information to the user.",
+        roles={"user": "<|user|>", "assistant": "<|assistant|>"},
+        seps=["<|end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        system_prefix_token_ids=[200022],  # <|system|>
+        stop_str=["<|endoftext|>", "<|end|>"],
+        stop_token_ids=[199999, 200020],  # <|endoftext|>, <|end|>
+    )
+)

python/mlc_llm/conversation_template/qwen2.py

@@ -0,0 +1,20 @@
+"""Qwen2 default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Same as chatml except system message, stop token, and stop string
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="qwen2",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message="You are a helpful assistant.",
+        roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"},
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|endoftext|>", "<|im_end|>"],
+        stop_token_ids=[151643, 151645],
+    )
+)

python/mlc_llm/conversation_template/qwen3.py

@@ -0,0 +1,26 @@
+"""Qwen3 conversation template.
+
+Matches Qwen2's ChatML structure but strips `<think>...</think>` blocks from
+historical assistant turns, mirroring Qwen3's official HF chat template. Small
+Qwen3 variants (e.g. 0.6B) otherwise emit `<|im_end|>` prematurely when their
+own thinking traces are echoed back in multi-turn context (see mlc-ai/mlc-llm#3482).
+"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="qwen3",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message="You are a helpful assistant.",
+        roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"},
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|endoftext|>", "<|im_end|>"],
+        stop_token_ids=[151643, 151645],
+        strip_reasoning_in_history=True,
+    )
+)

python/mlc_llm/conversation_template/qwen3_5.py

@@ -0,0 +1,45 @@
+"""Qwen3.5 conversation templates.
+
+qwen3_5: Thinking enabled — assistant prefix opens a <think> block for the model
+         to reason in before responding.
+qwen3_5_nothink: Thinking disabled — assistant prefix includes a closed empty
+                 <think> block so the model skips straight to responding.
+"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="qwen3_5",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message="You are a helpful assistant.",
+        roles={
+            "user": "<|im_start|>user",
+            "assistant": "<|im_start|>assistant\n<think>",
+        },
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|endoftext|>", "<|im_end|>"],
+        stop_token_ids=[248046, 248044],
+    )
+)
+
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="qwen3_5_nothink",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message="You are a helpful assistant.",
+        roles={
+            "user": "<|im_start|>user",
+            "assistant": "<|im_start|>assistant\n<think>\n\n</think>\n",
+        },
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|endoftext|>", "<|im_end|>"],
+        stop_token_ids=[248046, 248044],
+    )
+)

python/mlc_llm/conversation_template/redpajama.py

@@ -0,0 +1,20 @@
+"""RedPajama default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# RedPajama Chat
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="redpajama_chat",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "<human>", "assistant": "<bot>"},
+        seps=["\n"],
+        role_content_sep=": ",
+        role_empty_sep=":",
+        stop_str=["<human>"],
+        stop_token_ids=[0],
+    )
+)

python/mlc_llm/conversation_template/registry.py

@@ -0,0 +1,85 @@
+"""The conversation template registry and presets in MLC LLM"""
+
+from typing import ClassVar, Dict, Optional  # noqa: UP035
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+
+class ConvTemplateRegistry:
+    """Global conversation template registry for preset templates."""
+
+    _conv_templates: ClassVar[Dict[str, Conversation]] = {}  # noqa: UP006
+
+    @staticmethod
+    def register_conv_template(conv_template: Conversation, override: bool = False) -> None:
+        """Register a new conversation template in the global registry.
+        Using `override = True` to override the previously registered
+        template with the same name.
+        """
+        name = conv_template.name
+        if name is None:
+            raise ValueError("The template to register should have non-None name.")
+        if name in ConvTemplateRegistry._conv_templates and not override:
+            raise ValueError(
+                "The name of the template has been registered "
+                f"for {ConvTemplateRegistry._conv_templates[name].model_dump_json(by_alias=True)}"
+            )
+        ConvTemplateRegistry._conv_templates[name] = conv_template
+
+    @staticmethod
+    def get_conv_template(name: str) -> Optional[Conversation]:
+        """Return the conversation template specified by the given name,
+        or None if the template is not registered.
+        """
+        return ConvTemplateRegistry._conv_templates.get(name, None)
+
+
+# ChatML
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="chatml",
+        system_template=f"<|im_start|>system\n{MessagePlaceholders.SYSTEM.value}<|im_end|>\n",
+        system_message=(
+            "A conversation between a user and an LLM-based AI assistant. The "
+            "assistant gives helpful and honest answers."
+        ),
+        roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"},
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|im_end|>"],
+        stop_token_ids=[2],
+    )
+)
+
+# ChatML without a system prompt
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="chatml_nosystem",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"},
+        seps=["<|im_end|>\n"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|im_end|>"],
+        stop_token_ids=[2],
+    )
+)
+
+
+# Vanilla LM
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="LM",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "", "assistant": ""},
+        seps=[""],
+        role_content_sep="",
+        role_empty_sep="",
+        stop_str=[],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)

python/mlc_llm/conversation_template/rwkv.py

@@ -0,0 +1,24 @@
+"""RWKV default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# RWKV World
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="rwkv_world",
+        system_template=f"User: hi\n\nAssistant: {MessagePlaceholders.SYSTEM.value}",
+        system_message=(
+            "Hi. I am your assistant and I will provide expert full response "
+            "in full details. Please feel free to ask any question and I will "
+            "always answer it."
+        ),
+        roles={"user": "User", "assistant": "Assistant"},
+        seps=["\n\n"],
+        role_content_sep=": ",
+        role_empty_sep=": ",
+        stop_str=["\n\n"],
+        stop_token_ids=[0],
+    )
+)

python/mlc_llm/conversation_template/stablelm.py

@@ -0,0 +1,59 @@
+"""StableLM default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# StableLM Tuned Alpha
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="stablelm",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message=(
+            "<|SYSTEM|># StableLM Tuned (Alpha version)\n"
+            "- StableLM is a helpful and harmless open-source AI language model developed by "
+            "StabilityAI.\n"
+            "- StableLM is excited to be able to help the user, but will refuse to do "
+            "anything that could be considered harmful to the user.\n"
+            "- StableLM is more than just an information source, StableLM is also able to "
+            "write poetry, short stories, and make jokes.\n"
+            "- StableLM will refuse to participate in anything that could harm a human."
+        ),
+        roles={"user": "<|USER|>", "assistant": "<|ASSISTANT|>"},
+        seps=[""],
+        role_content_sep=": ",
+        role_empty_sep=": ",
+        stop_str=[""],
+        stop_token_ids=[50278, 50279, 50277, 1, 0],
+    )
+)
+
+# StableLM 3B
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="stablelm-3b",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "<|user|>", "assistant": "<|assistant|>"},
+        seps=["<|endoftext|>", "<|endoftext|>"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|endoftext|>"],
+        stop_token_ids=[0],
+    )
+)
+
+# StableLM-2
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="stablelm-2",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "<|user|>", "assistant": "<|assistant|>"},
+        seps=["<|endoftext|>", "<|endoftext|>"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["<|endoftext|>"],
+        stop_token_ids=[100257],
+    )
+)

python/mlc_llm/conversation_template/tinyllama.py

@@ -0,0 +1,20 @@
+"""Tiny Llama default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# TinyLlama v1.0
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="tinyllama_v1_0",
+        system_template=f"<|system|>\n{MessagePlaceholders.SYSTEM.value}</s>",
+        system_message="You are a helpful chatbot.",
+        roles={"user": "<|user|>", "assistant": "<|assistant|>"},
+        seps=["</s>"],
+        role_content_sep="\n",
+        role_empty_sep="\n",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+    )
+)

python/mlc_llm/conversation_template/wizardlm.py

@@ -0,0 +1,40 @@
+"""WiazrdLM and Coder default templates"""
+
+from mlc_llm.protocol.conversation_protocol import Conversation, MessagePlaceholders
+
+from .registry import ConvTemplateRegistry
+
+# Wizard LM 7B
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="wizardlm_7b",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message="",
+        roles={"user": "User", "assistant": "Response"},
+        seps=["###"],
+        role_content_sep=": ",
+        role_empty_sep=":",
+        stop_str=["###"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)
+
+# WizardCoder or WizardMath
+ConvTemplateRegistry.register_conv_template(
+    Conversation(
+        name="wizard_coder_or_math",
+        system_template=f"{MessagePlaceholders.SYSTEM.value}",
+        system_message=(
+            "Below is an instruction that describes a task. Write a response that appropriately "
+            "completes the request."
+        ),
+        roles={"user": "Instruction", "assistant": "Response"},
+        seps=["\n\n### ", "\n\n### "],
+        role_content_sep=":\n",
+        role_empty_sep=":\n",
+        stop_str=["</s>"],
+        stop_token_ids=[2],
+        system_prefix_token_ids=[1],
+    )
+)

python/mlc_llm/interface/calibrate.py

@@ -0,0 +1,172 @@
+"""Python entrypoint for calibration."""
+
+import asyncio
+import json
+import random
+from collections.abc import Mapping
+from typing import List, Optional, Tuple  # noqa: UP035
+
+import numpy as np
+import tqdm.asyncio
+import tvm
+from tvm.contrib import tvmjs
+
+from mlc_llm.serve.engine import AsyncMLCEngine, EngineConfig
+from mlc_llm.tokenizers import Tokenizer
+
+
+class CalibrationObserver:
+    """A singleton class to observe the calibration parameters.""" ""
+
+    instance: "CalibrationObserver" = None
+
+    params: Mapping[str, tvm.runtime.Tensor] = {}
+
+    @staticmethod
+    def get():
+        """Get the singleton instance of the class.""" ""
+        if CalibrationObserver.instance is None:
+            CalibrationObserver.instance = CalibrationObserver()
+        return CalibrationObserver.instance
+
+    @tvm.register_global_func("mlc_llm.calibration_observer")
+    @staticmethod
+    def callback(
+        name: str,
+        mode: str,
+        value: "tvm.runtime.Tensor",
+        out_value: "tvm.runtime.Tensor",
+    ):
+        """The callback function to update the saved calibration parameters."""
+        instance = CalibrationObserver.get()
+        if mode == "max":
+            reducer = np.maximum
+        else:
+            raise NotImplementedError(f"Unsupported calibration mode: {mode}")
+        if name in instance.params:
+            instance.params[name] = reducer(instance.params[name], value.numpy())
+        else:
+            instance.params[name] = value.numpy()
+        out_value.copyfrom(instance.params[name])
+
+    def save_params(self, output: str):
+        """Save the calibration parameters to the given output directory."""
+        tvmjs.dump_tensor_cache(
+            self.params,
+            output,
+            encode_format="f32-to-bf16",
+            meta_data=None,
+            show_progress=False,
+            update_if_exists=True,
+        )
+
+
+def sample_requests(
+    dataset_path: str,
+    num_requests: int,
+    tokenizer: Tokenizer,
+) -> List[Tuple[str, int, int]]:  # noqa: UP006
+    """Sample the requests from the given dataset."""
+    # Load the dataset.
+    with open(dataset_path, encoding="utf-8") as f:
+        dataset = json.load(f)
+
+    # Filter out the conversations with less than 2 turns.
+    dataset = [data for data in dataset if len(data["conversations"]) >= 2]
+    # Only keep the first two turns of each conversation.
+    dataset = [
+        (data["conversations"][0]["value"], data["conversations"][1]["value"]) for data in dataset
+    ]
+    prompts = [prompt for prompt, _ in dataset]
+    prompt_token_ids = tokenizer.encode_batch(prompts)
+    completions = [completion for _, completion in dataset]
+    completion_token_ids = tokenizer.encode_batch(completions)
+    tokenized_dataset: List[Tuple[str, List[int], int]] = []  # noqa: UP006
+    for i in range(len(dataset)):
+        output_len = len(completion_token_ids[i])
+        tokenized_dataset.append((prompts[i], prompt_token_ids[i], output_len))
+
+    # Filter out too long sequences.
+    filtered_dataset: List[Tuple[str, int, int]] = []  # noqa: UP006
+    for prompt, token_ids, output_len in tokenized_dataset:
+        prompt_len = len(token_ids)
+        if prompt_len < 4 or output_len < 4:
+            # Prune too short sequences.
+            continue
+        if prompt_len > 1024 or prompt_len + output_len > 2048:
+            # Prune too long sequences.
+            continue
+        filtered_dataset.append((prompt, prompt_len, output_len))
+
+    # Sample the requests.
+    sampled_requests = random.sample(filtered_dataset, num_requests)
+    return sampled_requests
+
+
+async def send_calibration_requests(
+    async_engine: AsyncMLCEngine,
+    sampled_requests: List[Tuple[str, int, int]],  # noqa: UP006
+    max_concurrent_requests: int,
+) -> None:
+    """Send the calibration requests to the engine."""
+    tasks = []
+
+    semaphore = asyncio.Semaphore(max_concurrent_requests)
+
+    async def generate_task(request_idx):
+        async with semaphore:
+            prompt, _, output_len = sampled_requests[request_idx]
+            await async_engine.chat.completions.create(
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=output_len,
+                request_id=str(request_idx),
+            )
+
+    for i in range(len(sampled_requests)):
+        task = asyncio.create_task(generate_task(i))
+        tasks.append(task)
+    await tqdm.asyncio.tqdm.gather(*tasks)
+
+
+def calibrate(
+    model: str,
+    device: str,
+    model_lib: Optional[str],
+    dataset: str,
+    output: str,
+    num_calibration_samples: int,
+    *,
+    seed: int,
+    max_num_sequence: Optional[int] = None,
+    max_total_sequence_length: Optional[int] = None,
+    prefill_chunk_size: Optional[int] = None,
+    max_history_size: Optional[int] = None,
+    gpu_memory_utilization: Optional[float] = None,
+) -> None:
+    """Calibrate the quantized model using the given dataset."""
+    random.seed(seed)
+    async_engine = AsyncMLCEngine(
+        model=model,
+        device=device,
+        model_lib=model_lib,
+        mode="server",
+        engine_config=EngineConfig(
+            max_num_sequence=max_history_size,
+            max_total_sequence_length=max_total_sequence_length,
+            prefill_chunk_size=prefill_chunk_size,
+            max_history_size=max_history_size,
+            gpu_memory_utilization=gpu_memory_utilization,
+        ),
+    )
+    sampled_requests = sample_requests(dataset, num_calibration_samples, async_engine.tokenizer)
+    asyncio.run(
+        send_calibration_requests(
+            async_engine,
+            sampled_requests,
+            max_concurrent_requests=max_num_sequence or 32,
+        )
+    )
+    async_engine.terminate()
+
+    calibrator = CalibrationObserver.get()
+    calibrator.save_params(output)

python/mlc_llm/interface/chat.py

@@ -0,0 +1,311 @@
+"""Python entrypoint of chat."""
+
+import dataclasses
+from typing import Any, Dict, List, Optional, Union  # noqa: UP035
+
+from prompt_toolkit import prompt as get_prompt
+from prompt_toolkit.key_binding import KeyBindings
+
+from mlc_llm.json_ffi import JSONFFIEngine
+from mlc_llm.protocol import openai_api_protocol
+from mlc_llm.serve.config import EngineConfig
+from mlc_llm.serve.engine import MLCEngine
+from mlc_llm.serve.engine_base import _query_engine_metrics
+from mlc_llm.support import argparse
+from mlc_llm.support.config import ConfigOverrideBase
+
+
+def _print_help_str():
+    help_str = """You can use the following special commands:
+  /help               print the special commands
+  /exit               quit the cli
+  /stats              print out stats of last request (token/sec)
+  /metrics            print out full engine metrics
+  /reset              restart a fresh chat
+  /set [overrides]    override settings in the generation config. For example,
+                      `/set temperature=0.5;top_p=0.8;seed=23;max_tokens=100;stop=str1,str2`
+                      Note: Separate stop words in the `stop` option with commas (,).
+  Multi-line input: Use escape+enter to start a new line.
+"""
+    print(help_str)
+
+
+def _set_up_key_bindings():
+    kb = KeyBindings()
+
+    @kb.add("escape", "enter")
+    def _(event):
+        event.current_buffer.insert_text("\n")
+
+    @kb.add("enter")
+    def _(event):
+        event.current_buffer.validate_and_handle()
+
+    return kb
+
+
+@dataclasses.dataclass
+class ChatCompletionOverride(ConfigOverrideBase):
+    """Flags for overriding chat completions."""
+
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    frequency_penalty: Optional[float] = None
+    presence_penalty: Optional[float] = None
+    max_tokens: Optional[int] = None
+    seed: Optional[int] = None
+    stop: Optional[Union[str, List[str]]] = None  # noqa: UP006
+
+    @staticmethod
+    def from_str(source: str) -> "ChatCompletionOverride":
+        """Parse model config override values from a string."""
+        parser = argparse.ArgumentParser(description="chat completion override values")
+        parser.add_argument("--temperature", type=float, default=None)
+        parser.add_argument("--top_p", type=float, default=None)
+        parser.add_argument("--frequency_penalty", type=float, default=None)
+        parser.add_argument("--presence_penalty", type=float, default=None)
+        parser.add_argument("--max_tokens", type=int, default=None)
+        parser.add_argument("--seed", type=int, default=None)
+        parser.add_argument("--stop", type=str, default=None)
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return ChatCompletionOverride(
+            temperature=results.temperature,
+            top_p=results.top_p,
+            frequency_penalty=results.frequency_penalty,
+            presence_penalty=results.presence_penalty,
+            max_tokens=results.max_tokens,
+            seed=results.seed,
+            stop=results.stop.split(",") if results.stop is not None else None,
+        )
+
+
+@dataclasses.dataclass
+class ModelConfigOverride(ConfigOverrideBase):
+    """Flags for overriding model config."""
+
+    context_window_size: Optional[int] = None
+    sliding_window_size: Optional[int] = None
+    prefill_chunk_size: Optional[int] = None
+    attention_sink_size: Optional[int] = None
+    tensor_parallel_shards: Optional[int] = None
+    pipeline_parallel_stages: Optional[int] = None
+    opt: Optional[str] = None
+
+    @staticmethod
+    def from_str(source: str) -> "ModelConfigOverride":
+        """Parse model config override values from a string."""
+        parser = argparse.ArgumentParser(description="model config override values")
+        parser.add_argument("--tensor_parallel_shards", type=int, default=None)
+        parser.add_argument("--pipeline_parallel_stages", type=int, default=None)
+        parser.add_argument("--opt", type=str, default=None)
+        parser.add_argument("--context_window_size", type=int, default=None)
+        parser.add_argument("--sliding_window_size", type=int, default=None)
+        parser.add_argument("--prefill_chunk_size", type=int, default=None)
+        parser.add_argument("--attention_sink_size", type=int, default=None)
+
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return ModelConfigOverride(
+            tensor_parallel_shards=results.tensor_parallel_shards,
+            pipeline_parallel_stages=results.pipeline_parallel_stages,
+            opt=results.opt,
+            context_window_size=results.context_window_size,
+            sliding_window_size=results.sliding_window_size,
+            prefill_chunk_size=results.prefill_chunk_size,
+            attention_sink_size=results.attention_sink_size,
+        )
+
+
+class ChatState:
+    """Simple helper class to manage chat state.
+
+    Chat state wraps around a  engine instance
+    and exposes the minimum set of tools to perform
+    interactive chat. It provides support for mlc_llm chat.
+    It also can be used to do interactive debugging
+    with different engine instance.
+
+    Examples
+    --------
+    .. code:: python
+
+        from openai import OpenAI
+        from mlc_llm import MLCEngine
+        from mlc_llm.serve import PopenServer
+        from mlc_llm.interface.chat import ChatState
+
+        def chat_with_engine(model):
+            # hookup with MLCEngine
+            ChatState(MLCEngine(model)).chat()
+
+        def chat_with_server(model):
+            # hookup with AsyncMLCEngine backed api server
+            with PopenServer(model) as server:
+                ChatState(
+                    OpenAI(base_url=server.openai_v1_base_url, api_key="None")
+                ).chat()
+    """
+
+    history: List[Dict[str, Any]]  # noqa: UP006
+    history_begin: int
+    # kwargs passed to completions
+    overrides: ChatCompletionOverride
+    # Underlying engine
+    engine: Union[JSONFFIEngine, MLCEngine]
+    last_finished_request_usage: Optional[openai_api_protocol.CompletionUsage]
+
+    def __init__(self, engine: Union[JSONFFIEngine, MLCEngine]):
+        self.engine = engine
+        self.history = []
+        self.history_window_begin = 0
+        self.overrides = ChatCompletionOverride()
+        # model is mainly used for compact reasons
+        self.model = "chat_model"
+        self.last_finished_request_usage = None
+
+    def slide_history(self):
+        """Slide history to fit into context window"""
+        history_window_size = len(self.history) - self.history_window_begin
+        assert history_window_size % 2 == 0
+        self.history_window_begin += ((history_window_size + 3) // 4) * 2
+
+    def process_system_prompts(self):
+        """Process system prompts"""
+        # TODO(mlc-team): possibly leverage debug option
+        # pass a simple prompt to warm up
+        for _ in self.engine.chat.completions.create(
+            messages=[{"role": "user", "content": ""}],
+            max_tokens=1,
+            model=self.model,
+            stream=True,
+        ):
+            pass
+
+    def generate(self, prompt: str):
+        """Run one generation with the prompt.
+
+        Parameters
+        ----------
+        prompt: str
+            The input prompt
+        """
+        self.history.append({"role": "user", "content": prompt})
+        output_text = ""
+        finish_reason_length = False
+        messages = self.history[self.history_window_begin :]
+
+        for response in self.engine.chat.completions.create(
+            messages=messages,
+            model=self.model,
+            stream=True,
+            stream_options={"include_usage": True},
+            **dataclasses.asdict(self.overrides),
+        ):
+            if response.usage is not None:
+                self.last_finished_request_usage = response.usage
+                continue
+            for choice in response.choices:
+                assert choice.delta.role == "assistant"
+                if isinstance(choice.delta.content, str):
+                    output_text += choice.delta.content
+                    print(choice.delta.content, end="", flush=True)
+                if choice.finish_reason == "length":
+                    finish_reason_length = True
+        if finish_reason_length:
+            print(" [output truncated due to context length limit...]")
+        # print additional \n when generation ends
+        print()
+        # record the history
+        self.history.append({"role": "assistant", "content": output_text})
+        if finish_reason_length:
+            self.slide_history()
+
+    def stats(self):
+        """Print statistics of the prefill and decode speed."""
+
+        def get_stats_text():
+            """Get text"""
+            if self.last_finished_request_usage is None:
+                return "N/A"
+            last_finished_request = self.last_finished_request_usage.extra
+            if last_finished_request is None:
+                return "N/A"
+            prefill_speed = last_finished_request.get("prefill_tokens_per_s", None)
+            decode_speed = last_finished_request.get("decode_tokens_per_s", None)
+            prefill_speed = f"{prefill_speed:.1f}" if prefill_speed is not None else "N/A"
+            decode_speed = f"{decode_speed:.1f}" if decode_speed is not None else "N/A"
+            return f"prefill: {prefill_speed} tok/s, decode: {decode_speed} tok/s"
+
+        print(get_stats_text(), flush=True)
+
+    def metrics(self):
+        """Print metrics as prometheus text"""
+        print(_query_engine_metrics(self.engine).prometheus_text(), flush=True)
+
+    def reset(self):
+        """Reset the chat history"""
+        self.history = []
+        self.history_window_begin = 0
+
+    def chat(self):
+        """Start an interactive chat session."""
+        _print_help_str()
+
+        self.process_system_prompts()
+        # Multi-line input support: set escape+enter as start a new line
+        kb = _set_up_key_bindings()
+
+        while True:
+            try:
+                prompt = get_prompt(
+                    ">>> ",
+                    key_bindings=kb,
+                    multiline=True,
+                )
+            except (KeyboardInterrupt, EOFError):
+                break
+            if prompt[:4] == "/set":
+                overrides = ChatCompletionOverride.from_str(prompt.split()[1])
+                for key, value in dataclasses.asdict(overrides).items():
+                    if value is not None:
+                        setattr(self.overrides, key, value)
+            elif prompt[:6] == "/stats":
+                self.stats()
+            elif prompt[:8] == "/metrics":
+                self.metrics()
+            elif prompt[:6] == "/reset":
+                self.reset()
+            elif prompt[:5] == "/exit":
+                break
+            elif prompt[:5] == "/help":
+                _print_help_str()
+            else:
+                self.generate(prompt)
+
+
+def chat(
+    model: str,
+    device: str,
+    model_lib: Optional[str],
+    overrides: ModelConfigOverride,
+):
+    """Chat cli entry"""
+    # By default we use JSONFFIEngine
+    engine = JSONFFIEngine(
+        model,
+        device,
+        model_lib=model_lib,
+        mode="interactive",
+        engine_config=EngineConfig(
+            max_single_sequence_length=overrides.context_window_size,
+            prefill_chunk_size=overrides.prefill_chunk_size,
+            sliding_window_size=overrides.sliding_window_size,
+            attention_sink_size=overrides.attention_sink_size,
+            tensor_parallel_shards=overrides.tensor_parallel_shards,
+            pipeline_parallel_stages=overrides.pipeline_parallel_stages,
+            opt=overrides.opt,
+        ),
+    )
+    try:
+        ChatState(engine).chat()
+    finally:
+        engine.terminate()

python/mlc_llm/interface/compile.py

@@ -0,0 +1,265 @@
+"""Python entrypoint of compilation."""
+
+import dataclasses
+from io import StringIO
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional, Tuple  # noqa: UP035
+
+from tvm import IRModule, relax, tirx
+from tvm.ir.transform import Pass, PassContext
+from tvm.relax.frontend import nn
+from tvm.target import Target
+
+from mlc_llm import compiler_pass as _  # noqa: F401
+from mlc_llm import op as op_ext
+from mlc_llm.cli.model_metadata import _report_memory_usage
+from mlc_llm.model import Model
+from mlc_llm.quantization import Quantization
+from mlc_llm.support import logging
+from mlc_llm.support.config import ConfigBase
+from mlc_llm.support.style import bold
+
+from .compiler_flags import ModelConfigOverride, OptimizationFlags
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class CompileArgs:
+    """Arguments to MLC LLM's compiler."""
+
+    config: Path
+    quantization: Quantization
+    model: Model
+    target: Target
+    opt: OptimizationFlags
+    build_func: Callable[[IRModule, "CompileArgs", Pass], None]
+    system_lib_prefix: str
+    output: Path
+    overrides: ModelConfigOverride
+    debug_dump: Optional[Path]
+
+    def __post_init__(self) -> None:
+        self.opt.update(self.target, self.quantization)
+
+    def display(self) -> None:
+        """Display the arguments to stdout."""
+        out = StringIO()
+        print(f"{bold('Compiling with arguments:')}", file=out)
+        print(f"  {bold('--config'):<25} {self.config}", file=out)
+        print(f"  {bold('--quantization'):<25} {self.quantization}", file=out)
+        print(f"  {bold('--model-type'):<25} {self.model.name}", file=out)
+        print(f"  {bold('--target'):<25} {self.target.export()}", file=out)
+        print(f"  {bold('--opt'):<25} {self.opt}", file=out)
+        print(f'  {bold("--system-lib-prefix"):<25} "{self.system_lib_prefix}"', file=out)
+        print(f"  {bold('--output'):<25} {self.output}", file=out)
+        print(f"  {bold('--overrides'):<25} {self.overrides}", file=out)
+        # As it's debug only, no need to display
+        # print(f"  {bold('--debug-dump'):<25} {self.debug_dump}", file=out)
+        print(out.getvalue().rstrip())
+
+
+def _apply_preproc_to_params_and_check_pipeline(
+    named_params: List[Tuple[str, nn.Parameter]],  # noqa: UP006
+    model_config,
+) -> Dict[str, tirx.PrimFunc]:  # noqa: UP006
+    extra_tirs: Dict[str, tirx.PrimFunc] = {}  # noqa: UP006
+    for name, param in named_params:
+        preprocs = param.attrs.get("preprocs", [])
+        shard_strategy = param.attrs.get("shard_strategy", None)
+        if shard_strategy is not None and model_config.tensor_parallel_shards > 1:
+            preprocs.append(
+                shard_strategy.gen_shard_info(
+                    shards=model_config.tensor_parallel_shards,
+                    weight=param,
+                )
+            )
+            if shard_strategy.name not in extra_tirs:
+                extra_tirs[shard_strategy.name] = shard_strategy.gen_tir(
+                    shards=model_config.tensor_parallel_shards,
+                    weight=param,
+                )
+        param.attrs["preprocs"] = preprocs
+
+        pipeline_parallel_stages = getattr(model_config, "pipeline_parallel_stages", 1)
+        if pipeline_parallel_stages != 1:
+            assert "pipeline_stages" in param.attrs, (
+                f'The pipeline stage is undefined for parameter "{name}" when the number '
+                f"of pipeline parallel stages is {pipeline_parallel_stages}"
+            )
+        param.attrs["pipeline_stages"] = (
+            [0]
+            if "pipeline_stages" not in param.attrs
+            else list(set(param.attrs["pipeline_stages"]))
+        )
+    return extra_tirs
+
+
+def _infer_kv_state_kind(model_type) -> str:
+    if "rwkv" in model_type:
+        return "rnn_state"
+    if "medusa" in model_type:
+        return "none"
+    if "qwen3_5" in model_type:
+        return "hybrid"
+    return "kv_cache"
+
+
+def _compile(args: CompileArgs, model_config: ConfigBase):
+    def _get_variable_bounds(model_config) -> Dict[str, int]:  # noqa: UP006
+        if hasattr(model_config, "sliding_window_size"):
+            return {
+                "rolling_cache_len": model_config.sliding_window_size,
+                "kv_seq_len": model_config.sliding_window_size + model_config.prefill_chunk_size,
+                "seq_len": model_config.prefill_chunk_size,
+                "batch_size": getattr(model_config, "max_batch_size", 1),
+            }
+        return {
+            "total_seq_len": model_config.context_window_size,
+            "seq_len": model_config.prefill_chunk_size,
+            "batch_size": getattr(model_config, "max_batch_size", 1),
+        }
+
+    def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]:  # noqa: UP006
+        return {
+            "name": name,
+            # Record dynamic shape as -1 (e.g. vocab_size)
+            "shape": [s if isinstance(s, int) else s.name for s in param.shape],
+            "dtype": str(param.dtype),
+            "preprocs": param.attrs["preprocs"],
+            "pipeline_stages": param.attrs.get("pipeline_stages", [0]),
+        }
+
+    logger.info("TOP LEVEL MODEL CONFIG BEFORE OVERRIDES: %s", str(model_config))
+    _kwargs = getattr(model_config, "kwargs", {})
+    model_config = args.overrides.apply(model_config)
+    with args.target:
+        op_ext.enable(
+            target=args.target,
+            flashinfer=args.opt.flashinfer,
+            faster_transformer=args.opt.faster_transformer,
+            cutlass=args.opt.cutlass,
+        )
+        # Step 1. Create the quantized model
+        logger.info("Creating model from: %s", model_config)
+        if (
+            args.quantization.kind == "ft-quant"
+            and hasattr(model_config, "tensor_parallel_shards")
+            and model_config.tensor_parallel_shards > 1
+        ):
+            raise NotImplementedError
+        if (
+            hasattr(args.quantization, "linear_weight_layout")
+            and args.quantization.linear_weight_layout == "KN"
+            and hasattr(model_config, "tensor_parallel_shards")
+            and model_config.tensor_parallel_shards > 1
+        ):
+            raise NotImplementedError(
+                "KN layout (q3f16_0 and q4f16_0) is not supported for tensor parallelism"
+            )
+        model, _ = args.model.quantize[args.quantization.kind](model_config, args.quantization)
+        # Step 2. Exporting the model to TVM
+        logger.info("Exporting the model to TVM compiler")
+        mod, named_params, ext_mods = model.export_tvm(
+            spec=model.get_default_spec(),
+            allow_extern=True,
+        )
+        # Step 3. Running relax compilation pipeline
+        logger.info("Running optimizations using TVM")
+        additional_tirs = _apply_preproc_to_params_and_check_pipeline(named_params, model_config)
+        variable_bounds = _get_variable_bounds(model_config)
+        cuda_graph_symbolic_capture_hints = {
+            "batch_decode": ["batch_size"],
+            "batch_decode_to_last_hidden_states": ["batch_size"],
+            "batch_verify": ["batch_size", "seq_len"],
+            "batch_verify_to_last_hidden_states": ["batch_size", "seq_len"],
+        }
+        avs = _kwargs.get("active_vocab_size", None)
+        if avs is not None and avs <= 0:
+            avs = None
+        metadata = {
+            "model_type": args.model.name,
+            "quantization": args.quantization.name,
+            "context_window_size": getattr(model_config, "context_window_size", -1),
+            "sliding_window_size": getattr(model_config, "sliding_window_size", -1),
+            "attention_sink_size": getattr(model_config, "attention_sink_size", -1),
+            "prefill_chunk_size": model_config.prefill_chunk_size,
+            "tensor_parallel_shards": model_config.tensor_parallel_shards,
+            "pipeline_parallel_stages": getattr(model_config, "pipeline_parallel_stages", 1),
+            "disaggregation": getattr(model_config, "disaggregation", False),
+            "kv_state_kind": _infer_kv_state_kind(args.model.name),
+            "max_batch_size": getattr(model_config, "max_batch_size", 1),
+            "active_vocab_size": avs,
+            "model_task": args.model.model_task,
+        }
+        if args.model.embedding_metadata:
+            metadata["embedding_metadata"] = dataclasses.asdict(args.model.embedding_metadata)
+        logger.info("Registering metadata: %s", metadata)
+        metadata["params"] = [_get_param_metadata(name, param) for name, param in named_params]
+        pass_config = {"relax.backend.use_cuda_graph": args.opt.cudagraph}
+        # TODO: Remove this workaround when the TVM CSE regression is fixed.
+        # Temporary workaround for TVM CSE regression that can produce
+        # dangling `cse_v*` vars during host codegen.
+        pass_config["tirx.disable_cse_tir"] = True
+
+        with PassContext(config=pass_config):
+            args.build_func(
+                mod,
+                args,
+                pipeline=relax.get_pipeline(
+                    "mlc_llm",
+                    target=args.target,
+                    flashinfer=args.opt.flashinfer,
+                    cublas_gemm=args.opt.cublas_gemm,
+                    faster_transformer=args.opt.faster_transformer,
+                    allreduce_strategy=args.opt.ipc_allreduce_strategy,
+                    variable_bounds=variable_bounds,
+                    cuda_graph_symbolic_capture_hints=cuda_graph_symbolic_capture_hints,
+                    additional_tirs=additional_tirs,
+                    ext_mods=ext_mods,
+                    metadata=metadata,
+                    debug_dump=args.debug_dump,
+                ),
+            )
+        _report_memory_usage(metadata=metadata, config=model_config)
+    logger.info("Generated: %s", bold(str(args.output)))
+
+
+def compile(
+    config: Dict[str, Any],  # noqa: UP006
+    quantization: Quantization,
+    model_type: Model,
+    target: Target,
+    opt: OptimizationFlags,
+    build_func: Callable[[IRModule, CompileArgs, Pass], None],
+    system_lib_prefix: str,
+    output: Path,
+    overrides: ModelConfigOverride,
+    debug_dump: Optional[Path] = None,
+):
+    """Compile a model given its configuration and quantization format to a specific target."""
+    avs = None
+    if "active_vocab_size" in config:
+        avs = config.pop("active_vocab_size")
+        logger.info("Active vocab size from input config: %s", str(avs))
+    if "model_config" in config:
+        model_config = config.pop("model_config")
+        model_config.update(config)
+        model_config = model_type.config.from_dict(model_config)
+    else:
+        model_config = model_type.config.from_dict(config)
+    model_config.kwargs = {"active_vocab_size": avs} if avs is not None else {}
+    args = CompileArgs(
+        model_config,
+        quantization,
+        model_type,
+        target,
+        opt,
+        build_func,
+        system_lib_prefix,
+        output,
+        overrides,
+        debug_dump,
+    )
+    args.display()
+    _compile(args, model_config)

python/mlc_llm/interface/compiler_flags.py

@@ -0,0 +1,227 @@
+"""Flags for overriding model config."""
+
+import dataclasses
+import enum
+from io import StringIO
+from typing import Optional
+
+from mlc_llm.support import argparse, logging
+from mlc_llm.support.config import ConfigOverrideBase
+
+logger = logging.getLogger(__name__)
+
+
+class IPCAllReduceStrategyType(enum.IntEnum):
+    """The all-reduce strategy."""
+
+    NONE = 0
+    ONESHOT = 1
+    TWOSHOT = 2
+    AUTO = 3
+
+
+@dataclasses.dataclass
+class OptimizationFlags:
+    """Optimization flags"""
+
+    flashinfer: bool = False
+    cublas_gemm: bool = False
+    faster_transformer: bool = False
+    cudagraph: bool = False
+    cutlass: bool = False
+    ipc_allreduce_strategy: IPCAllReduceStrategyType = IPCAllReduceStrategyType.NONE
+
+    def __repr__(self) -> str:
+        out = StringIO()
+        print(f"flashinfer={int(self.flashinfer)}", file=out, end="")
+        print(f";cublas_gemm={int(self.cublas_gemm)}", file=out, end="")
+        print(f";faster_transformer={int(self.faster_transformer)}", file=out, end="")
+        print(f";cudagraph={int(self.cudagraph)}", file=out, end="")
+        print(f";cutlass={int(self.cutlass)}", file=out, end="")
+        print(
+            f";ipc_allreduce_strategy={self.ipc_allreduce_strategy.name}",
+            file=out,
+            end="",
+        )
+        return out.getvalue().rstrip()
+
+    @staticmethod
+    def from_str(source: str) -> "OptimizationFlags":
+        """Parse optimization flags from a string."""
+
+        if source in OPT_FLAG_PRESET:
+            return OPT_FLAG_PRESET[source]
+
+        def boolean(value: str) -> bool:
+            if value == "0":
+                return False
+            if value == "1":
+                return True
+            raise ValueError(f"Invalid boolean value: {value}")
+
+        parser = argparse.ArgumentParser(description="optimization flags")
+        parser.add_argument("--flashinfer", type=boolean, default=True)
+        parser.add_argument("--cublas_gemm", type=boolean, default=False)
+        parser.add_argument("--faster_transformer", type=boolean, default=False)
+        parser.add_argument("--cudagraph", type=boolean, default=False)
+        parser.add_argument("--cutlass", type=boolean, default=False)
+        parser.add_argument(
+            "--ipc_allreduce_strategy",
+            type=str,
+            choices=["NONE", "ONESHOT", "TWOSHOT", "AUTO"],
+            default="NONE",
+        )
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return OptimizationFlags(
+            flashinfer=results.flashinfer,
+            cublas_gemm=results.cublas_gemm,
+            faster_transformer=results.faster_transformer,
+            cudagraph=results.cudagraph,
+            cutlass=results.cutlass,
+            ipc_allreduce_strategy=IPCAllReduceStrategyType[results.ipc_allreduce_strategy],
+        )
+
+    def update(self, target, quantization) -> None:
+        """Update optimization flags based on additional information."""
+
+        def _flashinfer(target) -> bool:
+            from mlc_llm.support.auto_target import (
+                detect_cuda_arch_list,
+            )
+
+            if not self.flashinfer:
+                return False
+            if target.kind.name != "cuda":
+                return False
+            arch_list = detect_cuda_arch_list(target)
+            for arch in arch_list:
+                if arch < 80:
+                    logger.warning("flashinfer is not supported on CUDA arch < 80")
+                    return False
+            return True
+
+        def _cublas_gemm(target, quantization) -> bool:
+            """correct cublas_gemm flag"""
+            if target.kind.name not in ["cuda", "rocm"]:
+                return False
+            if not (
+                quantization.name in ["q0f16", "q0bf16", "q0f32"]
+                or "e4m3" in quantization.name
+                or "e5m2" in quantization.name
+            ):
+                return False
+            return self.cublas_gemm
+
+        def _faster_transformer(target) -> bool:
+            """correct faster_transformer flag"""
+            if not target.kind.name == "cuda":
+                return False
+            return self.faster_transformer
+
+        def _cutlass(target) -> bool:
+            """correct cutlass flag"""
+            if not target.kind.name == "cuda":
+                return False
+            return self.cutlass
+
+        def _cudagraph(target) -> bool:
+            """correct cudagraph flag"""
+            if not target.kind.name == "cuda":
+                return False
+            return self.cudagraph
+
+        self.flashinfer = _flashinfer(target)
+        self.cublas_gemm = _cublas_gemm(target, quantization)
+        self.faster_transformer = _faster_transformer(target)
+        self.cutlass = _cutlass(target)
+        self.cudagraph = _cudagraph(target)
+
+
+@dataclasses.dataclass
+class ModelConfigOverride(ConfigOverrideBase):
+    """Flags for overriding model config."""
+
+    context_window_size: Optional[int] = None
+    sliding_window_size: Optional[int] = None
+    prefill_chunk_size: Optional[int] = None
+    attention_sink_size: Optional[int] = None
+    max_batch_size: Optional[int] = None
+    tensor_parallel_shards: Optional[int] = None
+    pipeline_parallel_stages: Optional[int] = None
+    disaggregation: Optional[bool] = None
+
+    def __repr__(self) -> str:
+        out = StringIO()
+        print(f"context_window_size={self.context_window_size}", file=out, end="")
+        print(f";sliding_window_size={self.sliding_window_size}", file=out, end="")
+        print(f";prefill_chunk_size={self.prefill_chunk_size}", file=out, end="")
+        print(f";attention_sink_size={self.attention_sink_size}", file=out, end="")
+        print(f";max_batch_size={self.max_batch_size}", file=out, end="")
+        print(f";tensor_parallel_shards={self.tensor_parallel_shards}", file=out, end="")
+        print(
+            f";pipeline_parallel_stages={self.pipeline_parallel_stages}",
+            file=out,
+            end="",
+        )
+        print(f";disaggregation={self.disaggregation}", file=out, end="")
+        return out.getvalue().rstrip()
+
+    @staticmethod
+    def from_str(source: str) -> "ModelConfigOverride":
+        """Parse model config override values from a string."""
+        parser = argparse.ArgumentParser(description="model config override values")
+        parser.add_argument("--context_window_size", type=int, default=None)
+        parser.add_argument("--sliding_window_size", type=int, default=None)
+        parser.add_argument("--prefill_chunk_size", type=int, default=None)
+        parser.add_argument("--attention_sink_size", type=int, default=None)
+        parser.add_argument("--max_batch_size", type=int, default=None)
+        parser.add_argument("--tensor_parallel_shards", type=int, default=None)
+        parser.add_argument("--pipeline_parallel_stages", type=int, default=None)
+        parser.add_argument(
+            "--disaggregation",
+            type=lambda x: str(x).lower() in ["true", "1", "yes", "True"],
+            default=None,
+        )
+        results = parser.parse_args([f"--{i}" for i in source.split(";") if i])
+        return ModelConfigOverride(
+            context_window_size=results.context_window_size,
+            sliding_window_size=results.sliding_window_size,
+            prefill_chunk_size=results.prefill_chunk_size,
+            attention_sink_size=results.attention_sink_size,
+            max_batch_size=results.max_batch_size,
+            tensor_parallel_shards=results.tensor_parallel_shards,
+            pipeline_parallel_stages=results.pipeline_parallel_stages,
+            disaggregation=results.disaggregation,
+        )
+
+
+OPT_FLAG_PRESET = {
+    "O0": OptimizationFlags(
+        flashinfer=False,
+        cublas_gemm=False,
+        cudagraph=False,
+    ),
+    "O1": OptimizationFlags(
+        flashinfer=False,
+        cublas_gemm=True,
+        faster_transformer=True,
+        cudagraph=False,
+        cutlass=True,
+    ),
+    "O2": OptimizationFlags(
+        flashinfer=True,
+        cublas_gemm=True,
+        faster_transformer=False,
+        cudagraph=True,
+        cutlass=True,
+        ipc_allreduce_strategy=IPCAllReduceStrategyType.NONE,
+    ),
+    "O3": OptimizationFlags(
+        flashinfer=True,
+        cublas_gemm=True,
+        faster_transformer=True,
+        cudagraph=True,
+        cutlass=True,
+        ipc_allreduce_strategy=IPCAllReduceStrategyType.AUTO,
+    ),
+}

python/mlc_llm/interface/convert_weight.py

@@ -0,0 +1,250 @@
+"""Python entrypoint of weight conversion."""
+
+import contextlib
+import dataclasses
+import math
+import os
+import tempfile
+from collections.abc import Iterator
+from io import StringIO
+from pathlib import Path
+from typing import Any, Dict, Optional, Tuple  # noqa: UP035
+
+from tvm import tirx
+from tvm.contrib import tvmjs
+from tvm.runtime import DataType, Device, Tensor
+from tvm.runtime import cpu as cpu_device
+from tvm.target import Target
+
+from mlc_llm.loader import LOADER
+from mlc_llm.model import Model
+from mlc_llm.quantization import Quantization
+from mlc_llm.support import logging, tqdm
+from mlc_llm.support.auto_weight import detect_weight
+from mlc_llm.support.preshard import apply_preshard
+from mlc_llm.support.style import bold, green
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class ConversionArgs:
+    """Arguments to MLC LLM's weight conversation and quantization flow."""
+
+    config: Path
+    quantization: Quantization
+    model: Model
+    device: Device
+    source: Path
+    source_format: str
+    output: Path
+    lora_adapter: Optional[Path] = None
+
+    def display(self) -> None:
+        """Display the arguments to stdout."""
+
+        def _device_to_str(device: Device) -> str:
+            return f"{Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]}:{device.index}"
+
+        out = StringIO()
+        print(f"{bold('Weight conversion with arguments:')}", file=out)
+        print(f"  {bold('--config'):<25} {self.config}", file=out)
+        print(f"  {bold('--quantization'):<25} {self.quantization}", file=out)
+        print(f"  {bold('--model-type'):<25} {self.model.name}", file=out)
+        print(f"  {bold('--device'):<25} {_device_to_str(self.device)}", file=out)
+        print(f"  {bold('--source'):<25} {self.source}", file=out)
+        print(f"  {bold('--source-format'):<25} {self.source_format}", file=out)
+        print(f"  {bold('--output'):<25} {self.output}", file=out)
+        if self.lora_adapter is not None:
+            print(f"  {bold('--lora-adapter'):<25} {self.lora_adapter}", file=out)
+        print(out.getvalue().rstrip())
+
+
+def _resolve_base_model_dir(source: Path) -> Path:
+    return source if source.is_dir() else source.parent
+
+
+@contextlib.contextmanager
+def _merge_lora_adapter_with_base_model(base_source: Path, lora_adapter: Path) -> Iterator[Path]:
+    base_model_dir = _resolve_base_model_dir(base_source)
+    if not base_model_dir.exists():
+        raise ValueError(f"Base model directory does not exist: {base_model_dir}")
+    if not lora_adapter.exists() or not lora_adapter.is_dir():
+        raise ValueError(f"LoRA adapter directory does not exist: {lora_adapter}")
+
+    try:
+        from peft import PeftModel
+        from transformers import AutoModelForCausalLM
+
+    except ImportError as err:
+        raise ImportError(
+            "`--lora-adapter` requires `peft` and `transformers` to be installed."
+        ) from err
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        merged_model_dir = Path(temp_dir) / "merged_model"
+        logger.info("Merging LoRA adapter %s into base model %s", lora_adapter, base_model_dir)
+
+        base_model = AutoModelForCausalLM.from_pretrained(
+            str(base_model_dir),
+            torch_dtype="auto",
+            trust_remote_code=False,
+            low_cpu_mem_usage=True,
+        )
+        merged_model = PeftModel.from_pretrained(
+            base_model, str(lora_adapter), is_trainable=False
+        ).merge_and_unload()
+        merged_model.save_pretrained(str(merged_model_dir), safe_serialization=True)
+        yield merged_model_dir
+
+
+def _convert_args(args: ConversionArgs) -> None:
+    pre_shards_num = os.getenv("MLC_INTERNAL_PRESHARD_NUM")
+    # model config & quantization config
+    model_config = args.model.config.from_file(args.config)
+    if (
+        args.quantization.kind == "ft-quant"
+        and hasattr(model_config, "tensor_parallel_shards")
+        and model_config.tensor_parallel_shards > 1
+    ):
+        raise NotImplementedError
+    if pre_shards_num is not None:
+        model_config.tensor_parallel_shards = int(pre_shards_num)
+    model, quantize_map = args.model.quantize[args.quantization.kind](
+        model_config, args.quantization
+    )
+    _, _named_params, _ = model.export_tvm(
+        spec=model.get_default_spec(),
+        allow_extern=True,
+    )
+    named_params = dict(_named_params)
+
+    if pre_shards_num is not None:
+        named_params, preshard_funcs = apply_preshard(named_params, int(pre_shards_num), args)
+    else:
+        preshard_funcs = None
+
+    def _check_param(name: str, param: Tensor):
+        nonlocal named_params
+        if name not in named_params:
+            raise ValueError(f"Parameter not found in model: {name}")
+        if name in param_names:
+            raise ValueError(f"Duplication: Parameter {name} already computed")
+
+        # Check shape (possibly dynamic)
+        def _check_shape(actual: tuple, expect: tuple):  # expect can have tirx.Var
+            if len(actual) != len(expect):
+                return False
+            for actual_i, expect_i in zip(actual, expect):
+                assert isinstance(expect_i, (int, tirx.Var))
+                if isinstance(expect_i, int) and actual_i != expect_i:
+                    return False
+            return True
+
+        expect_shape = named_params[name].shape
+        actual_shape = param.shape
+        if not _check_shape(actual_shape, expect_shape):
+            raise ValueError(
+                f"Parameter {name} has shape {param.shape}, but expected {expect_shape}"
+            )
+        # Check dtype
+        actual_dtype = param.dtype
+        expect_dtype = named_params[name].dtype
+        if actual_dtype != expect_dtype:
+            raise ValueError(
+                f"Parameter {name} has dtype {param.dtype}, but expected {expect_dtype}"
+            )
+        del named_params[name]
+
+    # load and quantize
+    param_names = set()
+    total_bytes = 0.0
+    total_params: int = 0
+
+    def _param_generator() -> Iterator[Tuple[str, Tensor]]:  # noqa: UP006
+        nonlocal total_params, total_bytes
+        with Target.from_device(args.device), tqdm.redirect():
+            loader = LOADER[args.source_format](
+                path=args.source,
+                extern_param_map=args.model.source[args.source_format](
+                    model_config, args.quantization
+                ),
+                quantize_param_map=quantize_map,
+            )
+            for name, param in loader.load(device=args.device, preshard_funcs=preshard_funcs):
+                _check_param(name, param)
+                param_names.add(name)
+                param = param.copyto(cpu_device())
+                total_bytes += math.prod(param.shape) * DataType(param.dtype).itemsize
+                yield name, param
+        total_params = loader.stats.total_param_num
+
+    def _metadata_callback() -> Dict[str, Any]:  # noqa: UP006
+        return {
+            "ParamSize": len(param_names),
+            "ParamBytes": total_bytes,
+            "BitsPerParam": total_bytes * 8.0 / total_params,
+        }
+
+    # dump to output directory
+    tvmjs.dump_tensor_cache(
+        _param_generator(),
+        str(args.output),
+        meta_data=_metadata_callback,
+        encode_format="f32-to-bf16",
+        show_progress=False,
+    )
+    if named_params:
+        raise ValueError(f"Parameter not found in source: {', '.join(named_params.keys())}")
+    # Log necessary statistics
+    logger.info(
+        "%s after quantization: %.3f GB",
+        green("Parameter size"),
+        total_bytes / (1024**3),
+    )
+    logger.info(f"%s: {total_params:,}", green("Total parameters"))
+    logger.info(
+        "%s: %.3f",
+        green("Bits per parameter"),
+        total_bytes * 8.0 / total_params,
+    )
+    logger.info("Saved to directory: %s", bold(str(args.output)))
+
+
+def convert_weight(
+    config: Path,
+    quantization: Quantization,
+    model: Model,
+    device: Device,
+    source: Path,
+    source_format: str,
+    output: Path,
+    lora_adapter: Optional[Path] = None,
+):
+    """MLC LLM's weight conversation and quantization flow."""
+    args = ConversionArgs(
+        config, quantization, model, device, source, source_format, output, lora_adapter
+    )
+
+    allowed_lora_source_formats = {"huggingface-safetensor", "huggingface-torch"}
+    if lora_adapter is not None and source_format not in allowed_lora_source_formats:
+        raise ValueError(
+            f"`--lora-adapter` only supports source formats: {sorted(allowed_lora_source_formats)}"
+        )
+
+    if lora_adapter is not None:
+        with _merge_lora_adapter_with_base_model(source, lora_adapter) as merged_model_dir:
+            merged_source, merged_source_format = detect_weight(
+                weight_path=merged_model_dir,
+                config_json_path=config,
+                weight_format="auto",
+            )
+            merged_args = dataclasses.replace(
+                args, source=merged_source, source_format=merged_source_format
+            )
+            merged_args.display()
+            _convert_args(merged_args)
+            return
+
+    args.display()
+    _convert_args(args)

python/mlc_llm/interface/gen_config.py

@@ -0,0 +1,359 @@
+"""Generator of mlc-chat-config.json and tokenizer configuration."""
+
+import dataclasses
+import json
+import re
+import shutil
+from dataclasses import asdict
+from pathlib import Path
+from typing import Optional
+
+from mlc_llm.conversation_template import ConvTemplateRegistry
+from mlc_llm.model import Model
+from mlc_llm.protocol.mlc_chat_config import MLCChatConfig
+from mlc_llm.quantization import Quantization
+from mlc_llm.support import convert_tiktoken, logging
+from mlc_llm.support.style import bold, green, red
+from mlc_llm.tokenizers import Tokenizer
+
+from .compiler_flags import ModelConfigOverride
+
+logger = logging.getLogger(__name__)
+
+FOUND = green("Found")
+NOT_FOUND = red("Not found")
+FAILED = red("Failed")
+
+
+def apply_system_defaults_for_missing_fields(mlc_chat_config: MLCChatConfig) -> None:
+    """Apply system default value."""
+    for key, value in mlc_chat_config.get_system_defaults_for_missing_fields().items():
+        setattr(mlc_chat_config, key, value)
+        logger.info("[System default] Setting %s: %s", bold(key), value)
+
+
+def check_string(s: str) -> bool:
+    """Check whether it's a string."""
+    s = s[1:] if s[0] == "b" else s
+    delimit = s[0]
+    if s[-1] != delimit or delimit not in ["'", '"']:
+        return False
+    for i in range(1, len(s) - 1):
+        if s[i] == delimit and s[i - 1] != "\\":
+            return False
+    return True
+
+
+def txt2rwkv_tokenizer(vocab: Path, out: Path) -> None:
+    """Generate tokenizer_model from RWKV vocab file."""
+    idx2token = {}
+
+    with vocab.open("r", encoding="utf-8") as f:
+        lines = f.readlines()
+
+    for line in lines:
+        idx = int(line[: line.index(" ")])
+        raw = line[line.index(" ") : line.rindex(" ")].strip()
+        if check_string(raw):
+            x = eval(raw)
+            x = x.encode("utf-8") if isinstance(x, str) else x
+            assert isinstance(x, bytes)
+            assert len(x) == int(line[line.rindex(" ") :])
+            idx2token[idx] = x
+        else:
+            raise ValueError("Unsupported vocab dictionary")
+
+    with (out / "tokenizer_model").open("wb") as f:
+        import msgpack
+
+        msgpack.pack(idx2token, f)
+
+
+def json2rwkv_tokenizer(vocab: Path, out: Path) -> None:
+    """Generate tokenizer_model from RWKV vocab file."""
+    idx2token = {}
+
+    with vocab.open("r", encoding="utf-8") as f:
+        data = json.load(f)
+        for key, value in data.items():
+            x = key.encode("utf-8") if isinstance(key, str) else key
+            assert isinstance(x, bytes)
+            idx2token[int(value)] = x
+
+    with (out / "tokenizer_model").open("wb") as f:
+        import msgpack
+
+        msgpack.pack(idx2token, f)
+
+
+def gen_config(
+    config: Path,
+    model: Model,
+    quantization: Quantization,
+    conv_template: str,
+    context_window_size: Optional[int],
+    sliding_window_size: Optional[int],
+    prefill_chunk_size: Optional[int],
+    attention_sink_size: Optional[int],
+    tensor_parallel_shards: Optional[int],
+    pipeline_parallel_stages: Optional[int],
+    disaggregation: Optional[bool],
+    max_batch_size: int,
+    output: Path,
+):
+    """Entrypoint of MLC Chat configuration generation."""
+    # Step 1. Initialize `mlc-chat-config.json` using `config.json`
+    conversation_reg = ConvTemplateRegistry.get_conv_template(conv_template)
+    if conversation_reg is None:
+        logger.warning(
+            "%s: Conversation template is not registered in ConvTemplateRegistry: %s",
+            red("Warning"),
+            conv_template,
+        )
+        conversation = conv_template
+    else:
+        conversation = conversation_reg.to_json_dict()
+
+    model_config = ModelConfigOverride(
+        context_window_size=context_window_size,
+        sliding_window_size=sliding_window_size,
+        prefill_chunk_size=prefill_chunk_size,
+        attention_sink_size=attention_sink_size,
+        max_batch_size=max_batch_size,
+        tensor_parallel_shards=tensor_parallel_shards,
+        pipeline_parallel_stages=pipeline_parallel_stages,
+        disaggregation=disaggregation,
+    ).apply(model.config.from_file(config))
+    mlc_chat_config = MLCChatConfig(
+        model_type=model.name,
+        quantization=quantization.name,
+        model_config=model_config.asdict(),
+        vocab_size=model_config.vocab_size,
+        active_vocab_size=getattr(model_config, "active_vocab_size", model_config.vocab_size),
+        context_window_size=getattr(model_config, "context_window_size", -1),
+        sliding_window_size=getattr(model_config, "sliding_window_size", -1),
+        prefill_chunk_size=model_config.prefill_chunk_size,
+        attention_sink_size=getattr(model_config, "attention_sink_size", -1),
+        tensor_parallel_shards=model_config.tensor_parallel_shards,
+        pipeline_parallel_stages=getattr(model_config, "pipeline_parallel_stages", 1),
+        disaggregation=getattr(model_config, "disaggregation", False),
+        conv_template=conversation,
+        model_task=model.model_task,
+        embedding_metadata=(
+            dataclasses.asdict(model.embedding_metadata) if model.embedding_metadata else None
+        ),
+    )
+    # Step 2. Load `generation_config.json` and `config.json` for text-generation related configs
+    for generation_config_filename in ["generation_config.json", "config.json"]:
+        generation_config = config.parent / generation_config_filename
+        if generation_config.exists():
+            with generation_config.open("r", encoding="utf-8") as in_file:
+                generation_config_json = json.load(in_file)
+            for key, value in generation_config_json.items():
+                if hasattr(mlc_chat_config, key) and getattr(mlc_chat_config, key) is None:
+                    setattr(mlc_chat_config, key, value)
+                    logger.info(
+                        "[%s] Setting %s: %s",
+                        generation_config_filename,
+                        bold(key),
+                        value,
+                    )
+        else:
+            logger.info("%s %s: %s", NOT_FOUND, generation_config_filename, generation_config)
+
+    # Step 3. Copy tokenizer configuration
+    # 3.1. Copy over the files and populate mlc_chat_config
+    for filename in TOKENIZER_FILES:
+        file = config.parent / filename
+        if file.exists():
+            mlc_chat_config.tokenizer_files.append(filename)
+            dest = output / filename
+            shutil.copy(file, dest)
+            logger.info("%s tokenizer config: %s. Copying to %s", FOUND, file, bold(str(dest)))
+        else:
+            logger.info("%s tokenizer config: %s", NOT_FOUND, file)
+    # 3.2. Generate `tokenizer_model` for rwkv if `rwkv_vocab_.*` is found
+    pattern = re.compile(r"rwkv_vocab_v\d{8}\.(json|txt)")
+    for item in config.parent.iterdir():
+        if item.is_file() and pattern.match(item.name):
+            logger.info(
+                "%s RWKV vocab file: %s. Genetating %s",
+                FOUND,
+                item,
+                bold("tokenizer_model"),
+            )
+            if item.name.endswith(".txt"):
+                txt2rwkv_tokenizer(item, output)
+            else:
+                json2rwkv_tokenizer(item, output)
+    # 3.3. If we have `tokenizer.model` but not `tokenizer.json`, try convert it to
+    # `tokenizer.json` with `transformers`.
+    tokenizer_json_file = config.parent / "tokenizer.json"
+    tokenizer_model_file = config.parent / "tokenizer.model"
+    if tokenizer_model_file.exists() and (not tokenizer_json_file.exists()):
+        logger.info(
+            "The model has `tokenizer.model` but not `tokenizer.json`. "
+            "It is always recommended to prefer JSON instead. "
+            "Attempting to convert using HuggingFace transformers library"
+        )
+        try:
+            from transformers import (
+                AutoTokenizer,
+            )
+
+            tokenizer_json_save_dest = output / "tokenizer.json"
+            fast_tokenizer = AutoTokenizer.from_pretrained(str(config.parent), use_fast=True)
+            fast_tokenizer.backend_tokenizer.save(str(tokenizer_json_save_dest))
+            mlc_chat_config.tokenizer_files.append("tokenizer.json")
+            logger.info(
+                "Successfully converted `tokenizer.model` to: %s",
+                tokenizer_json_save_dest,
+            )
+        except Exception:
+            logger.warning(
+                "Converting to `tokenizer.json` %s with the exception below. "
+                "Skipping the conversion.",
+                FAILED,
+                exc_info=True,
+            )
+    # 3.3. If we still don't have "tokenizer.json" at this point, try looking for "*.tiktoken" files
+    if (not tokenizer_json_file.exists()) and list(config.parent.glob("*.tiktoken")):
+        try:
+            logger.info(
+                "The model has tiktoken files but not `tokenizer.json`. "
+                "Attempting to convert from tiktoken files"
+            )
+            convert_tiktoken.convert_tiktoken(
+                str(config.parent), str(output), mlc_chat_config.context_window_size
+            )
+            mlc_chat_config.tokenizer_files.append("tokenizer.json")
+            mlc_chat_config.tokenizer_files.append("vocab.json")
+            mlc_chat_config.tokenizer_files.append("merges.txt")
+            mlc_chat_config.tokenizer_files.append("special_tokens_map.json")
+            logger.info("Succesfully converted from tiktoken files to: %s", str(output))
+        except Exception:
+            logger.exception("%s with the exception below. Skipping", FAILED)
+
+    # 3.4. Detect tokenizer info
+    mlc_chat_config.tokenizer_info = asdict(Tokenizer.detect_tokenizer_info(str(output)))
+    logger.info("Detected tokenizer info: %s", mlc_chat_config.tokenizer_info)
+
+    # 3.5. Ensure added_tokens do not have duplicated added_tokens, a mistake from model releaser
+    # that affects correctness of huggingface tokenizer.
+    # See https://huggingface.co/NousResearch/Hermes-2-Pro-Llama-3-8B/discussions/15.
+    if tokenizer_json_file.exists():
+        with open(tokenizer_json_file, encoding="utf-8") as f:
+            tokenizer_json = json.load(f)
+            if "added_tokens" in tokenizer_json:
+                appeared_content = set()
+                for added_token in tokenizer_json["added_tokens"]:
+                    content = added_token["content"]
+                    if content in appeared_content:
+                        logger.exception(
+                            "%s with incorrect tokenizer.json which has duplicated token %s. "
+                            "This affects correctness of huggingface tokenizer during runtime, "
+                            "please check your tokenizer.json to remove duplication manually.",
+                            FAILED,
+                            content,
+                        )
+                        raise ValueError("Duplicated vocab in tokenizer.json")
+                    appeared_content.add(content)
+
+    # Step 4. Load system default value
+    apply_system_defaults_for_missing_fields(mlc_chat_config)
+
+    # Step 5. Use HF tokenizer to detect active vocab size via len(tokenizer)
+    if tokenizer_json_file.exists():
+        try:
+            from transformers import (
+                AutoTokenizer,
+            )
+
+            hf_tokenizer = AutoTokenizer.from_pretrained(str(config.parent), use_fast=True)
+            active_vocab_size = len(hf_tokenizer)
+            if mlc_chat_config.active_vocab_size != active_vocab_size:
+                logger.info(
+                    "Overriding active_vocab_size from %d to %d using HF tokenizer",
+                    mlc_chat_config.active_vocab_size,
+                    active_vocab_size,
+                )
+                mlc_chat_config.active_vocab_size = active_vocab_size
+        except Exception:
+            logger.warning(
+                "Detecting active_vocab_size %s with the exception below. Skipping.",
+                FAILED,
+                exc_info=True,
+            )
+
+    # Step 5. Dump the configuration file to output directory
+    with (output / "mlc-chat-config.json").open("w", encoding="utf-8") as out_file:
+        json.dump(mlc_chat_config.model_dump(by_alias=True), out_file, indent=2)
+        logger.info("Dumping configuration file to: %s", bold(out_file.name))
+
+
+TOKENIZER_FILES = [
+    "tokenizer.model",
+    "tokenizer.json",
+    "vocab.json",
+    "merges.txt",
+    "added_tokens.json",
+    "tokenizer_config.json",
+]
+# FIXME: Copy RWKV tokenizer file
+
+CONV_TEMPLATES = {
+    "llama-4",
+    "llama-3",
+    "llama-3_1",
+    "chatml",
+    "chatml_nosystem",
+    "qwen2",
+    "open_hermes_mistral",
+    "neural_hermes_mistral",
+    "llama_default",
+    "llama-2",
+    "mistral_default",
+    "ministral3",
+    "ministral3_reasoning",
+    "gpt2",
+    "codellama_completion",
+    "codellama_instruct",
+    "redpajama_chat",
+    "rwkv_world",
+    "gorilla",
+    "gorilla-openfunctions-v2",
+    "dolly",
+    "oasst",
+    "stablelm",
+    "LM",
+    "stablelm-3b",
+    "gpt_bigcode",
+    "wizardlm_7b",
+    "wizard_coder_or_math",
+    "glm",
+    "phi-2",
+    "phi-3",
+    "phi-3-vision",
+    "phi-4",
+    "stablelm-2",
+    "gemma_instruction",
+    "gemma3_instruction",
+    "orion",
+    "llava",
+    "hermes2_pro_llama3",
+    "hermes3_llama-3_1",
+    "tinyllama_v1_0",
+    "aya-23",
+    "deepseek",
+    "deepseek_v2",
+    "deepseek_v3",
+    "deepseek_r1_qwen",
+    "deepseek_r1_llama",
+    "olmo",
+    "olmo2",
+    "nemotron",
+    "llm-jp",
+    "qwen3",
+    "qwen3_5",
+    "qwen3_5_nothink",
+}

python/mlc_llm/interface/help.py

@@ -0,0 +1,273 @@
+"""Help message for CLI arguments."""
+
+HELP = {
+    "config": (
+        """
+1) Path to a HuggingFace model directory that contains a `config.json` or
+2) Path to `config.json` in HuggingFace format, or
+3) The name of a pre-defined model architecture.
+
+A `config.json` file in HuggingFace format defines the model architecture, including the vocabulary
+size, the number of layers, the hidden size, number of attention heads, etc.
+Example: https://huggingface.co/codellama/CodeLlama-7b-hf/blob/main/config.json.
+
+A HuggingFace directory often contains a `config.json` which defines the model architecture,
+the non-quantized model weights in PyTorch or SafeTensor format, tokenizer configurations,
+as well as an optional `generation_config.json` provides additional default configuration for
+text generation.
+Example: https://huggingface.co/codellama/CodeLlama-7b-hf/tree/main.
+"""
+    ).strip(),
+    "quantization": """
+The quantization mode we use to compile. If unprovided, will infer from `model`.
+""".strip(),
+    "model": """
+A path to ``mlc-chat-config.json``, or an MLC model directory that contains `mlc-chat-config.json`.
+It can also be a link to a HF repository pointing to an MLC compiled model.
+""".strip(),
+    "model_lib": """
+The full path to the model library file to use (e.g. a ``.so`` file). If unspecified, we will use
+the provided ``model`` to search over possible paths. It the model lib is not found, it will be
+compiled in a JIT manner.
+""".strip(),
+    "model_type": """
+Model architecture such as "llama". If not set, it is inferred from `mlc-chat-config.json`.
+""".strip(),
+    "device_compile": """
+The GPU device to compile the model to. If not set, it is inferred from GPUs available locally.
+""".strip(),
+    "enable_subgroups": """
+Enable WebGPU subgroups in codegen. This only applies to WebGPU targets and will set
+supports_subgroups accordingly.
+""".strip(),
+    "device_quantize": """
+The device used to do quantization such as "cuda" or "cuda:0". Will detect from local available GPUs
+if not specified.
+""".strip(),
+    "device_deploy": """
+The device used to deploy the model such as "cuda" or "cuda:0". Will detect from local
+available GPUs if not specified.
+""".strip(),
+    "host": """
+The host LLVM triple to compile the model to. If not set, it is inferred from the local CPU and OS.
+Examples of the LLVM triple:
+1) iPhones: arm64-apple-ios;
+2) ARM64 Android phones: aarch64-linux-android;
+3) WebAssembly: wasm32-unknown-unknown-wasm;
+4) Windows: x86_64-pc-windows-msvc;
+5) ARM macOS: arm64-apple-darwin.
+""".strip(),
+    "opt": """
+Optimization flags. MLC LLM maintains a predefined set of optimization flags,
+denoted as O0, O1, O2, O3, where O0 means no optimization, O2 means majority of them,
+and O3 represents extreme optimization that could potentially break the system.
+Meanwhile, optimization flags could be explicitly specified via details knobs, e.g.
+--opt="cublas_gemm=1;cudagraph=0".
+""".strip(),
+    "system_lib_prefix": """
+Adding a prefix to all symbols exported. Similar to "objcopy --prefix-symbols".
+This is useful when compiling multiple models into a single library to avoid symbol
+conflicts. Different from objcopy, this takes no effect for shared library.
+""".strip(),
+    "context_window_size": """
+Option to provide the maximum sequence length supported by the model.
+This is usually explicitly shown as context length or context window in the model card.
+If this option is not set explicitly, by default,
+it will be determined by `context_window_size` or `max_position_embeddings` in `config.json`,
+and the latter is usually inaccurate for some models.
+""".strip(),
+    "output_compile": """
+The path to the output file. The suffix determines if the output file is a shared library or
+objects. Available suffixes:
+1) Linux: .so (shared), .tar (objects);
+2) macOS: .dylib (shared), .tar (objects);
+3) Windows: .dll (shared), .tar (objects);
+4) Android, iOS: .tar (objects);
+5) Web: .wasm (web assembly).
+""".strip(),
+    "source": """
+The path to original model weight, infer from `config` if missing.
+""".strip(),
+    "source_format": """
+The format of source model weight, infer from `config` if missing.
+""".strip(),
+    "output_quantize": """
+The output directory to save the quantized model weight. Will create `params_shard_*.bin` and
+`tensor-cache.json` in this directory.
+""".strip(),
+    "conv_template": """
+Conversation template. It depends on how the model is tuned. Use "LM" for vanilla base model
+""".strip(),
+    "output_gen_mlc_chat_config": """
+The output directory for generated configurations, including `mlc-chat-config.json` and tokenizer
+configuration.
+""".strip(),
+    "sliding_window_size": """
+(Experimental) The sliding window size in sliding window attention (SWA).
+This optional field overrides the `sliding_window_size` in config.json for
+those models that use SWA. Currently only useful when compiling Mistral.
+This flag subjects to future refactoring.
+""".strip(),
+    "prefill_chunk_size": """
+(Experimental) The chunk size during prefilling. By default,
+the chunk size is the same as sliding window or max sequence length.
+This flag subjects to future refactoring.
+""".strip(),
+    "attention_sink_size": """
+(Experimental) The number of stored sinks. Only supported on Mistral yet. By default,
+the number of sinks is 4. This flag subjects to future refactoring.
+""".strip(),
+    "max_batch_size": """
+The maximum allowed batch size set for the KV cache to concurrently support.
+""".strip(),
+    """tensor_parallel_shards""": """
+Number of shards to split the model into in tensor parallelism multi-gpu inference.
+""".strip(),
+    """pipeline_parallel_stages""": """
+Number of pipeline stages to split the model layers for pipeline parallelism.
+""".strip(),
+    """disaggregation""": """
+Whether enable disaggregation when compiling the model.
+""".strip(),
+    "overrides": """
+Model configuration override. Configurations to override `mlc-chat-config.json`. Supports
+`context_window_size`, `prefill_chunk_size`, `sliding_window_size`, `attention_sink_size`,
+`max_batch_size` and `tensor_parallel_shards`. Meanwhile, model config could be explicitly
+specified via details knobs, e.g. --overrides "context_window_size=1024;prefill_chunk_size=128".
+""".strip(),
+    "modelconfig_overrides": """
+Model configuration override. Supports overriding,
+`context_window_size`, `prefill_chunk_size`, `sliding_window_size`, `attention_sink_size`,
+`max_num_sequence` and `tensor_parallel_shards`. The overrides could be explicitly
+specified via details knobs, e.g. --overrides "context_window_size=1024;prefill_chunk_size=128".
+""".strip(),
+    "debug_dump": """
+Specifies the directory where the compiler will store its IRs for debugging purposes
+during various phases of compilation. By default, this is set to `None`, indicating
+that debug dumping is disabled.
+""".strip(),
+    "prompt": """
+The prompt of the text generation.
+""".strip(),
+    "generate_length": """
+The target length of the text generation.
+""".strip(),
+    "max_total_sequence_length_serve": """
+The KV cache total token capacity, i.e., the maximum total number of tokens that
+the KV cache support. This decides the GPU memory size that the KV cache consumes.
+If not specified, system will automatically estimate the maximum capacity based
+on the vRAM size on GPU.
+""".strip(),
+    "prefill_chunk_size_serve": """
+The maximum number of tokens the model passes for prefill each time.
+It should not exceed the prefill chunk size in model config.
+If not specified, this defaults to the prefill chunk size in model config.
+""".strip(),
+    "max_history_size_serve": """
+The maximum history length for rolling back the RNN state.
+If unspecified, the default value is 1.
+KV cache does not need this.
+""".strip(),
+    "enable_tracing_serve": """
+Enable Chrome Tracing for the server.
+After enabling, you can send POST request to the "debug/dump_event_trace" entrypoint
+to get the Chrome Trace. For example,
+"curl -X POST http://127.0.0.1:8000/debug/dump_event_trace -H "Content-Type: application/json" -d '{"model": "dist/llama"}'"
+""".strip(),  # noqa: E501
+    "mode_serve": """
+The engine mode in MLC LLM. We provide three preset modes: "local", "interactive" and "server".
+The default mode is "local".
+The choice of mode decides the values of "max_num_sequence", "max_total_seq_length" and
+"prefill_chunk_size" when they are not explicitly specified.
+1. Mode "local" refers to the local server deployment which has low request concurrency.
+   So the max batch size will be set to 4, and max total sequence length and prefill chunk size
+   are set to the context window size (or sliding window size) of the model.
+2. Mode "interactive" refers to the interactive use of server, which has at most 1 concurrent
+   request. So the max batch size will be set to 1, and max total sequence length and prefill
+   chunk size are set to the context window size (or sliding window size) of the model.
+3. Mode "server" refers to the large server use case which may handle many concurrent request
+   and want to use GPU memory as much as possible. In this mode, we will automatically infer
+   the largest possible max batch size and max total sequence length.
+You can manually specify arguments "max_num_sequence", "max_total_seq_length" and
+"prefill_chunk_size" via "--overrides" to override the automatic inferred values.
+For example: --overrides "max_num_sequence=32;max_total_seq_length=4096"
+""".strip(),
+    "additional_models_serve": """
+The model paths and (optional) model library paths of additional models (other than the main model).
+When engine is enabled with speculative decoding, additional models are needed.
+The way of specifying additional models is:
+"--additional-models model_path_1 model_path_2 ..." or
+"--additional-models model_path_1,model_lib_1 model_path_2 ...".
+When the model lib of a model is not given, JIT model compilation will be activated
+to compile the model automatically.
+""".strip(),
+    "gpu_memory_utilization_serve": """
+A number in (0, 1) denoting the fraction of GPU memory used by the server in total.
+It is used to infer to maximum possible KV cache capacity.
+When it is unspecified, it defaults to 0.85.
+Under mode "local" or "interactive", the actual memory usage may be significantly smaller than
+this number. Under mode "server", the actual memory usage may be slightly larger than this number.
+""".strip(),
+    "speculative_mode_serve": """
+The speculative decoding mode. Right now four options are supported:
+ - "disable", where speculative decoding is not enabled,
+ - "small_draft", denoting the normal speculative decoding (small draft) style,
+ - "eagle", denoting the eagle-style speculative decoding.
+ - "medusa", denoting the medusa-style speculative decoding.
+The default mode is "disable".
+""".strip(),
+    "spec_draft_length_serve": """
+The number of draft tokens to generate in speculative proposal.
+Being 0 means to enable adaptive speculative mode, where the draft length will be
+automatically adjusted based on engine state. The default values is 0.
+""".strip(),
+    "prefix_cache_mode_serve": """
+The prefix cache mode. Right now two options are supported:
+ - "disable", where prefix cache is not enabled,
+ - "radix", denoting the normal paged radix tree based prefix cache,
+The default mode is "radix".
+""".strip(),
+    "prefix_cache_max_num_recycling_seqs_serve": """
+The maximum number of sequences in prefix cache, default as max_batch_size.
+And set 0 to disable prefix cache, set -1 to have infinite capacity prefix cache.
+""".strip(),
+    "prefill_mode": """
+The prefill mode. "chunked" means the basic prefill with chunked input enabled. "hybrid" means the
+hybrid prefill or split-fuse, so that decode step will be converted into prefill.
+""".strip(),
+    "overrides_serve": """
+Overriding extra configurable fields of EngineConfig and model compilation config.
+Supporting fields that can be be overridden: "tensor_parallel_shards", "max_num_sequence",
+"max_total_seq_length", "prefill_chunk_size", "max_history_size", "gpu_memory_utilization",
+"spec_draft_length", "prefix_cache_max_num_recycling_seqs", "context_window_size",
+"sliding_window_size", "attention_sink_size".
+Please check out the documentation of EngineConfig in mlc_llm/serve/config.py for detailed docstring
+of each field.
+Example: --overrides "max_num_sequence=32;max_total_seq_length=4096;tensor_parallel_shards=2"
+""".strip(),
+    "config_package": """
+The path to "mlc-package-config.json" which is used for package build.
+See "https://github.com/mlc-ai/mlc-llm/blob/main/ios/MLCChat/mlc-package-config.json" as an example.
+""".strip(),
+    "mlc_llm_source_dir": """
+The source code path to MLC LLM.
+""".strip(),
+    "output_package": """
+The path of output directory for the package build outputs.
+""".strip(),
+    "calibration_dataset": """
+The path to the calibration dataset.
+    """.strip(),
+    "num_calibration_samples": """
+The number of samples used for calibration.
+    """.strip(),
+    "output_calibration": """
+The output directory to save the calibration params.
+    """.strip(),
+    "seed_calibrate": """
+The seed to sample the calibration dataset.""",
+    "pd_balance_factor": """
+How much prefill to move to decode engine. For example,
+0.1 means the last 10 percent tokens are prefilled by decode engine.
+    """.strip(),
+}

python/mlc_llm/interface/jit.py

@@ -0,0 +1,181 @@
+"""Just-in-time compilation of MLC-Chat models."""
+
+import dataclasses
+import hashlib
+import json
+import os
+import shlex
+import shutil
+import subprocess
+import sys
+import tempfile
+from pathlib import Path
+from typing import Any, Dict, Optional, Union  # noqa: UP035
+
+from tvm.runtime import Device
+
+from mlc_llm.model import MODELS
+from mlc_llm.support import logging
+from mlc_llm.support.auto_device import device2str
+from mlc_llm.support.constants import (
+    MLC_DSO_SUFFIX,
+    MLC_JIT_POLICY,
+    MLC_LLM_HOME,
+    MLC_TEMP_DIR,
+)
+from mlc_llm.support.style import blue, bold
+
+from .compiler_flags import ModelConfigOverride, OptimizationFlags
+
+logger = logging.getLogger(__name__)
+
+
+@dataclasses.dataclass
+class JITResult:
+    """The jit compilation result class."""
+
+    model_lib_path: str
+    system_lib_prefix: Optional[str] = None
+
+
+def log_jit_policy():
+    """log current jit policy"""
+    logger.info(
+        "%s = %s. Can be one of: ON, OFF, REDO, READONLY",
+        bold("MLC_JIT_POLICY"),
+        MLC_JIT_POLICY,
+    )
+
+
+def jit(
+    model_path: Path,
+    overrides: Dict[str, Any],  # noqa: UP006
+    device: Union[Device, str],
+    system_lib_prefix: Optional[str] = None,
+    *,
+    skip_log_jit_policy=False,
+) -> JITResult:
+    """Just-in-time compile a MLC-Chat model."""
+    # skip logging jit policy since when outside can hint once
+    if not skip_log_jit_policy:
+        log_jit_policy()
+
+    if MLC_JIT_POLICY == "OFF":
+        raise RuntimeError("JIT is disabled by MLC_JIT_POLICY=OFF")
+
+    with open(model_path / "mlc-chat-config.json", encoding="utf-8") as in_file:
+        mlc_chat_config = json.load(in_file)
+    model_type = mlc_chat_config.pop("model_type")
+    quantization = mlc_chat_config.pop("quantization")
+    lib_suffix = MLC_DSO_SUFFIX if device not in ["iphone", "macabi", "android"] else "tar"
+
+    def _get_optimization_flags() -> str:
+        opt = overrides.pop("opt", None)
+        if opt is None:
+            opt = "O2"
+        return repr(OptimizationFlags.from_str(opt))
+
+    def _get_overrides() -> str:
+        forbid_list = [
+            "context_window_size",
+            "sliding_window_size",
+            "attention_sink_size",
+        ]
+        result = []
+        for field in dataclasses.fields(ModelConfigOverride):
+            value = overrides.get(field.name, None)
+            if value is not None:
+                if field.name in forbid_list and value == -1:
+                    continue
+                result.append(f"{field.name}={value}")
+        return ";".join(result)
+
+    def _get_model_config() -> Dict[str, Any]:  # noqa: UP006
+        model_config = mlc_chat_config.pop("model_config")
+        model_config.update(mlc_chat_config)
+        for field in dataclasses.fields(ModelConfigOverride):
+            value = overrides.get(field.name, None)
+            if value is not None:
+                model_config[field.name] = value
+        return MODELS[model_type].config.from_dict(model_config).asdict()
+
+    def _run_jit(
+        opt: str,
+        overrides: str,
+        device: str,
+        system_lib_prefix: Optional[str],
+        dst: str,
+    ):
+        with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir:
+            dso_path = os.path.join(tmp_dir, f"lib.{lib_suffix}")
+            cmd = [
+                sys.executable,
+                "-m",
+                "mlc_llm",
+                "compile",
+                str(model_path),
+                "--opt",
+                opt,
+                "--overrides",
+                overrides,
+                "--device",
+                device,
+                "--output",
+                dso_path,
+            ]
+            if system_lib_prefix:
+                cmd += ["--system-lib-prefix", system_lib_prefix + "_"]
+            logger.info("Compiling using commands below:")
+            logger.info("%s", blue(shlex.join(cmd)))
+            subprocess.run(cmd, check=False, env=os.environ)
+            # note on windows: compilation can succeed but return code is still nonzero
+            # check whether file exists instead
+            if not os.path.isfile(dso_path):
+                raise RuntimeError("Cannot find compilation output, compilation failed")
+            shutil.move(dso_path, dst)
+            logger.info("Using compiled model lib: %s", bold(dst))
+
+    hash_key = {
+        "model_config": _get_model_config(),
+        "overrides": _get_overrides(),
+        "opt": _get_optimization_flags(),
+        "device": device2str(device) if isinstance(device, Device) else device,
+        "model_type": model_type,
+        "quantization": quantization,
+    }
+    if device in ["iphone", "macabi", "android"]:
+        if system_lib_prefix is None:
+            system_lib_hash_value = hashlib.md5(
+                json.dumps(
+                    hash_key,
+                    sort_keys=True,
+                    indent=2,
+                ).encode("utf-8")
+            ).hexdigest()
+            system_lib_prefix = f"{model_type}_{quantization}_{system_lib_hash_value}".replace(
+                "-", "_"
+            )
+        hash_key["system_lib_prefix"] = system_lib_prefix
+    hash_value = hashlib.md5(
+        json.dumps(
+            hash_key,
+            sort_keys=True,
+            indent=2,
+        ).encode("utf-8")
+    ).hexdigest()
+    dst = MLC_LLM_HOME / "model_lib" / f"{hash_value}.{lib_suffix}"
+    if dst.is_file() and MLC_JIT_POLICY in ["ON", "READONLY"]:
+        logger.info("Using cached model lib: %s", bold(str(dst)))
+        return JITResult(str(dst), system_lib_prefix)
+    if MLC_JIT_POLICY == "READONLY":
+        raise RuntimeError(
+            "No cached model lib found, and JIT is disabled by MLC_JIT_POLICY=READONLY"
+        )
+    _run_jit(
+        opt=hash_key["opt"],
+        overrides=hash_key["overrides"],
+        device=hash_key["device"],
+        system_lib_prefix=system_lib_prefix,
+        dst=str(dst),
+    )
+    return JITResult(str(dst), system_lib_prefix)

python/mlc_llm/interface/package.py

@@ -0,0 +1,402 @@
+"""Python entrypoint of package."""
+
+import dataclasses
+import json
+import os
+import shutil
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any, Dict, List, Literal  # noqa: UP035
+
+from mlc_llm.interface import jit
+from mlc_llm.support import download_cache, logging, style
+
+logging.enable_logging()
+logger = logging.getLogger(__name__)
+
+SUPPORTED_DEVICES = ["iphone", "macabi", "android"]
+
+
+def build_model_library(
+    package_config: Dict[str, Any],  # noqa: UP006
+    device: str,
+    bundle_dir: Path,
+    app_config_path: Path,
+) -> Dict[str, str]:  # noqa: UP006
+    """Build model libraries. Return the dictionary of "library prefix to lib path"."""
+    # - Create the bundle directory.
+    os.makedirs(bundle_dir, exist_ok=True)
+    # Clean up all the directories in `output/bundle`.
+    logger.info('Clean up all directories under "%s"', str(bundle_dir))
+    for content_path in bundle_dir.iterdir():
+        if content_path.is_dir():
+            shutil.rmtree(content_path)
+
+    # - Process each model, and prepare the app config.
+    app_config_model_list = []
+
+    model_entries = package_config.get("model_list", [])
+    if not isinstance(model_entries, list):
+        raise ValueError('The "model_list" in "mlc-package-config.json" is expected to be a list.')
+    model_lib_path_for_prepare_libs = package_config.get("model_lib_path_for_prepare_libs", {})
+    if not isinstance(model_lib_path_for_prepare_libs, dict):
+        raise ValueError(
+            'The "model_lib_path_for_prepare_libs" in "mlc-package-config.json" is expected to be '
+            "a dict."
+        )
+
+    jit.log_jit_policy()
+
+    for model_entry in package_config.get("model_list", []):
+        # - Parse model entry.
+        if not isinstance(model_entry, dict):
+            raise ValueError('The element of "model_list" is expected to be a dict.')
+        model = model_entry["model"]
+        model_id = model_entry["model_id"]
+        bundle_weight = model_entry.get("bundle_weight", False)
+        overrides = model_entry.get("overrides", {})
+        model_lib = model_entry.get("model_lib", None)
+
+        estimated_vram_bytes = model_entry["estimated_vram_bytes"]
+        if not isinstance(model, str):
+            raise ValueError('The value of "model" in "model_list" is expected to be a string.')
+        if not isinstance(model_id, str):
+            raise ValueError('The value of "model_id" in "model_list" is expected to be a string.')
+        if not isinstance(bundle_weight, bool):
+            raise ValueError(
+                'The value of "bundle_weight" in "model_list" is expected to be a boolean.'
+            )
+        if not isinstance(overrides, dict):
+            raise ValueError('The value of "overrides" in "model_list" is expected to be a dict.')
+        if model_lib is not None and not isinstance(model_lib, str):
+            raise ValueError('The value of "model_lib" in "model_list" is expected to be string.')
+
+        # - Load model config. Download happens when needed.
+        model_path = download_cache.get_or_download_model(model)
+
+        # - Jit compile if the model lib path is not specified.
+        model_lib_path = (
+            model_lib_path_for_prepare_libs.get(model_lib, None) if model_lib is not None else None
+        )
+        if model_lib_path is None:
+            if model_lib is None:
+                logger.info(
+                    'Model lib is not specified for model "%s". Now jit compile the model library.',
+                    model_id,
+                )
+            else:
+                logger.info(
+                    'Model lib path for "%s" is not specified in "model_lib_path_for_prepare_libs".'
+                    "Now jit compile the model library.",
+                    model_lib,
+                )
+            model_lib_path, model_lib = dataclasses.astuple(
+                jit.jit(
+                    model_path=model_path,
+                    overrides=overrides,
+                    device=device,
+                    system_lib_prefix=model_lib,
+                    skip_log_jit_policy=True,
+                )
+            )
+            assert model_lib is not None
+            model_lib_path_for_prepare_libs[model_lib] = model_lib_path
+
+        # - Set "model_url"/"model_path" and "model_id"
+        app_config_model_entry = {}
+        is_local_model = not model.startswith("HF://") and not model.startswith("https://")
+        app_config_model_entry["model_id"] = model_id
+        app_config_model_entry["model_lib"] = model_lib
+
+        # - Bundle weight
+        if is_local_model and not bundle_weight:
+            raise ValueError(
+                f'Model "{model}" in "model_list" is a local path.'
+                f'Please set \'"bundle_weight": true\' in the entry of model "{model}".'
+            )
+        if bundle_weight:
+            if not os.path.isfile(model_path / "tensor-cache.json"):
+                raise ValueError(
+                    f'Bundle weight is set for model "{model}". However, model weights are not'
+                    f'found under the directory "{model}". '
+                    + (
+                        "Please follow https://llm.mlc.ai/docs/compilation/convert_weights.html to "
+                        "convert model weights."
+                        if is_local_model
+                        else "Please report this issue to https://github.com/mlc-ai/mlc-llm/issues."
+                    )
+                )
+            # Overwrite the model weight directory in bundle.
+            bundle_model_weight_path = bundle_dir / model_id
+            logger.info(
+                "Bundle weight for %s, copy into %s",
+                style.bold(model_id),
+                style.bold(str(bundle_model_weight_path)),
+            )
+            if bundle_model_weight_path.exists():
+                shutil.rmtree(bundle_model_weight_path)
+            shutil.copytree(model_path, bundle_model_weight_path)
+        if bundle_weight and device in ["iphone", "macabi"]:
+            app_config_model_entry["model_path"] = model_id
+        else:
+            app_config_model_entry["model_url"] = model.replace("HF://", "https://huggingface.co/")
+
+        # - estimated_vram_bytes
+        app_config_model_entry["estimated_vram_bytes"] = estimated_vram_bytes
+
+        app_config_model_list.append(app_config_model_entry)
+
+    # - Dump "mlc-app-config.json".
+    app_config_json_str = json.dumps(
+        {"model_list": app_config_model_list},
+        indent=2,
+    )
+    with open(app_config_path, "w", encoding="utf-8") as file:
+        print(app_config_json_str, file=file)
+        logger.info(
+            'Dump the app config below to "%s":\n%s',
+            str(app_config_path),
+            style.green(app_config_json_str),
+        )
+    return model_lib_path_for_prepare_libs
+
+
+def validate_model_lib(
+    app_config_path: Path,
+    package_config_path: Path,
+    model_lib_path_for_prepare_libs: dict,
+    device: Literal["iphone", "macabi", "android"],
+    output: Path,
+) -> None:
+    """Validate the model lib prefixes of model libraries."""
+    if device == "android":
+        from tvm.support import ndk as cc
+    else:
+        from tvm.support import cc
+
+    with open(app_config_path, encoding="utf-8") as file:
+        app_config = json.load(file)
+
+    tar_list = []
+    model_set = set()
+
+    for model, model_lib_path in model_lib_path_for_prepare_libs.items():
+        model_lib_path = os.path.join(model_lib_path)
+        lib_path_valid = os.path.isfile(model_lib_path)
+        if not lib_path_valid:
+            raise RuntimeError(f"Cannot find file {model_lib_path} as an {device} model library")
+        tar_list.append(model_lib_path)
+        model_set.add(model)
+
+    os.makedirs(output / "lib", exist_ok=True)
+    if device in ["iphone", "macabi"]:
+        lib_name = "libmodel_iphone.a"
+    else:
+        lib_name = "libmodel_android.a"
+    lib_path = output / "lib" / lib_name
+
+    def _get_model_libs(lib_path: Path) -> List[str]:  # noqa: UP006
+        """Get the model lib prefixes in the given static lib path."""
+        global_symbol_map = cc.get_global_symbol_section_map(lib_path)
+        libs = []
+        suffix = "___tvm_ffi__library_bin"
+        for name, _ in global_symbol_map.items():
+            if name.endswith(suffix):
+                model_lib = name[: -len(suffix)]
+                if model_lib.startswith("_"):
+                    model_lib = model_lib[1:]
+                libs.append(model_lib)
+        return libs
+
+    cc.create_staticlib(lib_path, tar_list)
+    available_model_libs = _get_model_libs(lib_path)
+    logger.info("Creating lib from %s", str(tar_list))
+    logger.info("Validating the library %s", str(lib_path))
+    logger.info(
+        "List of available model libs packaged: %s,"
+        " if we have '-' in the model_lib string, it will be turned into '_'",
+        str(available_model_libs),
+    )
+    global_symbol_map = cc.get_global_symbol_section_map(lib_path)
+    error_happened = False
+
+    for item in app_config["model_list"]:
+        model_lib = item["model_lib"]
+        model_id = item["model_id"]
+        if model_lib not in model_set:
+            # NOTE: this cannot happen under new setting
+            # since if model_lib is not included, it will be jitted
+            raise RuntimeError(
+                f"ValidationError: model_lib={model_lib} specified for model_id={model_id} "
+                "is not included in model_lib_path_for_prepare_libs argument, "
+                "This will cause the specific model not being able to load, "
+                f"model_lib_path_for_prepare_libs={model_lib_path_for_prepare_libs}"
+            )
+
+        model_prefix_pattern = model_lib.replace("-", "_") + "___tvm_ffi__library_bin"
+        if (
+            model_prefix_pattern not in global_symbol_map
+            and "_" + model_prefix_pattern not in global_symbol_map
+        ):
+            # NOTE: no lazy format is ok since this is a slow pass
+            model_lib_path = model_lib_path_for_prepare_libs[model_lib]
+            log_msg = (
+                "ValidationError:\n"
+                f"\tmodel_lib {model_lib} requested in {str(app_config_path)}"
+                f" is not found in {str(lib_path)}\n"
+                f"\tspecifically the model_lib for {model_lib_path}.\n"
+                f"\tcurrent available model_libs in {str(lib_path)}: {available_model_libs}\n"
+                f"\tThis can happen when we manually specified model_lib_path_for_prepare_libs"
+                f" in {str(package_config_path)}\n"
+                f"\tConsider remove model_lib_path_for_prepare_libs (so library can be jitted)"
+                "or check the compile command"
+            )
+            logger.info(log_msg)
+            error_happened = True
+
+    if not error_happened:
+        logger.info(style.green("Validation pass"))
+    else:
+        logger.info(style.red("Validation failed"))
+        sys.exit(255)
+
+
+def build_android_binding(mlc_llm_source_dir: Path, output: Path) -> None:
+    """Build android binding in MLC LLM"""
+    mlc4j_path = mlc_llm_source_dir / "android" / "mlc4j"
+
+    # Move the model libraries to "build/lib/" for linking
+    os.makedirs(Path("build") / "lib", exist_ok=True)
+    src_path = str(output / "lib" / "libmodel_android.a")
+    dst_path = str(Path("build") / "lib" / "libmodel_android.a")
+    logger.info('Moving "%s" to "%s"', src_path, dst_path)
+    shutil.move(src_path, dst_path)
+
+    # Build mlc4j
+    logger.info("Building mlc4j")
+    subprocess.run([sys.executable, mlc4j_path / "prepare_libs.py"], check=True, env=os.environ)
+    # Copy built files back to output directory.
+    lib_path = output / "lib" / "mlc4j"
+    os.makedirs(lib_path, exist_ok=True)
+    logger.info('Clean up all directories under "%s"', str(lib_path))
+    for content_path in lib_path.iterdir():
+        if content_path.is_dir():
+            shutil.rmtree(content_path)
+
+    src_path = str(mlc4j_path / "src")
+    dst_path = str(lib_path / "src")
+    logger.info('Copying "%s" to "%s"', src_path, dst_path)
+    shutil.copytree(src_path, dst_path)
+
+    src_path = str(mlc4j_path / "build.gradle")
+    dst_path = str(lib_path / "build.gradle")
+    logger.info('Copying "%s" to "%s"', src_path, dst_path)
+    shutil.copy(src_path, dst_path)
+
+    src_path = str(Path("build") / "output")
+    dst_path = str(lib_path / "output")
+    logger.info('Copying "%s" to "%s"', src_path, dst_path)
+    shutil.copytree(src_path, dst_path)
+
+    os.makedirs(lib_path / "src" / "main" / "assets")
+    src_path = str(output / "bundle" / "mlc-app-config.json")
+    dst_path = str(lib_path / "src" / "main" / "assets" / "mlc-app-config.json")
+    logger.info('Moving "%s" to "%s"', src_path, dst_path)
+    shutil.move(src_path, dst_path)
+
+
+def build_iphone_binding(mlc_llm_source_dir: Path, output: Path) -> None:
+    """Build iOS binding in MLC LLM"""
+    # Build iphone binding
+    logger.info("Build iphone binding")
+    subprocess.run(
+        ["bash", mlc_llm_source_dir / "ios" / "prepare_libs.sh"],
+        check=True,
+        env=os.environ,
+    )
+
+    # Copy built libraries back to output directory.
+    for static_library in (Path("build") / "lib").iterdir():
+        dst_path = str(output / "lib" / static_library.name)
+        logger.info('Copying "%s" to "%s"', static_library, dst_path)
+        shutil.copy(static_library, dst_path)
+
+
+def build_macabi_binding(mlc_llm_source_dir: Path, output: Path) -> None:
+    """Build Mac Catalyst binding in MLC LLM"""
+    deployment_target = os.environ.get("MLC_MACABI_DEPLOYMENT_TARGET", "18.0")
+    macabi_arch = os.environ.get("MLC_MACABI_ARCH", "").strip() or "arm64"
+    logger.info("Build macabi binding (deployment target %s)", deployment_target)
+    cmd = [
+        "bash",
+        str(mlc_llm_source_dir / "ios" / "prepare_libs.sh"),
+        "--catalyst",
+        "--deployment-target",
+        deployment_target,
+    ]
+    if macabi_arch:
+        cmd += ["--arch", macabi_arch]
+    subprocess.run(cmd, check=True, env=os.environ)
+
+    # Copy built libraries back to output directory.
+    build_dir = Path(f"build-maccatalyst-{macabi_arch}")
+    for static_library in (build_dir / "lib").iterdir():
+        dst_path = str(output / "lib" / static_library.name)
+        logger.info('Copying "%s" to "%s"', static_library, dst_path)
+        shutil.copy(static_library, dst_path)
+
+
+def package(
+    package_config_path: Path,
+    mlc_llm_source_dir: Path,
+    output: Path,
+) -> None:
+    """Python entrypoint of package."""
+    logger.info('MLC LLM HOME: "%s"', mlc_llm_source_dir)
+
+    # - Read package config.
+    with open(package_config_path, encoding="utf-8") as file:
+        package_config = json.load(file)
+    if not isinstance(package_config, dict):
+        raise ValueError(
+            "The content of MLC package config is expected to be a dict with "
+            f'field "model_list". However, the content of "{package_config_path}" is not a dict.'
+        )
+
+    # - Read device.
+    if "device" not in package_config:
+        raise ValueError(f'JSON file "{package_config_path}" is required to have field "device".')
+    device = package_config["device"]
+    if device not in SUPPORTED_DEVICES:
+        raise ValueError(
+            f'The "device" field of JSON file {package_config_path} is expected to be one of '
+            f'{SUPPORTED_DEVICES}, while "{device}" is given in the JSON.'
+        )
+
+    bundle_dir = output / "bundle"
+    app_config_path = bundle_dir / "mlc-app-config.json"
+    # - Build model libraries.
+    model_lib_path_for_prepare_libs = build_model_library(
+        package_config, device, bundle_dir, app_config_path
+    )
+    # - Validate model libraries.
+    validate_model_lib(
+        app_config_path,
+        package_config_path,
+        model_lib_path_for_prepare_libs,
+        device,
+        output,
+    )
+
+    # - Copy model libraries
+    if device == "android":
+        build_android_binding(mlc_llm_source_dir, output)
+    elif device == "iphone":
+        build_iphone_binding(mlc_llm_source_dir, output)
+    elif device == "macabi":
+        build_macabi_binding(mlc_llm_source_dir, output)
+    else:
+        assert False, "Cannot reach here"
+
+    logger.info("All finished.")

python/mlc_llm/interface/router.py

@@ -0,0 +1,125 @@
+"""Python entrypoint of router."""
+
+from collections.abc import AsyncGenerator
+from http import HTTPStatus
+from typing import List, Literal, Optional, Type  # noqa: UP035
+
+import fastapi
+import uvicorn
+from fastapi.middleware.cors import CORSMiddleware
+
+from mlc_llm.protocol import error_protocol
+from mlc_llm.protocol.openai_api_protocol import CompletionLogProbs, CompletionRequest
+from mlc_llm.router import Router
+from mlc_llm.serve import engine_base, engine_utils
+
+
+def serve(
+    model: str,
+    model_lib: Optional[str],
+    router_host: str,
+    router_port: int,
+    endpoint_hosts: List[str],  # noqa: UP006
+    endpoint_ports: List[int],  # noqa: UP006
+    endpoint_num_gpus: List[int],  # noqa: UP006
+    enable_prefix_cache: bool,
+    router_mode: Literal["disagg", "round-robin"] = "round-robin",
+    pd_balance_factor: float = 0.0,
+    router_type: Type[Router] = Router,  # noqa: UP006
+):
+    """Start the router with the specified configuration."""
+    # 1. Instantiate router
+    router = router_type(
+        model=model,
+        model_lib=model_lib,
+        hosts=endpoint_hosts,
+        ports=endpoint_ports,
+        num_gpus=endpoint_num_gpus,
+        enable_prefix_cache=enable_prefix_cache,
+        router_mode=router_mode,
+        pd_balance_factor=pd_balance_factor,
+    )
+
+    router_app = fastapi.APIRouter()
+
+    @router_app.post("/v1/completions")
+    async def request_completion(request: CompletionRequest, raw_request: fastapi.Request):
+        """OpenAI-compatible completion API.
+        API reference: https://platform.openai.com/docs/api-reference/completions/create
+        """
+        if router is None:
+            return error_protocol.create_error_response(
+                HTTPStatus.BAD_REQUEST, message="Router is not initialized."
+            )
+        request_id = f"cmpl-{engine_utils.random_uuid()}"
+
+        # Streaming response.
+        if request.stream:
+            # We manually get the first response from generator to
+            # capture potential exceptions in this scope, rather then
+            # the StreamingResponse scope.
+            stream_generator = router.handle_completion(request, request_id)
+            first_response = await anext(  # noqa: F821
+                stream_generator
+            )
+
+            async def completion_stream_generator() -> AsyncGenerator[str, None]:
+                if isinstance(first_response, StopAsyncIteration):
+                    yield "data: [DONE]\n\n"
+                    return
+                yield f"data: {first_response.model_dump_json(by_alias=True)}\n\n"
+                async for response in stream_generator:
+                    yield f"data: {response.model_dump_json(by_alias=True)}\n\n"
+                yield "data: [DONE]\n\n"
+
+            return fastapi.responses.StreamingResponse(
+                completion_stream_generator(), media_type="text/event-stream"
+            )
+
+        # FIXME: Non-streaming response not fully implemented
+        request_final_usage = None
+        output_texts = [""] * request.n
+        finish_reasons: List[Optional[str]] = [None] * request.n  # noqa: UP006
+        logprob_results: List[Optional[CompletionLogProbs]] = [None] * request.n  # noqa: UP006
+
+        async for response in router.handle_completion(request, request_id):
+            if await raw_request.is_disconnected():
+                # In non-streaming cases, the engine will not be notified
+                # when the request is disconnected.
+                # Therefore, we check if it is disconnected each time,
+                # and explicitly return.
+                # Note that requesta abort is triggered when the async for and funciton scope ends.
+                return error_protocol.create_error_response(
+                    HTTPStatus.BAD_REQUEST, message="The request has disconnected"
+                )
+            # TODO(Charlie): This is copied from engine.py --
+            # why is it here? Non-streaming only has a single chunk right?
+            # this is the final chunk
+            # if response.usage is not None:
+            #     request_final_usage = response.usage
+            #     continue
+            for choice in response.choices:
+                output_texts[choice.index] += choice.text
+                if choice.finish_reason is not None and finish_reasons[choice.index] is None:
+                    finish_reasons[choice.index] = choice.finish_reason
+                if choice.logprobs is not None:
+                    logprob_results[choice.index] = choice.logprobs
+
+        assert all(finish_reason is not None for finish_reason in finish_reasons)
+        return engine_base.wrap_completion_response(
+            request_id=request_id,
+            model=request.model,
+            output_texts=output_texts,
+            finish_reasons=finish_reasons,
+            logprob_results=logprob_results,
+            usage=request_final_usage,
+        )
+
+    # 2. Set up app
+    app = fastapi.FastAPI()
+    app.add_middleware(CORSMiddleware)
+    app.include_router(router_app)
+    app.exception_handler(error_protocol.BadRequestError)(error_protocol.bad_request_error_handler)
+
+    # 3. Run
+    uvicorn.run(app, host=router_host, port=router_port, log_level="info")

python/mlc_llm/interface/serve.py

@@ -0,0 +1,131 @@
+"""Python entrypoint of serve."""
+
+from typing import Any, List, Literal, Optional, Tuple, Union  # noqa: UP035
+
+import fastapi
+import uvicorn
+from fastapi.middleware.cors import CORSMiddleware
+
+from mlc_llm.protocol import error_protocol
+from mlc_llm.serve import engine
+from mlc_llm.serve.embedding_engine import AsyncEmbeddingEngine
+from mlc_llm.serve.entrypoints import (
+    debug_entrypoints,
+    metrics_entrypoints,
+    microserving_entrypoints,
+    openai_entrypoints,
+)
+from mlc_llm.serve.server import ServerContext
+from mlc_llm.support import logging
+
+logger = logging.getLogger(__name__)
+
+
+def serve(
+    model: str,
+    device: str,
+    model_lib: Optional[str],
+    mode: Literal["local", "interactive", "server"],
+    enable_debug: bool,
+    additional_models: List[Union[str, Tuple[str, str]]],  # noqa: UP006
+    embedding_model: Optional[str],
+    embedding_model_lib: Optional[str],
+    tensor_parallel_shards: Optional[int],
+    pipeline_parallel_stages: Optional[int],
+    opt: Optional[str],
+    max_num_sequence: Optional[int],
+    max_total_sequence_length: Optional[int],
+    max_single_sequence_length: Optional[int],
+    prefill_chunk_size: Optional[int],
+    sliding_window_size: Optional[int],
+    attention_sink_size: Optional[int],
+    max_history_size: Optional[int],
+    gpu_memory_utilization: Optional[float],
+    speculative_mode: Literal["disable", "small_draft", "eagle", "medusa"],
+    spec_draft_length: Optional[int],
+    spec_tree_width: Optional[int],
+    prefix_cache_mode: Literal["disable", "radix"],
+    prefix_cache_max_num_recycling_seqs: Optional[int],
+    prefill_mode: Literal["hybrid", "chunked"],
+    enable_tracing: bool,
+    host: str,
+    port: int,
+    allow_credentials: bool,
+    allow_origins: Any,
+    allow_methods: Any,
+    allow_headers: Any,
+    api_key: Optional[str] = None,
+):
+    """Serve the model with the specified configuration."""
+    # Create engine and start the background loop
+    async_engine = engine.AsyncMLCEngine(
+        model=model,
+        device=device,
+        model_lib=model_lib,
+        mode=mode,
+        engine_config=engine.EngineConfig(
+            additional_models=additional_models,
+            tensor_parallel_shards=tensor_parallel_shards,
+            pipeline_parallel_stages=pipeline_parallel_stages,
+            opt=opt,
+            max_num_sequence=max_num_sequence,
+            max_total_sequence_length=max_total_sequence_length,
+            max_single_sequence_length=max_single_sequence_length,
+            prefill_chunk_size=prefill_chunk_size,
+            sliding_window_size=sliding_window_size,
+            attention_sink_size=attention_sink_size,
+            max_history_size=max_history_size,
+            gpu_memory_utilization=gpu_memory_utilization,
+            speculative_mode=speculative_mode,
+            spec_draft_length=spec_draft_length,
+            spec_tree_width=spec_tree_width,
+            prefix_cache_mode=prefix_cache_mode,
+            prefix_cache_max_num_recycling_seqs=prefix_cache_max_num_recycling_seqs,
+            prefill_mode=prefill_mode,
+        ),
+        enable_tracing=enable_tracing,
+    )
+
+    # Set up embedding model if specified
+    emb_engine = None
+    if embedding_model is not None:
+        if embedding_model_lib is None:
+            raise ValueError(
+                "--embedding-model-lib is required when --embedding-model is specified."
+            )
+        emb_engine = AsyncEmbeddingEngine(
+            model=embedding_model,
+            model_lib=embedding_model_lib,
+            device=device,
+        )
+        logger.info("Embedding model %s loaded successfully.", embedding_model)
+
+    with ServerContext() as server_context:
+        server_context.add_model(model, async_engine)
+        if emb_engine is not None:
+            server_context.add_embedding_engine(embedding_model, emb_engine)
+        server_context.api_key = api_key
+
+        app = fastapi.FastAPI()
+        app.add_middleware(
+            CORSMiddleware,
+            allow_credentials=allow_credentials,
+            allow_origins=allow_origins,
+            allow_methods=allow_methods,
+            allow_headers=allow_headers,
+        )
+
+        app.include_router(openai_entrypoints.app)
+        app.include_router(metrics_entrypoints.app)
+        app.include_router(microserving_entrypoints.app)
+
+        server_context.enable_debug = enable_debug
+
+        if enable_debug:
+            app.include_router(debug_entrypoints.app)
+            logger.info("Enable debug endpoint and debug_config in requests...")
+
+        app.exception_handler(error_protocol.BadRequestError)(
+            error_protocol.bad_request_error_handler
+        )
+        uvicorn.run(app, host=host, port=port, log_level="info")

python/mlc_llm/json_ffi/__init__.py

@@ -0,0 +1,8 @@
+"""JSON FFI is a pure string based interface of MLC LLM Engine.
+
+We build interfacing with JSON FFI for both testing purposes
+and internal use. For most python API usage, please use MLCEngine
+and MLCAsyncEngine
+"""
+
+from .engine import JSONFFIEngine

python/mlc_llm/json_ffi/engine.py

@@ -0,0 +1,295 @@
+import json
+import queue
+import threading
+from collections.abc import Iterator
+from typing import Any, Callable, Dict, List, Literal, Optional, Union  # noqa: UP035
+
+import tvm
+
+from mlc_llm.protocol import debug_protocol, openai_api_protocol
+from mlc_llm.serve import engine_utils
+from mlc_llm.serve.engine_base import (
+    EngineConfig,
+    EngineMetrics,
+    _check_engine_config,
+    _parse_models,
+    _process_model_args,
+    _query_engine_metrics,
+    detect_device,
+)
+from mlc_llm.tokenizers import Tokenizer
+
+
+class EngineState:
+    sync_queue: queue.Queue
+
+    def get_request_stream_callback(self) -> Callable[[str], None]:
+        # ChatCompletionStreamResponse
+
+        def _callback(chat_completion_stream_responses_json_str: str) -> None:
+            self._sync_request_stream_callback(chat_completion_stream_responses_json_str)
+
+        return _callback
+
+    def _sync_request_stream_callback(self, chat_completion_stream_responses_json_str: str) -> None:
+        # Put the delta outputs to the queue in the unblocking way.
+        self.sync_queue.put_nowait(chat_completion_stream_responses_json_str)
+
+    def handle_chat_completion(
+        self, ffi: dict, request_json_str: str, include_usage: bool, request_id: str
+    ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]:
+        """Helper class to handle chat completion
+
+        Note
+        ----
+        ffi is explicitly passed in to avoid cylic dependency
+        as ffi will capture EngineState
+        """
+        self.sync_queue = queue.Queue()
+
+        ffi["chat_completion"](request_json_str, request_id)
+
+        try:
+            last_chunk_arrived = False
+            while not last_chunk_arrived:
+                chat_completion_responses_json_str = self.sync_queue.get()
+                chat_completion_responses_list = json.loads(chat_completion_responses_json_str)
+                for chat_completion_response_json_dict in chat_completion_responses_list:
+                    chat_completion_response = (
+                        openai_api_protocol.ChatCompletionStreamResponse.model_validate(
+                            chat_completion_response_json_dict
+                        )
+                    )
+                    # the chunk with usage is always the last chunk
+                    if chat_completion_response.usage is not None:
+                        if include_usage:
+                            yield chat_completion_response
+                        last_chunk_arrived = True
+                        break
+                    yield chat_completion_response
+        except Exception as exception:
+            ffi["abort"](request_id)
+            raise exception
+
+
+class BackgroundLoops:
+    """Helper class to keep track of background loops"""
+
+    def __init__(self, ffi: dict):
+        self._ffi = ffi
+        # important: avoid self reference in closure
+        background_loop = self._ffi["run_background_loop"]
+        background_stream_back_loop = self._ffi["run_background_stream_back_loop"]
+
+        # Create the background engine-driving thread and start the loop.
+        self._background_loop_thread: threading.Thread = threading.Thread(target=background_loop)
+        self._background_stream_back_loop_thread: threading.Thread = threading.Thread(
+            target=background_stream_back_loop
+        )
+        self._background_loop_thread.start()
+        self._background_stream_back_loop_thread.start()
+        self._terminated = False
+
+    def __del__(self):
+        self.terminate()
+
+    def terminate(self):
+        if self._terminated:
+            return
+        self._terminated = True
+        self._ffi["exit_background_loop"]()
+        self._background_loop_thread.join()
+        self._background_stream_back_loop_thread.join()
+
+
+class Completions:
+    """Completions class to be compatible with OpenAI API"""
+
+    _ffi: dict
+    _state: EngineState
+    _background_loops: BackgroundLoops
+
+    def __init__(self, ffi: dict, state: EngineState, background_loops: BackgroundLoops):
+        self._ffi = ffi
+        self._state = state
+        self._background_loops = background_loops
+
+    def create(
+        self,
+        *,
+        messages: List[Dict[str, Any]],  # noqa: UP006
+        model: Optional[str] = None,
+        frequency_penalty: Optional[float] = None,
+        presence_penalty: Optional[float] = None,
+        logprobs: bool = False,
+        top_logprobs: int = 0,
+        logit_bias: Optional[Dict[int, float]] = None,  # noqa: UP006
+        max_tokens: Optional[int] = None,
+        n: int = 1,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,  # noqa: UP006
+        stream: bool = True,
+        stream_options: Optional[Dict[str, Any]] = None,  # noqa: UP006
+        temperature: Optional[float] = None,
+        top_p: Optional[float] = None,
+        tools: Optional[List[Dict[str, Any]]] = None,  # noqa: UP006
+        tool_choice: Optional[Union[Literal["none", "auto"], Dict]] = None,  # noqa: UP006
+        user: Optional[str] = None,
+        response_format: Optional[Dict[str, Any]] = None,  # noqa: UP006
+        request_id: Optional[str] = None,
+        extra_body: Optional[Dict[str, Any]] = None,  # noqa: UP006
+    ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]:
+        if request_id is None:
+            request_id = f"chatcmpl-{engine_utils.random_uuid()}"
+        debug_config = extra_body.get("debug_config", None) if extra_body is not None else None
+        if not stream:
+            raise ValueError("JSONFFIEngine only support stream=True")
+        request = openai_api_protocol.ChatCompletionRequest(
+            messages=[
+                openai_api_protocol.ChatCompletionMessage.model_validate(message)
+                for message in messages
+            ],
+            model=model,
+            frequency_penalty=frequency_penalty,
+            presence_penalty=presence_penalty,
+            logprobs=logprobs,
+            top_logprobs=top_logprobs,
+            logit_bias=logit_bias,
+            max_tokens=max_tokens,
+            n=n,
+            seed=seed,
+            stop=stop,
+            stream=stream,
+            stream_options=(
+                openai_api_protocol.StreamOptions.model_validate(stream_options)
+                if stream_options is not None
+                else None
+            ),
+            temperature=temperature,
+            top_p=top_p,
+            tools=(
+                [openai_api_protocol.ChatTool.model_validate(tool) for tool in tools]
+                if tools is not None
+                else None
+            ),
+            tool_choice=tool_choice,
+            user=user,
+            response_format=(
+                openai_api_protocol.RequestResponseFormat.model_validate(response_format)
+                if response_format is not None
+                else None
+            ),
+            debug_config=(
+                debug_protocol.DebugConfig.model_validate(debug_config)
+                if debug_config is not None
+                else None
+            ),
+        )
+        chatcmpl_generator = self._state.handle_chat_completion(
+            self._ffi,
+            request.model_dump_json(by_alias=True),
+            include_usage=(
+                request.stream_options is not None and request.stream_options.include_usage
+            ),
+            request_id=request_id,
+        )
+        for response in chatcmpl_generator:
+            yield response
+
+
+class Chat:
+    """Chat class to be compatible with OpenAI API"""
+
+    completions: Completions
+
+    def __init__(self, ffi: dict, state: EngineState, background_loops: BackgroundLoops):
+        self.completions = Completions(ffi, state, background_loops)
+
+
+class JSONFFIEngine:
+    chat: Chat
+
+    def __init__(
+        self,
+        model: str,
+        device: Union[str, tvm.runtime.Device] = "auto",
+        *,
+        model_lib: Optional[str] = None,
+        mode: Literal["local", "interactive", "server"] = "local",
+        engine_config: Optional[EngineConfig] = None,
+    ) -> None:
+        # - Check the fields fields of `engine_config`.
+        if engine_config is None:
+            engine_config = EngineConfig()
+        _check_engine_config(model, model_lib, mode, engine_config)
+
+        # - Initialize model loading info.
+        models = _parse_models(model, model_lib, engine_config.additional_models)
+        if isinstance(device, str):
+            device = detect_device(device)
+        assert isinstance(device, tvm.runtime.Device)
+        model_args = _process_model_args(models, device, engine_config)[0]
+
+        # - Load the raw model config into dict
+        for i, model_info in enumerate(models):
+            model_info.model_lib = model_args[i][1]
+
+        # - Initialize engine state and engine.
+        self._state = EngineState()
+        module = tvm.get_global_func("mlc.json_ffi.CreateJSONFFIEngine", allow_missing=False)()
+        self._ffi = {
+            key: module[key]
+            for key in [
+                "init_background_engine",
+                "reload",
+                "unload",
+                "reset",
+                "chat_completion",
+                "abort",
+                "run_background_loop",
+                "run_background_stream_back_loop",
+                "exit_background_loop",
+            ]
+        }
+        self.tokenizer = Tokenizer(model_args[0][0])
+        self._background_loops = BackgroundLoops(self._ffi)
+
+        engine_config.model = model_args[0][0]
+        engine_config.model_lib = model_args[0][1]
+        engine_config.additional_models = model_args[1:]
+        engine_config.mode = mode
+        self.engine_config = engine_config
+
+        self._ffi["init_background_engine"](
+            device.dlpack_device_type(),
+            device.index,
+            self._state.get_request_stream_callback(),
+        )
+        self._ffi["reload"](self.engine_config.asjson())
+
+        self.chat = Chat(self._ffi, self._state, self._background_loops)
+
+    def metrics(self) -> EngineMetrics:
+        """Get the engine metrics."""
+        return _query_engine_metrics(self)
+
+    def _raw_chat_completion(
+        self, request_json_str: str, include_usage: bool, request_id: str
+    ) -> Iterator[openai_api_protocol.ChatCompletionStreamResponse]:
+        """Raw chat completion API"""
+        return self._state.handle_chat_completion(
+            self._ffi, request_json_str, include_usage, request_id
+        )
+
+    def terminate(self):
+        """Explicitly terminate the engine"""
+        self._background_loops.terminate()
+
+    def _test_reload(self):
+        self._ffi["reload"](self.engine_config.asjson())
+
+    def _test_reset(self):
+        self._ffi["reset"]()
+
+    def _test_unload(self):
+        self._ffi["unload"]()

python/mlc_llm/libinfo.py

@@ -0,0 +1,71 @@
+"""Library information. This is a standalone file that can be used to get various info"""
+
+#! pylint: disable=protected-access
+import os
+import sys
+
+__version__ = "0.1.dev0"
+MLC_LIBRARY_PATH = os.environ.get("MLC_LIBRARY_PATH", None)
+
+
+def get_env_paths(env_var, splitter):
+    """Get path in env variable"""
+    if os.environ.get(env_var, None):
+        return [p.strip() for p in os.environ[env_var].split(splitter)]
+    return []
+
+
+def get_dll_directories():
+    """Get extra mlc llm dll directories"""
+    curr_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
+    source_dir = os.path.abspath(os.path.join(curr_dir, "..", ".."))
+    dll_path = [
+        curr_dir,
+        os.path.join(source_dir, "build"),
+        os.path.join(source_dir, "build", "Release"),
+    ]
+    if MLC_LIBRARY_PATH:
+        dll_path.append(MLC_LIBRARY_PATH)
+    if "CONDA_PREFIX" in os.environ:
+        dll_path.append(os.path.join(os.environ["CONDA_PREFIX"], "lib"))
+    if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"):
+        dll_path.extend(get_env_paths("LD_LIBRARY_PATH", ":"))
+    elif sys.platform.startswith("darwin"):
+        dll_path.extend(get_env_paths("DYLD_LIBRARY_PATH", ":"))
+    elif sys.platform.startswith("win32"):
+        dll_path.extend(get_env_paths("PATH", ";"))
+    return [os.path.abspath(p) for p in dll_path if os.path.isdir(p)]
+
+
+def find_lib_path(name, optional=False):
+    """Find mlc llm library
+
+    Parameters
+    ----------
+    name : str
+        The name of the library
+
+    optional: boolean
+        Whether the library is required
+    """
+    if sys.platform.startswith("linux") or sys.platform.startswith("freebsd"):
+        lib_name = f"lib{name}.so"
+    elif sys.platform.startswith("win32"):
+        lib_name = f"{name}.dll"
+    elif sys.platform.startswith("darwin"):
+        lib_name = f"lib{name}.dylib"
+    else:
+        lib_name = f"lib{name}.so"
+
+    dll_paths = get_dll_directories()
+    lib_dll_path = [os.path.join(p, lib_name) for p in dll_paths]
+    lib_found = [p for p in lib_dll_path if os.path.exists(p) and os.path.isfile(p)]
+    if not lib_found:
+        if not optional:
+            message = (
+                f"Cannot find libraries: {lib_name}\n"
+                + "List of candidates:\n"
+                + "\n".join(lib_dll_path)
+            )
+            raise RuntimeError(message)
+    return lib_found

python/mlc_llm/loader/__init__.py

@@ -0,0 +1,8 @@
+"""
+A subpackage of the compiler that represents mapping between external parameters, quantized
+parameters and parameters in MLC-defined models.
+"""
+
+from .huggingface_loader import HuggingFaceLoader
+from .loader import LOADER, Loader
+from .mapping import ExternMapping, QuantizeMapping