feat(vlm): add Cohere Compass / North-Micro-Vision (cohere_compass) VLM: Command-style parallel decoder with interleaved MRoPE plus deepstack vision tower #1354

Summary

model_type: "cohere_compass" (CohereLabs North-Micro-Vision-Instruct, 2.5B, generative image+text chat) is not detected by mlxcel. It is a generative VLM, not an embedding model: a 28-layer Command-style text decoder (parallel attention + MLP block, LayerNorm, tied embeddings, logit_scale) in which the sliding-window layers use an interleaved 3-axis MRoPE and the full-attention layers use no positional encoding, fed by a 27-block deepstack vision tower of exactly the Qwen3-VL shape (patch 16, temporal patch 2, spatial merge 2, interpolated 48x48 position grid, three deepstack mergers injected into the first three decoder layers). This issue adds the family by composing the existing Cohere2 decoder block, the Qwen3-VL vision encoder, and the Qwen-VL prompt/MRoPE machinery.

Current behavior

"cohere_compass" hits the _ => Err("Unsupported model type") arm at src/models/detection.rs:379; nothing under src/models/ or src/vision/ references it. The pieces exist separately:

• Parallel block + LayerNorm + sliding/global layer split with RoPE only on sliding layers: src/models/cohere2.rs:124-360 (Cohere2Attention::forward, Cohere2TransformerBlock::forward, logit_scale at src/models/cohere2.rs:473). Its RoPE is traditional/interleaved (src/models/cohere2.rs:166-170), which Compass does NOT use (see below).
• Interleaved MRoPE with mrope_section and the [THWTHW...TT] frequency layout: InterleavedMRoPE at src/models/qwen3_vl.rs:117-210.
• Deepstack vision encoder (Conv3d patch embed, bilinear-interpolated pos_embed, 2D vision RoPE, merger, deepstack_merger_list): src/vision/encoders/qwen3_vl.rs:36-64 (config) and :660-900 (Qwen3VLVisionEncoder::forward_with_grid).
• Deepstack injection into the first N decoder layers and per-sequence MRoPE state: src/models/qwen3_vl.rs:602-731 (set_deepstack_state), src/models/qwen_mrope_state.rs.
• Qwen-VL image preprocessing and <image_pad> expansion: src/vision/processors/qwen2_vl.rs:29-100 (Qwen2VLProcessor::new_with_norm), src/multimodal/qwen_vl.rs:20-60 (QwenVlmPromptInfo).
• Weight-prefix remap for both published layouts: rewrite_qwen3_vl_weight_key at src/loading/vlm.rs:756-783.
• {% generation %} in the chat template is already stripped by src/server/chat_template.rs:916-961.

Expected behavior

mlxcel generate -m models/North-Micro-Vision-Instruct-4bit --image photo.jpg -p "Describe this image." -n 64 describes the image; text-only prompts work without an image; /v1/chat/completions with an image_url part works through the server. Greedy ids match an oracle run of the same checkpoint.

Config (config.json)

Top level: model_type "cohere_compass", image_token_id 255031, video_token_id 255032, vision_start_token_id 255028, vision_end_token_id 255029, tie_word_embeddings true, eos_token_id 255001, pad_token_id 0, fusion_config: {"patch_embeddings": true} (informational; ignore), min_pixels / max_pixels null (fall back to preprocessor_config.json), text_config, vision_config, and quantization {group_size: 64, bits: 4, mode: affine} in the mlx conversion.

text_config (model_type "cohere_compass_text"):

vision_config (model_type "cohere_compass_vision"): depth 27, hidden_size 1152, intermediate_size 4304, num_heads 16, in_channels 3, patch_size 16, temporal_patch_size 2, spatial_merge_size 2, out_hidden_size 2048, num_position_embeddings 2304 (48x48 grid), deepstack_visual_indexes [8, 16, 24], hidden_act gelu_pytorch_tanh, use_rope default true. These are the Qwen3-VL vision config keys one for one; Qwen3VLVisionConfig (src/vision/encoders/qwen3_vl.rs:36) parses it unchanged.

Text decoder forward

h = embed_tokens(tokens)                                       # [B, L, 2048]
for i in 0..28:
    a = LayerNorm(h; layer_norm_eps)                           # weight only, no bias; computed in f32
    attn = Attn_i(a); mlp = SwiGLU_i(a)                        # BOTH read the same normed input
    h = h + attn + mlp                                         # parallel block
    if i < 3 and image present:
        h[visual positions] += deepstack_embeds[i]             # after layers 0, 1, 2 only
h = LayerNorm(h)
logits = embed_tokens.as_linear(h) * logit_scale               # 0.25

