fix(loading): Mistral3-VLM with mistral4 (MLA) text backbone fails to load (routes to Llama q_proj loader)

[inureyes]

Problem

Loading the Mistral Small 4 checkpoint mlx-community/Mistral-Small-4-119B-2603-4bit (a Mistral3ForConditionalGeneration VLM whose text_config.model_type is mistral4, i.e. an MLA + MoE text backbone with weights self_attn.q_a_proj / kv_a_proj_with_mqa / kv_b_proj) fails at load with: Failed to load text model: Weight not found: model.layers.0.self_attn.q_proj.weight.

Root cause

build_pixtral_family_context (src/loading/vlm_pixtral.rs:104-109) unconditionally parses the text config as models::llama3::ModelArgs and loads the text backbone via models::Llama3Model::from_weights (which expects standard attention q_proj). It never checks text_config.model_type == "mistral4", so an MLA text backbone (q_a_proj / kv_a_proj_with_mqa / kv_b_proj, no q_proj) cannot be found and the load fails. The non-VLM text-only path already handles this: load_llama_family_from_weights (src/loading/mod.rs:210-220) branches on is_mistral4_config(config) and routes to models::Mistral4Model::from_weights (after models::mistral4::sanitize_weights(weights, &args, "")). The VLM context builder is missing the equivalent branch. Note build_pixtral_family_context already strips the language_model. prefix via strip_language_model_prefix (line 100), so the weights are at the bare model.layers... namespace for both loaders.

Proposed fix

In build_pixtral_family_context, branch on is_mistral4_config(&full_config) (or text_config.model_type == "mistral4"). When true, parse the text config as models::mistral4::Mistral4Config, run models::mistral4::sanitize_weights(&mut weights, &args, ""), and build models::Mistral4Model::from_weights(&weights, &args) wrapped as LoadedModel::Mistral4, mirroring load_llama_family_from_weights. Keep the Llama path for non-mistral4 Mistral3 VLMs (Pixtral / standard Mistral text). Be careful that the Llama-specific config massaging in build_mistral_text_config (infer_llama_config_from_weights, apply_mistral_attention_head_override, which key off q_proj) does not corrupt the mistral4 config; for the mistral4 branch prefer the raw text_config with quantization inherited rather than the Llama-massaged value.

Impact

Mistral Small 4 (the only public mistral4 checkpoint) is a VLM, so this blocks loading it at all on the server/CLI. It also blocks real-model validation of the #421 dense-prefill-mask fix for mistral4 (the mask fix is in the mistral4 text backbone, which cannot currently load). Pre-existing (reproduces on main), independent of #421.

Acceptance criteria

• mlxcel-server -m <Mistral-Small-4-119B-2603-4bit> loads without the q_proj weight-not-found error and serves text generation.
• Text-only generation produces coherent output.
• Non-mistral4 Mistral3 VLMs (Pixtral / standard Mistral text backbones) still load unchanged.
• Ideally an image+text multimodal request also works (bonus; the projector/vision path is unchanged).

References

• Surfaced validating PR fix(models): size dense prefill masks from live_len for mistral4/nemotron_nas/qwen-vl under --max-kv-size (#421) #422 (fix(models): complete offset->live_len prefill-mask migration for dense-KVCache models under --max-kv-size (mistral4/nemotron_nas/qwen*-vl) #421). Checkpoint: https://huggingface.co/mlx-community/Mistral-Small-4-119B-2603-4bit