Adding New Model Architectures¶

This guide explains how to add a new model architecture to mlxcel. It is based on docs/MODEL_ADDITION_GUIDE.md (which is the canonical reference) and describes the concrete control surfaces you must touch.

Table of Contents¶

• Goals
• Before You Start
• Text Model Checklist
• VLM Checklist
• When to Create a New Family Module
• Model Implementation Pattern
• Weight Loading from SafeTensors
• Model Type Detection
• LoadedModel Dispatch
• Common Pitfalls
• Testing the New Model
• Shared Function and Execution Policy
• Dry-Run Workflow Validation
• Review Checklist
• Commit Strategy

Goals¶

• Keep new model additions predictable.
• Reuse existing control-plane helpers instead of adding new one-off branches.
• Add tests alongside the integration points that are easiest to regress.

Before You Start¶

1. Identify the reference implementation:
2. Text models: mlx_lm/models/ in ml-explore/mlx-lm

3. VLMs: mlx_vlm/models/ in Blaizzy/mlx-vlm

4. Classify the architecture:

5. Brand new model family → new file under src/models/
6. Format alias of an existing family → extend detection in detection.rs

7. VLM wrapper around an existing text model → add vision encoder/connector under src/vision/ and loader under src/loading/vlm_<family>.rs

8. Check whether an existing loader helper already matches: Start with src/model_metadata.rs: the converged registration table. Standard text models should extend that registration table first, because src/loading/config_backed.rs consumes the same source of truth.

9. Check for shared execution policy impact:

10. src/execution/runtime.rs for device/environment behavior
11. src/execution/sampling.rs for sampling defaults

Text Model Checklist¶

Follow these steps in order:

1. Add the implementation file under src/models/<family>.rs.

2. Register the module and re-export in src/models/mod.rs:

   pub mod my_new_model;
   pub use my_new_model::MyNewModel;
   

3. Add a ModelType variant in src/models/mod.rs:

   pub enum ModelType {
       // ... existing variants ...
       MyNewModel,
   }
   

4. Extend get_model_type() in src/models/detection.rs:

   "my_new_model" => Ok(ModelType::MyNewModel),
   

   Use shared helpers when the pattern fits:

   "my_family" => Ok(detect_text_or_vlm(&v, ModelType::MyFamily, ModelType::MyFamilyVLM)),
   

5. Add the LoadedModel variant in src/loaded_model.rs. Prefer extending the delegate_language_model! macro rather than adding a new repeated match arm.

6. Wire loading in src/loading/mod.rs. For standard config-backed models:

7. Update src/model_metadata.rs for kind, adapter support, and route selection.

8. If the model follows the standard text-model path, extend the shared registration surface instead of adding a parallel entry list in src/loading/config_backed.rs.

9. LoRA/adapter support (if applicable): verify load_model_from_weights() in src/loading/mod.rs. Non-standard adapter paths go in src/loading/special.rs.

VLM Checklist¶

1. Implement vision components (or reuse existing ones):
2. Vision encoder: src/vision/encoders/
3. Connector (bridges encoder output to text model input): src/vision/connectors/

4. Processor (image preprocessing): src/vision/processors/

5. Add VLM ModelType detection in src/models/detection.rs. For families with both text-only and VLM variants, use detect_text_or_vlm().

6. Add the loader entry in src/loading/vlm_<family>.rs (or the nearest peer family module). Register the directory entry point in try_load_vlm_model_from_dir() in src/loading/mod.rs.

7. Add LoadedModel capability helpers in src/loaded_model_capabilities.rs. Prefer extending VlmRuntimeRef over adding family-specific getters.

8. Reuse prompt helpers where possible:

9. Qwen-VL token insertion: src/multimodal/qwen_vl.rs
10. Generic image-token block expansion: src/multimodal/vlm_prompt.rs
11. Phi3V prompt tag handling: src/multimodal/phi3v_prompt.rs

When to Create a New Family Module¶

Do not create a new file by default. Create one when the family has a distinct control-plane identity.

Create a new loader family module when:

• config normalization is not a small variant of an existing family
• token defaults or weight-key remapping need dedicated tests
• the VLM wrapper uses a different prompt/runtime assembly path
• adding the logic inline would obscure an existing family boundary

Keep the model in an existing module when:

• it is primarily an alias or small config delta
• the same loader tests already express the policy
• the family is still recognizable after the change