Attention for layer i with layer_types[i]:

q, k, v = q_proj(a), k_proj(a), v_proj(a)      -> [B, 16, L, 128], [B, 8, L, 128], [B, 8, L, 128]
if layer_types[i] == "sliding_attention":
    (cos, sin) from the interleaved MRoPE at base 50000 (below), applied as
    q = q * cos + rotate_half_split(q) * sin;  k likewise     # rotate_half_split = concat(-x2, x1) on halves
    cache: RotatingKVCache(max_size = 4096); mask: causal restricted to the last 4096 keys
else:  # "full_attention"
    no rotation at all; cache: KVCache; mask: plain causal
out = o_proj(sdpa(q, k, v, scale = 128 ** -0.5))

Interleaved MRoPE (sliding layers): inv_freq[j] = 50000 ** (-2j / 128) for j in 0..64. Position ids are a [3, B, L] tensor (t, h, w). The frequency for index j uses axis t by default, axis h when j % 3 == 1 and j < 3 * 20 = 60, and axis w when j % 3 == 2 and j < 60 (so j = 60..63 stay on t). freqs[b, l, j] = pos[axis(j), b, l] * inv_freq[j]; cos/sin = cos/sin(concat(freqs, freqs)) (split layout). This is InterleavedMRoPE in src/models/qwen3_vl.rs:117 with mrope_section = [24, 20, 20] and base 50000; the only difference from Qwen3-VL is that it applies on sliding layers alone and the full layers get nothing.

Position ids (get_rope_index): text tokens advance all three axes by 1 per token. For each image (grid t, h, w in patches; merged grid h/2, w/2): the t * (h/2) * (w/2) image-pad tokens get t_idx, h_idx, w_idx from the merged grid plus the running text offset; the next text token starts at max(previous positions) + 1. rope_delta = max_position + 1 - L is cached per sequence for decode, where each new token has position offset + rope_delta on all three axes. This is the Qwen2-VL/Qwen3-VL rule already implemented by compute_position_ids_with_delta (src/models/qwen3_vl.rs:969) and qwen_mrope_state.rs; no video time scaling.

Vision tower forward

