compile.py

"""Python entrypoint of compilation."""

import dataclasses
from io import StringIO
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple  # noqa: UP035

from tvm import IRModule, relax, tirx
from tvm.ir.transform import Pass, PassContext
from tvm.relax.frontend import nn
from tvm.target import Target

from mlc_llm import compiler_pass as _  # noqa: F401
from mlc_llm import op as op_ext
from mlc_llm.cli.model_metadata import _report_memory_usage
from mlc_llm.model import Model
from mlc_llm.quantization import Quantization
from mlc_llm.support import logging
from mlc_llm.support.config import ConfigBase
from mlc_llm.support.runtime_lora import RuntimeLoRAConfig, validate_runtime_lora_scope
from mlc_llm.support.style import bold

from .compiler_flags import ModelConfigOverride, OptimizationFlags

logger = logging.getLogger(__name__)


@dataclasses.dataclass
class CompileArgs:
    """Arguments to MLC LLM's compiler."""

    config: Path
    quantization: Quantization
    model: Model
    target: Target
    opt: OptimizationFlags
    build_func: Callable[[IRModule, "CompileArgs", Pass], None]
    system_lib_prefix: str
    output: Path
    overrides: ModelConfigOverride
    debug_dump: Optional[Path]
    lora_adapter: Optional[Path]

    def __post_init__(self) -> None:
        self.opt.update(self.target, self.quantization)

    def display(self) -> None:
        """Display the arguments to stdout."""
        out = StringIO()
        print(f"{bold('Compiling with arguments:')}", file=out)
        print(f"  {bold('--config'):<25} {self.config}", file=out)
        print(f"  {bold('--quantization'):<25} {self.quantization}", file=out)
        print(f"  {bold('--model-type'):<25} {self.model.name}", file=out)
        print(f"  {bold('--target'):<25} {self.target.export()}", file=out)
        print(f"  {bold('--opt'):<25} {self.opt}", file=out)
        print(f'  {bold("--system-lib-prefix"):<25} "{self.system_lib_prefix}"', file=out)
        print(f"  {bold('--output'):<25} {self.output}", file=out)
        print(f"  {bold('--overrides'):<25} {self.overrides}", file=out)
        if self.lora_adapter is not None:
            print(f"  {bold('--lora-adapter'):<25} {self.lora_adapter}", file=out)
        # As it's debug only, no need to display
        # print(f"  {bold('--debug-dump'):<25} {self.debug_dump}", file=out)
        print(out.getvalue().rstrip())


def _apply_preproc_to_params_and_check_pipeline(
    named_params: List[Tuple[str, nn.Parameter]],  # noqa: UP006
    model_config,
) -> Dict[str, tirx.PrimFunc]:  # noqa: UP006
    extra_tirs: Dict[str, tirx.PrimFunc] = {}  # noqa: UP006
    for name, param in named_params:
        preprocs = param.attrs.get("preprocs", [])
        shard_strategy = param.attrs.get("shard_strategy", None)
        if shard_strategy is not None and model_config.tensor_parallel_shards > 1:
            preprocs.append(
                shard_strategy.gen_shard_info(
                    shards=model_config.tensor_parallel_shards,
                    weight=param,
                )
            )
            if shard_strategy.name not in extra_tirs:
                extra_tirs[shard_strategy.name] = shard_strategy.gen_tir(
                    shards=model_config.tensor_parallel_shards,
                    weight=param,
                )
        param.attrs["preprocs"] = preprocs

        pipeline_parallel_stages = getattr(model_config, "pipeline_parallel_stages", 1)
        if pipeline_parallel_stages != 1:
            assert "pipeline_stages" in param.attrs, (
                f'The pipeline stage is undefined for parameter "{name}" when the number '
                f"of pipeline parallel stages is {pipeline_parallel_stages}"
            )
        param.attrs["pipeline_stages"] = (
            [0]
            if "pipeline_stages" not in param.attrs
            else list(set(param.attrs["pipeline_stages"]))
        )
    return extra_tirs


def _infer_kv_state_kind(model_type) -> str:
    if "rwkv" in model_type:
        return "rnn_state"
    if "medusa" in model_type:
        return "none"
    if "qwen3_5" in model_type:
        return "hybrid"
    return "kv_cache"


