Skip to content

feat(models): port DeepSeek-VL2 (MoE text + SigLIP/SAM vision) #533

Description

@inureyes

Summary

Add DeepSeek-VL2 support to mlxcel. Checkpoints ship config.json with model_type: deepseek_vl_v2; mlxcel has no arm for it in src/models/detection.rs, so loading fails with "Unsupported model type".

DeepSeek-VL2 pairs a DeepSeek-MoE text backbone (the deepseek_v2 architecture family: MLA attention or plain MHA depending on config, softmax-gated routed experts plus shared experts) with a SigLIP-style ViT vision tower, a 2x2 space-to-depth MLP projector, and a global-plus-local tiling scheme that inserts learned newline and view-separator embeddings between tile rows and views.

Current state in mlxcel

  • src/models/detection.rs maps deepseek, deepseek_v2, deepseek_v3, deepseek_v32 (lines ~152-155) but not deepseek_vl_v2.
  • src/models/deepseek_v2.rs already implements the text backbone this model uses: MLA with optional q_lora_rank, YaRN rope scaling, dense-vs-MoE layer selection via is_moe_layer (first_k_dense_replace, moe_layer_freq), MoEGate (softmax scores, greedy top-k, routed_scaling_factor), SwitchGLU routed experts plus optional shared experts. Two gaps for this model:
    • No plain-MHA path. DeepSeek-VL2 text configs can set qk_nope_head_dim: 0, which switches every layer to standard multi-head attention (q_proj/k_proj/v_proj/o_proj, full-head rope). The smaller checkpoints in the family use this mode, so it is required.
    • MoEGate has no group-limited masking (n_group/topk_group). DeepSeek-VL2 configs set n_group: 1, topk_group: 1, which makes the masking a no-op, so the existing gate math already matches; the loader should still read the keys and fail loudly if a checkpoint sets n_group > 1.
  • src/models/switch_layers.rs provides shared SwitchLinear/SwitchGLU (gather_qmm / gather_mm) machinery; deepseek_v2.rs carries its own local copy that expects pre-fused model.layers.{i}.mlp.switch_mlp.* keys.
  • src/vision/encoders/siglip.rs is a SigLIP encoder with split q/k/v and patch_embedding.* naming. The DeepSeek-VL2 tower uses a different block layout (fused attn.qkv, blocks.{i} naming, a bare pos_embed parameter, no CLS token, trailing norm), so it needs either a key remap onto the existing encoder or a small dedicated encoder (see plan).
  • Multi-crop precedent: src/vision/processors/internvl.rs (aspect-ratio tile selection, crop loop, per-tile normalization) and src/multimodal/internvl_prompt.rs (placeholder expansion by per-image tile counts). DeepSeek-VL2's tiling rule differs in the details (candidate-resolution best fit, mean-color padding instead of resize-crop, global view first).
  • Canonical VLM registration points (all verified): ModelType enum plus the VLM list, display-name table, and macro list in src/models/mod.rs (~lines 266, 444, 598, 863), loader dispatch in src/loading/mod.rs (~line 187), loader module registration in src/loading/vlm.rs, LoadedModel variant plus macro arms in src/loaded_model.rs, VlmRuntimeRef dispatch in src/multimodal/vlm_runtime.rs, placeholder suppression in src/models/multimodal_placeholders.rs, docs entry in docs/supported-models.md.

Architecture

Shape notation below is always explicit named dimension order, e.g. (tiles, height, width, channels). Axis numbers are 0-based and refer to the shape written immediately beside them.

config.json layout

Top level: model_type: "deepseek_vl_v2", text config nested under language_config (accept text_config as an alias; some exports use it), vision_config, projector_config, and these scalar keys (defaults when absent):

key default meaning
image_token_index 100015 id of the <image> placeholder token
pad_id 100001 pad token id
select_layer -1 vision feature layer; only -1 (final normalized layer, all patch tokens) is used
vision_feature_select_strategy "default" no token dropping
tile_tag "2D" tile-formatting mode; only "2D" is required (error on "1D")
global_view_pos "head" global view placed before local tiles ("tail" reverses)
candidate_resolutions (required) list of [W, H] tile targets; also present in processor_config.json

language_config keys and defaults (the nested model_type is "deepseek_v2"): vocab_size 102400, hidden_size 1280, intermediate_size 6848, moe_intermediate_size 896, num_hidden_layers 30, num_attention_heads 32, num_key_value_heads 32 (falls back to num_attention_heads when null), n_shared_experts 2, n_routed_experts 64, routed_scaling_factor 1.0, kv_lora_rank 512, q_lora_rank 1536 (may be null), qk_rope_head_dim 64, v_head_dim 128, qk_nope_head_dim 128, topk_method "greedy", n_group 1, topk_group 1, num_experts_per_tok 6, moe_layer_freq 1, first_k_dense_replace 0, max_position_embeddings 2048, rms_norm_eps 1e-6, rope_theta 10000.0, rope_scaling null, attention_bias false, scoring_func "softmax".

Attention-mode rule: if qk_nope_head_dim == 0 the text layers use plain MHA; otherwise MLA exactly as in src/models/deepseek_v2.rs. Note that rope for this family is traditional (interleaved) by default, in both modes.

vision_config keys and defaults (nested model_type is "vision"; ignore its cls/params fields): layers 27, width 1152, intermediate_size 4304, num_attention_heads 16, image_size 384, patch_size 16, num_channels 3, layer_norm_eps 1e-6, mlp_ratio 3.7362.

projector_config keys and defaults: projector_type "downsample_mlp_gelu", input_dim 1152, n_embed 2048 (equals the text hidden_size in practice), depth 2, mlp_ratio 1, downsample_ratio 2, token_pooling false.

Text backbone

Identical to the in-tree DeepSeek-V2 block: RMSNorm pre-attention and pre-MLP, residuals, final model.norm, separate lm_head. MoE layers (where layer_idx >= first_k_dense_replace and layer_idx % moe_layer_freq == 0 and n_routed_experts is set) compute, with hidden states x of shape (batch, seq, hidden_size) and E = n_routed_experts:

gates  = x @ gate.weight^T          # gate.weight: (E, hidden_size); gates: (batch, seq, E)
scores = softmax(gates, axis=2)     # over the expert axis
# group-limited masking (no-op when n_group == 1):
#   reshape scores to (batch, seq, n_group, E / n_group);
#   group_score = max over axis 3;
#   zero the (n_group - topk_group) lowest-scoring groups;
#   reshape back to (batch, seq, E)
inds   = top-k(scores, k=num_experts_per_tok, axis=2)   # (batch, seq, k); no renormalization
w      = scores gathered at inds * routed_scaling_factor  # (batch, seq, k)
y      = sum over the k axis of w[..., None] * expert_outputs  # (batch, seq, hidden_size)
y      = y + shared_experts(x)      # shared MLP intermediate = moe_intermediate_size * n_shared_experts

The plain-MHA mode (when qk_nope_head_dim == 0): head_dim = hidden_size / num_attention_heads, scale head_dim^-0.5. Projections produce q: (batch, seq, num_attention_heads * head_dim), k/v: (batch, seq, num_key_value_heads * head_dim); reshape to (batch, seq, n_heads, head_dim) and transpose axes (0, 2, 1, 3) to (batch, n_heads, seq, head_dim) before attention, then invert the same transpose and reshape back to (batch, seq, hidden_size) for o_proj. Rope covers the full head_dim with traditional = true, base rope_theta, scale 1/factor only when rope_scaling.type == "linear"; GQA via num_key_value_heads.

Vision tower

