chore: import upstream snapshot with attribution

python/mlc_llm/support/__init__.py

@@ -0,0 +1,4 @@
+"""
+Common utilities used in the Python package. Do not import anything by default,
+as they may introduce unnecessary dependencies.
+"""

python/mlc_llm/support/argparse.py

@@ -0,0 +1,16 @@
+"""An enhanced argument parser for mlc-chat."""
+
+import argparse
+import sys
+
+
+class ArgumentParser(argparse.ArgumentParser):
+    """An enhanced argument parser for mlc-chat."""
+
+    def error(self, message):
+        """Overrides the behavior when erroring out"""
+        print("-" * 25 + " Usage " + "-" * 25)
+        self.print_help()
+        print("-" * 25 + " Error " + "-" * 25)
+        print(message, file=sys.stderr)
+        sys.exit(2)

python/mlc_llm/support/auto_config.py

@@ -0,0 +1,190 @@
+"""Help function for detecting the model configuration file `config.json`"""
+
+import json
+import tempfile
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from . import logging
+from .style import bold, green
+
+if TYPE_CHECKING:
+    from mlc_llm.model import Model
+    from mlc_llm.quantization import Quantization
+
+
+logger = logging.getLogger(__name__)
+
+FOUND = green("Found")
+
+
+def detect_mlc_chat_config(mlc_chat_config: str) -> Path:
+    """Detect and return the path that points to mlc-chat-config.json.
+    If `mlc_chat_config` is a directory, it looks for mlc-chat-config.json below it.
+
+    Parameters
+    ---------
+    mlc_chat_config : str
+        The path to `mlc-chat-config.json`, or the directory containing
+        `mlc-chat-config.json`.
+
+    Returns
+    -------
+    mlc_chat_config_json_path : pathlib.Path
+        The path points to mlc_chat_config.json.
+    """
+    from mlc_llm.model import MODEL_PRESETS
+
+    from .download_cache import download_and_cache_mlc_weights
+
+    if mlc_chat_config.startswith("HF://") or mlc_chat_config.startswith("http"):
+        mlc_chat_config_path = Path(download_and_cache_mlc_weights(model_url=mlc_chat_config))
+    elif isinstance(mlc_chat_config, str) and mlc_chat_config in MODEL_PRESETS:
+        logger.info("%s mlc preset model: %s", FOUND, mlc_chat_config)
+        content = MODEL_PRESETS[mlc_chat_config].copy()
+        content["model_preset_tag"] = mlc_chat_config
+        temp_file = tempfile.NamedTemporaryFile(
+            suffix=".json",
+            delete=False,
+        )
+        logger.info("Dumping config to: %s", temp_file.name)
+        mlc_chat_config_path = Path(temp_file.name)
+        with mlc_chat_config_path.open("w", encoding="utf-8") as mlc_chat_config_file:
+            json.dump(content, mlc_chat_config_file, indent=2)
+    else:
+        mlc_chat_config_path = Path(mlc_chat_config)
+    if not mlc_chat_config_path.exists():
+        raise ValueError(f"{mlc_chat_config_path} does not exist.")
+
+    if mlc_chat_config_path.is_dir():
+        # search mlc-chat-config.json under path
+        mlc_chat_config_json_path = mlc_chat_config_path / "mlc-chat-config.json"
+        if not mlc_chat_config_json_path.exists():
+            raise ValueError(f"Fail to find mlc-chat-config.json under {mlc_chat_config_path}.")
+    else:
+        mlc_chat_config_json_path = mlc_chat_config_path
+
+    logger.info("%s model configuration: %s", FOUND, mlc_chat_config_json_path)
+    return mlc_chat_config_json_path
+
+
+def detect_config(config: str) -> Path:
+    """Detect and return the path that points to config.json. If `config` is a directory,
+    it looks for config.json below it.
+
+    Parameters
+    ---------
+    config : str
+        The preset name of the model, or the path to `config.json`, or the directory containing
+        `config.json`.
+
+    Returns
+    -------
+    config_json_path : pathlib.Path
+        The path points to config.json.
+    """
+    from mlc_llm.model import MODEL_PRESETS
+
+    if isinstance(config, str) and config in MODEL_PRESETS:
+        logger.info("%s preset model: %s", FOUND, config)
+        content = MODEL_PRESETS[config].copy()
+        content["model_preset_tag"] = config
+        temp_file = tempfile.NamedTemporaryFile(
+            suffix=".json",
+            delete=False,
+        )
+        logger.info("Dumping config to: %s", temp_file.name)
+        config_path = Path(temp_file.name)
+        with config_path.open("w", encoding="utf-8") as config_file:
+            json.dump(content, config_file, indent=2)
+    else:
+        config_path = Path(config)
+    if not config_path.exists():
+        raise ValueError(f"{config_path} does not exist.")
+
+    if config_path.is_dir():
+        # search config.json under config path
+        config_json_path = config_path / "config.json"
+        if not config_json_path.exists():
+            raise ValueError(f"Fail to find config.json under {config_path}.")
+    else:
+        config_json_path = config_path
+
+    logger.info("%s model configuration: %s", FOUND, config_json_path)
+    return config_json_path
+
+
+def detect_model_type(model_type: str, config: Path) -> "Model":
+    """Detect the model type from the configuration file. If `model_type` is "auto", it will be
+    inferred from the configuration file. Otherwise, it will be used as the model type, and sanity
+    check will be performed.
+
+    Parameters
+    ----------
+    model_type : str
+        The model type, for example, "llama".
+
+    config : pathlib.Path
+        The path to config.json.
+
+    Returns
+    -------
+    model : mlc_llm.compiler.Model
+        The model type.
+    """
+
+    from mlc_llm.model import MODELS
+
+    if model_type == "auto":
+        with open(config, encoding="utf-8") as config_file:
+            cfg = json.load(config_file)
+        if "model_type" not in cfg and (
+            "model_config" not in cfg or "model_type" not in cfg["model_config"]
+        ):
+            raise ValueError(
+                f"'model_type' not found in: {config}. "
+                f"Please explicitly specify `--model-type` instead."
+            )
+        model_type = cfg["model_type"] if "model_type" in cfg else cfg["model_config"]["model_type"]
+    if model_type in ["mixformer-sequential"]:
+        model_type = "phi-msft"
+    logger.info("%s model type: %s. Use `--model-type` to override.", FOUND, bold(model_type))
+    if model_type not in MODELS:
+        raise ValueError(f"Unknown model type: {model_type}. Available ones: {list(MODELS.keys())}")
+    return MODELS[model_type]
+
+
+def detect_quantization(quantization_arg: str, config: Path) -> "Quantization":
+    """Detect the model quantization scheme from the configuration file or `--quantization`
+    argument. If `--quantization` is provided, it will override the value on the configuration
+    file.
+
+    Parameters
+    ----------
+    quantization_arg : str
+        The quantization scheme, for example, "q4f16_1".
+
+    config : pathlib.Path
+        The path to mlc-chat-config.json.
+
+    Returns
+    -------
+    quantization : mlc_llm.quantization.Quantization
+        The model quantization scheme.
+    """
+    from mlc_llm.quantization import (
+        QUANTIZATION,
+    )
+
+    with open(config, encoding="utf-8") as config_file:
+        cfg = json.load(config_file)
+    if quantization_arg is not None:
+        quantization = QUANTIZATION[quantization_arg]
+    elif "quantization" in cfg:
+        quantization = QUANTIZATION[cfg["quantization"]]
+    else:
+        raise ValueError(
+            f"'quantization' not found in: {config}. "
+            f"Please explicitly specify `--quantization` instead."
+        )
+    return quantization

python/mlc_llm/support/auto_device.py

@@ -0,0 +1,93 @@
+"""Automatic detection of the device available on the local machine."""
+
+import os
+import subprocess
+import sys
+from typing import Dict, Optional  # noqa: UP035
+
+import tvm
+from tvm.runtime import Device
+from tvm_ffi import DLDeviceType
+
+from . import logging
+from .style import bold, green, red
+
+FOUND = green("Found")
+NOT_FOUND = red("Not found")
+AUTO_DETECT_DEVICES = ["cuda", "rocm", "metal", "vulkan", "opencl", "cpu"]
+_RESULT_CACHE: Dict[str, bool] = {}  # noqa: UP006
+
+
+logger = logging.getLogger(__name__)
+
+
+def detect_device(device_hint: str) -> Optional[Device]:
+    """Detect locally available device from string hint."""
+    if device_hint == "auto":
+        device = None
+        for device_type in AUTO_DETECT_DEVICES:
+            cur_device = tvm.device(device_type=device_type, index=0)
+            if _device_exists(cur_device):
+                if device is None:
+                    device = cur_device
+        if device is None:
+            logger.info("%s: No available device detected", NOT_FOUND)
+            return None
+        logger.info("Using device: %s", bold(device2str(device)))
+        return device
+    try:
+        device = tvm.device(device_hint)
+    except Exception as err:
+        raise ValueError(f"Invalid device name: {device_hint}") from err
+    if not _device_exists(device):
+        raise ValueError(f"Device is not found on your local environment: {device_hint}")
+    return device
+
+
+def device2str(device: Device) -> str:
+    """Convert a TVM device object to string."""
+    return f"{tvm.runtime.Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]}:{device.index}"
+
+
+def _device_exists(device: Device) -> bool:
+    device_type = tvm.runtime.Device._DEVICE_TYPE_TO_NAME[device.dlpack_device_type()]
+    device_str = device2str(device)
+    if device_str in _RESULT_CACHE:
+        return _RESULT_CACHE[device_str]
+    cmd = [
+        sys.executable,
+        "-m",
+        "mlc_llm.cli.check_device",
+        device_type,
+    ]
+    prefix = "check_device:"
+    subproc_outputs = [
+        line[len(prefix) :].strip()
+        for line in subprocess.run(
+            cmd,
+            capture_output=True,
+            text=True,
+            check=False,
+            env=os.environ,
+        )
+        .stdout.strip()
+        .splitlines()
+        if line.startswith(prefix)
+    ]
+    if subproc_outputs:
+        if subproc_outputs[0]:
+            for i in subproc_outputs[0].split(","):
+                logger.info("%s device: %s:%s", FOUND, device_type, i)
+                _RESULT_CACHE[f"{device_type}:{i}"] = True
+                if device.dlpack_device_type() == DLDeviceType.kDLCPU:
+                    break
+    else:
+        logger.error(
+            "GPU device detection failed. Please report this issue with the output of command: %s",
+            " ".join(cmd),
+        )
+    if device_str in _RESULT_CACHE:
+        return _RESULT_CACHE[device_str]
+    logger.info("%s device: %s", NOT_FOUND, device_str)
+    _RESULT_CACHE[device_str] = False
+    return False

python/mlc_llm/support/auto_target.py

@@ -0,0 +1,554 @@
+"""Helper functions for target auto-detection."""
+
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple  # noqa: UP035
+
+from tvm import IRModule, relax
+from tvm.ir.transform import Pass
+from tvm.support import ndk, tar, xcode
+from tvm.target import Target
+from tvm_ffi import get_global_func, register_global_func
+
+from . import logging
+from .auto_device import AUTO_DETECT_DEVICES, detect_device, device2str
+from .constants import MLC_MULTI_ARCH
+from .style import bold, green, red
+
+if TYPE_CHECKING:
+    from mlc_llm.compiler.compile import CompileArgs
+
+
+logger = logging.getLogger(__name__)
+
+# TODO: add help message on how to specify the target manually
+HELP_MSG = """TBD"""
+FOUND = green("Found")
+NOT_FOUND = red("Not found")
+BuildFunc = Callable[[IRModule, "CompileArgs", Pass], None]
+
+
+def detect_target_and_host(
+    target_hint: str,
+    host_hint: str = "auto",
+    enable_subgroups: Optional[bool] = None,
+) -> Tuple[Target, BuildFunc]:  # noqa: UP006
+    """Detect the configuration for the target device and its host, for example, target GPU and
+    the host CPU.
+
+    Parameters
+    ----------
+    target_hint : str
+        The hint for the target device.
+
+    host_hint : str
+        The hint for the host CPU, default is "auto".
+    """
+    target, build_func = _detect_target_gpu(target_hint)
+    if target.host is None:
+        target = Target(target, host=_detect_target_host(host_hint))
+    target = _apply_webgpu_subgroups(target, enable_subgroups)
+    if target.kind.name == "cuda":
+        # Enable thrust for CUDA
+        target_dict = dict(target.export())
+        target_dict["libs"] = (
+            (target_dict["libs"] + ["thrust"]) if "libs" in target_dict else ["thrust"]
+        )
+        target = Target(target_dict)
+        _register_cuda_hook(target)
+    elif target.kind.name == "rocm":
+        target_dict = dict(target.export())
+        extra_libs = ["thrust", "rocblas", "miopen", "hipblas"]
+        target_dict["libs"] = (
+            (target_dict["libs"] + extra_libs) if "libs" in target_dict else extra_libs
+        )
+        target = Target(target_dict)
+    return target, build_func
+
+
+def _apply_webgpu_subgroups(target: Target, enable_subgroups: Optional[bool]) -> Target:
+    if not enable_subgroups:
+        return target
+    if target.kind.name != "webgpu":
+        logger.warning(
+            "--enable-subgroups is only supported for WebGPU targets; ignoring for %s",
+            target.kind.name,
+        )
+        return target
+    target_dict = dict(target.export())
+    target_dict["supports_subgroups"] = True
+    return Target(target_dict)
+
+
+def _detect_target_gpu(hint: str) -> Tuple[Target, BuildFunc]:  # noqa: UP006
+    if hint in ["iphone", "macabi", "android", "webgpu", "mali", "opencl"]:
+        hint += ":generic"
+    if hint == "auto" or hint in AUTO_DETECT_DEVICES:
+        target: Optional[Target] = None
+        device = detect_device(hint)
+        if device is not None:
+            device_str = device2str(device)
+            try:
+                target = Target.from_device(device)
+            except ValueError:
+                logger.info("%s: Cannot detect target from device: %s", NOT_FOUND, device_str)
+        if target is None:
+            raise ValueError(f"No target detected from device: {hint}. Please specify explicitly")
+        logger.info(
+            '%s configuration of target device "%s": %s',
+            FOUND,
+            bold(device_str),
+            target.export(),
+        )
+        return target, _build_default()
+    if hint in PRESET:
+        preset = PRESET[hint]
+        target = Target(preset["target"])
+        build = preset.get("build", _build_default)
+        return target, build()
+    if _is_device(hint):
+        logger.info("Detecting target device: %s", hint)
+        target = Target.from_device(hint)
+        logger.info("%s target: %s", FOUND, target.export())
+        return target, _build_default()
+    try:
+        logger.info("Try creating device target from string: %s", hint)
+        target = Target(hint)
+        logger.info("%s target: %s", FOUND, target.export())
+        return target, _build_default()
+    except Exception as err:
+        logger.info("%s: Failed to create target", NOT_FOUND)
+        raise ValueError(f"Invalid target: {hint}") from err
+
+
+def _detect_target_host(hint: str) -> Target:
+    """Detect the host CPU architecture."""
+    if hint == "auto":
+        target_triple = get_global_func("tvm.codegen.llvm.GetDefaultTargetTriple")()
+        target = Target.from_device("cpu")
+        logger.info("%s host LLVM triple: %s", FOUND, bold(target.attrs["mtriple"]))
+        logger.info("%s host LLVM CPU: %s", FOUND, bold(target.attrs["mcpu"]))
+        return target
+    target_triple = hint
+    logger.info("Using LLVM triple specified by --host: %s", bold(target_triple))
+    return Target({"kind": "llvm", "mtriple": target_triple})
+
+
+def _is_device(device: str):
+    if " " in device:
+        return False
+    if device.count(":") != 1:
+        return False
+    return True
+
+
+def _add_system_lib_prefix(mod: IRModule, prefix: str, is_system_lib: bool) -> IRModule:
+    if is_system_lib and prefix:
+        mod = mod.with_attrs({"system_lib_prefix": prefix})
+    elif is_system_lib:
+        logger.warning(
+            "%s is not specified when building a static library",
+            bold("--system-lib-prefix"),
+        )
+    elif prefix:
+        logger.warning(
+            "--system-lib-prefix is specified, but it will not take any effect "
+            "when building the shared library"
+        )
+    return mod
+
+
+def _build_metal_x86_64():
+    def build(mod: IRModule, args: "CompileArgs", pipeline=None):
+        output = args.output
+        mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=False)
+        assert output.suffix == ".dylib"
+        relax.build(
+            mod,
+            target=args.target,
+            relax_pipeline=pipeline,
+        ).export_library(
+            str(output),
+            fcompile=xcode.create_dylib,
+            sdk="macosx",
+            arch="x86_64",
+        )
+
+    return build
+
+
+def _build_iphone():
+    @register_global_func("tvm_callback_metal_compile", override=True)
+    def compile_metal(src, target):
+        libs = target.attrs.get("libs", None)
+        if libs:
+            return xcode.compile_metal(src, sdk=libs[0])
+        return xcode.compile_metal(src)
+
+    def build(mod: IRModule, args: "CompileArgs", pipeline=None):
+        output = args.output
+        mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True)
+        assert output.suffix == ".tar"
+        relax.build(
+            mod,
+            target=args.target,
+            relax_pipeline=pipeline,
+            system_lib=True,
+        ).export_library(
+            str(output),
+            fcompile=tar.tar,
+        )
+
+    return build
+
+
+def _build_android():
+    def build(mod: IRModule, args: "CompileArgs", pipeline=None):
+        output = args.output
+        mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True)
+        assert output.suffix == ".tar"
+        ex = relax.build(
+            mod,
+            target=args.target,
+            relax_pipeline=pipeline,
+            system_lib=True,
+        )
+        ex.export_library(
+            str(output),
+            fcompile=tar.tar,
+        )
+        if args.debug_dump is not None:
+            source = ex.mod.imports[0].imports[0].inspect_source()
+            with open(args.debug_dump / "kernel.cl", "w", encoding="utf-8") as f:
+                f.write(source)
+
+    return build
+
+
+def _build_android_so():
+    def build(mod: IRModule, args: "CompileArgs", pipeline=None):
+        output = args.output
+        mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=False)
+        assert output.suffix == ".so"
+        ex = relax.build(
+            mod,
+            target=args.target,
+            relax_pipeline=pipeline,
+            system_lib=False,
+        )
+        ex.export_library(
+            str(output),
+            fcompile=ndk.create_shared,
+        )
+        if args.debug_dump is not None:
+            source = ex.mod.imports[0].imports[0].inspect_source()
+            with open(args.debug_dump / "kernel.cl", "w", encoding="utf-8") as f:
+                f.write(source)
+
+    return build
+
+
+def _build_webgpu():
+    def build(mod: IRModule, args: "CompileArgs", pipeline=None):
+        output = args.output
+        mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True)
+        assert output.suffix == ".wasm"
+
+        # Try to locate `mlc_wasm_runtime.bc`
+        bc_path = None
+        bc_candidates = ["web/dist/wasm/mlc_wasm_runtime.bc"]
+        if os.environ.get("MLC_LLM_SOURCE_DIR", None):
+            mlc_source_home_dir = os.environ["MLC_LLM_SOURCE_DIR"]
+            bc_candidates.append(
+                os.path.join(mlc_source_home_dir, "web", "dist", "wasm", "mlc_wasm_runtime.bc")
+            )
+        error_info = (
+            "Cannot find library: mlc_wasm_runtime.bc\n"
+            + "Make sure you have run `./web/prep_emcc_deps.sh` and "
+            + "`export MLC_LLM_SOURCE_DIR=/path/to/mlc-llm` so that we can locate the file. "
+            + "We tried to look at candidate paths:\n"
+        )
+        for candidate in bc_candidates:
+            error_info += candidate + "\n"
+            if Path(candidate).exists():
+                bc_path = candidate
+        if not bc_path:
+            raise RuntimeError(error_info)
+
+        relax.build(
+            mod,
+            target=args.target,
+            relax_pipeline=pipeline,
+            system_lib=True,
+        ).export_library(
+            str(output),
+            libs=[bc_path],
+        )
+
+    return build
+
+
+def _build_mali():
+    def build(mod: IRModule, args: "CompileArgs", pipeline=None):
+        output = args.output
+        mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=True)
+        assert output.suffix == ".so"
+        mod = relax.build(
+            mod,
+            target=args.target,
+            relax_pipeline=pipeline,
+            system_lib=True,
+        )
+        if "TVM_NDK_CC" in os.environ:
+            mod.export_library(str(output), fcompile=ndk.create_shared)
+        else:
+            mod.export_library(str(output))
+
+    return build
+
+
+def _build_default():
+    def build(mod: IRModule, args: "CompileArgs", pipeline=None):
+        output = args.output
+        if output.suffix in [".tar", ".lib"]:
+            system_lib = True
+        elif output.suffix in [".so", ".dylib", ".dll"]:
+            system_lib = False
+        else:
+            logger.warning("Unknown output suffix: %s. Assuming shared library.", output.suffix)
+            system_lib = False
+        mod = _add_system_lib_prefix(mod, args.system_lib_prefix, is_system_lib=system_lib)
+        relax.build(
+            mod,
+            target=args.target,
+            relax_pipeline=pipeline,
+            system_lib=system_lib,
+        ).export_library(
+            str(output),
+        )
+
+    return build
+
+
+def detect_cuda_arch_list(target: Target) -> List[int]:  # noqa: UP006
+    """Detect the CUDA architecture list from the target."""
+
+    def convert_to_num(arch_str):
+        arch_num_str = "".join(filter(str.isdigit, arch_str))
+        assert arch_num_str, f"'{arch_str}' does not contain any digits"
+        return int(arch_num_str)
+
+    assert target.kind.name == "cuda", f"Expect target to be CUDA, but got {target}"
+    if MLC_MULTI_ARCH is not None:
+        multi_arch = [convert_to_num(x) for x in MLC_MULTI_ARCH.split(",")]
+    else:
+        assert target.attrs.get("arch", "").startswith("sm_")
+        multi_arch = [convert_to_num(target.attrs.get("arch")[3:])]
+    multi_arch = list(set(multi_arch))
+    return multi_arch
+
+
+def _register_cuda_hook(target: Target):
+    if MLC_MULTI_ARCH is None:
+        default_arch = target.attrs.get("arch", None)
+        logger.info("Generating code for CUDA architecture: %s", bold(default_arch))
+        logger.info(
+            "To produce multi-arch fatbin, set environment variable %s. "
+            "Example: MLC_MULTI_ARCH=70,72,75,80,86,87,89,90a",
+            bold("MLC_MULTI_ARCH"),
+        )
+        multi_arch = None
+    else:
+        logger.info("%s %s: %s", FOUND, bold("MLC_MULTI_ARCH"), MLC_MULTI_ARCH)
+        multi_arch = [x.strip() for x in MLC_MULTI_ARCH.split(",")]
+        logger.info("Generating code for CUDA architecture: %s", multi_arch)
+
+    @register_global_func("tvm_callback_cuda_compile", override=True)
+    def tvm_callback_cuda_compile(code):
+        """use nvcc to generate fatbin code for better optimization"""
+        from tvm.support import nvcc
+
+        if multi_arch is None:
+            ptx = nvcc.compile_cuda(code, target_format="fatbin", compiler="nvcc")
+        else:
+            arch = []
+            for compute_version in multi_arch:
+                arch += [
+                    "-gencode",
+                    f"arch=compute_{compute_version},code=sm_{compute_version}",
+                ]
+            ptx = nvcc.compile_cuda(code, target_format="fatbin", arch=arch, compiler="nvcc")
+        return ptx
+
+
+def detect_system_lib_prefix(
+    target_hint: str, prefix_hint: str, model_name: str, quantization: str
+) -> str:
+    """Detect the iOS / Android system lib prefix to identify the library needed to load the app.
+
+    Parameters
+    ----------
+    target_hint : str
+        The hint for the target device.
+
+    prefix_hint : str
+        The hint for the system lib prefix.
+    """
+    if prefix_hint == "auto" and (
+        target_hint.startswith("iphone")
+        or target_hint.startswith("macabi")
+        or target_hint.startswith("android")
+    ):
+        prefix = f"{model_name}_{quantization}_".replace("-", "_")
+        logger.warning(
+            "%s is automatically picked from the filename, %s, this allows us to use the filename "
+            "as the model_lib in android/iOS builds. Please avoid renaming the .tar file when "
+            "uploading the prebuilt.",
+            bold("--system-lib-prefix"),
+            bold(prefix),
+        )
+        return prefix
+    if target_hint not in ["iphone", "macabi", "android"]:
+        return ""
+    return prefix_hint
+
+
+_MACABI_ARCH = os.environ.get("MLC_MACABI_ARCH", "").strip() or "arm64"
+if _MACABI_ARCH not in ["arm64", "x86_64"]:
+    _MACABI_ARCH = "arm64"
+_MACABI_MTRIPLE = (
+    "x86_64-apple-ios18.0-macabi" if _MACABI_ARCH == "x86_64" else "arm64-apple-ios18.0-macabi"
+)
+
+PRESET = {
+    "iphone:generic": {
+        "target": {
+            "kind": "metal",
+            "max_threads_per_block": 256,
+            "max_shared_memory_per_block": 32768,
+            "thread_warp_size": 1,
+            "libs": ["iphoneos"],
+            "host": {
+                "kind": "llvm",
+                "mtriple": "arm64-apple-darwin",
+            },
+        },
+        "build": _build_iphone,
+    },
+    "macabi:generic": {
+        "target": {
+            "kind": "metal",
+            "max_threads_per_block": 256,
+            "max_shared_memory_per_block": 32768,
+            "thread_warp_size": 1,
+            "libs": ["macosx"],
+            "host": {
+                "kind": "llvm",
+                "mtriple": _MACABI_MTRIPLE,
+            },
+        },
+        "build": _build_iphone,
+    },
+    "android:generic": {
+        "target": {
+            "kind": "opencl",
+            "host": {
+                "kind": "llvm",
+                "mtriple": "aarch64-linux-android",
+            },
+        },
+        "build": _build_android,
+    },
+    "android:adreno": {
+        "target": {
+            "kind": "opencl",
+            "device": "adreno",
+            "max_threads_per_block": 512,
+            "host": {
+                "kind": "llvm",
+                "mtriple": "aarch64-linux-android",
+            },
+        },
+        "build": _build_android,
+    },
+    "android:adreno-so": {
+        "target": {
+            "kind": "opencl",
+            "device": "adreno",
+            "max_threads_per_block": 512,
+            "host": {
+                "kind": "llvm",
+                "mtriple": "aarch64-linux-android",
+            },
+        },
+        "build": _build_android_so,
+    },
+    "windows:adreno_x86": {
+        "target": {
+            "kind": "opencl",
+            "device": "adreno",
+            "max_threads_per_block": 512,
+            "host": {
+                "kind": "llvm",
+                "mtriple": "x86_64-pc-windows-msvc",
+            },
+        },
+    },
+    "metal:x86-64": {
+        "target": {
+            "kind": "metal",
+            "max_threads_per_block": 256,
+            "max_shared_memory_per_block": 32768,
+            "thread_warp_size": 1,
+        },
+        "build": _build_metal_x86_64,
+    },
+    "webgpu:generic": {
+        "target": {
+            "kind": "webgpu",
+            "host": {
+                "kind": "llvm",
+                "mtriple": "wasm32-unknown-unknown-wasm",
+            },
+        },
+        "build": _build_webgpu,
+    },
+    "opencl:generic": {
+        "target": {
+            "kind": "opencl",
+        },
+    },
+    "mali:generic": {
+        "target": {
+            "kind": "opencl",
+            "host": {
+                "kind": "llvm",
+                "mtriple": "aarch64-linux-gnu",
+            },
+        },
+        "build": _build_mali,
+    },
+    "metal:generic": {
+        "target": {
+            "kind": "metal",
+            "max_threads_per_block": 256,
+            "max_shared_memory_per_block": 32768,
+            "thread_warp_size": 1,
+        },
+    },
+    "vulkan:generic": {
+        "target": {
+            "kind": "vulkan",
+            "max_threads_per_block": 256,
+            "max_shared_memory_per_block": 32768,
+            "thread_warp_size": 1,
+            "supports_float16": 1,
+            "supports_int64": 1,
+            "supports_int16": 1,
+            "supports_int8": 1,
+            "supports_8bit_buffer": 1,
+            "supports_16bit_buffer": 1,
+            "supports_storage_buffer_storage_class": 1,
+        },
+    },
+}