Identical to Qwen3-VL: pixel_values [N, 3*2*16*16 = 1536] -> Conv3d patch embed (patch_embed.proj, bias true) -> add bilinear-interpolated pos_embed (48x48 grid resampled to each image's h x w, reordered into merge-window order) -> 27 blocks of LayerNorm(1e-6) -> qkv (bias) -> 2D vision RoPE (head_dim 72; 18 frequencies per spatial axis at theta 10000, the row-half and column-half tables concatenated to 36 and then doubled to cover the full 72 dims, split layout, computed in f32) -> SDPA per image (cu_seqlens) -> proj, LayerNorm(1e-6) -> linear_fc1 -> GELU(tanh) -> linear_fc2; after blocks 8, 16, 24 a deepstack_merger_list[k] (post-shuffle LayerNorm on 1152*4, linear_fc1, GELU, linear_fc2 -> 2048) produces the k-th deepstack embedding; the final merger (LayerNorm on 1152 before the 2x2 shuffle, linear_fc1, GELU, linear_fc2 -> 2048) produces the image embeddings that replace the image-pad tokens. Qwen3VLVisionEncoder::forward_with_grid (src/vision/encoders/qwen3_vl.rs:817) returns exactly (features, deepstack_features).

Weight keys

Two layouts are published; the existing rewrite_qwen3_vl_weight_key maps both onto one:

original (CohereLabs):        model.language_model.*,   model.visual.*,        lm_head.* (absent, tied)
mlx conversion:               language_model.model.*,   vision_tower.*
target after remap:           model.* (text),           vision_tower.*

Text keys under model.: embed_tokens.weight, layers.{i}.input_layernorm.weight (no bias), layers.{i}.self_attn.{q,k,v,o}_proj.weight (+ scales/biases), layers.{i}.mlp.{gate,up,down}_proj.weight, norm.weight. There is no post_attention_layernorm. Vision keys under vision_tower.: patch_embed.proj.{weight,bias} (5-D Conv3d weight; transpose [O, C, T, P, P] -> [O, T, P, P, C] when the second dim equals in_channels, as the Qwen3-VL loader already does), pos_embed.weight [2304, 1152], blocks.{i}.{norm1,norm2}.{weight,bias}, blocks.{i}.attn.{qkv,proj}.{weight,bias}, blocks.{i}.mlp.{linear_fc1,linear_fc2}.{weight,bias}, merger.{norm,linear_fc1,linear_fc2}.*, deepstack_merger_list.{0,1,2}.{norm,linear_fc1,linear_fc2}.*. Sanitize: drop any position_ids key; drop lm_head.* when tied; otherwise the Qwen3-VL remap.

Preprocessing and prompt

preprocessor_config.json: patch_size 16, merge_size 2, temporal_patch_size 2, image_mean / image_std [0.5, 0.5, 0.5], size.shortest_edge 65536 (min pixels), size.longest_edge 16777216 (max pixels), resample bicubic, do_convert_rgb true. Resize with the Qwen smart_resize (factor patch_size * merge_size = 32, aspect ratio capped at 200), normalize, duplicate along the temporal axis (t = 1 images become 2 frames), flatten patches in merge-window order, image_grid_thw = [1, h/16, w/16]. Use Qwen2VLProcessor::new_with_norm(16, 2, 2, [0.5; 3], [0.5; 3]) with the min/max pixel bounds above.

Chat template (chat_template.jinja): {{ bos_token }} then turns wrapped as <|START_OF_TURN_TOKEN|><|USER_TOKEN|>...<|END_OF_TURN_TOKEN|>, assistant turns as <|START_OF_TURN_TOKEN|><|CHATBOT_TOKEN|>; every image content part renders as the literal <|VISION_START|><|IMAGE_PAD|><|VISION_END|> (ids 255028, 255031, 255029) and the runtime must expand the single <|IMAGE_PAD|> into t * (h/2) * (w/2) copies (the insert_qwen_vl_image_tokens rule). The template contains {% generation %} blocks. Stop on <|END_OF_TURN_TOKEN|> (255001).

Implementation plan

1. src/models/cohere_compass.rs (new, text decoder): CompassTextConfig (serde, with the per-layer-type rope_parameters reader: a layer type maps to None (no RoPE) when its entry is JSON null or absent while other layer types are present; otherwise to {theta, mrope_section: Option<[usize; 3]>}), CompassAttention (q/k/v/o via FusedQKVLinear::from_weights_separate, Option<InterleavedMRoPE>, Option<window>), CompassBlock (single LayerNorm from mlxcel_core::layers::LayerNorm without bias, parallel residual as in Cohere2TransformerBlock::forward at src/models/cohere2.rs:349), CompassLanguageModel (embed_tokens: UnifiedEmbedding, tied head, logit_scale, make_caches = Rotating(4096) for sliding layers else KV, deepstack injection + MRoPE state copied from Qwen3VLLanguageModel in src/models/qwen3_vl.rs:602-760, forward_with_embeddings). With mask == None and L > 1 build the causal and window masks locally (maskless-prefill contract).
2. Vision: reuse vision::encoders::qwen3_vl::Qwen3VLVisionEncoder directly; add cohere_compass_vision as an accepted model_type if the encoder checks it.
3. src/vision/cohere_compass.rs (new): CohereCompassModel { vision_tower, text_model, image_token_id, vision_start_token_id, processor } implementing the QwenVlm trait in src/multimodal/qwen_vl.rs (prompt_info, input_embeddings, input_embeddings_with_cache, MRoPE bind/take/install), modeled on src/vision/qwen3_vl.rs:36-140.
4. Loader src/loading/vlm_cohere_compass.rs (or a branch in the Qwen3-VL loader): parse top-level config, text_config, vision_config, quantization; remap keys with remap_qwen3_vl_weights(raw, false); build the processor from preprocessor_config.json; ModelType::CohereCompassVLM => load_cohere_compass_vlm in src/loading/mod.rs.
5. Detection: "cohere_compass" => Ok(ModelType::CohereCompassVLM) in src/models/detection.rs with cohere_compass_model_type_is_detected; ModelType/LoadedModel/model_metadata.rs rows with kind: Vlm; arch-string arm for TP dispatch totality.
6. Server: add the family to the Qwen-VL multimodal dispatch tables (src/multimodal/batched_dispatch.rs, src/server/ image handling) so /v1/chat/completions image parts route through QwenVlmPromptInfo.
7. docs/supported-models.md: entry under "Vision-language and multimodal variants".

Validation

(a) Unit tests in src/models/cohere_compass_tests.rs and src/vision/cohere_compass_tests.rs:

• rope_parameters_null_means_no_rope: the published rope_parameters block yields Some(theta 50000, section [24,20,20]) for sliding layers and None for full layers; a config with only top-level rope_theta and rope_on_all_layers: true yields RoPE on every layer.
• interleaved_selector_matches_section: for section [24, 20, 20] and 64 frequencies, axis is h exactly at {1, 4, ..., 58}, w at {2, 5, ..., 59}, t elsewhere.
• split_rope_layout: q = [1, 2, 3, 4], cos = 0, sin = 1 gives [-3, -4, 1, 2] (split), not [-2, 1, -4, 3] (interleave).
• parallel_block_has_no_post_norm: loading a synthetic checkpoint with a post_attention_layernorm key fails or ignores it; forward equals x + attn(n) + mlp(n) with the same n.
• logit_scale_applied: logits of a tiny model scale by 0.25.
• deepstack_injects_first_three_layers_only.
• sliding_layers_get_rotating_cache: 28 caches, rotating except at {3, 7, 11, 15, 19, 23, 27}.
• text_prefill_is_causal: 4 layers (3 sliding + 1 full), 96-token prefill with mask == None; first-48 logits equal those of a 48-token prefill within 1e-4.
• image_pad_expands_to_merged_grid: a 448x448 image gives grid [1, 28, 28] and 196 image tokens between the start/end tokens.

(b) Real checkpoint: mlx-community/North-Micro-Vision-Instruct-4bit (about 1.5 GB; CohereLabs/North-Micro-Vision-Instruct bf16 about 5 GB is the unquantized original).

./target/release/mlxcel download mlx-community/North-Micro-Vision-Instruct-4bit
./target/release/mlxcel generate -m models/North-Micro-Vision-Instruct-4bit -p "Summarize the history of the Apollo program in five sentences." -n 64 --temp 0
./target/release/mlxcel generate -m models/North-Micro-Vision-Instruct-4bit --image tests/fixtures/<a real photo>.jpg -p "Describe this image in detail." -n 64 --temp 0

Acceptance: text-only greedy ids match an oracle run of the same checkpoint through at least the first 32 tokens on three prompts (one at least 64 tokens long); the image run names the salient objects of the photo and matches the oracle's first 16 greedy tokens on the same image and prompt; a two-image prompt produces position ids whose maximum equals the oracle's (rope_delta parity); the server /v1/chat/completions path with an image_url part returns the same text as the CLI.

Acceptance criteria

• src/models/cohere_compass.rs implements the parallel LayerNorm block, per-layer-type RoPE (none on full layers), split rotate-half, sliding/full cache split, deepstack injection and logit_scale. Correction: the sliding layers do NOT use Qwen3-VL's InterleavedMRoPE. mrope_interleaved is in transformers' ignore_keys_at_rope_validation and is never read; CohereCompassRotaryEmbedding pre-permutes inv_freq and assigns axes in contiguous [H, W, T] sections. Measured against that class, this implementation is within 2.2e-07 in cos and the interleave specified here is off by 1.96. Implemented in src/models/cohere_compass_rope.rs. The cache split is also masks-plus-window over a dense KVCache rather than RotatingKVCache, which is what cohere2.rs / olmo3.rs do and what make_caches() -> Vec<KVCache> allows.
• src/vision/cohere_compass.rs wires the Qwen3-VL vision encoder and the Qwen-VL prompt/MRoPE machinery for this family; both weight-prefix layouts load (every validation run below was done on the mlx conversion and on a float32 checkpoint written in the CohereLabs layout)
• generates for text-only and image prompts via CLI and server. Text is token-exact against a transformers 5.18.0.dev0 oracle on the same weights: 32/32 identical tokens read from server logprobs, on both checkpoints. Image greedy ids are not: the prompt, the 196/245 image-token counts and the MRoPE positions all match bit for bit, but generation flips to the oracle's rank-2 token at step 2 (0.238 nats) for one image and step 0 (0.866 nats) for two. TF32, both vision activations, DeepStack order/depth and quantization are ruled out; the residual is in the shared Qwen3-VL tower and is filed as Qwen3-VL vision tower is not token-exact against transformers on the image path #1738.
• mlxcel arch lists Cohere Compass / North-Micro-Vision and mlxcel list shows the checkpoint (mlxcel list enumerates local snapshots rather than architectures, so arch is the family listing)
• docs/supported-models.md updated
• detection table in src/models/detection.rs updated with cohere_compass_model_type_is_detected
• cargo test --workspace --profile test-fast --features metal,accelerate: not runnable on this host (Linux + CUDA, no Metal, no Accelerate); unrun. cargo test --profile test-fast --features cuda --lib cohere_compass passes 23/23.
• cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass (green in CI on PR feat(vlm): add the Cohere Compass (North-Micro-Vision) VLM #1735)

Out of scope

• Video input (<|VIDEO_PAD|>, video_token_id 255032, video_preprocessor_config.json, timestamp-separated frames).
• The CohereCompassTextForSequenceClassification reranker head (score_shift_a/b, pooling); no published checkpoint uses it.
• Tensor-parallel sharding.