A no-CLS ViT over each tile. The tower consumes channels-last input (tiles, image_size, image_size, channels); if the processor emits channels-first (tiles, channels, image_size, image_size), permute axes (0, 2, 3, 1) at the tower boundary (precedent: src/vision/kimi_vl.rs line ~147).

  • patch embed: Conv2d(num_channels -> width, kernel = stride = patch_size) producing (tiles, gh, gw, width) with gh = gw = floor(image_size / patch_size); flatten axes 1 and 2 (row-major over (gh, gw)) to (tiles, N, width), N = gh * gw
  • add learned pos_embed of shape (1, N, width), broadcast over axis 0
  • layers pre-norm blocks: x = x + attn(norm1(x)); x = x + mlp(norm2(x))
    • attention: fused qkv Linear(width -> 3*width, bias) produces (tiles, N, 3 * width); split into three equal (tiles, N, width) chunks along axis 2, ordered q, k, v; reshape each to (tiles, N, num_attention_heads, head_dim) and transpose axes (0, 2, 1, 3) to (tiles, num_attention_heads, N, head_dim); scaled dot-product attention with scale head_dim^-0.5, no mask; transpose back (0, 2, 1, 3) and reshape to (tiles, N, width); output proj Linear(width -> width, bias)
    • mlp: fc1 (width -> intermediate_size, bias) -> GELU (tanh approximation) -> fc2 (bias)
    • block LayerNorms use layer_norm_eps (1e-6)
  • final LayerNorm norm (eps 1e-5) over all tokens; the output (tiles, N, width) is the feature the projector consumes
  • the checkpoint also contains an attention-pooling head (attn_pool.*: latent, q, kv, proj, norm, mlp.fc1/fc2) that is not used for VLM inference; drop those weights at load

Projector (downsample_mlp_gelu)

Input (tiles, N, input_dim), g = sqrt(N), ds = downsample_ratio:

  1. reshape to (tiles, g, g, input_dim) (token axis unflattened row-major into height-then-width); zero-pad axes 1 and 2 at the high end up to g_pad, the next multiple of ds
  2. space-to-depth: for each non-overlapping ds x ds block, scanned row-major (outer loop over height, inner over width), concatenate the block's values into one feature vector ordered channel-outermost: flattened index c * ds * ds + dy * ds + dx, where c is the channel and (dy, dx) the position inside the block. Result: (tiles, N', input_dim * ds * ds) with N' = (g_pad / ds)^2. This feature ordering must match exactly; the first projector Linear was trained against it.
  3. Linear(input_dim * ds * ds -> n_embed * mlp_ratio) -> GELU -> Linear(n_embed * mlp_ratio -> n_embed); with depth > 2, additional GELU+Linear pairs sit in the middle

Output: (tiles, N', n_embed), h' = w' = sqrt(N') = ceil(floor(image_size / patch_size) / ds), token order row-major over (h', w').

Tile-sequence assembly (tile_tag "2D")