python/mlc_llm/support/auto_weight.py

@@ -0,0 +1,178 @@
+"""Help functions for detecting weight paths and weight formats."""
+
+import json
+from pathlib import Path
+from typing import List, Optional, Tuple  # noqa: UP035
+
+from . import logging
+from .style import bold, green, red
+
+logger = logging.getLogger(__name__)
+
+FOUND = green("Found")
+NOT_FOUND = red("Not found")
+
+
+def detect_weight(
+    weight_path: Path,
+    config_json_path: Path,
+    weight_format: str = "auto",
+) -> Tuple[Path, str]:  # noqa: UP006
+    """Detect the weight directory, and detect the weight format.
+
+    Parameters
+    ---------
+    weight_path : pathlib.Path
+        The path to weight files. If `weight_path` is not None, check if it exists. Otherwise, find
+        `weight_path` in `config.json` or use the same directory as `config.json`.
+
+    config_json_path: pathlib.Path
+        The path to `config.json`.
+
+    weight_format : str
+        The hint for the weight format. If it is "auto", guess the weight format.
+        Otherwise, check the weights are in that format.
+        Available weight formats:
+            - auto (guess the weight format)
+            - huggingface-torch (validate via checking pytorch_model.bin.index.json)
+            - huggingface-safetensor (validate via checking model.safetensors.index.json)
+            - awq
+            - ggml
+            - gguf
+
+    Returns
+    -------
+    weight_config_path : pathlib.Path
+        The path that points to the weights config file or the weights directory.
+
+    weight_format : str
+        The valid weight format.
+    """
+    if weight_path is None:
+        assert config_json_path is not None and config_json_path.exists(), (
+            "Please provide config.json path."
+        )
+
+        # 1. Find the weight_path in config.json
+        with open(config_json_path, encoding="utf-8") as i_f:
+            config = json.load(i_f)
+        if "weight_path" in config:
+            weight_path = Path(config["weight_path"])
+            logger.info('Found "weight_path" in config.json: %s', weight_path)
+            if not weight_path.exists():
+                raise ValueError(f"weight_path doesn't exist: {weight_path}")
+        else:
+            # 2. Find the weights file in the same directory as config.json
+            weight_path = config_json_path.parent
+    else:
+        if not weight_path.exists():
+            raise ValueError(f"weight_path doesn't exist: {weight_path}")
+
+    logger.info("Finding weights in: %s", weight_path)
+
+    # check weight format
+    # weight_format = "auto", guess the weight format.
+    # otherwise, check the weight format is valid.
+    if weight_format == "auto":
+        return _guess_weight_format(weight_path)
+
+    if weight_format not in AVAILABLE_WEIGHT_FORMAT:
+        raise ValueError(
+            f"Available weight format list: {AVAILABLE_WEIGHT_FORMAT}, but got {weight_format}"
+        )
+    if weight_format in CHECK_FORMAT_METHODS:
+        check_func = CHECK_FORMAT_METHODS[weight_format]
+        weight_config_path = check_func(weight_path)
+        if not weight_config_path:
+            raise ValueError(f"The weight is not in {weight_format} format.")
+    else:
+        weight_config_path = weight_path
+    return weight_config_path, weight_format
+
+
+def _guess_weight_format(weight_path: Path) -> Tuple[Path, str]:  # noqa: UP006
+    possible_formats: List[Tuple[Path, str]] = []  # noqa: UP006
+    for weight_format, check_func in CHECK_FORMAT_METHODS.items():
+        weight_config_path = check_func(weight_path)
+        if weight_config_path:
+            possible_formats.append((weight_config_path, weight_format))
+
+    if len(possible_formats) == 0:
+        raise ValueError(
+            "Fail to detect source weight format. "
+            "Use `--source-format` to explicitly specify the format."
+        )
+
+    weight_config_path, selected_format = possible_formats[0]
+    logger.info(
+        "Using source weight configuration: %s. Use `--source` to override.",
+        bold(str(weight_config_path)),
+    )
+    logger.info(
+        "Using source weight format: %s. Use `--source-format` to override.",
+        bold(selected_format),
+    )
+    return weight_config_path, selected_format
+
+
+def _check_pytorch(weight_path: Path) -> Optional[Path]:
+    pytorch_json_path = weight_path / "pytorch_model.bin.index.json"
+    if pytorch_json_path.exists():
+        logger.info(
+            "%s source weight format: huggingface-torch. Source configuration: %s",
+            FOUND,
+            pytorch_json_path,
+        )
+        return pytorch_json_path
+
+    pytorch_file_path = weight_path / "pytorch_model.bin"
+    if pytorch_file_path.exists():
+        logger.info(
+            "%s source weight format: huggingface-torch. Source configuration: %s",
+            FOUND,
+            pytorch_file_path,
+        )
+        return pytorch_file_path
+
+    logger.info("%s Huggingface PyTorch", NOT_FOUND)
+    return None
+
+
+def _check_safetensor(weight_path: Path) -> Optional[Path]:
+    safetensor_json_path = weight_path / "model.safetensors.index.json"
+    if safetensor_json_path.exists():
+        logger.info(
+            "%s source weight format: huggingface-safetensor. Source configuration: %s",
+            FOUND,
+            safetensor_json_path,
+        )
+        return safetensor_json_path
+
+    safetensor_file_path = weight_path / "model.safetensors"
+    if safetensor_file_path.exists():
+        from safetensors.torch import (
+            load_file,
+        )
+
+        weights = load_file(safetensor_file_path, device="cpu")
+        weight_map = {key: "model.safetensors" for key in weights}
+        with open(safetensor_json_path, "w", encoding="utf-8") as file:
+            json.dump({"weight_map": weight_map}, file, indent=2)
+        logger.info(
+            "%s source weight format: huggingface-safetensor. Source configuration: %s",
+            FOUND,
+            safetensor_json_path,
+        )
+        return safetensor_json_path
+
+    logger.info("%s Huggingface Safetensor", NOT_FOUND)
+    return None
+
+
+CHECK_FORMAT_METHODS = {
+    "huggingface-torch": _check_pytorch,
+    "huggingface-safetensor": _check_safetensor,
+}
+
+# "ggml", "gguf" are not supported yet.
+AVAILABLE_WEIGHT_FORMAT = ["huggingface-torch", "huggingface-safetensor", "awq"]