If you are unsure, extend the existing family module first and split only when the test file or router starts to lose a clear boundary.

Model Implementation Pattern¶

All model implementations use mlxcel-core for MLX C++ bindings via the cxx bridge. The LanguageModel trait is defined in src/lib/mlxcel-core/src/generate.rs:

pub trait LanguageModel {
    /// Forward pass: returns logits of shape [batch, seq_len, vocab_size]
    fn forward(
        &self,
        input_ids: &MlxArray,
        caches: &mut [KVCache],
        mask: Option<&MlxArray>,
    ) -> UniquePtr<MlxArray>;

    /// Create per-layer KV caches
    fn make_caches(&self) -> Vec<KVCache>;

    /// Number of transformer/SSM layers
    fn num_layers(&self) -> usize;

    /// EOS token IDs for this model
    fn eos_token_ids(&self) -> Vec<i32>;

    /// Whether tile-aligned padded prefill is safe (default: true).
    /// Recurrent hybrids (NemotronH, Jamba, Mamba) return false because
    /// pad positions corrupt their conv/SSM state.
    fn supports_padded_prefill(&self) -> bool { true }

    /// Drop `excess` pad positions from model-owned state after a padded
    /// prefill. The default has nothing to trim; families that keep their
    /// own caches and still answer `supports_padded_prefill() == true` must
    /// override it.
    fn trim_state(&self, seq: Option<SequenceId>, excess: i32) -> Result<(), String>;

    /// Whether the scheduler's decode lookahead can be rewound from this
    /// model's own per-sequence state (default: false, synchronous decode).
    fn supports_decode_lookahead_rewind(&self) -> bool { false }

    /// Forward with optional pre-computed embeddings (for VLM prefill).
    /// The default ignores the embeddings and calls `forward`.
    fn forward_with_embeddings(
        &self,
        input_ids: &MlxArray,
        input_embeddings: Option<&MlxArray>,
        caches: &mut [KVCache],
        mask: Option<&MlxArray>,
    ) -> UniquePtr<MlxArray>;
}

The trait has more methods with defaults (chunked prefill, batched forward, sequence-addressed forwards); the excerpt shows the ones a new model most often touches. Attention layers do not choose between dense and paged attention themselves: they call cache.attend(&q, k, v, scale, mask) on their KVCache (or mlxcel_core::cache::attend_batched for a batched step), and the cache dispatches to the dense, rotating, chunked, or paged implementation it holds. A family with model-owned caches implements trim_state to take padded prefill, and supports_decode_lookahead_rewind plus rewind_decode_appends to take the server's lookahead decode; without them the family stays on unpadded prefill and synchronous decode.

Decode runs on one batch-native engine (mlxcel_core::engine::Engine) for the CLI and the server alike, and a single sequence is a batch of one, so a new model decodes through forward_batched / forward_batched_with_ids even when only one request is active. supports_batching() defaults to true. batched_decode_kind() defaults to BatchedDecodeKind::PerRowFallback, which loops forward over the rows; override forward_batched to share one graph across rows and return BatchedDecodeKind::Native, which the server logs at startup as decode path: native batched forward. Per-row fallback works but aggregate throughput does not scale with --parallel. A family whose state lives inside the model (recurrent, linear-attention, or model-owned caches) keys that state by SequenceId (sequence_state_layout, prepare_sequence_state, release_sequence_state_by_id, and forward_batched_with_ids, which receives the ids) and can then also report supports_batching() == true. Return false only for a runtime that serves a single request at a time, such as a seq2seq or duplex-speech worker; the server then clamps max_batch_size to 1.

A minimal text model implementation:

use crate::ffi::{self, MlxArray, MlxStream};
use crate::layers::KVCache;
use cxx::UniquePtr;

pub struct MyNewModel {
    // Weight tensors held as UniquePtr<MlxArray>
    embed_tokens: UniquePtr<MlxArray>,
    layers: Vec<MyTransformerLayer>,
    norm: UniquePtr<MlxArray>,
    lm_head: UniquePtr<MlxArray>,
    n_layers: usize,
}

impl MyNewModel {
    /// Load from a weights map (SafeTensors key → MlxArray)
    pub fn from_weights(
        config: &MyModelConfig,
        weights: &std::collections::HashMap<String, UniquePtr<MlxArray>>,
    ) -> anyhow::Result<Self> {
        // Extract tensors from the weights map
        let embed_tokens = weights
            .get("model.embed_tokens.weight")
            .ok_or_else(|| anyhow::anyhow!("missing embed_tokens"))?;
        // ... load all layers ...
        Ok(Self { /* ... */ })
    }
}