Two learned vectors of shape (n_embed,) live at the checkpoint top level: image_newline and the view separator (stored under the misspelled key view_seperator). Per image with tile grid (tw, th) (width count first, from images_spatial_crop), operating on that image's projector output slice of shape (1 + tw*th, N', D) where D = n_embed and tile 0 is the global view:

  • global view: take (N', D), reshape to (h', w', D); broadcast image_newline to (h', 1, D) and concatenate along axis 1 -> (h', w' + 1, D); reshape to (h' * (w' + 1), D)
  • local views: take (th * tw, h' * w', D), reshape to (th, tw, h', w', D), transpose axes (0, 2, 1, 3, 4) -> (th, h', tw, w', D), reshape to (th * h', tw * w', D); broadcast image_newline to (th * h', 1, D) and concatenate along axis 1 -> (th * h', tw * w' + 1, D); reshape to ((th * h') * (tw * w' + 1), D)
  • reshape view_separator to (1, D); concatenate along axis 0 in the order [global, view_separator, local] (global_view_pos == "head"; "tail" gives [local, view_separator, global])

The assembled rows overwrite the token embeddings (seq, D) at exactly the positions holding image_token_index (the images_seq_mask positions), in order. Text-only requests bypass the vision path entirely.

Processor

Configuration comes from processor_config.json: candidate_resolutions, patch_size, downsample_ratio, image_mean (0.5, 0.5, 0.5), image_std (0.5, 0.5, 0.5), normalize true, image_token "<image>", pad_token "<|▁pad▁|>". image_size = candidate_resolutions[0][0]. Note the processor's patch_size/image_size govern token counts and must agree with the vision config.

Per image:

  1. Cropping is enabled only when the request carries at most 2 images; otherwise every image gets the 1x1 grid.
  2. With cropping, pick (best_W, best_H) from candidate_resolutions: for each candidate compute scale = min(W/w0, H/h0), effective = min(floor(w0*scale) * floor(h0*scale), w0*h0), wasted = W*H - effective; choose max effective, ties broken by min wasted. Without cropping, best = (image_size, image_size).
  3. Global view: aspect-preserving resize of the full image into (image_size, image_size) with centered padding, pad color = mean color (mid gray).
  4. Local views: pad the image the same way into (best_W, best_H), then crop non-overlapping image_size x image_size tiles row-major (outer loop over height, inner over width).
  5. Normalize each tile: divide by 255, then (x - mean) / std with mean/std broadcast along the channel axis. Tile order per image: global first, then locals. Emit the stack in a defined layout, either (tiles, channels, image_size, image_size) plus the (0, 2, 3, 1) permute at the tower boundary, or (tiles, image_size, image_size, channels) directly.
  6. images_spatial_crop entry: [tw, th] = [best_W / image_size, best_H / image_size] (width count first).
  7. Placeholder expansion: each <image> occurrence in the prompt is replaced by image_token_index repeated h'*(w'+1) + 1 + (th*h')*(tw*w'+1) times, with h' = w' = ceil(floor(image_size / patch_size) / downsample_ratio). Prepend BOS; do not append EOS for generation. Examples: a 384px tower with patch 14 gives h' = 14, so a 1x1 grid costs 210 + 1 + 210 = 421 tokens and a 2x1 grid costs 210 + 1 + 14*29 = 617; patch 16 gives h' = 12 and 156 + 1 + 156 = 313 for 1x1. Derive from the actual checkpoint config, never hardcode.

Chat format, when tokenizer_config.json ships no chat template: user turns render as <|User|>: {content}\n\n, assistant turns as <|Assistant|> {content}\n\n, and the generation prompt is <|Assistant|>:. <|User|>/<|Assistant|> are additional special tokens; grounding tokens (<|ref|>, <|/ref|>, <|det|>, <|/det|>, <|grounding|>) exist in the tokenizer but need no special handling here.

Safetensors weight keys

  • text: language.model.embed_tokens.weight, language.model.layers.{i}.self_attn.*, language.model.layers.{i}.{input_layernorm,post_attention_layernorm}.weight, language.model.norm.weight, language.lm_head.weight
    • MLA attention: q_proj (or q_a_proj/q_a_layernorm/q_b_proj when q_lora_rank is set), kv_a_proj_with_mqa, kv_a_layernorm, kv_b_proj, o_proj
    • MHA attention (qk_nope_head_dim == 0): q_proj, k_proj, v_proj, o_proj
  • dense MLP layers: language.model.layers.{i}.mlp.{gate_proj,up_proj,down_proj}.weight
  • MoE layers: ...mlp.gate.weight of shape (n_routed_experts, hidden_size), per-expert ...mlp.experts.{e}.{gate_proj,up_proj,down_proj}.weight, shared ...mlp.shared_experts.{gate_proj,up_proj,down_proj}.weight
  • vision: vision.patch_embed.proj.{weight,bias}, vision.pos_embed of shape (1, N, width), vision.blocks.{i}.norm1.{weight,bias}, vision.blocks.{i}.attn.qkv.{weight,bias}, vision.blocks.{i}.attn.proj.{weight,bias}, vision.blocks.{i}.norm2.{weight,bias}, vision.blocks.{i}.mlp.{fc1,fc2}.{weight,bias}, vision.norm.{weight,bias}, plus droppable vision.attn_pool.*
  • projector: projector.layers.0.{weight,bias}, projector.layers.2.{weight,bias} (indices shift with depth)
  • top level: image_newline and view_seperator (checkpoint key keeps the typo), each of shape (n_embed,)

