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

Code before:
        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]:
        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": param.dtype,
            "preprocs": param.attrs["preprocs"],
            "pipeline_stages": param.attrs.get("pipeline_stages", [0]),
        }
    logger.info("TOP LEVEL MODEL CONFIG BEFORE OVERRIDES: %s", str(model_config))
    _kwargs = getattr(model_config, "kwargs", {})
    model_config = args.overrides.apply(model_config)
    with args.target:
        op_ext.enable(
            target=args.target,
            flashinfer=args.opt.flashinfer,
            faster_transformer=args.opt.faster_transformer,
            cutlass=args.opt.cutlass,
        )
        # Step 1. Create the quantized model
        logger.info("Creating model from: %s", model_config)
        if (
            args.quantization.kind == "ft-quant"
            and hasattr(model_config, "tensor_parallel_shards")
            and model_config.tensor_parallel_shards > 1  # type: ignore
        ):
            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  # type: ignore
        ):
            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

Review comment (suggestion):
Just a note here, i noticed that this override wipes out any `kwargs` in the original model_config. This PR isn't the place to address it probably, but I just wanted to call it out.

Code after:
        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]:
        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": param.dtype,
            "preprocs": param.attrs["preprocs"],
            "pipeline_stages": param.attrs.get("pipeline_stages", [0]),
        }

    logger.info("TOP LEVEL MODEL CONFIG BEFORE OVERRIDES: %s", str(model_config))
    _kwargs = getattr(model_config, "kwargs", {})
    model_config = args.overrides.apply(model_config)
    with args.target:
        op_ext.enable(
            target=args.target,
            flashinfer=args.opt.flashinfer,
            faster_transformer=args.opt.faster_transformer,
            cutlass=args.opt.cutlass,
        )
        # Step 1. Create the quantized model
        logger.info("Creating model from: %s", model_config)
        if (
            args.quantization.kind == "ft-quant"
            and hasattr(model_config, "tensor_parallel_shards")
            and model_config.tensor_parallel_shards > 1  # type: ignore
        ):
            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  # type: ignore
        ):
            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)