impl LanguageModel for MyNewModel {
    fn forward(
        &self,
        input_ids: &MlxArray,
        caches: &mut [KVCache],
        mask: Option<&MlxArray>,
    ) -> UniquePtr<MlxArray> {
        // 1. Embed input tokens
        let mut h = ffi::take(input_ids, &self.embed_tokens, 0);
        // 2. Run through transformer layers
        for (i, layer) in self.layers.iter().enumerate() {
            h = layer.forward(&h, &mut caches[i], mask);
        }
        // 3. Apply final norm
        h = ffi::rms_norm(&h, &self.norm, 1e-5);
        // 4. Apply LM head
        ffi::linear(&h, &self.lm_head, None)
    }

    fn make_caches(&self) -> Vec<KVCache> {
        (0..self.n_layers).map(|_| KVCache::new()).collect()
    }

    fn num_layers(&self) -> usize {
        self.n_layers
    }

    fn eos_token_ids(&self) -> Vec<i32> {
        vec![2] // </s> token, model-specific
    }
}

Weight Loading from SafeTensors¶

Model weights are distributed as SafeTensors files (.safetensors). The loading pipeline in src/loading/ reads these files and builds the weight map that is passed to from_weights().

The weight map is a HashMap<String, UniquePtr<MlxArray>> where the key is the full weight name as found in the SafeTensors file (e.g., "model.layers.0.self_attn.q_proj.weight").

Key helpers:

• src/models/sanitize.rs: load_and_sanitize_weights(): loads all .safetensors files from a directory, applying key normalization and tied embedding handling.
• src/models/sanitize.rs: sanitize_tied_embeddings(): copies the embedding table to the LM head when the model uses tied embeddings.
• src/models/sanitize.rs: sanitize_config_json(): normalizes known non-standard JSON patterns in config.json before parsing.

Weights read through mlxcel_core::weights (load_weights_from_dir and its filtered variants) are registered with the weight materialization scope that every model-load entry point opens, so a new family inherits eager checkpoint upload at load without code of its own. Tensors the loader drops and modality-tower tensors stay lazy. Do not call eval on weights in a family loader to force residency.

// Loading pattern used by most standard text models
let weights = load_and_sanitize_weights(&model_path)?;
let config: MyModelConfig = serde_json::from_str(&config_str)?;
let model = MyNewModel::from_weights(&config, &weights)?;

Weight Key Remapping¶

When the weight file uses different key names than the implementation expects, add a remapping step in the loader:

let weights: HashMap<String, _> = raw_weights
    .into_iter()
    .map(|(k, v)| {
        let new_key = k
            .replace("transformer.h.", "model.layers.")
            .replace("ln_f", "model.norm");
        (new_key, v)
    })
    .collect();

Model Type Detection¶

The detection function in src/models/detection.rs reads config.json from the model directory and matches model_type:

pub fn get_model_type(model_path: &Path) -> Result<ModelType> {
    let config: serde_json::Value = /* read and parse config.json */;
    let model_type = config["model_type"].as_str()?;
    match model_type {
        "my_new_model" => Ok(ModelType::MyNewModel),
        // ...
    }
}

If the same model_type string can refer to both a text-only and VLM variant, use the detect_text_or_vlm() helper:

"my_family" => Ok(detect_text_or_vlm(
    &config,
    ModelType::MyFamily,
    ModelType::MyFamilyVLM,
)),

This helper returns the VLM variant when config.json declares a vision tower (a non-empty vision_config, without language_model_only: true) and the text-only variant otherwise. detect_text_or_vlm_by_weights also scans the checkpoint's weight names and picks the text-only variant when no vision-tower weights ship, so a vision-stripped checkpoint loads as its text model.

Add a unit test in src/models/detection_tests.rs to cover the new mapping.

LoadedModel Dispatch¶

src/loaded_model.rs contains the LoadedModel enum that holds one variant per supported model type. The delegate_language_model! macro generates the dispatch code for the LanguageModel trait methods:

delegate_language_model! {
    MyNewModel => LoadedModel::MyNewModel(m) => m,
    // other variants...
}

For VLMs, also add the capability wiring in src/loaded_model_capabilities.rs:

impl VlmRuntimeRef<'_> {
    pub fn encode_images(&self, images: &[ImageData]) -> Result<UniquePtr<MlxArray>> {
        match self {
            Self::MyFamilyVLM(m) => m.encode_images(images),
            // ...
        }
    }
}

Common Pitfalls¶

Regression Hotspots¶

Most regressions appear around registration, routing, or capability dispatch rather than in the decoder math itself. Check these files explicitly:

• src/models/mod.rs
• missing module export or ModelType variant
• src/models/detection.rs
• missing aliases in get_model_type()
• text/VLM misclassification when vision_config is present
• src/loading/mod.rs
• divergence between load_model() and load_model_from_weights()
• adding a standard config-backed model as a one-off special case
• adding a VLM directly into load_model() instead of try_load_vlm_model_from_dir()
• src/model_metadata.rs
• forgetting text/VLM kind, adapter support, route policy, or shared standard-text registration before wiring loaders
• src/loading/config_backed.rs
• bypassing the shared registration surface
• forgetting wrapper constructors for models such as Llama4, Gemma3, or Ministral3
• src/loading/nonstandard.rs
• leaving directory-only loader families in src/loading/mod.rs instead of the non-standard registry
• src/loading/special.rs
• adding adapter/owned-weight special handling inline instead of the special-weight registry
• forgetting Qwen3.5 text-config normalization or owned-weight sanitization before construction
• src/loading/vlm.rs and sibling family modules under src/loading/
• wrong default token IDs
• missing top-level quantization inheritance
• incorrect weight-key remapping between text and vision towers
• src/loaded_model.rs and src/loaded_model_capabilities.rs
• missing dispatch arm for a new variant
• missing capability wiring in VlmRuntimeRef
• updating the all-model dispatch macro but forgetting the multimodal switchboard

Keep src/loading/mod.rs focused on route selection. If a new model family adds substantial construction logic, prefer a dedicated sibling module and call it from the router instead of growing load_model() or load_model_from_weights() inline.

For very large model families, extract internal helper hotspots into a focused sibling helper module when the code changes for different reasons than the main decoder stack. Current examples include src/models/gemma3n_helpers.rs, src/models/llama4_helpers.rs, and the VLM family modules under src/loading/.

Testing the New Model¶

Detection Tests¶

Add unit tests in src/models/detection_tests.rs that verify the model_type string maps to the correct ModelType variant:

#[test]
fn test_my_new_model_detection() {
    let config = serde_json::json!({"model_type": "my_new_model"});
    // write config.json to a temp dir and call get_model_type()
}

Sanitization Tests¶

If the model requires weight key remapping or config normalization, add tests in src/models/sanitize_tests.rs.

Loader Tests¶

For loader normalization or token-default logic, add tests near src/loading/tests.rs or a dedicated src/loading/<family>_tests.rs.

Smoke Test (requires a local model)¶

If you have the model weights locally, run a smoke test through the CLI:

mlxcel generate -m ./models/my-new-model -p "Hello, world!" -n 20

Or via the server:

mlxcel-server --model ./models/my-new-model --port 8080 &
curl http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"model": "my-new-model", "messages": [{"role": "user", "content": "Hello"}]}'

Running Ignored Tests¶