python/mlc_llm/support/config.py

@@ -0,0 +1,111 @@
+"""
+A common base class for configuration. A configuration could be initialized from its constructor,
+a JSON string or a JSON file, and irrelevant fields during initialization are automatically moved
+to the `kwargs` field.
+
+Take model configuration as an example: it is usually a JSON file in HuggingFace that contains
+the model's hyperparameters. For instance, Vicuna-13b-v1.5-16k contains the following
+[JSON file](https://huggingface.co/lmsys/vicuna-13b-v1.5-16k/blob/main/config.json).
+The base class allows us to load the configuration from this JSON file, moving irrelevant fields
+into `kwargs`, such as `transformers_version` and `use_cache`.
+"""
+
+import dataclasses
+import json
+from pathlib import Path
+from typing import Any, Dict, Type, TypeVar  # noqa: UP035
+
+from . import logging
+from .style import bold, red
+
+logger = logging.getLogger(__name__)
+
+ConfigClass = TypeVar("ConfigClass", bound="ConfigBase")
+
+
+@dataclasses.dataclass
+class ConfigBase:
+    """Base class for configurations, providing a common interface for loading configs from a
+    JSON file or a dict. It requires the subclasses to be dataclasses, and has an `kwargs` field
+    that stores the extra fields that are not defined in the dataclass.
+    """
+
+    @classmethod
+    def from_dict(cls: Type[ConfigClass], source: Dict[str, Any]) -> ConfigClass:  # noqa: UP006
+        """Create a config object from a dictionary.
+
+        Parameters
+        ----------
+        source : Dict[str, Any]
+            Source to create config from, usually loaded from `config.json` in HuggingFace style.
+
+        Returns
+        -------
+        cfg : ConfigClass
+            An instance of the config object.
+        """
+        field_names = [field.name for field in dataclasses.fields(cls)]
+        fields = {k: v for k, v in source.items() if k in field_names}
+        kwargs = {k: v for k, v in source.items() if k not in field_names}
+        return cls(**fields, kwargs=kwargs)
+
+    @classmethod
+    def from_file(cls: Type[ConfigClass], source: Path) -> ConfigClass:  # noqa: UP006
+        """Create a config object from a file.
+
+        Parameters
+        ----------
+        cfg_cls : Type[ConfigClass]
+            The config class to create, for example, LlamaConfig.
+
+        source : pathlib.Path
+            Path to the source file, usually `config.json` in HuggingFace repo.
+
+        Returns
+        -------
+        cfg : ConfigClass
+            An instance of the config object.
+        """
+        with source.open("r", encoding="utf-8") as in_file:
+            return cls.from_dict(json.load(in_file))
+
+    def asdict(self):
+        """Convert the config object to a dictionary.
+
+        Returns
+        -------
+        Dict[str, Any]
+            A dictionary representation of the config object.
+        """
+        result = dataclasses.asdict(self)
+        result.pop("kwargs")
+        return result
+
+
+class ConfigOverrideBase:
+    """Base class for ConfigOverride, providing a common interface for overriding configs.
+    It requires the subclasses to be dataclasses.
+    """
+
+    def apply(self, config):
+        """Apply the overrides to the given config."""
+        updated = config.asdict()
+        for field in dataclasses.fields(self):
+            key = field.name
+            value = getattr(self, key)
+            if value is None:
+                continue
+            if key not in updated:
+                logger.warning(
+                    "%s: Cannot override %s, because %s does not have this field",
+                    red("Warning"),
+                    bold(key),
+                    bold(type(config).__name__),
+                )
+            else:
+                logger.info(f"Overriding {bold(key)} from {updated[key]} to {value}")
+                updated[key] = value
+        return type(config).from_dict(updated)
+
+
+__all__ = ["ConfigBase", "ConfigOverrideBase"]

