Implementation:Mlc ai Mlc llm Mistral Loader

Overview

The Mistral Loader module (python/mlc_llm/model/mistral/mistral_loader.py) defines parameter mappings for converting Mistral model weights from HuggingFace and AWQ formats into MLC LLM's internal representation. The loader follows the same structural pattern as the Llama loader, providing both a huggingface function for standard weights and an awq function for pre-quantized AWQ weights.

Location

• File: python/mlc_llm/model/mistral/mistral_loader.py
• Lines: 166
• Module: mlc_llm.model.mistral

Function: huggingface

def huggingface(model_config: MistralConfig, quantization: Quantization) -> ExternMapping:

Returns a parameter mapping from MLC LLM parameter names to HuggingFace PyTorch parameter names for the Mistral architecture.

Parameters:

Parameter	Type	Description
model_config	MistralConfig	The configuration of the Mistral model.
quantization	Quantization	The quantization configuration.

Initialization

model = MistralForCasualLM(model_config)
if quantization is not None:
    model.to(quantization.model_dtype)
_, _named_params, _ = model.export_tvm(
    spec=model.get_default_spec(),
    allow_extern=True,
)
named_parameters = dict(_named_params)

Note that the class is named MistralForCasualLM (with a typo -- "Casual" instead of "Causal"), which is preserved from the model definition.

Per-Layer Mappings

For each of the num_hidden_layers transformer layers:

QKV Projection Fusion

Q, K, V projection weights are concatenated along axis 0:

attn = f"model.layers.{i}.self_attn"
mlc_name = f"{attn}.qkv_proj.weight"
mapping.add_mapping(
    mlc_name,
    [
        f"{attn}.q_proj.weight",
        f"{attn}.k_proj.weight",
        f"{attn}.v_proj.weight",
    ],
    functools.partial(
        lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
        dtype=mlc_param.dtype,
    ),
)

MLP Gate-Up Fusion

Gate and up projection weights are concatenated along axis 0:

mlp = f"model.layers.{i}.mlp"
mlc_name = f"{mlp}.gate_up_proj.weight"
mapping.add_mapping(
    mlc_name,
    [
        f"{mlp}.gate_proj.weight",
        f"{mlp}.up_proj.weight",
    ],
    functools.partial(
        lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype),
        dtype=mlc_param.dtype,
    ),
)

Unused Parameters

mapping.add_unused(f"{attn}.rotary_emb.inv_freq")

Identity Fallback

Remaining parameters are mapped with a simple dtype cast:

for mlc_name, mlc_param in named_parameters.items():
    if mlc_name not in mapping.param_map:
        mapping.add_mapping(
            mlc_name,
            [mlc_name],
            functools.partial(
                lambda x, dtype: x.astype(dtype),
                dtype=mlc_param.dtype,
            ),
        )

Function: awq

def awq(model_config: MistralConfig, quantization: Quantization) -> ExternMapping:

Returns a parameter mapping from MLC LLM parameter names to AWQ pre-quantized parameter names.

Initialization

model, _ = awq_quant(model_config, quantization)
_, _named_params = model.export_tvm(
    spec=model.get_default_spec(),
    allow_extern=True,
)
named_parameters = dict(_named_params)

Note that the Mistral AWQ path uses a two-element return from export_tvm (without the third element), unlike the HuggingFace path which unpacks three elements.

Quantized Parameter Fusion

For each layer, three AWQ-specific suffixes (qweight, qzeros, scales) are mapped for both QKV attention and gate-up MLP projections. Unlike the HuggingFace path, AWQ concatenation uses axis=0 for both operations:

for quantize_suffix in ["qweight", "qzeros", "scales"]:
    mlc_name = f"{attn}.qkv_proj.{quantize_suffix}"
    mapping.add_mapping(
        mlc_name,
        [
            f"{attn}.q_proj.{quantize_suffix}",
            f"{attn}.k_proj.{quantize_suffix}",
            f"{attn}.v_proj.{quantize_suffix}",
        ],
        functools.partial(
            lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
            dtype=mlc_param.dtype,
        ),
    )

The MLP gate-up AWQ mapping follows the same pattern:

for quantize_suffix in ["qweight", "qzeros", "scales"]:
    mlc_name = f"{mlp}.gate_up_proj.{quantize_suffix}"
    mapping.add_mapping(
        mlc_name,
        [
            f"{mlp}.gate_proj.{quantize_suffix}",
            f"{mlp}.up_proj.{quantize_suffix}",
        ],
        functools.partial(
            lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype),
            dtype=mlc_param.dtype,
        ),
    )

Comparison with Llama Loader

The Mistral loader is structurally very similar to the Llama loader but differs in one key aspect:

Aspect	Llama AWQ	Mistral AWQ
Concatenation axis	axis=1 (due to transposed AWQ GEMM)	axis=0
export_tvm return	3 elements	2 elements (no third return value)

Dependencies

• functools -- for functools.partial
• numpy -- for array concatenation and dtype casting
• mlc_llm.loader.ExternMapping -- the core mapping data structure
• mlc_llm.quantization.Quantization -- quantization configuration
• .mistral_model.MistralConfig, .mistral_model.MistralForCasualLM -- Mistral model definitions
• .mistral_quantization.awq_quant -- AWQ quantization utility