Implementation plan

  1. src/models/mod.rs: add ModelType::DeepSeekVL2 to the enum (~line 266), the VLM model-type list (~line 444), the display-name table (~line 598: ("DeepSeek-VL2", "DeepSeek VLM")), and the macro model-type list (~line 863).
  2. src/models/detection.rs: add the "deepseek_vl_v2" => Ok(ModelType::DeepSeekVL2) arm next to the existing deepseek arms; cover in src/models/detection_tests.rs.
  3. Text backbone (src/models/deepseek_v2.rs): add the plain-MHA attention path selected by qk_nope_head_dim == 0 (new loader branch plus attention struct, full-head traditional rope); add serde fields for topk_method, n_group, topk_group, scoring_func and return a clear error when a config requires unimplemented gate behavior (n_group > 1 or topk_method/scoring_func other than greedy/softmax).
  4. src/vision/encoders/deepseek_vl_v2.rs: dedicated tower for the fused-qkv, blocks.{i}, pos_embed, trailing-norm layout described above (the existing src/vision/encoders/siglip.rs expects split q/k/v and different key names; a dedicated encoder mirrors how src/vision/encoders/paddleocr_vl.rs handles its own SigLIP variant). Register in src/vision/encoders/mod.rs.
  5. src/vision/connectors/deepseek_vl_v2.rs: the downsample_mlp_gelu projector (space-to-depth plus Linear/GELU stack, with the exact channel-outermost feature ordering from the Architecture section); error on projector_type values other than downsample_mlp_gelu and on token_pooling: true. Register in src/vision/connectors/mod.rs.
  6. src/vision/processors/deepseek_vl_v2.rs: candidate-resolution selection, mean-color padding, global-plus-local tiling, normalization, images_spatial_crop bookkeeping, and per-image token-count computation. Register in src/vision/processors/mod.rs. Tiling loop precedent: src/vision/processors/internvl.rs.
  7. src/vision/deepseek_vl_v2.rs: VLM wrapper struct holding text model, tower, projector, processor, image_newline/view_separator vectors, and the config scalars (image_token_index, tile_tag, global_view_pos); implements tile-sequence assembly and embedding scatter, exposing the same surface as src/vision/paddleocr_vl.rs (input_embeddings, LanguageModel plumbing). Register in src/vision/mod.rs.
  8. src/loading/vlm_deepseek_vl_v2.rs: parse language_config (alias text_config), vision_config, projector_config; run the sanitize/remap below; build the wrapper. Wire as a module in src/loading/vlm.rs (mirroring the vlm_paddleocr.rs registration at lines ~56-79) and dispatch from src/loading/mod.rs (~line 187). Add the LoadedModel::DeepSeekVL2 variant and macro arms in src/loaded_model.rs.
  9. src/multimodal/: expand <image> placeholders with the per-image tile-dependent counts (model-specific deepseek_vl_v2_prompt.rs following src/multimodal/internvl_prompt.rs, or the generic src/multimodal/vlm_prompt.rs if its expansion fits); add the VlmRuntimeRef variant and dispatch arm in src/multimodal/vlm_runtime.rs; register in src/multimodal/mod.rs.
  10. src/models/multimodal_placeholders.rs: report image_token_index so the placeholder id is suppressed from text output.
  11. docs/supported-models.md: add the entry in the established format.

Weight loading and sanitize

Do the remap in the loader (per-model fn, shared helpers in src/models/sanitize.rs where applicable):

  • language.model.* -> model.* and language.lm_head.* -> lm_head.* so DeepSeekV2Model::from_weights sees its expected names; keep vision.* and projector.* namespaces for the tower and connector.
  • Stack per-expert weights: model.layers.{i}.mlp.experts.{e}.{gate_proj,up_proj,down_proj}.{weight,scales,biases} -> model.layers.{i}.mlp.switch_mlp.{...}, stacking the n_routed_experts tensors along a new axis 0. Unquantized source shapes: gate/up (moe_intermediate_size, hidden_size) each, down (hidden_size, moe_intermediate_size); stacked results (n_routed_experts, moe_intermediate_size, hidden_size) and (n_routed_experts, hidden_size, moe_intermediate_size). Precedent: the stacking block in src/models/deepseek_v3.rs (~line 1056) and the .mlp.experts. -> .mlp.switch_mlp. rename in src/loading/vlm.rs (~line 568). Accept checkpoints that already ship fused switch_mlp keys.
  • Every required permutation, with source and target shapes:
    • vision.patch_embed.proj.weight: the checkpoint stores (out_channels=width, in_channels=num_channels, kH=patch_size, kW=patch_size); mlxcel's conv2d consumes channels-last weights, so permute axes (0, 2, 3, 1) to (out_channels, kH, kW, in_channels). Precedents: transpose_axes(&w, &[0, 2, 3, 1]) in src/loading/vlm_kimi_vl.rs (~line 229) and src/vision/encoders/minicpmo.rs (~line 317). Skip when the tensor is already 4-D channels-last (detect by shape, as those loaders do).
    • Pixel input: if the processor emits (tiles, channels, image_size, image_size), permute axes (0, 2, 3, 1) to (tiles, image_size, image_size, channels) before the conv (precedent: src/vision/kimi_vl.rs ~line 147).
    • All Linear weights (q_proj, qkv, fc1, projector layers, experts, ...): stored (out_features, in_features) and consumed as-is by UnifiedLinear/SwitchLinear; no permute. The fused vision qkv.weight (3 * width, width) is ordered q, k, v along axis 0; keep it fused and split the activation, or split the weight along axis 0 into three (width, width) blocks in that order.
    • vision.pos_embed (1, N, width): consumed as-is (broadcast over axis 0); if a variant ships (N, width), insert axis 0.
  • Normalize view_seperator -> the canonical internal name; both image_newline and the separator are (n_embed,) vectors, no permute.
  • Drop vision.attn_pool.* and any position_ids buffers.
  • Quantized exports: honor top-level quantization {group_size, bits}; handle mixed exports where the vision tower or projector stays unquantized by deciding per key (the UnifiedLinear/SwitchLinear loaders already support this pattern).
  • No tie_word_embeddings in this family's configs: load the separate lm_head.

Validation and acceptance criteria

  • Co-located unit tests: processor (candidate-resolution selection incl. tie-break, tile grid and crop order, per-image token-count formula, the 3-plus-images no-cropping rule), projector shapes and space-to-depth feature ordering (incl. the odd-grid padding path), tile-sequence assembly order (global/newline/separator/local layout and the local-view axis transpose), MoE gate top-k, detection arm.
  • Synthetic parity test tests/deepseek_vl_v2_parity.rs following the tests/internvl_parity.rs / tests/paddleocr_vl_parity.rs pattern: model-type detection, text-only forward with finite logits, image forward with finite logits, tiling assertions on a fixture image.
  • Real-checkpoint smoke test: load a real deepseek_vl_v2 checkpoint and generate through mlxcel's server/CLI runtime, both text-only and with an image; verify the placeholder expansion count matches the number of scattered feature rows and that output is coherent.
  • Known failure mode: synthetic parity can pass while real checkpoints break on weight-key layout (the language. prefix, per-expert vs fused MoE keys, the view_seperator spelling), config defaults (qk_nope_head_dim == 0 MHA mode, first_k_dense_replace, q_lora_rank: null), tokenizer/chat-template details, and MoE fusion. Real-checkpoint load plus generate is the completion bar.
  • Completion requires integration into the working runtime (detection -> loader -> server/CLI generation), not standalone modules.

Out of scope

  • Grounding post-processing (<|ref|>/<|det|> box parsing) beyond the tokens passing through as text.
  • The "1D" tile-tag path and token_pooling/identity/linear/mlp_gelu projector variants (error clearly if encountered).
  • Video inputs and batch/collate training paths.
  • Group-limited MoE routing for n_group > 1 (not used by this family; error clearly).

Effort: high

Metadata

Metadata

Assignees

Labels

area:modelsModel architectures, weights, loading, metadatapriority:mediumMedium prioritystatus:readyReady to be worked ontype: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