python/mlc_llm/support/constants.py

@@ -0,0 +1,90 @@
+"""Environment variables used by the MLC LLM."""
+
+import os
+import sys
+from pathlib import Path
+from typing import List  # noqa: UP035
+
+MLC_CHAT_CONFIG_VERSION = "0.1.0"
+
+
+def _check():
+    if MLC_JIT_POLICY not in ["ON", "OFF", "REDO", "READONLY"]:
+        raise ValueError(
+            'Invalid MLC_JIT_POLICY. It has to be one of "ON", "OFF", "REDO", "READONLY"'
+            f"but got {MLC_JIT_POLICY}."
+        )
+
+    if MLC_DOWNLOAD_CACHE_POLICY not in ["ON", "OFF", "REDO", "READONLY"]:
+        raise ValueError(
+            "Invalid MLC_AUTO_DOWNLOAD_POLICY. "
+            'It has to be one of "ON", "OFF", "REDO", "READONLY"'
+            f"but got {MLC_DOWNLOAD_CACHE_POLICY}."
+        )
+
+
+def _get_cache_dir() -> Path:
+    if "MLC_LLM_HOME" in os.environ:
+        result = Path(os.environ["MLC_LLM_HOME"])
+    elif sys.platform == "win32":
+        result = Path(os.environ["LOCALAPPDATA"])
+        result = result / "mlc_llm"
+    elif os.getenv("XDG_CACHE_HOME", None) is not None:
+        result = Path(os.getenv("XDG_CACHE_HOME"))
+        result = result / "mlc_llm"
+    else:
+        result = Path(os.path.expanduser("~/.cache"))
+        result = result / "mlc_llm"
+    result.mkdir(parents=True, exist_ok=True)
+    if not result.is_dir():
+        raise ValueError(
+            f"The default cache directory is not a directory: {result}. "
+            "Use environment variable MLC_LLM_HOME to specify a valid cache directory."
+        )
+    (result / "model_weights").mkdir(parents=True, exist_ok=True)
+    (result / "model_lib").mkdir(parents=True, exist_ok=True)
+    return result
+
+
+def _get_dso_suffix() -> str:
+    if "MLC_DSO_SUFFIX" in os.environ:
+        return os.environ["MLC_DSO_SUFFIX"]
+    if sys.platform == "win32":
+        return "dll"
+    if sys.platform == "darwin":
+        return "dylib"
+    return "so"
+
+
+def _get_test_model_path() -> List[Path]:  # noqa: UP006
+    paths = []
+    if "MLC_LLM_TEST_MODEL_PATH" in os.environ:
+        paths += [Path(p) for p in os.environ["MLC_LLM_TEST_MODEL_PATH"].split(os.pathsep)]
+    # by default, we reuse the cache dir via mlc_llm chat
+    # note that we do not auto download for testcase
+    # to avoid networking dependencies
+    base_list = ["hf"]
+    paths += [_get_cache_dir() / "model_weights" / base / "mlc-ai" for base in base_list] + [
+        Path(os.path.abspath(os.path.curdir)),
+        Path(os.path.abspath(os.path.curdir)) / "dist",
+    ]
+    return paths
+
+
+def _get_read_only_weight_caches() -> List[Path]:  # noqa: UP006
+    if "MLC_LLM_READONLY_WEIGHT_CACHE" in os.environ:
+        return [Path(p) for p in os.environ["MLC_LLM_READONLY_WEIGHT_CACHE"].split(os.pathsep)]
+    return []
+
+
+MLC_TEMP_DIR = os.getenv("MLC_TEMP_DIR", None)
+MLC_MULTI_ARCH = os.environ.get("MLC_MULTI_ARCH", None)
+MLC_JIT_POLICY = os.environ.get("MLC_JIT_POLICY", "ON")
+MLC_DSO_SUFFIX = _get_dso_suffix()
+MLC_TEST_MODEL_PATH: List[Path] = _get_test_model_path()  # noqa: UP006
+
+MLC_DOWNLOAD_CACHE_POLICY = os.environ.get("MLC_DOWNLOAD_CACHE_POLICY", "ON")
+MLC_LLM_HOME: Path = _get_cache_dir()
+MLC_LLM_READONLY_WEIGHT_CACHE = _get_read_only_weight_caches()
+
+_check()

python/mlc_llm/support/convert_tiktoken.py