def _compile(args: CompileArgs, model_config: ConfigBase):
    def _get_variable_bounds(model_config) -> Dict[str, int]:  # noqa: UP006
        if hasattr(model_config, "sliding_window_size"):
            return {
                "rolling_cache_len": model_config.sliding_window_size,
                "kv_seq_len": model_config.sliding_window_size + model_config.prefill_chunk_size,
                "seq_len": model_config.prefill_chunk_size,
                "batch_size": getattr(model_config, "max_batch_size", 1),
            }
        return {
            "total_seq_len": model_config.context_window_size,
            "seq_len": model_config.prefill_chunk_size,
            "batch_size": getattr(model_config, "max_batch_size", 1),
        }

    def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]:  # noqa: UP006
        return {
            "name": name,
            # Record dynamic shape as -1 (e.g. vocab_size)
            "shape": [s if isinstance(s, int) else s.name for s in param.shape],
            "dtype": str(param.dtype),
            "preprocs": param.attrs["preprocs"],
            "pipeline_stages": param.attrs.get("pipeline_stages", [0]),
        }

    logger.info("TOP LEVEL MODEL CONFIG BEFORE OVERRIDES: %s", str(model_config))
    _kwargs = getattr(model_config, "kwargs", {})
    model_config = args.overrides.apply(model_config)
    runtime_lora = getattr(model_config, "runtime_lora", None)
    if runtime_lora is not None:
        validate_runtime_lora_scope(
            model_name=args.model.name,
            quantization_name=args.quantization.name,
            tensor_parallel_shards=model_config.tensor_parallel_shards,
        )
    with args.target:
        op_ext.enable(
            target=args.target,
            flashinfer=args.opt.flashinfer,
            faster_transformer=args.opt.faster_transformer,
            cutlass=args.opt.cutlass,
        )
        # Step 1. Create the quantized model
        logger.info("Creating model from: %s", model_config)
        if (
            args.quantization.kind == "ft-quant"
            and hasattr(model_config, "tensor_parallel_shards")
            and model_config.tensor_parallel_shards > 1
        ):
            raise NotImplementedError
        if (
            hasattr(args.quantization, "linear_weight_layout")
            and args.quantization.linear_weight_layout == "KN"
            and hasattr(model_config, "tensor_parallel_shards")
            and model_config.tensor_parallel_shards > 1
        ):
            raise NotImplementedError(
                "KN layout (q3f16_0 and q4f16_0) is not supported for tensor parallelism"
            )
        model, _ = args.model.quantize[args.quantization.kind](model_config, args.quantization)
        # Step 2. Exporting the model to TVM
        logger.info("Exporting the model to TVM compiler")
        mod, named_params, ext_mods = model.export_tvm(
            spec=model.get_default_spec(),
            allow_extern=True,
        )
        # Step 3. Running relax compilation pipeline
        logger.info("Running optimizations using TVM")
        additional_tirs = _apply_preproc_to_params_and_check_pipeline(named_params, model_config)
        variable_bounds = _get_variable_bounds(model_config)
        cuda_graph_symbolic_capture_hints = {
            "batch_decode": ["batch_size"],
            "batch_decode_to_last_hidden_states": ["batch_size"],
            "batch_verify": ["batch_size", "seq_len"],
            "batch_verify_to_last_hidden_states": ["batch_size", "seq_len"],
        }
        avs = _kwargs.get("active_vocab_size", None)
        if avs is not None and avs <= 0:
            avs = None
        metadata = {
            "model_type": args.model.name,
            "quantization": args.quantization.name,
            "context_window_size": getattr(model_config, "context_window_size", -1),
            "sliding_window_size": getattr(model_config, "sliding_window_size", -1),
            "attention_sink_size": getattr(model_config, "attention_sink_size", -1),
            "prefill_chunk_size": model_config.prefill_chunk_size,
            "tensor_parallel_shards": model_config.tensor_parallel_shards,
            "pipeline_parallel_stages": getattr(model_config, "pipeline_parallel_stages", 1),
            "disaggregation": getattr(model_config, "disaggregation", False),
            "kv_state_kind": _infer_kv_state_kind(args.model.name),
            "max_batch_size": getattr(model_config, "max_batch_size", 1),
            "active_vocab_size": avs,
            "model_task": args.model.model_task,
        }
        if runtime_lora is not None:
            metadata["runtime_lora"] = runtime_lora.asdict()
        if args.model.embedding_metadata:
            metadata["embedding_metadata"] = dataclasses.asdict(args.model.embedding_metadata)
        logger.info("Registering metadata: %s", metadata)
        metadata["params"] = [_get_param_metadata(name, param) for name, param in named_params]
        pass_config = {"relax.backend.use_cuda_graph": args.opt.cudagraph}
        # TODO: Remove this workaround when the TVM CSE regression is fixed.
        # Temporary workaround for TVM CSE regression that can produce
        # dangling `cse_v*` vars during host codegen.
        pass_config["tirx.disable_cse_tir"] = True

        with PassContext(config=pass_config):
            args.build_func(
                mod,
                args,
                pipeline=relax.get_pipeline(
                    "mlc_llm",
                    target=args.target,
                    flashinfer=args.opt.flashinfer,
                    cublas_gemm=args.opt.cublas_gemm,
                    faster_transformer=args.opt.faster_transformer,
                    allreduce_strategy=args.opt.ipc_allreduce_strategy,
                    variable_bounds=variable_bounds,
                    cuda_graph_symbolic_capture_hints=cuda_graph_symbolic_capture_hints,
                    additional_tirs=additional_tirs,
                    ext_mods=ext_mods,
                    metadata=metadata,
                    debug_dump=args.debug_dump,
                ),
            )
        _report_memory_usage(metadata=metadata, config=model_config)
    logger.info("Generated: %s", bold(str(args.output)))


def compile(
    config: Dict[str, Any],  # noqa: UP006
    quantization: Quantization,
    model_type: Model,
    target: Target,
    opt: OptimizationFlags,
    build_func: Callable[[IRModule, CompileArgs, Pass], None],
    system_lib_prefix: str,
    output: Path,
    overrides: ModelConfigOverride,
    debug_dump: Optional[Path] = None,
    lora_adapter: Optional[Path] = None,
):
    """Compile a model given its configuration and quantization format to a specific target."""
    config = dict(config)
    if "model_config" in config:
        config["model_config"] = dict(config["model_config"])
    if lora_adapter is not None:
        config["runtime_lora"] = RuntimeLoRAConfig.from_peft_directory(lora_adapter).asdict()

    avs = None
    if "active_vocab_size" in config:
        avs = config.pop("active_vocab_size")
        logger.info("Active vocab size from input config: %s", str(avs))
    if "model_config" in config:
        model_config = config.pop("model_config")
        model_config.update(config)
        model_config = model_type.config.from_dict(model_config)
    else:
        model_config = model_type.config.from_dict(config)
    model_config.kwargs = {"active_vocab_size": avs} if avs is not None else {}
    args = CompileArgs(
        model_config,
        quantization,
        model_type,
        target,
        opt,
        build_func,
        system_lib_prefix,
        output,
        overrides,
        debug_dump,
        lora_adapter,
    )
    args.display()
    _compile(args, model_config)