tracking: VLM-reference model families needing architecture investigation before a port
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
## Summary
A model-coverage audit surfaced the model families below as not supported by mlxcel. Unlike the families recorded as out of scope in #524, these are plausible port candidates: most are autoregressive text or vision-language models shaped like mlxcel's existing runtimes. What is missing is confirmation against a concrete checkpoint (exact `model_type`, config key set, weight naming, and whether the family is worth porting), so they are grouped here for triage instead of being filed as standalone port issues.
Today each of these fails at load in `src/models/detection.rs::get_model_type` with `Unsupported model type: `, since no match arm exists for their `model_type` strings.
Each candidate below records: an architecture note, the `model_type` string, the natural in-tree reuse slot (paths verified against the current tree), what is unconfirmed, and candidate-specific triage steps. The shared workflow is at the bottom. Check a box when the candidate has either a filed port issue (link it) or a recorded not-planned disposition.
- [ ] `jina_vlm`
- [ ] `laguna`
- [x] `bonsai` (triaged: diffusion image generation, out of scope; rationale below)
- [ ] `zaya1_vl`
- [ ] `hrm_text`
- [ ] `multi_modality`
- [ ] `falcon_ocr`
- [ ] `falcon_perception`
## Candidates
### `jina_vlm`
Autoregressive VLM. A SigLIP-so400m-class ViT tower (hidden 1152, 27 layers, patch 14, image size 378, tanh-GELU, no CLS token) supplies features taken from two intermediate layers, concatenated, spatially pooled 2x2 through a latent cross-attention block, and projected by a SwiGLU projector into the text stream; projected image features are scattered into the embedding sequence at image-token positions. The text decoder is Qwen2-class: fused QKV projection, GQA (16 query heads, 8 KV heads, head_dim 128), per-head Q/K RMSNorm, SwiGLU MLP with fused gate-up weights, RoPE theta 1e6, no biases, vocabulary 151936 plus a separate 128-entry extension embedding table for added special tokens (im_start/im_end/im_patch/im_col markers).
- `model_type`: `jina_vlm` (text and vision sub-configs carry the same string; the processor family tag is `jvlm`).
- In-tree reuse: `src/vision/encoders/siglip.rs` for the tower; the Molmo runtime (`src/models/molmo.rs`, `src/vision/molmo_vl.rs`, `src/vision/encoders/molmo.rs`) implements the same connector recipe (multi-layer ViT feature concat plus pooling cross-attention plus pad embeddings); `src/models/qwen2.rs` for the text block, with the per-head Q/K norm structure available in `src/models/qwen3.rs`.
- Unconfirmed: which public checkpoint ships this; the config schema is nested (block_config with attn/ffn/lnorm sub-blocks, vl_connector_config, keys like `n_layers`/`n_kv_heads`/`vit_layers`/`pooling_h`) rather than flat HF-style keys, which needs confirming against a real `config.json`; handling of the split base-plus-extension embedding.
- Triage: locate the checkpoint; confirm `model_type` and the nested config keys; diff the connector against the Molmo connector and the text block against `qwen2.rs`; size the loader work for the fused QKV and split-embedding weight layout.
### `laguna`
Text-only MoE decoder with gated attention. Attention supports per-layer variable head counts, optional QKV and output biases, per-head Q/K RMSNorm, a learned softplus output gate (per-head or per-channel), optional learned attention sinks, and mixed `full_attention`/`sliding_attention` layer types with separate rotary parameters for the sliding-window layers. The MoE block routes through a sigmoid router carrying a score-correction bias (aux-loss-free load balancing), top-k probability normalization, a routed scaling factor, and optional router-logit softcapping, over stacked routed experts plus one shared expert; some layers stay dense via `mlp_only_layers`/`decoder_sparse_step`. Checkpoints may ship packed sub-byte quantized tensors (`weight_packed`/`weight_scale`/`weight_shape`) that a loader must unpack into quantized-matmul scales and biases, with the router kept at higher precision.
- `model_type`: `laguna` (supplied by the checkpoint; the family imposes no other alias).
- In-tree reuse: nearest analogs are `src/models/qwen3_next.rs` (sigmoid-gated attention output, hybrid layer types), `src/models/gpt_oss.rs` (per-head attention sinks plus sliding windows), `src/models/deepseek_v3.rs` (sigmoid router with `e_score_correction_bias` and a shared expert), and `src/models/switch_layers.rs` (stacked per-expert weights).
- Unconfirmed: the vendor/product identity and whether a public checkpoint exists at all; whether mlxcel's loaders can consume the packed quantization format or need an unpack pass.
- Triage: identify a checkpoint (this is the gating step; without one the candidate stays parked); confirm `model_type` and the gating/sink/MoE config keys; size the delta as a composition of the four modules above; check the packed-tensor unpack against the unified quantized loaders.
### `bonsai` (triaged: out of scope)
Not a language or vision-language model. Bonsai is a distilled few-step text-to-image generator: a latent flow-matching (rectified-flow) DiT transformer with ternary/2-bit quantized weights, a bundled prompt text encoder, a VAE for latent decode, and an Euler flow-matching scheduler running about 4 denoising steps, with optional classifier-free guidance. Running it requires a native sub-4-bit quantized-matmul kernel. Known checkpoint: `prism-ml/bonsai-image-ternary-4B-mlx-2bit`, a multi-component diffusion directory rather than a single `config.json` checkpoint.
- Family tag: `bonsai` (no top-level `model_type` in the usual sense).
- Disposition: out of scope for mlxcel's text-generation runtime for the same reasons as the image-generation families recorded in #524 (diffusion sampling loop, VAE, and an image output path, none of which exist in mlxcel). No port issue; box checked as triaged.
### `zaya1_vl`
Autoregressive VLM pairing a Qwen2.5-VL-style ViT tower (windowed plus full attention blocks, RMSNorm, rotary position embedding, 2x2 patch merger) with a heavily customized MoE text backbone. The attention is convolution-augmented: Q/K projections are packed and passed through grouped causal depthwise 1-D convolutions over the sequence, value streams are computed from time-shifted hidden states, Q/K are L2-normalized with a learned per-KV-head temperature, and rotary covers only half the head dimension; this requires a dual per-layer cache (standard KV plus convolution/shift state). Routing uses a down-projection to a small router space followed by a two-layer MLP, with router-state carryover across layers and a mixture-of-depths skip expert that passes tokens through unchanged. Per-sublayer learned residual scaling, tied embeddings, and per-expert vision LoRA adapters applied only at image-token positions round out the customization.
- `model_type`: `zaya1_vl` (text sub-config `zaya1_vl`; the vision sub-config ships as `qwen2_5_vl`).
- In-tree reuse: `src/vision/encoders/qwen2_5_vl.rs` covers the tower essentially as-is. The text backbone has no in-tree analog; the closest structural precedent for a mixed-state cache is the hybrid cache handling in `src/models/qwen3_next.rs`.
- Unconfirmed: public checkpoint availability and size; whether the custom attention, router carryover, and skip-expert mechanics justify what is the largest port delta on this list.
- Triage: locate the checkpoint; confirm `model_type` and the custom config keys (mixture-of-depths and router-carryover flags, convolution kernel sizes, `num_query_groups`); prototype the dual cache against the existing cache abstractions before scoping; expect the convolution weights to need axis-order handling at load (see the standing layout check below).
### `hrm_text`
Text-only hierarchical recurrent language model. Two Llama-style decoder stacks, a low-level module and a high-level module (RMSNorm, SwiGLU, RoPE, GQA, plus a sigmoid gate on the attention output and a fused gate+Q+K+V projection), are replayed in a fixed recurrence: for each high-level cycle, the low module iterates several times over the sum of low and high states, then the high module updates once; the effective depth is computed as `num_layers_per_stack * H_cycles * (L_cycles + 1)` rather than read from a flat layer count. Supports prefix-LM masking (bidirectional attention over the prompt, causal after it) and scales embeddings by the inverse initializer range. The vocabulary size (151808) indicates a Qwen-family tokenizer.
- `model_type`: `hrm_text`.
- In-tree reuse: the individual blocks map onto `src/models/qwen3.rs` / `src/models/llama3.rs` structure, but mlxcel's generation engine assumes a flat, non-shared layer list; the weight-shared recurrence, computed depth, and prefix-LM mask have no in-tree slot and would need engine-level accommodation.
- Unconfirmed: which concrete checkpoint exists; whether a recurrence-based reasoning model is in scope for a serving runtime at all (value unproven).
- Triage: locate a checkpoint; confirm `model_type`, the cycle-count config keys, and the fused-projection weight layout; decide scope on value first, since the engine accommodation dominates the cost.
### `multi_modality`
Identity confirmed: this is DeepSeek-VL, first generation. A plain Llama-style text decoder (SwiGLU, RMSNorm, RoPE with optional linear scaling, GQA, no biases, vocabulary 102400) behind a vision stack that is either a single SigLIP tower or a hybrid dual tower combining SigLIP (low resolution, semantic) with a SAM-style high-resolution encoder, merged by an MLP projector (plain GELU MLP, or a split high/low variant for the hybrid tower). Each image placeholder token expands to 576 image tokens whose projected features are scattered into the embedding sequence.
- `model_type`: `multi_modality` (checkpoint-supplied; the text sub-config carries `llama`).
- In-tree reuse: `src/models/llama3.rs` for the text backbone and `src/vision/encoders/siglip.rs` for the low-resolution tower; the SAM-style high-resolution encoder and the split hybrid projector are new. #533 (DeepSeek-VL2 port) covers the successor family and would share the SigLIP-plus-SAM vision work.
- Unconfirmed: value, since DeepSeek-VL2 supersedes it; which public variants use the hybrid tower versus SigLIP alone.
- Triage: decide whether to fold this into #533 as a stretch goal (shared vision stack) or record not-planned as superseded; if kept, confirm the projector variant per checkpoint size before scoping.
### `falcon_ocr`
Compact (~300M) early-fusion OCR model; checkpoint `tiiuae/Falcon-OCR`. There is no vision encoder: images are cut into 16x16 patches and linearly projected straight into the token stream, and a single decoder processes image and text under a hybrid mask (bidirectional within each image block, causal over text). The decoder blocks are Llama-derived with custom pieces: fused QKV, per-head Q/K RMSNorm, learned per-head attention sinks, a squared-ReLU gated MLP, and a 3-D rotary scheme that concatenates 1-D temporal rotary with a 2-D per-head spatial rotary applied at image positions. Weight names follow a `tok_embeddings`/`layers.N.attention`/`feed_forward`/`output` scheme rather than HF-style names, the fused MLP weight ships interleaved and must be split at load, and the model disallows chunked prefill.
- `model_type`: `falcon_ocr` (same string in the sub-configs; the config is also expressible through raw `dim`/`n_layers`/`n_heads`/`n_kv_heads`/`ffn_dim` keys).
- In-tree reuse: per-head attention sinks exist in `src/models/gpt_oss.rs`; the 2-D/3-D position-index machinery in `src/models/qwen_mrope_state.rs` is the nearest analog for the spatial rotary; placeholder expansion lives in `src/models/multimodal_placeholders.rs`. The early-fusion patch projector, hybrid mask, and squared-ReLU MLP are new but small.
- Unconfirmed: the exact config key set on the shipped checkpoint (raw keys versus HF-style), the tokenizer and prompt contract, and whether the no-chunked-prefill constraint conflicts with mlxcel's prefill path.
- Triage: pull `tiiuae/Falcon-OCR` config and weights index; confirm `model_type`, key scheme, and the interleaved fused-weight layout (named axis order) against the notes above; verify prefill compatibility; size the delta as a new small runtime borrowing sinks and position-index code.
### `falcon_perception`
Same early-fusion decoder family as `falcon_ocr` (patch projector, hybrid image/text mask, Q/K norm, attention sinks, squared-ReLU MLP, temporal plus spatial rotary) with detection and segmentation heads attached: Fourier coordinate and size encoders, a box-decoder MLP that turns the hidden states at dedicated coordinate/size tokens into center-xy and height-width values, and on the larger variant a segmentation decoder with a learned cross-attention upsampler producing masks. Decoding is autoregressive; special coord/size/seg token ids trigger the auxiliary heads during generation. Checkpoints: `tiiuae/Falcon-Perception` (~0.6B, detection plus segmentation) and `tiiuae/Falcon-Perception-300M` (detection only).
- `model_type`: `falcon_perception`.
- In-tree reuse: nearly everything is shared with a `falcon_ocr` port, so sequencing matters: port `falcon_ocr` first, then this adds the heads. Box and mask outputs need an output contract beyond chat completions, the same open question recorded for the grounding family in #524; box-only output could follow the `mlxcel detect` precedent (`src/commands/detect.rs`).
- Unconfirmed: whether box/mask outputs can be usefully surfaced through mlxcel's current APIs; the effort of the mask upsampler.
- Triage: triage after `falcon_ocr`; confirm the head config keys (`coord_token_id`/`size_token_id`/`seg_token_id`, segmentation flags) on the real checkpoints; decide the output contract before filing a port issue.
## Triage workflow
To graduate a candidate:
1. Locate a real public checkpoint and record the repo id here.
2. Confirm `config.json::model_type` and the actual config key set against the notes above. mlxcel dispatches on `model_type` in `src/models/detection.rs::get_model_type`; anything unlisted errors with `Unsupported model type`.
3. Size the port delta against the named in-tree module(s). Standing layout check: checkpoints serialize 4-D convolution weights in `(out_channels, in_channels, kernel_h, kernel_w)` order, while mlxcel's MLX runtime runs channels-last convolutions expecting `(out_channels, kernel_h, kernel_w, in_channels)`; fused or packed tensors (QKV, gate-up, interleaved MLP weights, packed quantization) carry their own axis and interleave conventions. Check every convolution, patch embedding, and fused tensor for layout deltas against the named in-tree module and plan the load-time permutation or split (in-tree precedent: `src/vision/detection/rt_detr_v2/sanitize.rs`, which also guards against double-transposing weights already stored in MLX layout).
4. Decide in scope or not.
5. In scope: file a dedicated port issue following the format used by #533 through #547 (Summary naming the `model_type` and the current `Unsupported model type` failure, What it is, In-tree reuse, Scope covering the `ModelType` variant plus `detection.rs` arm, encoder/processor/connector work, registry and metadata tables, `docs/supported-models.md`, tests plus real-checkpoint validation, Effort). Link the new issue next to the candidate's checkbox and check the box.
6. Out of scope: record a one-line rationale next to the checkbox and check it. Task-specific non-generative vision families belong with the decision record in #524.
## Non-goals
- Implementing any port inside this issue; every confirmed candidate gets its own port issue.
- Re-triaging families already covered by dedicated port issues (#533 through #547) or by the out-of-scope decision record (#524).
- Building non-text output paths (boxes, masks, images); if a candidate needs one, that infrastructure is scoped separately before the port issue is filed.
Contributor guide
Research direction
Start at src/models/detection.rs::get_model_type and use the candidate-specific reuse paths and checkpoint details in this issue to verify model_type, config keys, and weight layouts against concrete checkpoints. Compare candidates with the referenced model and vision implementations, then record either a filed port issue with its link or a not-planned disposition for every unchecked family.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- rust
- Domain
- ai, machine-learning
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100