@@ -0,0 +1,171 @@
+"""
+Adapted from https://gist.github.com/xenova/a452a6474428de0182b17605a98631ee
+Generator of mlc-chat-config.json and tokenizer configuration.
+"""
+
+# isort: off
+import json
+import os
+from typing import Dict, List, Optional  # noqa: UP035
+
+
+def bpe(
+    mergeable_ranks: Dict[bytes, int],  # noqa: UP006
+    token: bytes,
+    max_rank: Optional[int] = None,
+) -> List[bytes]:  # noqa: UP006
+    """Adapted from https://github.com/openai/tiktoken/issues/60#issuecomment-1499977960"""
+    parts = [bytes([b]) for b in token]
+    while True:
+        min_idx = None
+        min_rank = None
+        for i, pair in enumerate(zip(parts[:-1], parts[1:])):
+            rank = mergeable_ranks.get(pair[0] + pair[1])
+            if rank is not None and (min_rank is None or rank < min_rank):
+                min_idx = i
+                min_rank = rank
+        if min_rank is None or (max_rank is not None and min_rank >= max_rank):
+            break
+        assert min_idx is not None
+        parts = [*parts[:min_idx], parts[min_idx] + parts[min_idx + 1], *parts[min_idx + 2 :]]
+    return parts
+
+
+def generate_vocab_and_merges(encoder, mergeable_ranks):
+    """Generate vocab and merges in huggingface tokenizers format"""
+
+    from transformers.models.gpt2.tokenization_gpt2 import (
+        bytes_to_unicode,
+    )
+
+    byte_encoder = bytes_to_unicode()
+
+    def token_bytes_to_string(b):
+        """Convert a token from bytes to a string"""
+        return "".join([byte_encoder[ord(char)] for char in b.decode("latin-1")])
+
+    merges = []
+    vocab = {}
+    for token, rank in mergeable_ranks.items():
+        vocab[token_bytes_to_string(token)] = rank
+
+        if len(token) == 1:
+            continue
+        merged = tuple(bpe(mergeable_ranks, token, max_rank=rank))
+        assert len(merged) == 2
+
+        merges.append(" ".join(map(token_bytes_to_string, merged)))
+
+    # Also add special tokens
+    vocab.update(encoder._special_tokens)
+
+    return vocab, merges
+
+
+def convert_tiktoken(model_path, output_dir, context_window_size=None):
+    """Convert tiktoken tokenizers to huggingface tokenizers style"""
+    try:
+        from transformers import AutoTokenizer
+    except ImportError:
+        raise ImportError(
+            'Converting tiktoken tokenizer requires the "transformers" package.'
+            'Please install the "transformers" package to convert toktoken tokenizer'
+        )
+
+    tiktoken_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    encoder = tiktoken_tokenizer.tokenizer
+
+    vocab, merges = generate_vocab_and_merges(encoder, tiktoken_tokenizer.get_vocab())
+
+    added_tokens = [
+        {
+            "id": id,
+            "content": content,
+            "single_word": False,
+            "lstrip": False,
+            "rstrip": False,
+            "normalized": False,
+            "special": True,
+        }
+        for content, id in encoder._special_tokens.items()
+    ]
+
+    tokenizer_template = {
+        "version": "1.0",
+        "truncation": None,
+        "padding": None,
+        "added_tokens": added_tokens,
+        "normalizer": None,
+        "pre_tokenizer": {
+            "type": "ByteLevel",
+            "add_prefix_space": False,
+            "trim_offsets": True,
+            "use_regex": True,
+        },
+        "post_processor": {
+            "type": "ByteLevel",
+            "add_prefix_space": True,
+            "trim_offsets": False,
+            "use_regex": True,
+        },
+        "decoder": {
+            "type": "ByteLevel",
+            "add_prefix_space": True,
+            "trim_offsets": True,
+            "use_regex": True,
+        },
+        "model": {
+            "type": "BPE",
+            "dropout": None,
+            "unk_token": None,
+            "continuing_subword_prefix": "",
+            "end_of_word_suffix": "",
+            "fuse_unk": False,
+            "byte_fallback": False,
+            "vocab": vocab,
+            "merges": merges,
+        },
+    }
+
+    tokenizer_config_template = {
+        "add_prefix_space": False,
+        "bos_token": "<|endoftext|>",
+        "clean_up_tokenization_spaces": True,
+        "eos_token": "<|endoftext|>",
+        "unk_token": "<|endoftext|>",
+    }
+
+    tokenizer_name = type(tiktoken_tokenizer).__name__
+
+    tokenizer_config_template["tokenizer_class"] = tokenizer_name
+    if context_window_size:
+        tokenizer_config_template["model_max_length"] = context_window_size
+    tokenizer_config_template = dict(sorted(tokenizer_config_template.items(), key=lambda x: x[0]))
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    # Save to files
+    with open(os.path.join(output_dir, "vocab.json"), "w", encoding="utf-8") as fp:
+        json.dump(vocab, fp, indent=2, ensure_ascii=False)
+
+    with open(os.path.join(output_dir, "tokenizer.json"), "w", encoding="utf-8") as fp:
+        json.dump(tokenizer_template, fp, indent=2, ensure_ascii=False)
+
+    with open(os.path.join(output_dir, "tokenizer_config.json"), "w", encoding="utf-8") as fp:
+        json.dump(tokenizer_config_template, fp, indent=2, ensure_ascii=False)
+
+    with open(os.path.join(output_dir, "special_tokens_map.json"), "w", encoding="utf-8") as fp:
+        json.dump(
+            {
+                "bos_token": "<|endoftext|>",
+                "eos_token": "<|endoftext|>",
+                "unk_token": "<|endoftext|>",
+            },
+            fp,
+            indent=2,
+            ensure_ascii=False,
+        )
+
+    with open(os.path.join(output_dir, "merges.txt"), "w", encoding="utf-8") as fp:
+        fp.write("#version: 0.2\n")
+        fp.write("\n".join(merges))

python/mlc_llm/support/download_cache.py

