chore: import upstream snapshot with attribution

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