diff --git a/docs/supported-models.md b/docs/supported-models.md index a55081d1..6d464f09 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -83,6 +83,7 @@ Implemented VLM variants include: - GLM-4V (`glm4v`): GLM-4V ViT vision tower (3D patch embedding, bilinear-resampled learned position embeddings, Conv2d spatial downsample, SwiGLU patch merger) plus a GLM-4 text backbone driven by sectioned even/odd MRoPE. Reuses the shared Qwen-VL image processor and prompt/token plumbing. - GLM-4V MoE (`glm4v_moe`): GLM-4.5V-class variant reusing the GLM-4V ViT vision tower with a GLM-4 MoE text backbone (grouped `noaux_tc` routing, shared experts, `first_k_dense_replace` dense layers) driven by sectioned half-split MRoPE. Reuses the GLM-4 MoE machinery and the shared Qwen-VL runtime. - Granite Vision (`granite_vision`, or `llava_next` with a `granite` text config): IBM's document VLM. A SigLIP vision tower with four intermediate feature taps (concatenated on the channel axis) feeds a 2-layer GELU projector; a learned `image_newline` embedding, LLaVA-Next AnyRes multi-tile preprocessing (`image_grid_pinpoints`), and a dense Granite text backbone (embedding / attention / residual / logits multipliers) complete the model. Both `config.json` spellings route to the same loader. Best for document understanding, charts, and tables. +- Granite 4 Vision (`granite4_vision`): IBM's document VLM with multi-depth visual injection. A SigLIP tower feeds eight window-QFormer projectors (a QFormer with self+cross attention per `4x4`/`8x8` window, with deepstack mean-pool and spatial strided-offset query downsamplers) whose packed outputs are added into the residual stream at eight different depths of a Granite 4 hybrid (`granitemoehybrid`) text backbone during prefill, rather than merged once. Reuses the shared AnyRes tiling and the four Granite scalar multipliers. Best for document understanding. - GLM-OCR (`glm_ocr`): document-OCR sibling of GLM-4V. A 24-block ViT (3D patch embedding, per-head q/k RMSNorm on the packed `cu_seqlens` attention, 2D vision RoPE, Conv2d spatial downsample, SwiGLU patch merger) feeds a 16-layer GLM-4 text decoder driven by full-width even/odd MRoPE (`rope_parameters` with `mrope_section [16, 24, 24]`, `partial_rotary_factor 1.0`). The tower has no learned position embedding or post-conv norm, and the loader drops the next-n prediction (MTP) layer. Patches are reordered from the processor's raster order into spatial-merge-window order so the rotary, downsample, and merged-token grid stay spatially aligned (OCR reads scrambled patches wrong). Best for plain text, tables, and formula recognition. - Youtu-VL - Kimi-VL and Kimi-VL 2.5 (`kimi_vl` / `kimi_k25`): MoonViT native-resolution vision encoder (Conv2d patch embedding, learned plus bicubically-interpolated 2D position embedding, a shared 2D rotary embedding, block-diagonal cross-image attention, and `spatial_merge_size` patch merging) feeding a `LayerNorm -> Linear -> GELU -> Linear` connector into a DeepSeek-V3-style MoE text backbone. Detected, loaded, and served end to end: the safetensors directory loader wires the MoonViT tower, the connector, and the DeepSeek-V3 MoE backbone; the native-resolution processor patchifies each image; and the runtime expands each `<|media_pad|>` placeholder into `(h/merge) * (w/merge)` tokens before the merged vision features replace them. Image path only; the Kimi-VL 2.5 3D MoonViT video (image plus video patch embedding) path is tracked as a separate follow-up. diff --git a/src/commands/generate_vlm.rs b/src/commands/generate_vlm.rs index 25df68d0..86e70d09 100644 --- a/src/commands/generate_vlm.rs +++ b/src/commands/generate_vlm.rs @@ -218,6 +218,15 @@ fn print_preparation_summary(summary: VlmPreparationSummary) { image_blocks, total_image_tokens ); } + VlmPreparationSummary::Granite4Vision { + image_blocks, + total_image_tokens, + } => { + println!( + "Granite 4 Vision: expanded {} placeholder(s) ({} total image tokens)", + image_blocks, total_image_tokens + ); + } VlmPreparationSummary::ImageBlocks(stats) => match stats.action { ImageTokenBlockAction::Expanded { existing_image_count, diff --git a/src/distributed/tensor_parallel/inference.rs b/src/distributed/tensor_parallel/inference.rs index 32fe470b..db4826a7 100644 --- a/src/distributed/tensor_parallel/inference.rs +++ b/src/distributed/tensor_parallel/inference.rs @@ -258,6 +258,9 @@ fn fallback_architecture(model_type: ModelType) -> &'static str { // Granite Vision's text backbone is Granite; TP is refused for VLM-kind // models earlier, this keeps the dispatch table total. ModelType::GraniteVisionVLM => "granite", + // Granite 4 Vision's text backbone is granitemoehybrid; TP is refused for + // VLM-kind models earlier, this keeps the dispatch table total. + ModelType::Granite4VisionVLM => "granitemoehybrid", // Youtu-VL is not currently supported by tensor-parallel inference; // we return a placeholder architecture string here so the planner // does not panic on the dispatch table lookup. The actual loader diff --git a/src/loaded_model.rs b/src/loaded_model.rs index 49c4ef3c..30b14ee4 100644 --- a/src/loaded_model.rs +++ b/src/loaded_model.rs @@ -78,6 +78,7 @@ pub enum LoadedModel { MllamaVLM(vision::MllamaVLModel), LlavaVLM(vision::VisionLanguageModel), GraniteVisionVLM(vision::GraniteVisionVLModel), + Granite4VisionVLM(vision::Granite4VisionVLModel), Qwen2VL(vision::Qwen2VLModel), Qwen25VL(vision::Qwen25VLModel), Qwen3VL(vision::Qwen3VLModel), @@ -210,6 +211,7 @@ macro_rules! delegate_language_model { LoadedModel::MllamaVLM(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::LlavaVLM(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::GraniteVisionVLM(inner) => LanguageModel::$method(inner, $($arg),*), + LoadedModel::Granite4VisionVLM(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Qwen2VL(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Qwen25VL(inner) => LanguageModel::$method(inner, $($arg),*), LoadedModel::Qwen3VL(inner) => LanguageModel::$method(inner, $($arg),*), diff --git a/src/loaded_model_capabilities.rs b/src/loaded_model_capabilities.rs index 8b93d2c2..a3f3b61f 100644 --- a/src/loaded_model_capabilities.rs +++ b/src/loaded_model_capabilities.rs @@ -61,6 +61,7 @@ pub enum VlmRuntimeRef<'a> { /// LFM2-VL (lfm2_vl) runtime. Lfm2Vl(&'a vision::Lfm2VlModel), GraniteVision(&'a vision::GraniteVisionVLModel), + Granite4Vision(&'a vision::Granite4VisionVLModel), Standard(&'a vision::VisionModule), } @@ -175,6 +176,7 @@ impl LoadedModel { Self::Llama4VLM(vlm) => Some(VlmRuntimeRef::Standard(&vlm.vision)), Self::LlavaVLM(vlm) => Some(VlmRuntimeRef::Standard(&vlm.vision)), Self::GraniteVisionVLM(model) => Some(VlmRuntimeRef::GraniteVision(model)), + Self::Granite4VisionVLM(model) => Some(VlmRuntimeRef::Granite4Vision(model)), _ => None, } } diff --git a/src/loading/mod.rs b/src/loading/mod.rs index d6f9d9f3..d4c5671f 100644 --- a/src/loading/mod.rs +++ b/src/loading/mod.rs @@ -176,6 +176,7 @@ fn try_load_vlm_model_from_dir( ModelType::Gemma4Unified => Some(load_gemma4_unified(model_path)?), ModelType::LlavaVLM => Some(load_llava_vlm(model_path)?), ModelType::GraniteVisionVLM => Some(load_granite_vision_vlm(model_path)?), + ModelType::Granite4VisionVLM => Some(load_granite4_vision_vlm(model_path)?), ModelType::LlavaBunnyVLM => Some(load_llava_bunny_vlm(model_path)?), ModelType::AyaVisionVLM => Some(load_aya_vision_vlm(model_path)?), ModelType::PaliGemmaVLM => Some(load_paligemma_vlm(model_path)?), diff --git a/src/loading/vlm.rs b/src/loading/vlm.rs index bca16e91..cd379e76 100644 --- a/src/loading/vlm.rs +++ b/src/loading/vlm.rs @@ -41,6 +41,8 @@ use models::sanitize_config_json; mod gemma; #[path = "vlm_gemma_unified.rs"] mod gemma_unified; +#[path = "vlm_granite4_vision.rs"] +mod granite4_vision; #[path = "vlm_granite_vision.rs"] mod granite_vision; #[path = "vlm_idefics2.rs"] @@ -75,6 +77,7 @@ mod youtu_vl_loader; pub(crate) use gemma::{load_gemma3_vlm, load_gemma3n_vlm, load_gemma4_vlm}; pub(crate) use gemma_unified::load_gemma4_unified; pub(crate) use granite_vision::load_granite_vision_vlm; +pub(crate) use granite4_vision::load_granite4_vision_vlm; pub(crate) use idefics2::load_idefics2_vlm; pub(crate) use internvl::load_internvl_vlm; pub(crate) use kimi_vl_loader::load_kimi_vl_vlm; diff --git a/src/loading/vlm_granite4_vision.rs b/src/loading/vlm_granite4_vision.rs new file mode 100644 index 00000000..80f5e465 --- /dev/null +++ b/src/loading/vlm_granite4_vision.rs @@ -0,0 +1,254 @@ +// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +//! Granite 4 Vision loader (SigLIP + 8 window-QFormer projectors + Granite 4 +//! hybrid text backbone with multi-depth injection). + +use anyhow::Result; +use std::path::Path; + +use crate::LoadedModel; +use crate::models; +use crate::vision; +use crate::vision::connectors::granite4_vision::{Downsampler, WindowQFormerProjector}; + +use super::{load_vlm_weights_common, read_sanitized_vlm_config, strip_language_model_prefix}; + +/// Default LLaVA-Next `image_grid_pinpoints` for Granite 4 Vision. +fn default_pinpoints() -> Vec<(i32, i32)> { + let mut p: Vec<(i32, i32)> = Vec::new(); + for w in (384..=3840).step_by(384) { + p.push((384, w)); + } + for w in (384..=1920).step_by(384) { + p.push((768, w)); + } + for w in (384..=1152).step_by(384) { + p.push((1152, w)); + } + for h in [1536, 1920] { + for w in [384, 768] { + p.push((h, w)); + } + } + for h in [2304, 2688, 3072, 3456, 3840] { + p.push((h, 384)); + } + p +} + +fn parse_pinpoints(full_config: &serde_json::Value) -> Vec<(i32, i32)> { + full_config + .get("image_grid_pinpoints") + .and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|pair| { + let p = pair.as_array()?; + Some((p.first()?.as_i64()? as i32, p.get(1)?.as_i64()? as i32)) + }) + .collect::>() + }) + .filter(|v| !v.is_empty()) + .unwrap_or_else(default_pinpoints) +} + +/// `downsample_rate` `"q/w"` -> `(q, w)`; defaults to `(4, 8)`. +fn parse_downsample_rate(full_config: &serde_json::Value) -> (i32, i32) { + full_config + .get("downsample_rate") + .and_then(|v| v.as_str()) + .and_then(|s| { + let mut it = s.split('/'); + let q = it.next()?.trim().parse::().ok()?; + let w = it.next()?.trim().parse::().ok()?; + Some((q, w)) + }) + .unwrap_or((4, 8)) +} + +fn parse_i32_pairs(v: Option<&serde_json::Value>) -> Vec<[i32; 2]> { + v.and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|pair| { + let p = pair.as_array()?; + Some([p.first()?.as_i64()? as i32, p.get(1)?.as_i64()? as i32]) + }) + .collect() + }) + .unwrap_or_default() +} + +fn parse_i32_list(v: Option<&serde_json::Value>) -> Vec { + v.and_then(|v| v.as_array()) + .map(|arr| { + arr.iter() + .filter_map(|x| x.as_i64().map(|i| i as i32)) + .collect() + }) + .unwrap_or_default() +} + +/// Load a Granite 4 Vision VLM. +pub(crate) fn load_granite4_vision_vlm(model_path: &Path) -> Result { + let (_config_str, full_config) = read_sanitized_vlm_config(model_path)?; + + // Text backbone config: inherit top-level quantization, parse as granitemoehybrid. + let mut text_config_val = full_config + .get("text_config") + .cloned() + .ok_or_else(|| anyhow::anyhow!("Missing text_config in config.json"))?; + if let Some(obj) = text_config_val.as_object_mut() + && obj.get("quantization").is_none() + && let Some(q) = full_config.get("quantization") + { + obj.insert("quantization".to_string(), q.clone()); + } + let text_args: models::granitemoehybrid::ModelArgs = serde_json::from_value(text_config_val) + .map_err(|e| anyhow::anyhow!("Failed to parse Granite 4 text config: {}", e))?; + + let vision_config: vision::config::VisionConfig = serde_json::from_value( + full_config + .get("vision_config") + .cloned() + .ok_or_else(|| anyhow::anyhow!("Missing vision_config in config.json"))?, + ) + .map_err(|e| anyhow::anyhow!("Failed to parse Granite 4 vision config: {}", e))?; + + let group_size = text_args.group_size(); + let bits = text_args.bits(); + let (q, w) = parse_downsample_rate(&full_config); + + let weights = strip_language_model_prefix(load_vlm_weights_common(model_path, None)?); + + // Vision tower and 8 window-QFormer projectors read by ref before the text + // backbone consumes the weight map. The granite4 SigLIP tower uses the + // sigmoid `GELU(approx="fast")` in its MLP (unlike Granite Vision #539, which + // uses the tanh variant), matching the reference vision encoder. + let vision_tower = + vision::encoders::siglip::SigLipVisionModel::from_weights_with_quant_and_gelu( + &weights, + &vision_config, + "vision_tower.vision_model", + group_size, + bits, + true, // use_fast_gelu (sigmoid GELU) + ) + .map_err(|e| anyhow::anyhow!("Failed to load Granite 4 vision tower: {}", e))?; + + let spatial_offsets = [(0, 0), (0, 1), (1, 0), (1, 1)]; + let mut projectors = Vec::with_capacity(8); + for i in 0..4 { + projectors.push( + WindowQFormerProjector::from_weights( + &weights, + &format!("layerwise_projectors.{i}"), + Downsampler::MeanPool, + q, + w, + group_size, + bits, + ) + .map_err(|e| anyhow::anyhow!("Failed to load layerwise projector {i}: {}", e))?, + ); + } + for (i, (row_off, col_off)) in spatial_offsets.iter().enumerate() { + projectors.push( + WindowQFormerProjector::from_weights( + &weights, + &format!("spatial_projectors.{i}"), + Downsampler::Strided { + row_off: *row_off, + col_off: *col_off, + }, + q, + w, + group_size, + bits, + ) + .map_err(|e| anyhow::anyhow!("Failed to load spatial projector {i}: {}", e))?, + ); + } + + let image_newline = weights + .get("image_newline") + .map(|x| mlxcel_core::copy(x)) + .ok_or_else(|| anyhow::anyhow!("Missing image_newline weight"))?; + + // Stream routing from config. + let deepstack_map = parse_i32_pairs(full_config.get("deepstack_layer_map")); + let spatial_targets = parse_i32_list(full_config.get("spatial_target_layers")); + let spatial_vision_layer = full_config + .get("spatial_vision_layer") + .and_then(|v| v.as_i64()) + .unwrap_or(-1) as i32; + let use_spatial = full_config + .get("use_spatial_sampling") + .and_then(|v| v.as_bool()) + .unwrap_or(true); + + // Collect the deepstack vision taps; the spatial tap reuses one of them when + // present, otherwise it is appended. + let mut taps: Vec = deepstack_map.iter().map(|p| p[0]).collect(); + let spatial_tap_idx = match taps.iter().position(|&t| t == spatial_vision_layer) { + Some(idx) => idx, + None => { + taps.push(spatial_vision_layer); + taps.len() - 1 + } + }; + + let mut stream_specs: Vec<(usize, usize, usize)> = Vec::new(); + for (i, pair) in deepstack_map.iter().enumerate() { + stream_specs.push((i, i, pair[1] as usize)); + } + if use_spatial { + for (i, &layer) in spatial_targets.iter().enumerate() { + stream_specs.push((4 + i, spatial_tap_idx, layer as usize)); + } + } + + let image_token_index = full_config + .get("image_token_index") + .and_then(|v| v.as_i64()) + .unwrap_or(100352) as i32; + let pinpoints = parse_pinpoints(&full_config); + let image_size = vision_config.image_size as u32; + let patch_grid = (vision_config.image_size / vision_config.patch_size) as i32; // 24 + let feature_side = patch_grid / (w / q); // 24 / 2 = 12 + let base_tokens = feature_side * feature_side; // 144 + let processor = vision::processors::anyres::AnyResProcessor::new(pinpoints, image_size); + + // Text backbone consumes the weight map last. + let eos_token_id = text_args.eos_token_ids().first().copied().unwrap_or(100257); + let text_model = models::GraniteMoeHybridModel::from_weights(text_args, weights) + .map_err(|e| anyhow::anyhow!("Failed to load Granite 4 text model: {}", e))?; + + let vlm = vision::Granite4VisionVLModel::new( + text_model, + vision_tower, + projectors, + image_newline, + processor, + image_token_index, + taps, + stream_specs, + feature_side, + base_tokens, + eos_token_id, + ); + + Ok(LoadedModel::Granite4VisionVLM(vlm)) +} diff --git a/src/model_metadata.rs b/src/model_metadata.rs index 1c3e50ae..02ee0697 100644 --- a/src/model_metadata.rs +++ b/src/model_metadata.rs @@ -115,6 +115,7 @@ macro_rules! for_each_model_registration { Gemma4Unified => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("Gemma4 Unified cannot be loaded with LoRA adapters yet") }; LlavaVLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("LLaVA VLM cannot be loaded with LoRA adapters yet") }; GraniteVisionVLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("Granite Vision VLM cannot be loaded with LoRA adapters yet") }; + Granite4VisionVLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("Granite 4 Vision VLM cannot be loaded with LoRA adapters yet") }; LlavaBunnyVLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("LLaVA VLM cannot be loaded with LoRA adapters yet") }; AyaVisionVLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("Aya Vision VLM cannot be loaded with LoRA adapters yet") }; PaliGemmaVLM => { kind: Vlm, directory: Vlm, weight: None, adapter: Some("PaliGemma VLM cannot be loaded with LoRA adapters yet") }; diff --git a/src/models/detection.rs b/src/models/detection.rs index 1597c3f0..4564c329 100644 --- a/src/models/detection.rs +++ b/src/models/detection.rs @@ -238,6 +238,7 @@ pub fn get_model_type(model_path: &Path) -> Result { "moondream3" => Ok(ModelType::Moondream3VLM), "moondream2" | "moondream1" => Ok(ModelType::Moondream2VLM), "granite_vision" => Ok(ModelType::GraniteVisionVLM), + "granite4_vision" => Ok(ModelType::Granite4VisionVLM), "llava" | "llava_next" => { // The original IBM Granite Vision checkpoint ships as `llava_next` // with a `granite` text backbone; route it to the Granite VLM. diff --git a/src/models/granitemoehybrid.rs b/src/models/granitemoehybrid.rs index fe0b4158..7db58a9a 100644 --- a/src/models/granitemoehybrid.rs +++ b/src/models/granitemoehybrid.rs @@ -821,6 +821,11 @@ impl GraniteMoeHybridSharedMLP { enum FeedForward { Dense(GraniteMoeHybridMLP), + /// Dense feed-forward kept in the fused `shared_mlp` form (a single + /// `input_linear` producing gate+up). Used for quantized dense checkpoints, + /// where the fused quantized `input_linear` cannot be split into separate + /// `mlp.gate_proj`/`up_proj` at load. + DenseShared(GraniteMoeHybridSharedMLP), Moe { moe: GraniteMoeHybridMoE, shared: GraniteMoeHybridSharedMLP, @@ -831,6 +836,7 @@ impl FeedForward { fn forward(&self, x: &MlxArray) -> UniquePtr { match self { FeedForward::Dense(mlp) => mlp.forward(x), + FeedForward::DenseShared(shared) => shared.forward(x), FeedForward::Moe { moe, shared } => { let moe_out = moe.forward(x); let shared_out = shared.forward(x); @@ -946,6 +952,83 @@ impl GraniteMoeHybridLayerCache { // Granite 4.x model. +/// Multi-depth visual injection for Granite 4 Vision: a per-target packed +/// feature tensor `(total_image_tokens, hidden)` added at image-token positions +/// right before the named decoder layer during prefill. +pub struct HybridInjection<'a> { + /// `(1, seq)` mask (int/bool) that is nonzero at image-token positions. + pub visual_pos_mask: &'a MlxArray, + /// `(layer_index, features)` pairs; `features` is `(total_image_tokens, hidden)`. + pub targets: Vec<(usize, UniquePtr)>, +} + +/// Scatter-add `visual_embeds` `(n_img, hidden)` onto `h` `(1, seq, hidden)` at +/// the nonzero positions of `visual_pos_mask` `(1, seq)`, in order. Batch-1 only +/// (VLM prefill is unpadded batch-1). `pub(crate)` for the injection regression +/// test (guards the `slice_update`-is-functional footgun). +pub(crate) fn inject_at_positions( + h: &MlxArray, + visual_pos_mask: &MlxArray, + visual_embeds: &MlxArray, +) -> UniquePtr { + let h_shape = mlxcel_core::array_shape(h); + if h_shape[0] != 1 { + return mlxcel_core::copy(h); + } + let mask_1d = mlxcel_core::slice(visual_pos_mask, &[0, 0], &[1, h_shape[1]]); + let mask_1d = mlxcel_core::squeeze_axis(&mask_1d, 0); + let mask_i32 = mlxcel_core::astype(&mask_1d, mlxcel_core::dtype::INT32); + mlxcel_core::eval(&mask_i32); + + let seq_len = h_shape[1] as usize; + let mut positions = Vec::new(); + for i in 0..seq_len { + let val = mlxcel_core::slice(&mask_i32, &[i as i32], &[i as i32 + 1]); + mlxcel_core::eval(&val); + if mlxcel_core::item_i32(&val) != 0 { + positions.push(i as i32); + } + } + if positions.is_empty() { + return mlxcel_core::copy(h); + } + + let hidden = h_shape[2]; + let batch_h = mlxcel_core::slice(h, &[0, 0, 0], &[1, h_shape[1], hidden]); + let batch_h = mlxcel_core::squeeze_axis(&batch_h, 0); + + let idx_arr = mlxcel_core::from_slice_i32(&positions, &[positions.len() as i32]); + let current = mlxcel_core::take(&batch_h, &idx_arr, 0); + let n_img = positions.len() as i32; + let embeds = mlxcel_core::slice(visual_embeds, &[0, 0], &[n_img, hidden]); + let embeds = mlxcel_core::astype(&embeds, mlxcel_core::array_dtype(&batch_h)); + let updated = mlxcel_core::add(¤t, &embeds); + + // `slice_update` is functional (returns a new array); the result MUST be + // captured or the write is lost. Contiguous image runs (the single-image + // case) collapse to one slice_update; disjoint runs fall back to per-row. + let mut result = mlxcel_core::copy(&batch_h); + let contiguous = positions + .iter() + .enumerate() + .all(|(i, &p)| p == positions[0] + i as i32); + if contiguous { + let start = positions[0]; + result = + mlxcel_core::slice_update(&result, &updated, &[start, 0], &[start + n_img, hidden]); + } else { + for (local_idx, &pos) in positions.iter().enumerate() { + let val = mlxcel_core::slice( + &updated, + &[local_idx as i32, 0], + &[local_idx as i32 + 1, hidden], + ); + result = mlxcel_core::slice_update(&result, &val, &[pos, 0], &[pos + 1, hidden]); + } + } + mlxcel_core::expand_dims(&result, 0) +} + pub struct GraniteMoeHybridModel { config: ModelArgs, embed_tokens: UnifiedEmbedding, @@ -974,8 +1057,50 @@ impl GraniteMoeHybridModel { ) -> UniquePtr { // h = embed_tokens(x) * embedding_multiplier. let h = self.embed_tokens.forward(inputs); - let mut h = mlxcel_core::multiply_scalar(&h, self.embedding_multiplier); + let h = mlxcel_core::multiply_scalar(&h, self.embedding_multiplier); + self.run_from_embedded(h, caches, None) + } + + /// Raw token-embedding lookup, no `embedding_multiplier`. Used by the + /// Granite 4 Vision wrapper, which zeroes image slots then relies on + /// multi-depth injection ([`HybridInjection`]) for image content. + pub fn input_embeddings(&self, input_ids: &MlxArray) -> UniquePtr { + self.embed_tokens.forward(input_ids) + } + /// Forward from a pre-computed (un-scaled) embedding stream, applying + /// `embedding_multiplier` here, with optional multi-depth visual injection + /// (Granite 4 Vision). Routes through the model's own sequence state so + /// recurrent conv/SSM state and KV offsets advance exactly once. + pub(crate) fn forward_embeds_with_injection( + &self, + inputs_embeds: &MlxArray, + injection: Option<&HybridInjection>, + seq_id: Option, + ) -> UniquePtr { + let run = |internal: &mut [GraniteMoeHybridLayerCache]| { + let h = mlxcel_core::multiply_scalar(inputs_embeds, self.embedding_multiplier); + self.run_from_embedded(h, internal, injection) + }; + match seq_id { + Some(_) => self.sequence_state.with_or_create_sequence_state( + seq_id, + || GraniteMoeHybridModel::make_caches(self), + run, + ), + None => self.sequence_state.with_sequence_state(None, run), + } + } + + /// Run the decoder stack + norm + head from an already-`embedding_multiplier` + /// scaled hidden state. Injects visual features at masked positions *before* + /// each targeted layer when `injection` is present. + fn run_from_embedded( + &self, + mut h: UniquePtr, + caches: &mut [GraniteMoeHybridLayerCache], + injection: Option<&HybridInjection>, + ) -> UniquePtr { let shape = mlxcel_core::array_shape(&h); let seq_len = shape[1]; @@ -993,7 +1118,16 @@ impl GraniteMoeHybridModel { None }; - for (layer, cache) in self.layers.iter().zip(caches.iter_mut()) { + for (idx, (layer, cache)) in self.layers.iter().zip(caches.iter_mut()).enumerate() { + // Multi-depth visual injection happens before the targeted layer + // (layer 0 injection lands right after the embedding scale). + if let Some(inj) = injection { + for (target_layer, feats) in inj.targets.iter() { + if *target_layer == idx { + h = inject_at_positions(&h, inj.visual_pos_mask, feats); + } + } + } let mask = if layer.is_attention() { attn_mask.as_deref() } else { @@ -1227,6 +1361,26 @@ fn build_feed_forward( shared_intermediate: args.shared_intermediate_size as i32, }; Ok(FeedForward::Moe { moe, shared }) + } else if weights.contains_key(&format!("{prefix}.shared_mlp.input_linear.weight")) + && !weights.contains_key(&format!("{prefix}.mlp.gate_proj.weight")) + { + // Dense checkpoint whose fused `shared_mlp.input_linear` was left intact + // by sanitize (quantized: the fused weight cannot be split). Use it as-is. + Ok(FeedForward::DenseShared(GraniteMoeHybridSharedMLP { + input_linear: UnifiedLinear::from_weights( + weights, + &format!("{prefix}.shared_mlp.input_linear"), + gs, + bits, + )?, + output_linear: UnifiedLinear::from_weights( + weights, + &format!("{prefix}.shared_mlp.output_linear"), + gs, + bits, + )?, + shared_intermediate: args.shared_intermediate_size as i32, + })) } else { let mlp_prefix = format!("{prefix}.mlp"); Ok(FeedForward::Dense(GraniteMoeHybridMLP { @@ -1320,9 +1474,14 @@ fn sanitize_weights(mut weights: WeightMap, args: &ModelArgs) -> WeightMap { if weights.contains_key(&format!("{layer_prefix}.mlp.gate_proj.weight")) { continue; } + // Quantized fused `input_linear` cannot be split (slicing a quantized + // tensor is unsound); leave the whole fused `shared_mlp` intact so + // `build_feed_forward` loads it via `FeedForward::DenseShared`. + if weights.contains_key(&format!("{layer_prefix}.shared_mlp.input_linear.scales")) { + continue; + } let input_key = format!("{layer_prefix}.shared_mlp.input_linear.weight"); if weights.contains_key(&input_key) - && !weights.contains_key(&format!("{layer_prefix}.shared_mlp.input_linear.scales")) && let Some(input_weight) = weights.remove(&input_key) { let shape = mlxcel_core::array_shape(&input_weight); @@ -1413,6 +1572,43 @@ impl LanguageModel for GraniteMoeHybridModel { ) } + fn forward_with_embeddings( + &self, + input_ids: &MlxArray, + input_embeddings: Option<&MlxArray>, + _caches: &mut [KVCache], + _mask: Option<&MlxArray>, + ) -> UniquePtr { + match input_embeddings { + Some(embeds) => self.forward_embeds_with_injection(embeds, None, None), + None => self.sequence_state.with_sequence_state(None, |internal| { + self.forward_with_caches(input_ids, internal) + }), + } + } + + fn forward_with_embeddings_and_sequence_id( + &self, + input_ids: &MlxArray, + input_embeddings: Option<&MlxArray>, + seq_id: Option, + _caches: &mut [KVCache], + _mask: Option<&MlxArray>, + ) -> UniquePtr { + match input_embeddings { + Some(embeds) => self.forward_embeds_with_injection(embeds, None, seq_id), + None => self.sequence_state.with_or_create_sequence_state( + seq_id, + || GraniteMoeHybridModel::make_caches(self), + |internal| self.forward_with_caches(input_ids, internal), + ), + } + } + + fn embed_tokens(&self, input_ids: &MlxArray) -> Option> { + Some(self.input_embeddings(input_ids)) + } + fn supports_snapshot_reuse(&self) -> bool { true } diff --git a/src/models/granitemoehybrid_tests.rs b/src/models/granitemoehybrid_tests.rs index 04003d1c..d0e9003d 100644 --- a/src/models/granitemoehybrid_tests.rs +++ b/src/models/granitemoehybrid_tests.rs @@ -197,3 +197,50 @@ fn granite_hybrid_quantization_read_from_config() { assert_eq!(bf16.group_size(), 64); assert_eq!(bf16.bits(), 4); } + +/// Regression for the `slice_update`-is-functional footgun: `inject_at_positions` +/// MUST write the visual features into `h` at the image-token positions. A +/// dropped `slice_update` return leaves the image slots zero and silently makes +/// the VLM blind (Granite 4 Vision #541). Runs on the default device. +#[test] +fn inject_at_positions_writes_features_into_hidden_state() { + use super::granitemoehybrid::inject_at_positions; + + let (seq, hidden) = (5i32, 4i32); + // Image slots start zeroed, exactly as `get_input_embeddings` produces them. + let h = mlxcel_core::zeros(&[1, seq, hidden], mlxcel_core::dtype::FLOAT32); + // Image tokens occupy positions 1, 2, 3 (a contiguous single-image run). + let mask = mlxcel_core::from_slice_i32(&[0, 1, 1, 1, 0], &[1, seq]); + // Three feature rows, each a distinct non-zero constant. + let feats = mlxcel_core::from_slice_f32( + &[1.0, 1.0, 1.0, 1.0, 2.0, 2.0, 2.0, 2.0, 3.0, 3.0, 3.0, 3.0], + &[3, hidden], + ); + + let out = inject_at_positions(&h, &mask, &feats); + mlxcel_core::eval(&out); + let val = |pos: i32, col: i32| { + let c = mlxcel_core::slice(&out, &[0, pos, col], &[1, pos + 1, col + 1]); + mlxcel_core::eval(&c); + mlxcel_core::item_f32(&c) + }; + + // Non-image positions stay zero; image positions carry their feature row. + assert_eq!(val(0, 0), 0.0, "non-image position must be untouched"); + assert_eq!( + val(1, 0), + 1.0, + "image position 1 must receive feature row 0" + ); + assert_eq!( + val(2, 2), + 2.0, + "image position 2 must receive feature row 1" + ); + assert_eq!( + val(3, 3), + 3.0, + "image position 3 must receive feature row 2" + ); + assert_eq!(val(4, 0), 0.0, "non-image position must be untouched"); +} diff --git a/src/models/mod.rs b/src/models/mod.rs index a4c8885c..3d0267f4 100644 --- a/src/models/mod.rs +++ b/src/models/mod.rs @@ -233,61 +233,62 @@ pub use kokoro::KokoroModel; #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ModelType { // Standard Transformer models - Llama, // Llama 1/2/3, Mistral - Llama4, // Llama 4 (MoE) - Llama4VLM, // Llama 4 VLM (vision-language) - MllamaVLM, // Llama 3.2 Vision (mllama): tiled ViT + gated cross-attention - Qwen2, // Qwen 2/2.5 - Qwen3, // Qwen 3 - Qwen3Moe, // Qwen 3 MoE - Qwen3Next, // Qwen 3 with GatedDeltaNet - Qwen35, // Qwen 3.5 Hybrid (Transformer + GatedDeltaNet) - Qwen35VLM, // Qwen 3.5 VLM (Qwen3-VL vision + Qwen3.5 hybrid text) - Qwen35Moe, // Qwen 3.5 MoE Hybrid - Qwen35MoeVLM, // Qwen 3.5 MoE VLM - Gemma, // Gemma 1 - Gemma2, // Gemma 2 - Gemma3, // Gemma 3 (text-only) - Gemma4, // Gemma 4 text-only route - DiffusionGemma, // DiffusionGemma (block-diffusion on the Gemma 4 MoE backbone) - Gemma3VLM, // Gemma 3 VLM (vision-language) - Gemma4VLM, // Gemma 4 VLM (vision-language) - Gemma4Unified, // Gemma 4 Unified (encoder-free text + vision + audio) - LlavaVLM, // LLaVA (CLIP/SigLIP + Llama/Qwen2) - GraniteVisionVLM, // Granite Vision (SigLIP multi-tap + Granite text, AnyRes) - LlavaBunnyVLM, // LLaVA-Bunny (SigLIP + Qwen2) - AyaVisionVLM, // Aya Vision (SigLIP + Cohere2) - PaliGemmaVLM, // PaliGemma (SigLIP + Gemma) - PixtralVLM, // Pixtral (ViT w/ 2D RoPE + Mistral) - Mistral3VLM, // Mistral 3 VLM (Pixtral ViT + PatchMerger + Mistral) - Qwen2VL, // Qwen2-VL (custom ViT + Qwen2 w/ MRoPE) - Qwen25VL, // Qwen2.5-VL (windowed ViT + Qwen2 w/ MRoPE) - Qwen3VL, // Qwen3-VL (ViT + interleaved MRoPE + DeepStack) - Qwen3VLMoe, // Qwen3-VL-MoE (Qwen3-VL + MoE text backbone) - PaddleOcrVL, // PaddleOCR-VL (NaViT vision + ERNIE-4.5 w/ MRoPE) - Glm4v, // GLM-4V (GLM-4V ViT + GLM-4 text w/ sectioned MRoPE) - Glm4vMoe, // GLM-4V MoE (GLM-4V ViT + GLM-4 MoE text w/ MRoPE) - GlmOcr, // GLM-OCR (GLM-OCR ViT + GLM-4 text w/ full-width MRoPE) - YoutuVLM, // Youtu-VL (SigLIP2 windowed-attn + DeepSeek-V3-style MLA) - InternVLChatVLM, // InternVL (internvl_chat): InternViT + pixel-shuffle mlp1 + Qwen2 text - SmolVLM, // SmolVLM/SmolVLM2 (smolvlm): SigLIP + pixel-shuffle connector + SmolLM2 text - Idefics2, // Idefics2 (idefics2): SigLIP + perceiver-resampler connector + Mistral text - MiniCPMOVLM, // MiniCPM-o (dynamic SigLIP + resampler + Qwen3-VL text) - MiniCPMV46VLM, // MiniCPM-V 4.6 (SigLIP + VitMerger + Merger + Qwen3.5 text) - Moondream3VLM, // Moondream3 (custom ViT + custom text decoder, query/caption image path) - Moondream2VLM, // Moondream2 (SigLIP-style ViT + Phi text decoder + crop tiling) - Gemma3n, // Gemma 3n (text-only) - Gemma3nVLM, // Gemma 3n VLM (MobileNetV5 + Gemma3n) - Phi, // Phi 1/2 - Phi3, // Phi 3 - Phi4MMVLM, // Phi-4 Multimodal (SigLIP2 NaFlex + Phi4 text, image path only) - Phi4SigLipVLM, // Phi-4 reasoning vision (SigLIP2 NaFlex + Phi3-style text) - Phi3VLM, // Phi 3.5 Vision (CLIP + Phi3) - MolmoVLM, // Molmo v1 (CLIP ViT + attention pooling + OLMo-style text) - Molmo2VLM, // Molmo2 (custom ViT + attention pooling + Molmo2 text) - MolmoPointVLM, // Molmo-Point (custom ViT + point prediction + Molmo2 text) - Phi3Small, // Phi 3 Small - PhiMoe, // Phi MoE + Llama, // Llama 1/2/3, Mistral + Llama4, // Llama 4 (MoE) + Llama4VLM, // Llama 4 VLM (vision-language) + MllamaVLM, // Llama 3.2 Vision (mllama): tiled ViT + gated cross-attention + Qwen2, // Qwen 2/2.5 + Qwen3, // Qwen 3 + Qwen3Moe, // Qwen 3 MoE + Qwen3Next, // Qwen 3 with GatedDeltaNet + Qwen35, // Qwen 3.5 Hybrid (Transformer + GatedDeltaNet) + Qwen35VLM, // Qwen 3.5 VLM (Qwen3-VL vision + Qwen3.5 hybrid text) + Qwen35Moe, // Qwen 3.5 MoE Hybrid + Qwen35MoeVLM, // Qwen 3.5 MoE VLM + Gemma, // Gemma 1 + Gemma2, // Gemma 2 + Gemma3, // Gemma 3 (text-only) + Gemma4, // Gemma 4 text-only route + DiffusionGemma, // DiffusionGemma (block-diffusion on the Gemma 4 MoE backbone) + Gemma3VLM, // Gemma 3 VLM (vision-language) + Gemma4VLM, // Gemma 4 VLM (vision-language) + Gemma4Unified, // Gemma 4 Unified (encoder-free text + vision + audio) + LlavaVLM, // LLaVA (CLIP/SigLIP + Llama/Qwen2) + GraniteVisionVLM, // Granite Vision (SigLIP multi-tap + Granite text, AnyRes) + Granite4VisionVLM, // Granite 4 Vision (SigLIP + window-QFormer + Granite-4 hybrid) + LlavaBunnyVLM, // LLaVA-Bunny (SigLIP + Qwen2) + AyaVisionVLM, // Aya Vision (SigLIP + Cohere2) + PaliGemmaVLM, // PaliGemma (SigLIP + Gemma) + PixtralVLM, // Pixtral (ViT w/ 2D RoPE + Mistral) + Mistral3VLM, // Mistral 3 VLM (Pixtral ViT + PatchMerger + Mistral) + Qwen2VL, // Qwen2-VL (custom ViT + Qwen2 w/ MRoPE) + Qwen25VL, // Qwen2.5-VL (windowed ViT + Qwen2 w/ MRoPE) + Qwen3VL, // Qwen3-VL (ViT + interleaved MRoPE + DeepStack) + Qwen3VLMoe, // Qwen3-VL-MoE (Qwen3-VL + MoE text backbone) + PaddleOcrVL, // PaddleOCR-VL (NaViT vision + ERNIE-4.5 w/ MRoPE) + Glm4v, // GLM-4V (GLM-4V ViT + GLM-4 text w/ sectioned MRoPE) + Glm4vMoe, // GLM-4V MoE (GLM-4V ViT + GLM-4 MoE text w/ MRoPE) + GlmOcr, // GLM-OCR (GLM-OCR ViT + GLM-4 text w/ full-width MRoPE) + YoutuVLM, // Youtu-VL (SigLIP2 windowed-attn + DeepSeek-V3-style MLA) + InternVLChatVLM, // InternVL (internvl_chat): InternViT + pixel-shuffle mlp1 + Qwen2 text + SmolVLM, // SmolVLM/SmolVLM2 (smolvlm): SigLIP + pixel-shuffle connector + SmolLM2 text + Idefics2, // Idefics2 (idefics2): SigLIP + perceiver-resampler connector + Mistral text + MiniCPMOVLM, // MiniCPM-o (dynamic SigLIP + resampler + Qwen3-VL text) + MiniCPMV46VLM, // MiniCPM-V 4.6 (SigLIP + VitMerger + Merger + Qwen3.5 text) + Moondream3VLM, // Moondream3 (custom ViT + custom text decoder, query/caption image path) + Moondream2VLM, // Moondream2 (SigLIP-style ViT + Phi text decoder + crop tiling) + Gemma3n, // Gemma 3n (text-only) + Gemma3nVLM, // Gemma 3n VLM (MobileNetV5 + Gemma3n) + Phi, // Phi 1/2 + Phi3, // Phi 3 + Phi4MMVLM, // Phi-4 Multimodal (SigLIP2 NaFlex + Phi4 text, image path only) + Phi4SigLipVLM, // Phi-4 reasoning vision (SigLIP2 NaFlex + Phi3-style text) + Phi3VLM, // Phi 3.5 Vision (CLIP + Phi3) + MolmoVLM, // Molmo v1 (CLIP ViT + attention pooling + OLMo-style text) + Molmo2VLM, // Molmo2 (custom ViT + attention pooling + Molmo2 text) + MolmoPointVLM, // Molmo-Point (custom ViT + point prediction + Molmo2 text) + Phi3Small, // Phi 3 Small + PhiMoe, // Phi MoE // MoE models GptOss, @@ -436,6 +437,7 @@ pub const ALL_MODEL_TYPES: &[ModelType] = &[ ModelType::Gemma4Unified, ModelType::LlavaVLM, ModelType::GraniteVisionVLM, + ModelType::Granite4VisionVLM, ModelType::LlavaBunnyVLM, ModelType::AyaVisionVLM, ModelType::PaliGemmaVLM, @@ -764,6 +766,10 @@ impl ModelType { // ----- Other VLM (cross-family vision-language stacks) ----- ModelType::LlavaVLM => ("LLaVA (CLIP/SigLIP + Llama/Qwen2)", "Other VLM"), ModelType::GraniteVisionVLM => ("Granite Vision (SigLIP + Granite)", "Granite VLM"), + ModelType::Granite4VisionVLM => ( + "Granite 4 Vision (SigLIP + Granite 4 hybrid)", + "Granite VLM", + ), ModelType::LlavaBunnyVLM => ("LLaVA-Bunny (SigLIP + Qwen2)", "Other VLM"), ModelType::InternVLChatVLM => { ("InternVL (InternViT + pixel-shuffle + Qwen2)", "Other VLM") @@ -861,6 +867,7 @@ mod metadata_tests { Gemma4Unified, LlavaVLM, GraniteVisionVLM, + Granite4VisionVLM, LlavaBunnyVLM, AyaVisionVLM, PaliGemmaVLM, diff --git a/src/multimodal/vlm_runtime.rs b/src/multimodal/vlm_runtime.rs index 6ef3f861..7e160822 100644 --- a/src/multimodal/vlm_runtime.rs +++ b/src/multimodal/vlm_runtime.rs @@ -149,6 +149,11 @@ pub enum VlmPreparationSummary { image_blocks: usize, total_image_tokens: i32, }, + /// Granite 4 Vision expanded each `` into `num_image_tokens` copies. + Granite4Vision { + image_blocks: usize, + total_image_tokens: i32, + }, /// Kimi-VL expanded each `<|media_pad|>` into `(h/merge)*(w/merge)` /// media-placeholder tokens. KimiVL { @@ -1032,6 +1037,35 @@ where preparation, })) } + VlmRuntimeRef::Granite4Vision(granite4) => { + // AnyRes tiling at grid side 12 (patch 16, downsample /2); each image + // packs into `144 + H*(W+1)` tokens per injected stream. + let (pixel_values, infos) = granite4.processor.preprocess_with_tiles(images); + + let preparation = + crate::multimodal::granite_vision_prompt::insert_granite_vision_image_tokens( + prompt_tokens, + &infos, + granite4.image_token_index, + granite4.feature_side, + granite4.base_tokens, + ) + .map(|stats| VlmPreparationSummary::Granite4Vision { + image_blocks: stats.image_blocks, + total_image_tokens: stats.total_image_tokens, + }); + + let _ = active_caches; + let _ = image_cache_keys; + + let input_ids_arr = prompt_ids_array(prompt_tokens); + let embeddings = granite4.get_input_embeddings(&input_ids_arr, &pixel_values, &infos); + + Ok(Some(PreparedVlmEmbeddings { + embeddings, + preparation, + })) + } VlmRuntimeRef::KimiVL(kimi) => { // MoonViT native-resolution preprocessing: each image is patchified // into [num_patches, C, p, p] plus its (h, w) patch grid. diff --git a/src/server/chat_template.rs b/src/server/chat_template.rs index 512b059b..1668d7a9 100644 --- a/src/server/chat_template.rs +++ b/src/server/chat_template.rs @@ -682,6 +682,18 @@ fn preprocess_template(template: String) -> String { out = out.replace(v, ""); } } + // Some chat templates (e.g. Granite 4 Vision) write the image marker as a + // multi-line string literal — `"\n"` split across two source lines — + // which minijinja renders as empty/whitespace instead of the marker, so the + // per-model image-token expansion never finds a placeholder and mis-places + // the image. Normalize such literals to an escaped single-line form. + for marker in ["", "