Skip to content

tracking: VLM-reference model families needing architecture investigation before a port #548

Description

@inureyes

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: <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
  • 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.

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. feat(models): port DeepSeek-VL2 (MoE text + SigLIP/SAM vision) #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 feat(models): port DeepSeek-VL2 (MoE text + SigLIP/SAM vision) #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 tracking: non-generative vision models present in the VLM reference but out of mlxcel's LLM/VLM-generation scope #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 feat(models): port DeepSeek-VL2 (MoE text + SigLIP/SAM vision) #533 through feat(models): port GLM-OCR (ViT + GLM-4 text OCR) #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 tracking: non-generative vision models present in the VLM reference but out of mlxcel's LLM/VLM-generation scope #524.

Non-goals

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:modelsModel architectures, weights, loading, metadatapriority:lowLow prioritystatus:investigationFeasibility spike / under investigationtype:enhancementNew features, capabilities, or significant additions

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions