Implementation:Mlc ai Mlc llm Ministral3 Loader

Overview

The Ministral3 Loader module (python/mlc_llm/model/ministral3/ministral3_loader.py) defines the parameter mapping for converting Ministral3 (Mistral 3 conditional generation) model weights from HuggingFace format to MLC LLM's internal representation. This is the most complex loader in the codebase, as it supports block-scale FP8 quantization in addition to standard weight loading, and handles both text-only and vision-language model configurations.

Location

• File: python/mlc_llm/model/ministral3/ministral3_loader.py
• Lines: 265
• Module: mlc_llm.model.ministral3

Helper Function: _dequantize_block_scale_weight

def _dequantize_block_scale_weight(
    weight: np.ndarray, weight_scale: np.ndarray, block_size: Tuple[int, int]
) -> np.ndarray:

Reconstructs float weights from FP8 block-scale storage by multiplying each block of the quantized weight by its corresponding scale factor. This is used internally to verify or transform FP8 block-quantized weights.

Parameters:

Parameter	Type	Description
weight	np.ndarray	The FP8 quantized weight matrix.
weight_scale	np.ndarray	A 2D scale matrix with one scale per block.
block_size	Tuple[int, int]	The (rows, cols) dimensions of each quantization block.

The function iterates over blocks defined by block_size and applies element-wise multiplication:

for i in range(num_row_blocks):
    row_start = i * block_rows
    row_end = min(row_start + block_rows, rows)
    scale_row = weight_scale[i]
    for j in range(num_col_blocks):
        col_start = j * block_cols
        col_end = min(col_start + block_cols, cols)
        out[row_start:row_end, col_start:col_end] = (
            weight[row_start:row_end, col_start:col_end] * scale_row[j]
        )

Function: huggingface

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

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

Initialization and BlockScale Handling

The function begins by instantiating the model and optionally applying block-scale quantization:

model = Mistral3ForConditionalGeneration(model_config)
if quantization is not None:
    model.to(quantization.model_dtype)
if isinstance(quantization, BlockScaleQuantize):
    model = quantization.quantize_model(model, QuantizeMapping({}, {}), "")
    if model_config.weight_block_size is None:
        raise ValueError(
            "The input Ministral 3 model is not fp8 block quantized. "
            "Thus BlockScaleQuantize is not supported."
        )

A validation check ensures that non-BlockScale quantization is not applied to FP8 block-quantized models:

if (
    not isinstance(quantization, BlockScaleQuantize)
    and model_config.weight_block_size is not None
):
    raise ValueError(
        "The input Ministral 3 model is fp8 block quantized. "
        "Please use BlockScaleQuantize for the model."
    )

Prefix Handling

The loader strips the language_model. prefix from MLC parameter names and conditionally adds it back for HuggingFace names when a vision configuration is present:

if any(name.startswith("language_model.") for name in raw_params):
    named_parameters = {
        name.replace("language_model.", "", 1): value for name, value in raw_params.items()
    }

hf_prefix = ""
if "vision_config" in model_config.kwargs:
    hf_prefix = "language_model."

Core Helper: add_weight_and_scale_mapping

This is the central function that handles mapping for both the weight itself and its associated block-scale quantization artifacts:

def add_weight_and_scale_mapping(
    weight_mlc_name: str,
    weight_hf_names: List[str],
    weight_transform_func: Callable,
    activation_transform_func: Optional[Callable] = None,
):

For each weight parameter, it:

1. Adds the primary weight mapping using the provided transform function.
2. If BlockScaleQuantize is active, adds a mapping for the _scale_inv suffix parameter (the inverse scale for block-scale quantization).
3. If BlockScaleQuantize is active, adds a mapping for the .activation_scale parameter.

The scale transform functions include sophisticated shape handling to deal with mismatches between HuggingFace and MLC scale tensor shapes, including broadcasting scalar scales and reshaping 1D scales to 2D:

def _weight_scale_transform(*arrays, dtype: str, _transform=weight_transform_func):
    processed = []
    for arr in arrays:
        arr_np = np.asarray(arr)
        if arr_np.ndim == 0:
            arr_np = arr_np.reshape((1,))
        processed.append(arr_np)
    result = _transform(*processed, dtype=dtype)
    result = np.asarray(result, dtype=dtype)
    if result.shape == expected_weight_scale_shape:
        return result
    if result.shape == ():
        return np.full(expected_weight_scale_shape, result.item(), dtype=dtype)
    # ... additional shape handling

Shared Activation Scale Transform

A factory function that creates activation scale transforms which verify that all source activation scales are identical before using the first:

def make_shared_activation_transform(target_name: str):
    def func(first: np.ndarray, *rest: np.ndarray, dtype: str):
        for _, arr in enumerate(rest, start=1):
            if not np.allclose(arr, first):
                raise ValueError(
                    f"Activation scales for {target_name} must be identical between "
                    "concatenated sources."
                )
        return first.astype(dtype)
    return func

This is necessary because when Q, K, V projections are fused, their activation scales must be consistent.

Per-Layer Mappings

For each layer, the following mappings are registered:

QKV Projection Fusion

attn = f"model.layers.{i}.self_attn"
mlc_name = f"{attn}.qkv_proj.weight"
proj_sources = [hf(f"{attn}.{proj}.weight") for proj in ["q_proj", "k_proj", "v_proj"]]
add_weight_and_scale_mapping(
    mlc_name,
    proj_sources,
    lambda q, k, v, dtype: np.concatenate([q, k, v], axis=0).astype(dtype),
    activation_transform_func=make_shared_activation_transform(
        f"{mlc_name}_activation_scale"
    ),
)

MLP Gate-Up Fusion

mlp = f"model.layers.{i}.mlp"
mlc_name = f"{mlp}.gate_up_proj.weight"
gate_sources = [hf(f"{mlp}.{proj}.weight") for proj in ["gate_proj", "up_proj"]]
add_weight_and_scale_mapping(
    mlc_name,
    gate_sources,
    lambda gate, up, dtype: np.concatenate([gate, up], axis=0).astype(dtype),
    activation_transform_func=make_shared_activation_transform(
        f"{mlc_name}_activation_scale"
    ),
)

Individual Linear Layers

The output projection and down projection are mapped individually with identity transforms:

for linear_name in [f"{attn}.o_proj.weight", f"{mlp}.down_proj.weight"]:
    add_weight_and_scale_mapping(
        linear_name,
        [hf(linear_name)],
        identity_transform,
    )

Unused Parameters

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

Identity Fallback

Remaining parameters are mapped with a simple dtype cast, applying the HuggingFace prefix:

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

Key Design Decisions

• Block-scale FP8 support: This loader is unique in supporting BlockScaleQuantize, which requires mapping not only the weight but also the per-block scale inverse and activation scale tensors.
• Validation guards: Two complementary checks ensure that block-scale quantized models use BlockScaleQuantize and vice versa.
• Shape broadcasting: The scale transform functions handle various shape mismatches (scalar, 1D, 2D) that arise from differences between HuggingFace and MLC storage conventions.
• Activation scale consistency: When fusing Q/K/V projections, the loader validates that all source activation scales are identical, raising an error if they differ.

Dependencies

• functools -- for functools.partial
• numpy -- for array concatenation, broadcasting, and dtype casting
• mlc_llm.loader.ExternMapping, mlc_llm.loader.QuantizeMapping -- mapping data structures
• mlc_llm.quantization.BlockScaleQuantize, mlc_llm.quantization.Quantization -- quantization types
• .ministral3_model.Ministral3Config, .ministral3_model.Mistral3ForConditionalGeneration -- model definitions