@@ -0,0 +1,237 @@
+"""Common utilities for downloading files from HuggingFace or other URLs online."""
+
+import concurrent.futures as cf
+import hashlib
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import List, Optional, Tuple  # noqa: UP035
+
+import requests
+
+from . import logging, tqdm
+from .constants import (
+    MLC_DOWNLOAD_CACHE_POLICY,
+    MLC_LLM_HOME,
+    MLC_LLM_READONLY_WEIGHT_CACHE,
+    MLC_TEMP_DIR,
+)
+from .style import bold
+
+logger = logging.getLogger(__name__)
+
+
+def log_download_cache_policy():
+    """log current download policy"""
+    logger.info(
+        "%s = %s. Can be one of: ON, OFF, REDO, READONLY",
+        bold("MLC_DOWNLOAD_CACHE_POLICY"),
+        MLC_DOWNLOAD_CACHE_POLICY,
+    )
+
+
+def _ensure_directory_not_exist(path: Path, force_redo: bool) -> None:
+    if path.exists():
+        if force_redo:
+            logger.info("Deleting existing directory: %s", path)
+            shutil.rmtree(path)
+        else:
+            raise ValueError(f"Directory already exists: {path}")
+    else:
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+
+def git_clone(url: str, destination: Path, ignore_lfs: bool) -> None:
+    """Clone a git repository into a directory."""
+    repo_name = ".tmp"
+    command = ["git", "clone", url, repo_name]
+    _ensure_directory_not_exist(destination, force_redo=False)
+    try:
+        env = os.environ.copy()
+        env["GIT_LFS_SKIP_SMUDGE"] = "1"
+        with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir:
+            logger.info("[Git] Cloning %s to %s", bold(url), destination)
+            subprocess.run(
+                command,
+                env=env,
+                cwd=tmp_dir,
+                check=True,
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+            git_dir = os.path.join(tmp_dir, repo_name)
+            if not ignore_lfs:
+                git_lfs_pull(Path(git_dir))
+            shutil.move(git_dir, str(destination))
+    except subprocess.CalledProcessError as error:
+        raise ValueError(
+            f"Git clone failed with return code {error.returncode}: {error.stderr}. "
+            f"The command was: {command}"
+        ) from error
+
+
+def git_lfs_pull(repo_dir: Path, ignore_extensions: Optional[List[str]] = None) -> None:  # noqa: UP006
+    """Pull files with Git LFS."""
+    filenames = (
+        subprocess.check_output(
+            ["git", "-C", str(repo_dir), "lfs", "ls-files", "-n"],
+            stderr=subprocess.STDOUT,
+        )
+        .decode("utf-8")
+        .splitlines()
+    )
+    if ignore_extensions is not None:
+        filenames = [
+            filename
+            for filename in filenames
+            if not any(filename.endswith(extension) for extension in ignore_extensions)
+        ]
+    logger.info("[Git LFS] Downloading %d files with Git LFS: %s", len(filenames), filenames)
+    with tqdm.redirect():
+        for file in tqdm.tqdm(filenames):
+            logger.info("[Git LFS] Downloading %s", file)
+            subprocess.check_output(
+                ["git", "-C", str(repo_dir), "lfs", "pull", "--include", file],
+                stderr=subprocess.STDOUT,
+            )
+
+
+def download_file(
+    url: str,
+    destination: Path,
+    md5sum: Optional[str],
+) -> Tuple[str, Path]:  # noqa: UP006
+    """Download a file from a URL to a destination file."""
+    with requests.get(url, stream=True, timeout=30) as response:
+        response.raise_for_status()
+        with destination.open("wb") as file:
+            for chunk in response.iter_content(chunk_size=8192):
+                file.write(chunk)
+    if md5sum is not None:
+        hash_md5 = hashlib.md5()
+        with destination.open("rb") as file:
+            for chunk in iter(lambda: file.read(8192), b""):
+                hash_md5.update(chunk)
+        file_md5 = hash_md5.hexdigest()
+        if file_md5 != md5sum:
+            raise ValueError(
+                f"MD5 checksum mismatch for downloaded file: {destination}. "
+                f"Expected {md5sum}, got {file_md5}"
+            )
+    return url, destination
+
+
+def download_and_cache_mlc_weights(
+    model_url: str,
+    num_processes: int = 4,
+    force_redo: Optional[bool] = None,
+) -> Path:
+    """Download weights for a model from the HuggingFace Git LFS repo."""
+    log_download_cache_policy()
+    if MLC_DOWNLOAD_CACHE_POLICY == "OFF":
+        raise RuntimeError(f"Cannot download {model_url} as MLC_DOWNLOAD_CACHE_POLICY=OFF")
+
+    prefixes, mlc_prefix = ["HF://", "https://huggingface.co/"], ""
+    mlc_prefix = next(p for p in prefixes if model_url.startswith(p))
+    assert mlc_prefix
+
+    git_url_template = "https://huggingface.co/{user}/{repo}"
+    bin_url_template = "https://huggingface.co/{user}/{repo}/resolve/main/{record_name}"
+
+    if model_url.count("/") != 1 + mlc_prefix.count("/") or not model_url.startswith(mlc_prefix):
+        raise ValueError(f"Invalid model URL: {model_url}")
+    user, repo = model_url[len(mlc_prefix) :].split("/")
+    domain = "hf"
+
+    readonly_cache_dirs = []
+    for base in MLC_LLM_READONLY_WEIGHT_CACHE:
+        cache_dir = base / domain / user / repo
+        readonly_cache_dirs.append(str(cache_dir))
+        if (cache_dir / "mlc-chat-config.json").is_file():
+            logger.info("Use cached weight: %s", bold(str(cache_dir)))
+            return cache_dir
+
+    if force_redo is None:
+        force_redo = MLC_DOWNLOAD_CACHE_POLICY == "REDO"
+
+    git_dir = MLC_LLM_HOME / "model_weights" / domain / user / repo
+    readonly_cache_dirs.append(str(git_dir))
+
+    try:
+        _ensure_directory_not_exist(git_dir, force_redo=force_redo)
+    except ValueError:
+        logger.info("Weights already downloaded: %s", bold(str(git_dir)))
+        return git_dir
+
+    if MLC_DOWNLOAD_CACHE_POLICY == "READONLY":
+        raise RuntimeError(
+            f"Cannot find cache for {model_url}, "
+            "cannot proceed to download as MLC_DOWNLOAD_CACHE_POLICY=READONLY, "
+            "please check settings MLC_LLM_READONLY_WEIGHT_CACHE, "
+            f"local path candidates: {readonly_cache_dirs}"
+        )
+
+    with tempfile.TemporaryDirectory(dir=MLC_TEMP_DIR) as tmp_dir_prefix:
+        tmp_dir = Path(tmp_dir_prefix) / "tmp"
+        git_url = git_url_template.format(user=user, repo=repo)
+        git_clone(git_url, tmp_dir, ignore_lfs=True)
+        git_lfs_pull(tmp_dir, ignore_extensions=[".bin"])
+        shutil.rmtree(tmp_dir / ".git", ignore_errors=True)
+        with (tmp_dir / "tensor-cache.json").open(encoding="utf-8") as in_file:
+            param_metadata = json.load(in_file)["records"]
+        with cf.ProcessPoolExecutor(max_workers=num_processes) as executor:
+            futures = []
+            for record in param_metadata:
+                record_name = record["dataPath"]
+                file_url = bin_url_template.format(user=user, repo=repo, record_name=record_name)
+                file_dest = tmp_dir / record_name
+                file_md5 = record.get("md5sum", None)
+                futures.append(executor.submit(download_file, file_url, file_dest, file_md5))
+            with tqdm.redirect():
+                for future in tqdm.tqdm(cf.as_completed(futures), total=len(futures)):
+                    file_url, file_dest = future.result()
+                    logger.info("Downloaded %s to %s", file_url, file_dest)
+        logger.info("Moving %s to %s", tmp_dir, bold(str(git_dir)))
+        shutil.move(str(tmp_dir), str(git_dir))
+    return git_dir
+
+
+def get_or_download_model(model: str) -> Path:
+    """Use user-provided argument ``model`` to get model_path
+
+    We define "valid" as having an ``mlc-chat-config.json`` right under the folder.
+
+    Parameters
+    ----------
+    model : str
+        User's input; may a path or url
+
+    Returns
+    ------
+    model_path : Path
+        A "valid" path to model folder, with
+        ``(model_path / "mlc-chat-config.json").is_file`` being True
+
+    Note
+    ----
+    This function may perform additional download and caching
+
+    Raises
+    ------
+    FileNotFoundError: if we cannot find a valid `model_path`.
+    """
+    if model.startswith("HF://"):
+        logger.info("Downloading model from HuggingFace: %s", model)
+        model_path = download_and_cache_mlc_weights(model)
+    else:
+        model_path = Path(model)
+
+    if not model_path.is_dir():
+        raise FileNotFoundError(f"Cannot find model {model}, directory does not exist")
+    mlc_config_path = model_path / "mlc-chat-config.json"
+    if mlc_config_path.is_file():
+        return model_path
+    raise FileNotFoundError(f"Cannot find {str(mlc_config_path)} in the model directory provided")

python/mlc_llm/support/logging.py

@@ -0,0 +1,24 @@
+"""
+Logging support for MLC. It derives from Python's logging module, and in the future,
+it can be easily replaced by other logging modules such as structlog.
+"""
+
+import logging
+import os
+
+
+def enable_logging():
+    """Enable MLC's default logging format"""
+    if os.getenv("MLC_UNSET_LOGGING"):
+        return
+    logging.basicConfig(
+        level=logging.INFO,
+        style="{",
+        datefmt="%Y-%m-%d %H:%M:%S",
+        format="[{asctime}] {levelname} {filename}:{lineno}: {message}",
+    )
+
+
+def getLogger(name: str):
+    """Get a logger according to the given name"""
+    return logging.getLogger(name)

python/mlc_llm/support/max_thread_check.py

@@ -0,0 +1,38 @@
+"""Helper functions for checking max num thread."""
+
+from tvm.target import Target
+
+
+def get_max_num_threads_per_block(target: Target) -> int:
+    """
+    max(max_num_threads, max_threads_per_block); if latter does not exist, return max_num_threads.
+    We add this method since some targets have both fields and `max_threads_per_block` is larger.
+    """
+    max_num_threads = target.attrs.get("max_num_threads")
+    max_threads_per_block = target.attrs.get("max_threads_per_block", None)
+    if max_threads_per_block is None:
+        return max_num_threads
+    return max(max_num_threads, max_threads_per_block)
+
+
+def check_thread_limits(target: Target, bdx: int, bdy: int, bdz: int, gdz: int):
+    """
+    Check whether max num threads exceeded given a target.
+
+    Parameters
+    ----------
+    bdx: threadIdx.x
+    bdy: threadIdx.y
+    bdz: threadIdx.z
+    gdz: blockIdx.z
+    """
+    max_num_threads_per_block = get_max_num_threads_per_block(target)
+
+    assert bdx * bdy * bdz <= max_num_threads_per_block, (
+        f"{target.kind} max num threads exceeded: {bdx}*{bdy}*{bdz}>{max_num_threads_per_block}"
+    )
+
+    if target.kind.name == "webgpu":
+        # https://gpuweb.github.io/gpuweb/#dom-supported-limits-maxcomputeworkgroupsizez
+        assert bdz <= 64, f"webgpu's threadIdx.z cannot exceed 64, but got bdz={bdz}"
+        assert gdz == 1, f"webgpu's blockIdx.z should be 1, but got gdz={gdz}"

python/mlc_llm/support/preshard.py

@@ -0,0 +1,123 @@
+"""Functions for pre-sharding weights"""
+
+import logging
+from collections.abc import Sequence
+from typing import Any, Callable, Dict, Tuple  # noqa: UP035
+
+from tvm import IRModule, relax
+from tvm.relax.frontend import nn
+from tvm.runtime import Device, Tensor
+from tvm.s_tir import dlight as dl
+from tvm.target import Target
+
+logger = logging.getLogger("preshard")
+
+
+def _sharded_param_name(param_name, worker_id):
+    return f"{param_name}_shard-{worker_id}"
+
+
+def _create_shard_func(bb: relax.BlockBuilder, param: nn.Parameter, tensor_parallel_shards: int):
+    shard_strategy = param.attrs.get("shard_strategy", None)
+    # generate tirx shard function
+    tir_func = shard_strategy.gen_tir(shards=tensor_parallel_shards, weight=param)
+    tir_func = tir_func.with_attr("global_symbol", f"{shard_strategy.name}_tir")
+    # add tirx shard function to the IRModule
+    tir_gvar = bb.add_func(tir_func, func_name=f"{shard_strategy.name}_tir")
+    # create relax function that
+    #     1. shard weight with tirx shard function, result: [num_shards, *sharded_weight_shape]
+    #     2. split the sharded weight along dim 0, result: num_shards * [1, *sharded_weight_shape]
+    #     3. squeeze the 0th-dim of all shards, result: num_shards * [*sharded_weight_shape]
+    weight_shape = param.shape
+    weight_shape[shard_strategy.dim] = weight_shape[shard_strategy.dim] * tensor_parallel_shards
+    sharded_weight_shape = [tensor_parallel_shards, *param.shape]
+    weight_var = relax.Var("weight", relax.TensorType(weight_shape, param.dtype))
+    with bb.function(name=shard_strategy.name, params=[weight_var]):
+        with bb.dataflow():
+            lv0 = bb.emit(
+                relax.call_tir(
+                    tir_gvar,
+                    weight_var,
+                    out_ty=relax.TensorType(sharded_weight_shape, param.dtype),
+                )
+            )
+            lv1 = bb.emit(relax.op.split(lv0, indices_or_sections=tensor_parallel_shards, axis=0))
+            output_vars = []
+            for i in range(tensor_parallel_shards):
+                lvi = bb.emit(relax.TupleGetItem(lv1, i))
+                squeezed_lvi = bb.emit(relax.op.squeeze(lvi, 0))
+                output_vars.append(squeezed_lvi)
+            gv = bb.emit_output(output_vars)
+        bb.emit_func_output(gv)
+
+
+def _compile_shard_funcs(mod: IRModule, device: Device):
+    target = Target.from_device(device)
+    with target:
+        mod = relax.transform.LegalizeOps()(mod)
+        mod = dl.ApplyDefaultSchedule(
+            dl.gpu.Matmul(),
+            dl.gpu.GEMV(),
+            dl.gpu.Reduction(),
+            dl.gpu.GeneralReduction(),
+            dl.gpu.Fallback(),
+        )(mod)
+    ex = relax.build(mod, target=target)
+    vm = relax.VirtualMachine(ex, device)
+    return vm
+
+
+def apply_preshard(
+    named_params: Dict[str, nn.Parameter],  # noqa: UP006
+    tensor_parallel_shards: int,
+    args: Any,
+) -> Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]]]:  # noqa: UP006
+    """Apply pre-sharding to the named parameters.
+
+    Parameters
+    ----------
+    named_params : Dict[str, nn.Parameter]
+        The named parameters of the model. If the model is quantized, the named parameters should
+        the state dictionary of the quantized model.
+    tensor_parallel_shards : int
+        The number of tensor parallel shards.
+    args : Any
+        The parsed arguments of weight conversion.
+
+    Returns
+    -------
+    Tuple[Dict[str, nn.Parameter], Dict[str, Callable[[Tensor], Sequence[Tensor]]]
+        The updated named parameters and the mapping from parameter name to the shard function.
+    """
+    bb = relax.BlockBuilder()
+    param_to_shard_func = {}
+    shard_func_names = set()
+    new_named_params: Dict[str, nn.Parameter] = {}  # noqa: UP006
+    has_shard_strategy = False
+    for name, param in named_params.items():
+        shard_strategy = param.attrs.get("shard_strategy", None)
+        if shard_strategy is not None:
+            has_shard_strategy = True
+            for i in range(tensor_parallel_shards):
+                new_named_params[_sharded_param_name(name, i)] = param
+            # create shard functions
+            param_to_shard_func[name] = shard_strategy.name
+            if shard_strategy.name not in shard_func_names:
+                _create_shard_func(bb, param, tensor_parallel_shards)
+                shard_func_names.add(shard_strategy.name)
+        else:
+            new_named_params[name] = param
+
+    if not has_shard_strategy:
+        logger.warning(
+            "No parameters with 'shard_strategy' found."
+            "At least one parameter must have a 'shard_strategy' for presharding. "
+            "The model will continue to convert weights in a non-presharded manner."
+        )
+
+    mod = bb.finalize()
+    vm = _compile_shard_funcs(mod, args.device)
+
+    for name in param_to_shard_func:
+        param_to_shard_func[name] = vm[param_to_shard_func[name]]
+    return new_named_params, param_to_shard_func