For tests that require actual model weights (marked #[ignore]):

cargo test --test model_loading -- --ignored --test-threads=1

Broader Test Expectations¶

Add tests in the same slice as the model/control-plane change:

• model detection helpers: tests near src/models/detection.rs
• sanitization helpers: tests near src/models/sanitize.rs
• loader normalization or token-default logic: tests near src/loading/tests.rs or the relevant src/loading/vlm*_tests.rs
• shared vision merge contracts: tests near src/vision/merge_tests.rs
• prompt/token expansion logic: dedicated helper tests such as src/multimodal/qwen_vl_tests.rs, src/multimodal/vlm_prompt_tests.rs, or src/multimodal/phi3v_prompt_tests.rs
• runtime validation: scripts/run_quality_gate.sh plus at least one local smoke test when a matching model exists

If the slice touched MLX-heavy ignored helper tests, run them explicitly with --ignored --test-threads=1.

Shared Function and Execution Policy¶

When touching shared functions used by multiple model families, update the usage comments described in AGENTS.md, especially under:

• src/lib/mlxcel-core/src/layers.rs
• src/lib/mlxcel-core/src/utils.rs

Those comments act as the retest list for future shared-function changes.

Keep entry-point policy in the shared execution layer when behavior must be identical across CLI, server, and future frontends:

• src/execution/runtime.rs
• environment-driven device selection (MLXCEL_DEVICE)
• GPU wired-memory limit setup
• src/execution/sampling.rs
• centralized SamplingConfig assembly from resolved request defaults
• shared greedy vs non-greedy branching

If a new frontend or request type needs different defaults, resolve those defaults at the edge and keep the final conversion in src/execution/.

Keep CLI-only prompt formatting and terminal output behavior in src/commands/generate.rs instead of moving it into shared loading or server modules.

For server-only boot behavior, keep startup policy in src/server/startup.rs instead of growing src/server/mod.rs:

• API key / chat-template resolution precedence
• startup-time normalization of CLI-compatible flags
• warmup behavior
• Unix-socket vs TCP binding

Keep shared server types in focused modules:

• src/server/config.rs for request/default configuration structs
• src/server/state.rs for AppState and metrics containers
• src/server/model_provider.rs for the public request/response channel API
• src/server/model_worker.rs for the long-lived worker thread, VLM request prep, and decode state

Keep server edge adapters out of route files once more than one endpoint needs the same behavior:

• src/server/chat_request.rs for OpenAI chat message flattening and prompt fallback assembly
• src/server/request_options.rs for request-default merging into ServerGenerateOptions
• src/server/media.rs for data: / file:// image-source parsing
• src/server/streaming.rs for shared SSE channel and [DONE] emission helpers

Dry-Run Workflow Validation¶

Use these dry runs before coding to verify that the architecture still reduces ambiguity when adding new model support.

Dry Run A: Standard Text Model¶

Assume a new text model follows the existing config-backed loading path.

Required surfaces:

1. src/models/<family>.rs
2. src/models/mod.rs
3. src/models/detection.rs
4. src/model_metadata.rs through the converged registration surface plus static_model_descriptor() / model_load_policy()
5. src/loading/config_backed.rs only if shared config-backed loading behavior itself must change
6. src/loaded_model.rs
7. src/loaded_model_capabilities.rs only if the family changes multimodal capability exposure
8. tests near src/models/detection_tests.rs, src/models/sanitize_tests.rs, and src/loading/tests.rs

What should not happen:

• no new one-off construction branch inside load_model()
• no direct CLI or server changes unless user-visible behavior changes
• no family-specific getter added to LoadedModel if an existing capability is enough

This keeps route selection centralized, adapter support declared in one policy surface, and the expected edit list short enough to review before coding.

Dry Run B: New VLM Family¶

Assume a new VLM family needs its own token defaults and weight-key remapping.

Required surfaces:

1. text model and/or VLM wrapper under src/models/ or src/vision/
2. src/models/mod.rs
3. src/models/detection.rs
4. src/loading/vlm_<family>.rs
5. src/loading/vlm.rs
6. src/model_metadata.rs through static_model_descriptor() / model_load_policy()
7. src/loaded_model.rs
8. src/loaded_model_capabilities.rs through VlmRuntimeRef
9. src/multimodal/ only if prompt/runtime preparation is truly new
10. tests near src/loading/vlm_<family>_tests.rs and any new multimodal helper test file

What should not happen:

• no concrete model-type checks added to CLI or server request paths
• no family-specific loading logic added directly to src/loading/mod.rs
• no duplicated prompt-rewrite logic across CLI and server

Review Checklist¶

Before opening a PR for a new model or VLM family, confirm:

1. loading policy was updated through src/model_metadata.rs
2. LoadedModel wiring stayed inside src/loaded_model.rs and src/loaded_model_capabilities.rs
3. CLI and server still depend on shared helpers rather than the concrete model type
4. unit tests cover the new policy or normalization logic
5. at least one smoke test exists when a local model is available
6. loaders and models log through tracing, not println!
7. attention goes through KVCache::attend / cache::attend_batched, and a family with model-owned caches either implements trim_state or answers supports_padded_prefill() == false
8. batched decode is decided on purpose: the family reports BatchedDecodeKind::Native, or its per-row fallback is acceptable, and per-sequence state is keyed by SequenceId

Commit Strategy¶

Prefer small checkpoints that isolate one control-plane surface:

• model detection
• loader normalization
• prompt preparation
• runtime initialization

This keeps regressions searchable and makes future model additions easier to compare against previous slices.