lablup / lablup/mlxcel

feat(mage_vl): port Mage-VL 4B (Mage-ViT with 4:6:6 3D RoPE + 2x2 patch merger + plain Qwen3 decoder)

Open
#1,365 0 comments 0 reactions 0 assignees View on GitHub
arch:dense area:cli area:docs area:inference area:models modelsize:small modeltype:vlm priority:medium status:ready type:enhancement
Dominant language
Rust
Stars
467
Forks
54
Avg merge
4h 25m
Merged PRs (30d)
310

Description

## Summary

Mage-VL (`microsoft/Mage-VL`, `model_type: "mage_vl"`, `architectures: ["MageVLForConditionalGeneration"]`) is a 4B "codec-native" streaming VLM. For still images the codec part does not apply: the image path is an ordinary Qwen2-VL-style patch processor (patch 16, 2x2 merge, temporal patch 1, CLIP mean/std, dynamic resolution) feeding a new 24-layer "Mage-ViT" tower with an interleaved 3D rotary embedding split 4:6:6 over (t, h, w), a LayerNorm / MLP patch merger, and a stock Qwen3-4B-Instruct-2507 decoder with plain 1-D RoPE (no MRoPE). The codec-native video path (hardware-codec-derived frame groups and motion vectors through the checkpoint's `codec` preprocessor block) has no torch-free implementation anywhere and is out of scope; image input is the deliverable. mlxcel has no `mage_vl` arm.

## Current behavior

- `src/models/detection.rs` has no `mage_vl` arm.
- Closest tower: `src/vision/encoders/qwen3_vl.rs` (2D rotary, window-free) and `src/vision/encoders/qwen2_vl.rs`. Neither has the interleaved `rotate_half`, the `concat([freqs, freqs])` lane pairing, the 4:6:6 split, or the `arange(size)/size` exponent that Mage-ViT uses; a new encoder file is required.
- Closest processor: `src/vision/processors/qwen2_vl.rs:29-240` (`Qwen2VLProcessor { patch_size, temporal_patch_size, spatial_merge_size, min_pixels, max_pixels, mean, std }`, `preprocess_with_grid` emitting `[rows, C*P*P]` in merge-block order). It already supports `temporal_patch_size = 1` and custom mean/std (`new_with_norm`).
- Decoder: `src/models/qwen3.rs` (`Qwen3Model`, `forward_with_embeddings` at line 854, `embed_tokens` at line 864). Its `ModelArgs.rope_theta` defaults to 10000 (`src/models/qwen3.rs:49-50, 68-70`) and there is no `rope_parameters` lifting in `qwen3.rs`; the only place that lifts a nested `rope_parameters.rope_theta` today is the Qwen3.5 / GLM-OCR VLM loader (`src/loading/vlm_qwen.rs:422-428, 552`).
- The hub config stores the decoder's rope under `text_config.rope_parameters: {"rope_theta": 5000000, "rope_type": "default"}` and has no flat `rope_theta`. Loading it through `serde` as-is would silently run the decoder at base 10000: fluent at short lengths, wrong at long ones.

## Expected behavior

`mlxcel generate -m models/Mage-VL-8bit --image photo.jpg -p "Describe this image." -n 64` resizes the image with the Qwen2-VL `smart_resize` rule (factor 32 = patch 16 x merge 2, `min_pixels 3136`, `max_pixels 4000000`), patchifies into `[N, 768]` rows (`3*16*16`) in 2x2 block order, runs Mage-ViT, merges each 4-patch group into one 2560-wide token, scatters the `N/4` tokens onto the `<|image_pad|>` (151655) positions of the chat-templated prompt, and decodes with Qwen3.

### config.json

Top level: `model_type "mage_vl"`, `image_token_id 151655`, `video_token_id 151656`, `vision_start_token_id 151652`, `vision_end_token_id 151653`, `bos_token_id 151643`, `eos_token_id 151645`, `tie_word_embeddings false`.

`text_config`: `model_type "qwen3"`, `hidden_size 2560`, `num_hidden_layers 36`, `intermediate_size 9728`, `num_attention_heads 32`, `num_key_value_heads 8`, `head_dim 128`, `rms_norm_eps 1e-6`, `vocab_size 151936`, `max_position_embeddings 262144`, `rope_parameters {rope_theta 5000000, rope_type "default"}`, `attention_bias false`, `tie_word_embeddings false`.

`vision_config`: `model_type "mage_vl_vision"`, `hidden_size 1024`, `num_hidden_layers 24`, `num_attention_heads 16` (head_dim 64), `intermediate_size 4096`, `patch_size 16`, `image_size 448` (informational), `num_channels 3`, `layer_norm_eps 1e-6`, `layer_norm_type "layer_norm"`, `hidden_act "gelu"` (exact erf), `rope_theta 10000.0`, `spatial_merge_size 2`, `temporal_patch_size 1`, `out_hidden_size 2560`, `frame_windows_size 4`, `use_head false`, `use_patch_position_encoding false`, `max_position_embeddings 8192`.

`preprocessor_config.json`: `image_processor_type "Qwen2VLImageProcessor"`, `patch_size 16`, `merge_size 2`, `temporal_patch_size 1`, `min_pixels 3136`, `max_pixels 4000000`, `image_mean [0.48145466, 0.4578275, 0.40821073]`, `image_std [0.26862954, 0.26130258, 0.27577711]`, `resample 3` (bicubic), plus a `codec` block that is video-only and must be ignored.

Chat template (`chat_template.jinja`): Qwen2-VL style, `<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>{text}<|im_end|>\n<|im_start|>assistant\n`. Each `<|image_pad|>` expands to `t*h*w/4` copies, where `(t,h,w)` is that image's patch grid.

### Weight keys (hub layout)

```
model.visual.embeddings.patch_embedding.weight # [1024, 3, 16, 16] torch, no bias
model.visual.layernorm_pre.{weight,bias}
model.visual.encoder.layers.{0..23}.layer_norm1.{weight,bias}
model.visual.encoder.layers.N.self_attn.qkv.{weight,bias} # [3072, 1024]; split is contiguous q|k|v, not head-interleaved
model.visual.encoder.layers.N.self_attn.proj.{weight,bias}
model.visual.encoder.layers.N.layer_norm2.{weight,bias}
model.visual.encoder.layers.N.mlp.{fc1,fc2}.{weight,bias} # 1024 -> 4096 -> 1024, exact GELU
model.visual.merger.ln_q.{weight,bias} # LayerNorm(1024), eps = layer_norm_eps
model.visual.merger.mlp.0.{weight,bias} # Linear 4096 -> 4096
model.visual.merger.mlp.2.{weight,bias} # Linear 4096 -> 2560
model.language_model.embed_tokens.weight, model.language_model.layers.N.*, model.language_model.norm.weight
lm_head.weight
```

Converted MLX checkpoints may already store `vision_tower.*` / `language_model.model.*` / `language_model.lm_head.*` and a channels-last patch kernel; accept both.

### Mage-ViT forward

Input: patch rows `x: [L, 768]` (one image, `L = h*w`, rows in 2x2 block order) and per-row positions `pos: [L, 3]` of `(t, hi, wi)` generated in the same block order:

```
for t in 0..T: for hb in 0..h/2: for wb in 0..w/2: for dh in 0..2: for dw in 0..2:
pos.push((t, hb*2 + dh, wb*2 + dw))
```

1. `x = x.reshape(L, 3, 16, 16).transpose(0, 2, 3, 1)` then a 16x16 stride-16 conv with no bias (equivalently a Linear over the 768 vector after the same channel reorder) -> `[L, 1024]`, add batch dim -> `[1, L, 1024]`.
2. Rotary frequencies with `head_dim = 64`, `half = 32`, `unit = half / 16 = 2`, `t_size = 8`, `h_size = 12`, `w_size = 12`; `inv_t[i] = base^(-i/8)` for `i < 8`, `inv_h[j] = base^(-j/12)`, `inv_w[k] = base^(-k/12)` (note the exponent is `i/size`, not `2i/size`). `freqs = concat(t*inv_t, hi*inv_h, wi*inv_w)` -> `[L, 32]`, then `freqs = concat(freqs, freqs)` -> `[L, 64]` (plain concatenation, so lanes `2i` and `2i+1` carry different frequencies).
3. `x = LayerNorm(layernorm_pre)(x)`.
4. 24 pre-norm layers: `h = x + proj(attn(LN1(x)))`; `x = h + fc2(gelu_exact(fc1(LN2(h))))`. Attention: `qkv = Linear(x).reshape(1, L, 3, 16, 64)` -> q, k, v each `[1, 16, L, 64]`; apply rotary in f32 as `q' = q*cos(freqs) + rotate_half(q)*sin(freqs)` with the INTERLEAVED `rotate_half((x1,x2,x3,x4,...)) = (-x2, x1, -x4, x3, ...)`; scale `64^-0.5`; mask: block-diagonal over windows of `frame_windows_size = 4` frames, which for a single image (`t = 1`) is one block, so no mask.
5. No post-layernorm (`use_head false`).
6. Merger: `y = LN(ln_q)(x).reshape(L/4, 4096)`; `y = Linear(4096,4096)`; `gelu_exact`; `Linear(4096, 2560)` -> `[L/4, 2560]`. `use_patch_position_encoding` is false, so no `pos_emb_h/w`.

Output rows map 1:1, in order, onto the `<|image_pad|>` positions (or `<|video_pad|>` 151656 when a video path is added later).

### Decoder

`Qwen3Model` with `rope_theta = 5e6` lifted from `rope_parameters`, `tie_word_embeddings false`, EOS `[151645]` plus `<|endoftext|>` 151643.

## Implementation plan

1. **Detection.** `src/models/detection.rs`: `"mage_vl" => Ok(ModelType::MageVLM)`. `src/models/mod.rs`: enum, supported list, description `("Mage-VL (Mage-ViT 3D-RoPE + 2x2 merger + Qwen3)", "Other VLM")`. Test in `src/models/detection_tests.rs`.

2. **Config (`src/vision/mage_vl_config.rs`, new).** `MageVlVisionConfig` with the fields above and serde defaults matching them; `MageVlConfig { text_config: qwen3::ModelArgs, vision_config, image_token_id (151655), video_token_id (151656), vision_start_token_id, vision_end_token_id, eos_token_id }`. Add `fn lift_rope_parameters(text_config: &mut serde_json::Value)` in the loader that copies `rope_parameters.rope_theta` into `rope_theta` when the flat key is absent (and maps a non-default `rope_type` into `rope_scaling` with a `type` key). Unit test: a config with only `rope_parameters.rope_theta = 5e6` yields `ModelArgs.rope_theta == 5e6`.

3. **Encoder (`src/vision/encoders/mage_vl.rs`, new).** `MageVlVisionEncoder { patch_weight, layernorm_pre: LayerNorm, layers: Vec, merger: MageVlPatchMerger, inv_freq_t/h/w: Vec, num_heads, head_dim, layer_norm_eps }`. Provide `fn positions_from_grid(grid_thw: &[(i32,i32,i32)], merge: usize) -> Vec<[i32; 3]>` (pure Rust, tested) and `fn forward(&self, patches: &MlxArray, grid_thw: &[(i32,i32,i32)]) -> UniquePtr` returning `[sum(t*h*w)/4, 2560]`. Implement rotary with `mlxcel_core` primitives: build `cos`/`sin` as `[1, L, 64]` f32 from the computed `freqs`, and `rotate_half` as `reshape(.., L, 32, 2)` -> stack(`-odd`, `even`) -> reshape; do not use `mlxcel_core::fast_rope` (its rotation convention is split-half or traditional, neither matches). Use `mlxcel_core::scaled_dot_product_attention` with `mask = None` for `t = 1`; build the block-diagonal additive mask only when `t > frame_windows_size` (video, future). Load weights from either prefix (`model.visual.` or `vision_tower.`) and transpose `patch_embedding.weight` with `[0,2,3,1]` only when `shape[1] == 3`. Tests in `src/vision/encoders/mage_vl_tests.rs`: `positions_follow_2x2_block_order` (for grid `(1,4,4)` the first eight rows are `(0,0,0),(0,0,1),(0,1,0),(0,1,1),(0,0,2),(0,0,3),(0,1,2),(0,1,3)`), `every_group_of_four_shares_one_merge_cell`, `positions_cover_grid_exactly_once` for `(2,8,6)`, `rotate_half_is_interleaved` (input `[1,2,3,4]` gives `[-2,1,-4,3]`), `inv_freq_sizes_are_8_12_12_for_head_dim_64`, `forward_output_rows_equal_patches_over_four` on a random-weight 2-layer config.

4. **Processor.** Reuse `Qwen2VLProcessor::new_with_norm(16, 1, 2, clip_mean, clip_std)` with `min_pixels`/`max_pixels` from `preprocessor_config.json`; its `preprocess_with_grid` already emits `[rows, 768]` rows in merge-block order with `temporal_patch_size = 1`. Add a test `qwen2_vl_processor_temporal_patch_1_row_count` asserting `rows == t*h*w` (not `2*t*h*w`) when `temporal_patch_size == 1`, since the existing callers all use 2.

5. **Runtime (`src/vision/mage_vl.rs`, new).** `MageVlModel { text_model: Qwen3Model, vision: MageVlVisionEncoder, processor: Qwen2VLProcessor, image_token_id, video_token_id, vision_start_token_id, vision_end_token_id, eos_token_ids }`; `get_input_embeddings(input_ids, pixel_values, grid_thw)` = `merge::merge_llava(image_token_id, &features, &text_model.embed_tokens(ids), input_ids)` (positions where `id == image_token_id || id == video_token_id`). `LanguageModel` by delegation to `Qwen3Model` (batching and paged decode capabilities pass through unchanged; the text model is a plain Qwen3, so `supports_batching` stays true).

6. **Loader (`src/loading/vlm_mage_vl.rs`, new).** Read config; lift rope parameters; inherit top-level `quantization` into `text_config`; remap keys: `model.language_model.X -> model.X`, `lm_head.X` unchanged, `model.visual.X -> vision_tower.X`, and accept already-converted `language_model.model.X` / `language_model.lm_head.X` via `strip_language_model_prefix`; build `Qwen3Model::from_weights`, the encoder with prefix `vision_tower`, and the processor. bf16 -> f16 conversion through `load_vlm_weights_common`. Register in `src/loading/vlm.rs`, dispatch in `src/loading/mod.rs`.

7. **Dispatch.** `LoadedModel::MageVLM` (`src/loaded_model.rs`), `VlmRuntimeRef::MageVl` (`src/loaded_model_capabilities.rs`), and a `vlm_runtime.rs` arm: render the chat template (the checkpoint ships `chat_template.jinja`; the CLI path can reuse the Qwen2-VL prompt builder that emits `<|vision_start|><|image_pad|><|vision_end|>` per image, see the Qwen runtime in `src/multimodal/qwen_vl.rs`), call `processor.preprocess_with_grid`, expand each `<|image_pad|>` to `t*h*w/4` copies (the Qwen expansion helper in `src/multimodal/qwen_vl.rs` already computes `pixels / (patch*merge)^2`; reuse it with patch 16 and merge 2), build embeddings, and report `VlmPreparationSummary::MageVl { images, image_tokens }` (print in `src/commands/generate_vlm.rs`). Do NOT route through `QwenVlRuntime` / MRoPE: positions for the decoder are the ordinary 1-D sequence positions.

8. **Output suppression.** `[151652, 151653, 151655, 151656]`.

9. **Docs.** `docs/supported-models.md`: add Mage-VL (image input only; codec-native video listed as unsupported).

## Validation

(a) Unit tests listed in steps 2 to 4 plus `detection_tests.rs`.

(b) Real checkpoint: `mlx-community/Mage-VL-8bit` (about 5 GB) or `mlx-community/Mage-VL-OptiQ-4bit`; the bf16 original `microsoft/Mage-VL` is 9.5 GB.

```
./target/release/mlxcel generate -m models/Mage-VL-8bit --image /tmp/solid_red.png -p "What color is this image? Answer with one word." -n 16
./target/release/mlxcel generate -m models/Mage-VL-8bit --image tests/fixtures/test_image.png -p "Describe this image." -n 64
```

Acceptance: finite logits; "red" for the solid image; a 224x224 image yields a `(1, 14, 14)` grid and 49 image tokens; greedy token-exact match for the first 32 tokens against the checkpoint's own `transformers` implementation (CPU bf16) on `tests/fixtures/test_image.png`; a long-context sanity check (a 1500-token prompt plus image) stays fluent, which is the symptom that separates `rope_theta = 5e6` from the 10000 default.

## Acceptance criteria

- [ ] `mage_vl` resolves to `ModelType::MageVLM`
- [ ] Decoder RoPE base is read from `text_config.rope_parameters.rope_theta` (unit test)
- [ ] Mage-ViT reproduces the interleaved rotation, the `concat(freqs, freqs)` pairing, the 4:6:6 split, and the block-order positions (unit tests)
- [ ] Image tokens per image equal `t*h*w/4`; features land on `<|image_pad|>` positions in order
- [ ] Solid-color answer correct and 32-token greedy parity on the real checkpoint
- [ ] docs/supported-models.md updated
- [ ] detection table in src/models/detection.rs updated with a test
- [ ] cargo test --workspace --profile test-fast --features metal,accelerate passes
- [ ] cargo clippy --workspace --all-targets -- -D warnings and cargo fmt --all -- --check pass

## Out of scope

- Video input, both the codec-native path (HEVC/DCVC frame groups described by the `codec` block of `preprocessor_config.json`) and a plain frame sampler. The encoder's `frame_windows_size` block-diagonal mask is specified above so a later video path only needs the processor.
- `use_patch_position_encoding = true` (absolute `pos_emb_h/w` tables in the merger): reject at load.

Contributor guide

Open the contributing guide

Research direction

Start with the detection and model registration points in src/models/detection.rs, src/models/mod.rs, and src/models/detection_tests.rs, then read the Qwen2-VL processor, Qwen3 decoder, and existing VLM loaders. Implement the planned config, encoder, runtime, loader, and dispatch files, using the named Mage-VL tests to validate ordering and rotary behavior. Done means the listed tests pass and the image-only mlxcel generate command produces a response from the Mage-VL checkpoint.

Written by the indexing model from the issue text.

Assessment

Tech stack
rust
Domain
ai, computer-vision
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.