python/mlc_llm/support/random.py

@@ -0,0 +1,17 @@
+"""Utility functions for random number generation."""
+
+import sys
+
+
+def set_global_random_seed(seed):
+    """Set global random seed for python, numpy, torch and tvm."""
+    if "numpy" in sys.modules:
+        sys.modules["numpy"].random.seed(seed)
+    if "torch" in sys.modules:
+        sys.modules["torch"].manual_seed(seed)
+    if "random" in sys.modules:
+        sys.modules["random"].seed(seed)
+    if "tvm" in sys.modules:
+        set_seed = sys.modules["tvm"].get_global_func("mlc.random.set_seed")
+        if set_seed:
+            set_seed(seed)

python/mlc_llm/support/style.py

@@ -0,0 +1,62 @@
+"""Printing styles."""
+
+from enum import Enum
+
+
+class Styles(Enum):
+    """Predefined set of styles to be used.
+
+    Reference:
+    - https://en.wikipedia.org/wiki/ANSI_escape_code#3-bit_and_4-bit
+    - https://stackoverflow.com/a/17303428
+    """
+
+    RED = "\033[91m"
+    GREEN = "\033[92m"
+    YELLOW = "\033[93m"
+    BLUE = "\033[94m"
+    PURPLE = "\033[95m"
+    CYAN = "\033[96m"
+    BOLD = "\033[1m"
+    UNDERLINE = "\033[4m"
+    END = "\033[0m"
+
+
+def red(text: str) -> str:
+    """Return red text."""
+    return f"{Styles.RED.value}{text}{Styles.END.value}"
+
+
+def green(text: str) -> str:
+    """Return green text."""
+    return f"{Styles.GREEN.value}{text}{Styles.END.value}"
+
+
+def yellow(text: str) -> str:
+    """Return yellow text."""
+    return f"{Styles.YELLOW.value}{text}{Styles.END.value}"
+
+
+def blue(text: str) -> str:
+    """Return blue text."""
+    return f"{Styles.BLUE.value}{text}{Styles.END.value}"
+
+
+def purple(text: str) -> str:
+    """Return purple text."""
+    return f"{Styles.PURPLE.value}{text}{Styles.END.value}"
+
+
+def cyan(text: str) -> str:
+    """Return cyan text."""
+    return f"{Styles.CYAN.value}{text}{Styles.END.value}"
+
+
+def bold(text: str) -> str:
+    """Return bold text."""
+    return f"{Styles.BOLD.value}{text}{Styles.END.value}"
+
+
+def underline(text: str) -> str:
+    """Return underlined text."""
+    return f"{Styles.UNDERLINE.value}{text}{Styles.END.value}"

python/mlc_llm/support/tensor_parallel.py

@@ -0,0 +1,118 @@
+"""Sharding operators for tensor parallelism."""
+
+import dataclasses
+from contextlib import contextmanager
+from typing import Any, Dict, List, Optional  # noqa: UP035
+
+from tvm import te, tirx, topi
+from tvm.relax.frontend import nn
+
+
+@dataclasses.dataclass
+class ShardSingleDim:
+    """
+    Shard a tensor by a single dimension.
+
+
+    Parameters
+    ----------
+    name : str
+        The name of the shard func
+
+    dim : int
+        The dimension to shard
+
+    segs : Optional[List[int]]
+        The length of segments along `dim`. Default to None. If specified,
+        shard a tensor by its "segmented" dimension, where each segment has a different length
+        and sharded evenly on each worker.
+
+    """
+
+    name: str
+    dim: int
+    segs: Optional[List[int]] = None  # noqa: UP006
+
+    def gen_tir(self, shards: int, weight: nn.Tensor) -> tirx.PrimFunc:
+        """Generate a TIR function that shards the weight tensor by its rows."""
+        shape = weight.shape
+        segs = self.segs or [shape[self.dim]]
+        assert sum(segs) == shape[self.dim]
+        # NOTE: we use int64 to prevent int32 overflow
+        shape = [tirx.IntImm("int64", v) for v in shape]
+        segs = [tirx.IntImm("int64", v) for v in segs]
+        w = te.placeholder(
+            [tirx.IntImm("int64", v) for v in self._compute_in_shape(shards, weight)],
+            weight.dtype,
+            name="w",
+        )
+        ws: List[te.Tensor] = []  # noqa: UP006
+        offset = 0
+        for idx, sub_seg in enumerate(segs):
+            ws.append(
+                topi.transpose(
+                    topi.reshape(
+                        te.compute(
+                            (
+                                *shape[: self.dim],
+                                sub_seg * shards,
+                                *shape[self.dim + 1 :],
+                            ),
+                            lambda *idx: w[
+                                (
+                                    *idx[: self.dim],
+                                    idx[self.dim] + offset,
+                                    *idx[self.dim + 1 :],
+                                )
+                            ],
+                            name=f"w_{idx}",
+                        ),
+                        (
+                            *shape[: self.dim],
+                            tirx.IntImm("int64", shards),
+                            sub_seg,
+                            *shape[self.dim + 1 :],
+                        ),
+                    ),
+                    [self.dim, *range(self.dim), *range(self.dim + 1, len(shape) + 1)],
+                )
+            )
+            offset += sub_seg * shards
+        o = topi.concatenate(ws, axis=1 + self.dim)
+        func = te.create_prim_func([w, o])
+        return func
+
+    def gen_shard_info(self, shards: int, weight: nn.Tensor) -> Dict[str, Any]:  # noqa: UP006
+        """Generate shard info for this sharding strategy."""
+        return {
+            "func_name": self.name,
+            "in_shape": self._compute_in_shape(shards, weight),
+            "out_shape": (shards, *weight.shape),
+            "out_dtype": str(weight.dtype),
+        }
+
+    def _compute_in_shape(self, shards: int, weight: nn.Tensor) -> List[int]:  # noqa: UP006
+        """Compute the weight shape before sharding."""
+        shape = weight.shape
+        return [*shape[: self.dim], shape[self.dim] * shards, *shape[self.dim + 1 :]]
+
+
+@contextmanager
+def shard_bias(linear: nn.Linear, tensor_parallel_shards: int):
+    """
+    A context manager to shard the bias of a linear into `tensor_parallel_shards` shards.
+
+
+    Parameters
+    ----------
+    linear : nn.Linear
+        The linear layer whose bias would be sharded.
+
+    tensor_parallel_shards : int
+        The number of shards.
+    """
+    original_bias = linear.bias
+    if tensor_parallel_shards > 1:
+        linear.bias = linear.bias / tensor_parallel_shards
+    yield
+    linear.bias = original_bias

python/mlc_llm/support/tqdm.py

@@ -0,0 +1,39 @@
+"""Utils to better use tqdm"""
+
+import contextlib
+import inspect
+import io
+
+from tqdm import tqdm
+from tqdm.contrib.logging import logging_redirect_tqdm as _redirect_logging
+
+
+@contextlib.contextmanager
+def _redirect_print():
+    old_print = print
+
+    def new_print(*args, **kwargs):
+        with io.StringIO() as output:
+            kwargs["file"] = output
+            kwargs["end"] = ""
+            old_print(*args, **kwargs)
+            content = output.getvalue()
+        tqdm.write(content)
+
+    try:
+        inspect.builtins.print = new_print
+        yield
+    finally:
+        inspect.builtins.print = old_print
+
+
+@contextlib.contextmanager
+def redirect():
+    """Redirect tqdm output to logging and print."""
+
+    with _redirect_logging():
+        with _redirect_print():
+            yield
+
+
+__all__ = ["redirect", "tqdm"]