Repository: mlc-ai/mlc-llm
File: python/mlc_chat/interface/compile.py (Python)

Code before:
def _compile(args: CompileArgs, model_config: ConfigBase):
    def _get_variable_bounds(model_config) -> Dict[str, int]:
        if hasattr(model_config, "sliding_window_size"):
            return {
                "seq_len": model_config.prefill_chunk_size,
                "rolling_cache_len": model_config.sliding_window_size,
                "kv_seq_len": model_config.sliding_window_size + model_config.prefill_chunk_size,
            }
        return {
            "seq_len": model_config.prefill_chunk_size,
            "total_seq_len": model_config.context_window_size,
        }

    def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]:
        return {
            "name": name,
            "shape": list(param.shape),
            "dtype": param.dtype,
            "preprocs": param.attrs["preprocs"],
        }

    def _find_kv_cache_bytes(model: nn.Module, model_config) -> int:
        all_kv_cache = nn.core._attribute_finder(  # pylint: disable=protected-access
            model,
            prefix="",
            condition_yield=lambda x: isinstance(x, nn.KVCache),
        )
        result = 0
        for _, kv_cache in all_kv_cache:
            result += math.prod(kv_cache.unit_shape) * np.dtype(kv_cache.dtype).itemsize
        if getattr(model_config, "sliding_window_size", -1) > 0:
            window_size = model_config.sliding_window_size
        elif getattr(model_config, "context_window_size", -1) > 0:
            window_size = model_config.context_window_size
        else:
            window_size = 0
        return result * window_size

    args.overrides.apply(model_config)
    with args.target:
        op_ext.enable(
            target=args.target,
            flashinfer=args.opt.flashinfer,
        )
        # Step 1. Create the quantized model
        logger.info("Creating model from: %s", args.config)
        model, _ = args.model.quantize[args.quantization.kind](model_config, args.quantization)
        kv_cache_bytes = _find_kv_cache_bytes(model, model_config)
        # Step 2. Exporting the model to TVM Unity
        logger.info("Exporting the model to TVM Unity compiler")
        mod, named_params, ext_mods = model.export_tvm(

Review comment (performance):
@MasterJH5574 This for now only assumes `nn.KVCache`, and we will need to update this method to support serving cases where KVCache can be more complicated.

Code after:
def _compile(args: CompileArgs, model_config: ConfigBase):
    def _get_variable_bounds(model_config) -> Dict[str, int]:
        if hasattr(model_config, "sliding_window_size"):
            return {
                "seq_len": model_config.prefill_chunk_size,
                "rolling_cache_len": model_config.sliding_window_size,
                "kv_seq_len": model_config.sliding_window_size + model_config.prefill_chunk_size,
            }
        return {
            "seq_len": model_config.prefill_chunk_size,
            "total_seq_len": model_config.context_window_size,
        }

    def _get_param_metadata(name: str, param: nn.Parameter) -> Dict[str, Any]:
        return {
            "name": name,
            "shape": list(param.shape),
            "dtype": param.dtype,
            "preprocs": param.attrs["preprocs"],
        }

    def _find_kv_cache_bytes(model: nn.Module, model_config) -> int:
        all_kv_cache = nn.core._attribute_finder(  # pylint: disable=protected-access
            model,
            prefix="",
            condition_yield=lambda x: isinstance(x, nn.KVCache),
        )
        result = 0
        for _, kv_cache in all_kv_cache:
            result += math.prod(kv_cache.unit_shape) * np.dtype(kv_cache.dtype).itemsize
        if getattr(model_config, "sliding_window_size", -1) > 0:
            window_size = model_config.sliding_window_size
        elif getattr(model_config, "context_window_size", -1) > 0:
            window_size = model_config.context_window_size
        else:
            window_size = 0
        return result * window_size

    model_config = args.overrides.apply(model_config)
    with args.target:
        op_ext.enable(
            target=args.target,
            flashinfer=args.opt.flashinfer,
        )
        # Step 1. Create the quantized model
        logger.info("Creating model from: %s", model_config)
        model, _ = args.model.quantize[args.quantization.kind](model_config, args.quantization)
        kv_cache_bytes = _find_kv_cache_bytes(model, model_config)
        # Step 2. Exporting the model to TVM Unity
        logger.info("Exporting the model to TVM Unity compiler")
        mod, named_params, ext_mods = model.export_tvm(