From d249a9e90eda9c8588b4f093293be6c08f3bf9ec Mon Sep 17 00:00:00 2001 From: artem-spector Date: Tue, 21 Apr 2026 15:43:39 +0300 Subject: [PATCH] Add Granite 4.1 Vision as built-in multimodal model (#40282) Signed-off-by: Artem Spector Signed-off-by: artemspector Co-authored-by: artemspector Co-authored-by: Claude Opus 4.6 --- docs/models/supported_models.md | 1 + .../vision_language_multi_image.py | 33 + .../multimodal/generation/test_common.py | 37 + tests/models/registry.py | 4 + vllm/model_executor/models/granite4_vision.py | 930 ++++++++++++++++++ vllm/model_executor/models/registry.py | 4 + vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 2 + .../configs/granite4_vision.py | 94 ++ .../transformers_utils/processors/__init__.py | 2 + .../processors/granite4_vision.py | 78 ++ 11 files changed, 1186 insertions(+) create mode 100644 vllm/model_executor/models/granite4_vision.py create mode 100644 vllm/transformers_utils/configs/granite4_vision.py create mode 100644 vllm/transformers_utils/processors/granite4_vision.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 746980b8fe..f1ba46b707 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -560,6 +560,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + IE+ + VE+ | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ | | `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + IE+ + VE+ | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ | | `GlmOcrForConditionalGeneration` | GLM-OCR | T + IE+ | `zai-org/GLM-OCR`, etc. | ✅︎ | ✅︎ | +| `Granite4VisionForConditionalGeneration` | Granite 4 Vision | T + IE+ | `ibm-granite/granite-4.1-3b-vision`, etc. | ✅︎ | ✅︎ | | `GraniteSpeechForConditionalGeneration` | Granite Speech | T + A | `ibm-granite/granite-speech-3.3-8b` | ✅︎ | ✅︎ | | `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision-Instruct-3B | T + I+ + V+ | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | | | | `HCXVisionV2ForCausalLM` | HyperCLOVAX-SEED-Think-32B | T + I+ + V+ | `naver-hyperclovax/HyperCLOVAX-SEED-Think-32B` | | | diff --git a/examples/offline_inference/vision_language_multi_image.py b/examples/offline_inference/vision_language_multi_image.py index 1a3b7fd954..1b68a23b3b 100755 --- a/examples/offline_inference/vision_language_multi_image.py +++ b/examples/offline_inference/vision_language_multi_image.py @@ -310,6 +310,38 @@ def load_gemma3(question: str, image_urls: list[str]) -> ModelRequestData: ) +def load_granite4_vision(question: str, image_urls: list[str]) -> ModelRequestData: + model_name = "ibm-granite/granite-vision-4.1-4b" + engine_args = EngineArgs( + model=model_name, + max_model_len=4096, + max_num_seqs=16, + limit_mm_per_prompt={"image": len(image_urls)}, + ) + + placeholders = [{"type": "image", "image": url} for url in image_urls] + messages = [ + { + "role": "user", + "content": [ + *placeholders, + {"type": "text", "text": question}, + ], + } + ] + + processor = AutoProcessor.from_pretrained(model_name) + prompt = processor.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + + return ModelRequestData( + engine_args=engine_args, + prompt=prompt, + image_data=[fetch_image(url) for url in image_urls], + ) + + def load_h2ovl(question: str, image_urls: list[str]) -> ModelRequestData: model_name = "h2oai/h2ovl-mississippi-800m" @@ -1487,6 +1519,7 @@ model_example_map = { "deepseek_ocr": load_deepseek_ocr, "exaone4_5": load_exaone4_5, "gemma3": load_gemma3, + "granite4_vision": load_granite4_vision, "h2ovl_chat": load_h2ovl, "hunyuan_vl": load_hunyuan_vl, "hyperclovax_seed_vision": load_hyperclovax_seed_vision, diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 1147ccef35..29680bafd5 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -86,6 +86,29 @@ COMMON_BROADCAST_SETTINGS = { # which cases would be selected and deselected by pytest. In general, # this is a good idea for checking your command first, since tests are slow. + +def _granite4_vision_vllm_to_hf_output(vllm_output, model): + """Post-processor for granite4_vision vLLM output. + + Self-contained to avoid calling AutoConfig/AutoTokenizer without + trust_remote_code (needed while the model is not in upstream HF). + """ + output_ids, output_str, out_logprobs = vllm_output + mm_token_id = 100352 + hf_output_ids = [ + token_id + for idx, token_id in enumerate(output_ids) + if token_id != mm_token_id or idx == 0 or output_ids[idx - 1] != mm_token_id + ] + hf_output_str = ( + output_str[1:] if output_str and output_str[0] == " " else output_str + ) + eos_token_id = 100257 + if hf_output_ids and hf_output_ids[-1] == eos_token_id: + hf_output_str = hf_output_str + "<|end_of_text|>" + return hf_output_ids, hf_output_str, out_logprobs + + VLM_TEST_SETTINGS = { #### Core tests to always run in the CI "llava": VLMTestInfo( @@ -492,6 +515,20 @@ VLM_TEST_SETTINGS = { auto_cls=AutoModelForImageTextToText, marks=[large_gpu_mark(min_gb=32)], ), + "granite4_vision": VLMTestInfo( + models=["ibm-granite/granite-vision-4.1-4b"], + test_type=(VLMTestType.IMAGE), + prompt_formatter=lambda img_prompt: f"<|user|>\n{img_prompt}\n<|assistant|>\n", + max_model_len=8192, + auto_cls=AutoModelForImageTextToText, + vllm_output_post_proc=_granite4_vision_vllm_to_hf_output, + image_size_factors=[(1.0,)], + vllm_runner_kwargs={ + "enable_lora": True, + "max_lora_rank": 256, + "default_mm_loras": {"image": "ibm-granite/granite-vision-4.1-4b"}, + }, + ), "h2ovl": VLMTestInfo( models=[ "h2oai/h2ovl-mississippi-800m", diff --git a/tests/models/registry.py b/tests/models/registry.py index f5968438cb..ea1f019056 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -899,6 +899,10 @@ _MULTIMODAL_EXAMPLE_MODELS = { "zai-org/GLM-ASR-Nano-2512", min_transformers_version="5.0.0", ), + "Granite4VisionForConditionalGeneration": _HfExamplesInfo( + "ibm-granite/granite-vision-4.1-4b", + is_available_online=False, + ), "GraniteVision": _HfExamplesInfo("ibm-granite/granite-vision-3.3-2b"), "GraniteSpeechForConditionalGeneration": _HfExamplesInfo( "ibm-granite/granite-speech-3.3-2b", diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py new file mode 100644 index 0000000000..147f02eced --- /dev/null +++ b/vllm/model_executor/models/granite4_vision.py @@ -0,0 +1,930 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""vLLM implementation of Granite 4 Vision. + +Uses GraniteForCausalLM as the language backbone with SigLIP vision encoder +and deepstack feature injection via WindowQFormer projectors. + +LoRA support: use --enable-lora --default-mm-loras for LM-only LoRA adapters. +""" + +import math +from collections.abc import Iterable, Mapping +from fractions import Fraction +from itertools import islice + +import torch +import torch.nn as nn +from transformers import BatchFeature +from transformers.models.blip_2.configuration_blip_2 import Blip2QFormerConfig +from transformers.models.llava_next.modeling_llava_next import ( + get_anyres_image_grid_shape, + image_size_to_num_patches, + unpad_image, +) + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import CacheConfig, VllmConfig +from vllm.distributed.parallel_state import get_pp_group +from vllm.logger import init_logger +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.granite import GraniteForCausalLM, GraniteModel +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.llava import LlavaDummyInputsBuilder +from vllm.model_executor.models.llava_next import ( + BaseLlavaNextMultiModalProcessor, + LlavaNextImageEmbeddingInputs, + LlavaNextImageInputs, + LlavaNextImagePixelInputs, + LlavaNextProcessingInfo, +) +from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.model_executor.models.siglip import SiglipVisionModel +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + PPMissingLayer, + WeightsMapper, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFieldConfig +from vllm.sequence import IntermediateTensors + +from .blip2 import Blip2QFormerModel + +logger = init_logger(__name__) + + +# --------------------------------------------------------------------------- +# Downsampler modules (translated from HF downsampling.py) +# --------------------------------------------------------------------------- + + +class InterpolateDownsampler: + """Spatial downsampling via area interpolation.""" + + def __init__(self, config, mode="area"): + self.orig_image_side = ( + config.vision_config.image_size // config.vision_config.patch_size + ) + self.new_image_side = int( + self.orig_image_side * Fraction(config.downsample_rate) + ) + self.mode = mode + + def __call__(self, image_features: torch.Tensor) -> torch.Tensor: + batch_size, _, dim = image_features.size() + up_shape = [batch_size, self.orig_image_side, self.orig_image_side, dim] + large = image_features.view(up_shape).permute(0, 3, 1, 2) + small = torch.nn.functional.interpolate( + large, + size=(self.new_image_side, self.new_image_side), + mode=self.mode, + ) + return small.permute(0, 2, 3, 1).flatten(1, 2) + + +class SpatialOffsetDownsampler: + """Sample one position from each 2x2 block (offset 0-3 = TL/TR/BL/BR).""" + + def __init__(self, config, offset: int = 0): + self.orig_image_side = ( + config.vision_config.image_size // config.vision_config.patch_size + ) + self.new_image_side = self.orig_image_side // 2 + offsets = [(0, 0), (0, 1), (1, 0), (1, 1)] + self.offset_h, self.offset_w = offsets[offset] + + def __call__(self, image_features: torch.Tensor) -> torch.Tensor: + B, _, C = image_features.shape + features_2d = image_features.reshape( + B, self.orig_image_side, self.orig_image_side, C + ) + n = self.new_image_side + blocks = features_2d.reshape(B, n, 2, n, 2, C) + sampled = blocks[:, :, self.offset_h, :, self.offset_w, :] + return sampled.reshape(B, -1, C) + + +class WindowQFormerDownsampler(nn.Module): + """Window-based QFormer downsampler (matches HF downsampling.py exactly).""" + + def __init__( + self, + config, + quant_config: QuantizationConfig | None = None, + cache_config: CacheConfig | None = None, + spatial_offset: int | None = None, + prefix: str = "", + ): + super().__init__() + llm_hidden_size = config.text_config.hidden_size + vision_hidden_size = config.vision_config.hidden_size + + self.dropout = nn.Dropout(config.projector_dropout) + + if spatial_offset is not None: + self.downsampler = SpatialOffsetDownsampler(config, offset=spatial_offset) + else: + self.downsampler = InterpolateDownsampler(config) + + qformer_config = Blip2QFormerConfig( + hidden_size=vision_hidden_size, + num_attention_heads=vision_hidden_size // 64, + intermediate_size=3072, + num_hidden_layers=1, + encoder_hidden_size=vision_hidden_size, + cross_attention_frequency=1, + max_position_embeddings=2048, + use_qformer_text_input=False, + ) + self.qformer = Blip2QFormerModel( + qformer_config, + quant_config=quant_config, + cache_config=cache_config, + prefix=f"{prefix}.qformer", + ) + + self.image_side = ( + config.vision_config.image_size // config.vision_config.patch_size + ) + q, w = config.downsample_rate.split("/") + self.query_side, self.window_side = int(q), int(w) + self.query_length = self.query_side**2 + + embed_std = 1 / math.sqrt(vision_hidden_size) + self.norm = nn.LayerNorm(vision_hidden_size, eps=1e-6) + self.query = nn.Parameter( + torch.randn(1, self.query_length, vision_hidden_size) * embed_std + ) + self.image_positions = nn.Parameter( + torch.randn(1, self.window_side**2, vision_hidden_size) * embed_std + ) + self.out_linear = nn.Linear(vision_hidden_size, llm_hidden_size, bias=True) + + def _win(self, x: torch.Tensor, side: int, win: int) -> torch.Tensor: + """(B, side*side, C) → (B*n*n, win*win, C) where n=side//win.""" + B, _, C = x.shape + n = side // win + return ( + x.view(B, side, side, C) + .view(B, n, win, n, win, C) + .transpose(2, 3) + .flatten(0, 2) + .flatten(1, 2) + ) + + def _unwin(self, xw: torch.Tensor, n: int, win: int) -> torch.Tensor: + """(B*n*n, win*win, C) → (B, (n*win)^2, C).""" + Bnn, _, C = xw.shape + B = Bnn // (n * n) + side = n * win + return ( + xw.view(B, n, n, win, win, C) + .transpose(2, 3) + .contiguous() + .view(B, side, side, C) + .flatten(1, 2) + ) + + def forward(self, image_features: torch.Tensor) -> torch.Tensor: + B, HW, C = image_features.shape + assert self.image_side * self.image_side == HW + n = self.image_side // self.window_side + + image_features = self.norm(image_features) + enc = self._win(image_features, self.image_side, self.window_side) + + downsampled = self.downsampler(image_features) + new_side = n * self.query_side + downsampled_w = self._win(downsampled, new_side, self.query_side) + + query_embeds = self.query + downsampled_w + encoder_embeds = self.dropout(enc + self.image_positions) + out_w = self.qformer( + query_embeds=query_embeds, + encoder_hidden_states=encoder_embeds, + ) + + out = self._unwin(out_w, n=n, win=self.query_side) + out = self.dropout(out) + return self.out_linear(out) + + +# --------------------------------------------------------------------------- +# LLM subclasses with deepstack injection in the layer loop +# --------------------------------------------------------------------------- + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": 0, + "intermediate_tensors": 0, + "inputs_embeds": 0, + "deepstack_input_embeds": 0, + } +) +class Granite4VisionLLMModel(GraniteModel): + """GraniteModel with deepstack feature injection in the layer loop.""" + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + deepstack_input_embeds: IntermediateTensors | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + hidden_states = hidden_states * self.config.embedding_multiplier + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + # Recover deepstack features forwarded from the previous PP rank. + if deepstack_input_embeds is None: + ds_keys = [ + k for k in intermediate_tensors.tensors if k.startswith("ds_") + ] + if ds_keys: + deepstack_input_embeds = IntermediateTensors( + {k: intermediate_tensors[k] for k in ds_keys} + ) + + for layer_idx, layer in islice( + enumerate(self.layers), self.start_layer, self.end_layer + ): + if deepstack_input_embeds is not None: + key = f"ds_{layer_idx}" + if key in deepstack_input_embeds.tensors: + feat = deepstack_input_embeds[key] + # Resize to match hidden_states in case of CUDA graph padding + num_tokens = hidden_states.size(0) + buf_len = feat.shape[0] + if buf_len != num_tokens: + feat = torch.nn.functional.pad( + feat[:num_tokens], + (0, 0, 0, max(0, num_tokens - buf_len)), + ) + hidden_states = hidden_states + feat + hidden_states = layer(positions, hidden_states) + + if not get_pp_group().is_last_rank: + # Forward hidden_states and any deepstack features for later ranks. + it = {"hidden_states": hidden_states} + if deepstack_input_embeds is not None: + remaining = { + k: v + for k, v in deepstack_input_embeds.tensors.items() + if int(k.split("_")[1]) >= self.end_layer + } + it.update(remaining) + return IntermediateTensors(it) + + hidden_states = self.norm(hidden_states) + return hidden_states + + +class Granite4VisionLLMForCausalLM(GraniteForCausalLM): + """GraniteForCausalLM backed by Granite4VisionLLMModel.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + nn.Module.__init__(self) + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + self.model = Granite4VisionLLMModel( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + if config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + logit_scale = getattr(config, "logit_scale", 1.0) + if hasattr(config, "logits_scaling"): + logit_scale /= config.logits_scaling + self.logits_processor = LogitsProcessor( + config.vocab_size, scale=logit_scale + ) + else: + self.lm_head = PPMissingLayer() + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors: + tensors = super().make_empty_intermediate_tensors(batch_size, dtype, device) + # Include deepstack buffers so non-first PP ranks receive them. + # _ds_layer_indices is set directly on this instance by the outer model. + for llm_layer in getattr(self, "_ds_layer_indices", []): + tensors.tensors[f"ds_{llm_layer}"] = torch.zeros( + (batch_size, self.config.hidden_size), dtype=dtype, device=device + ) + return tensors + + +# --------------------------------------------------------------------------- +# Processing info / processor (reuses LlavaNext patterns) +# --------------------------------------------------------------------------- + + +class Granite4VisionProcessingInfo(LlavaNextProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config() + + def get_hf_processor(self, **kwargs): + return self.ctx.get_hf_processor(**kwargs) + + def get_num_image_tokens( + self, + *, + image_width: int, + image_height: int, + ) -> int: + hf_config = self.get_hf_config() + vision_encoder_info = self.get_vision_encoder_info() + + # After QFormer downsampling, patch grid is scaled by downsample_rate + ds_rate = Fraction(hf_config.downsample_rate) + patch_grid = vision_encoder_info.get_patch_grid_length() # 24 for 384/16 + downsampled_grid = int(patch_grid * ds_rate) # 12 for rate 4/8 + + # Base feature: downsampled_grid^2 + base_feature_size = downsampled_grid * downsampled_grid + + num_patch_height, num_patch_width = get_anyres_image_grid_shape( + image_size=(image_height, image_width), + grid_pinpoints=hf_config.image_grid_pinpoints, + patch_size=vision_encoder_info.get_image_size(), + ) + + ( + unpadded_feature_size, + newline_feature_size, + ) = self._get_num_unpadded_features( + original_height=image_height, + original_width=image_width, + npatches=downsampled_grid, + num_patch_height=num_patch_height, + num_patch_width=num_patch_width, + ) + + return unpadded_feature_size + newline_feature_size + base_feature_size + + +class Granite4VisionMultiModalProcessor( + BaseLlavaNextMultiModalProcessor[Granite4VisionProcessingInfo] +): + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + return dict( + pixel_values=MultiModalFieldConfig.batched("image"), + image_sizes=MultiModalFieldConfig.batched("image"), + ) + + +# --------------------------------------------------------------------------- +# Top-level model +# --------------------------------------------------------------------------- + + +@MULTIMODAL_REGISTRY.register_processor( + Granite4VisionMultiModalProcessor, + info=Granite4VisionProcessingInfo, + dummy_inputs=LlavaDummyInputsBuilder, +) +class Granite4VisionForConditionalGeneration( + nn.Module, SupportsLoRA, SupportsMultiModal, SupportsPP +): + """vLLM implementation of Granite 4 Vision. + + Architecture: + - SigLIP vision tower -> WindowQFormerDownsampler projectors + - Deepstack: 4 vision layers projected and injected at 4 LLM layers + - Spatial: 4 offset groups from last vision layer injected at 4 more LLM layers + - Granite language backbone with embedding_multiplier + - logits_scaling via LogitsProcessor + + The outer model runs the LLM layer loop directly (like HF does) to inject + deepstack features. This avoids wrapping the inner model and keeps weight + loading simple. + + LoRA support: + - Full merge: --hf-overrides '{"adapter_path": "path/to/lora"}' merges + LM-only LoRA deltas at load time (W += scaling * B @ A). + - Native LoRA: --enable-lora --default-mm-loras '{"image": "path/to/lora"}' + lets vLLM runtime serve LM LoRA per-request. + Both modes expect a LM-only adapter (no modules_to_save). + """ + + # LoRA class attributes (matches GraniteForCausalLM) + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + embedding_modules = {} + + # Weight mapping: HF checkpoint -> vLLM parameter names + # HF: model.language_model.layers.0... + # vLLM: language_model.model.layers.0... + # (because GraniteForCausalLM.model = GraniteModel) + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "model.language_model.": "language_model.model.", + "model.layerwise_projectors.": "layerwise_projectors.", + "model.spatial_projectors.": "spatial_projectors.", + "model.image_newline": "image_newline", + "model.vision_tower.": "vision_tower.", + "lm_head.": "language_model.lm_head.", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return "" + raise ValueError(f"Only image modality is supported, got {modality}") + + def get_mm_mapping(self) -> MultiModelKeys: + return MultiModelKeys.from_string_field( + language_model="language_model", + connector=["layerwise_projectors", "spatial_projectors"], + tower_model="vision_tower", + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.vllm_config = vllm_config + + # ----- Vision tower + projectors (marked as tower) ----- + with self._mark_tower_model(vllm_config, "image"): + # Do NOT use init_vision_tower_for_llava here — it truncates the + # encoder to vision_feature_layer depth. Deepstack needs ALL hidden + # states (deepstack_layer_map uses negative indices into the full + # encoder output list). + self.vision_tower = SiglipVisionModel( + config.vision_config, + quant_config=quant_config, + require_post_norm=False, + prefix=maybe_prefix(prefix, "vision_tower"), + ) + + # image_newline parameter + if config.use_image_newline_parameter: + self.image_newline = nn.Parameter( + torch.empty(config.text_config.hidden_size) + ) + else: + self.image_newline = None + + cache_config = vllm_config.cache_config + + # Deepstack projectors: one per (vision_layer, llm_layer) pair + self.layerwise_projectors = nn.ModuleList( + [ + WindowQFormerDownsampler( + config, + quant_config=quant_config, + cache_config=cache_config, + prefix=maybe_prefix(prefix, f"layerwise_projectors.{i}"), + ) + for i in range(len(config.deepstack_layer_map)) + ] + ) + + # Spatial projectors: 4 offset groups + self.spatial_projectors = None + if config.use_spatial_sampling: + self.spatial_projectors = nn.ModuleList( + [ + WindowQFormerDownsampler( + config, + quant_config=quant_config, + cache_config=cache_config, + spatial_offset=i, + prefix=maybe_prefix(prefix, f"spatial_projectors.{i}"), + ) + for i in range(4) + ] + ) + + # ----- Language model (marked as LM) ----- + with self._mark_language_model(vllm_config): + self.language_model = Granite4VisionLLMForCausalLM( + vllm_config=vllm_config.with_hf_config(config.text_config), + prefix=maybe_prefix(prefix, "language_model"), + ) + + self.make_empty_intermediate_tensors = ( + self.language_model.make_empty_intermediate_tensors + ) + + # Store config values we need + self._deepstack_layer_map = config.deepstack_layer_map # [[-19, 9], ...] + self._use_spatial_sampling = getattr(config, "use_spatial_sampling", False) + self._spatial_vision_layer = getattr(config, "spatial_vision_layer", -1) + self._spatial_target_layers = getattr(config, "spatial_target_layers", []) + self._vision_feature_select_strategy = getattr( + config, "vision_feature_select_strategy", "full" + ) + self._downsample_rate = Fraction(config.downsample_rate) + + # Ordered list of LLM layer indices for each deepstack level. + # Pre-populated from config so it's available during CUDA graph capture + # (before any embed_multimodal call). + self._ds_layer_indices: list[int] = [ + llm_layer for _, llm_layer in config.deepstack_layer_map + ] + list(getattr(config, "spatial_target_layers", [])) + + # Share ds_layer_indices with the LLM causal model so + # make_empty_intermediate_tensors includes the correct keys + # (its self.config is text_config, no deepstack_layer_map). + self.language_model._ds_layer_indices = self._ds_layer_indices + + # Pre-allocated persistent GPU buffers for deepstack features. + # Written via .copy_() in embed_input_ids(), read by forward() via a + # slice. Because the buffer address is fixed, CUDA graph replay sees + # the updated values written just before each prefill. + # Shape: (max_num_batched_tokens, lm_hidden_size) per level. + n_layerwise = len(config.deepstack_layer_map) + n_spatial = len(getattr(config, "spatial_target_layers", [])) + num_ds_levels = n_layerwise + n_spatial + lm_hidden = config.text_config.hidden_size + max_tokens = vllm_config.scheduler_config.max_num_batched_tokens + # Allocated on CPU first; moved to GPU in embed_input_ids on first use. + self._ds_buffers: list[torch.Tensor] = [ + torch.zeros(max_tokens, lm_hidden) for _ in range(num_ds_levels) + ] + self._ds_num_tokens: int = 0 # tokens written in last embed_input_ids call + + # ----- Vision feature extraction ----- + + def _get_vision_hidden_states( + self, pixel_values: torch.Tensor + ) -> list[torch.Tensor]: + """Run vision tower and return all hidden states (including input embeddings). + + Uses SiglipEncoder's built-in return_all_hidden_states support. + Returns list[Tensor] where index 0 = embeddings, index i = after layer i-1. + """ + vt = self.vision_tower + vm = vt.vision_model if hasattr(vt, "vision_model") else vt + + hidden_states = vm.embeddings(pixel_values) + all_hidden_states = vm.encoder( + inputs_embeds=hidden_states, + return_all_hidden_states=True, + ) + return all_hidden_states + + def _pack_and_unpad_image_features( + self, + image_features: list[torch.Tensor] | tuple[torch.Tensor, ...], + image_sizes: torch.Tensor, + ) -> list[torch.Tensor]: + """Reshape, unpad, and pack image features. + + Matches HF Granite4VisionModel.pack_and_unpad_image_features exactly. + """ + config = self.config + ds_rate = self._downsample_rate + new_image_features = [] + + for image_idx, image_feature in enumerate(image_features): + if image_feature.shape[0] > 1: + # Multi-patch: first is base, rest are high-res + base_image_feature = image_feature[0] + image_feature = image_feature[1:] + + height = width = ( + config.vision_config.image_size // config.vision_config.patch_size + ) + # After QFormer downsampling + height = int(height * ds_rate) + width = int(width * ds_rate) + + num_patch_height, num_patch_width = get_anyres_image_grid_shape( + image_sizes[image_idx], + config.image_grid_pinpoints, + config.vision_config.image_size, + ) + + image_feature = image_feature.view( + num_patch_height, num_patch_width, height, width, -1 + ) + image_feature = ( + image_feature.permute(4, 0, 2, 1, 3) + .contiguous() + .flatten(1, 2) + .flatten(2, 3) + ) + image_feature = unpad_image(image_feature, image_sizes[image_idx]) + + if self.image_newline is not None: + image_feature = torch.cat( + ( + image_feature, + self.image_newline[:, None, None] + .expand(*image_feature.shape[:-1], 1) + .to(image_feature.device, image_feature.dtype), + ), + dim=-1, + ) + + image_feature = image_feature.flatten(1, 2).transpose(0, 1) + image_feature = torch.cat((base_image_feature, image_feature), dim=0) + else: + image_feature = image_feature[0] + if self.image_newline is not None: + image_feature = torch.cat( + (image_feature, self.image_newline[None].to(image_feature)), + dim=0, + ) + + new_image_features.append(image_feature) + + return new_image_features + + def _get_all_layer_features( + self, + pixel_values: torch.Tensor, + image_sizes: torch.Tensor, + ) -> tuple[list[int], list[torch.Tensor]]: + """Extract deepstack + spatial features for all levels. + + Returns: + llm_layer_indices: ordered list of target LLM layer indices + per_image_packed: one tensor per image, shape + (num_tokens_i, lm_hidden_size * num_levels), + all levels packed on dim=-1. + + Packing on dim=-1 means the framework's token-level slicing for + chunked prefill preserves all levels intact. + """ + select_strategy = self._vision_feature_select_strategy + + image_num_patches = [ + image_size_to_num_patches( + image_size=imsize, + grid_pinpoints=self.config.image_grid_pinpoints, + patch_size=self.config.vision_config.image_size, + ) + for imsize in image_sizes + ] + + if pixel_values.dim() == 5: + pixel_values = torch.cat( + [pv[:np_] for pv, np_ in zip(pixel_values, image_num_patches)], + dim=0, + ) + + all_hidden_states = self._get_vision_hidden_states(pixel_values) + + # Collect per-level: (llm_layer, [per_image_tensor, ...]) + levels: list[tuple[int, list[torch.Tensor]]] = [] + + for proj_idx, (vision_layer, llm_layer) in enumerate(self._deepstack_layer_map): + selected = all_hidden_states[vision_layer] + if select_strategy == "default": + selected = selected[:, 1:] + projected = self.layerwise_projectors[proj_idx](selected) + per_image = self._pack_and_unpad_image_features( + torch.split(projected, image_num_patches, dim=0), image_sizes + ) + levels.append((llm_layer, per_image)) + + if self._use_spatial_sampling and self.spatial_projectors is not None: + spatial_hidden = all_hidden_states[self._spatial_vision_layer] + if select_strategy == "default": + spatial_hidden = spatial_hidden[:, 1:] + for group_idx, llm_layer in enumerate(self._spatial_target_layers): + projected = self.spatial_projectors[group_idx](spatial_hidden) + per_image = self._pack_and_unpad_image_features( + torch.split(projected, image_num_patches, dim=0), image_sizes + ) + levels.append((llm_layer, per_image)) + + llm_layer_indices = [llm_layer for llm_layer, _ in levels] + num_images = len(image_sizes) + per_image_packed = [ + torch.cat([levels[lvl][1][img] for lvl in range(len(levels))], dim=-1) + for img in range(num_images) + ] + + return llm_layer_indices, per_image_packed + + # ----- Multimodal interface ----- + + def _parse_and_validate_image_input( + self, **kwargs: object + ) -> LlavaNextImageInputs | None: + pixel_values = kwargs.pop("pixel_values", None) + image_sizes = kwargs.pop("image_sizes", None) + image_embeds = kwargs.pop("image_embeds", None) + + if pixel_values is None and image_embeds is None: + return None + + if pixel_values is not None: + expected_h = expected_w = self.config.vision_config.image_size + return LlavaNextImagePixelInputs( + type="pixel_values", + pixel_values=pixel_values, + image_sizes=image_sizes, + resolve_bindings={"h": expected_h, "w": expected_w}, + ) + + if image_embeds is not None: + return LlavaNextImageEmbeddingInputs( + type="image_embeds", + data=image_embeds, + ) + + raise AssertionError("Unreachable") + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: + """Run vision tower and return per-image packed feature tensors. + + Each returned tensor has shape (num_tokens_i, lm_hidden_size * num_levels) + with all deepstack levels packed on dim=-1. The framework caches these + tensors and slices along dim=0 for chunked prefill — all levels survive + intact because slicing is token-wise, not feature-wise. + + embed_input_ids() splits the packed tensor back into per-level buffers. + """ + image_input = self._parse_and_validate_image_input(**kwargs) + if image_input is None: + return [] + + if image_input["type"] == "image_embeds": + return [image_input["data"]] + + pixel_values = image_input["pixel_values"] + image_sizes = image_input.get("image_sizes") + + if isinstance(pixel_values, list): + pixel_values = torch.cat(pixel_values, dim=0) + + llm_layer_indices, per_image_packed = self._get_all_layer_features( + pixel_values, image_sizes + ) + self._ds_layer_indices = llm_layer_indices + return per_image_packed + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = True, + ) -> torch.Tensor: + """Merge text and vision embeddings, apply embedding_multiplier. + + HF flow: + 1. inputs_embeds = embed_tokens(input_ids) + 2. inputs_embeds.masked_fill(vision_mask, 0.0) + 3. hidden_states = inputs_embeds * embedding_multiplier + 4. layer loop injects deepstack features at target layers + + multimodal_embeddings contains packed tensors from embed_multimodal(): + shape (num_tokens_i, lm_hidden_size * num_levels). We split on dim=-1 + to get per-level features, build batch-sized buffers (zero at text + positions), and store in self._ds_features for forward(). + """ + lm_inner = self.language_model.model + + has_vision = ( + multimodal_embeddings is not None + and is_multimodal is not None + and len(multimodal_embeddings) > 0 + and is_multimodal.any() + ) + + if not has_vision: + self._ds_num_tokens = 0 + embeds = lm_inner.embed_input_ids(input_ids) + return embeds * lm_inner.config.embedding_multiplier + + # 1. Text embeddings + text_embeds = lm_inner.embed_input_ids(input_ids) + + # 2. Zero image positions (matches HF masked_fill(vision_mask, 0.0)) + text_embeds[is_multimodal] = 0.0 + + # 3. Apply embedding_multiplier + inputs_embeds = text_embeds * lm_inner.config.embedding_multiplier + + # 4. Split packed tensors into per-level features and build buffers. + # multimodal_embeddings is a list of per-image packed tensors + # (possibly a chunk slice from the framework's encoder cache). + # Concatenate along token dim → (total_mm_tokens, lm_h * num_levels). + N, lm_h = inputs_embeds.shape + all_packed = torch.cat( + [t.to(dtype=inputs_embeds.dtype) for t in multimodal_embeddings], + dim=0, + ) + level_features = all_packed.split(lm_h, dim=-1) # num_levels tensors + + # Ensure persistent buffers are on the right device/dtype (first call). + buf0 = self._ds_buffers[0] + if buf0.device != inputs_embeds.device or buf0.dtype != inputs_embeds.dtype: + self._ds_buffers = [ + b.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype) + for b in self._ds_buffers + ] + + for level_idx in range(len(self._ds_layer_indices)): + target = self._ds_buffers[level_idx][:N] + target.zero_() + target[is_multimodal] = level_features[level_idx] + + self._ds_num_tokens = N + return inputs_embeds + + # ----- Forward ----- + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + + # Build IntermediateTensors from pre-allocated persistent buffers. + # Always pass deepstack when inputs_embeds is non-None (prefill path), + # including during CUDA graph capture (buffers are zero → no-op injection). + # This ensures the graph captures the injection code path. + if ( + inputs_embeds is not None + and get_pp_group().is_first_rank + and self._ds_layer_indices + ): + ds: IntermediateTensors | None = IntermediateTensors( + { + f"ds_{llm_layer}": self._ds_buffers[lvl] + for lvl, llm_layer in enumerate(self._ds_layer_indices) + } + ) + else: + ds = None + + hidden_states = self.language_model.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + deepstack_input_embeds=ds, + ) + + # Clear buffers after use so stale features don't leak into the next request. + if ( + inputs_embeds is not None + and get_pp_group().is_first_rank + and self._ds_num_tokens > 0 + ): + n = self._ds_num_tokens + for buf in self._ds_buffers: + buf[:n].zero_() + self._ds_num_tokens = 0 + + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + # GraniteForCausalLM.compute_logits uses + # LogitsProcessor(scale=1/logits_scaling) + return self.language_model.compute_logits(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index b3cebd5100..4ba774a9fe 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -407,6 +407,10 @@ _MULTIMODAL_MODELS = { "granite_speech", "GraniteSpeechForConditionalGeneration", ), + "Granite4VisionForConditionalGeneration": ( + "granite4_vision", + "Granite4VisionForConditionalGeneration", + ), "H2OVLChatModel": ("h2ovl", "H2OVLChatModel"), "HunYuanVLForConditionalGeneration": ( "hunyuan_vision", diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index d909c4a266..cf2676a8f7 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -92,6 +92,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict( flex_olmo="FlexOlmoConfig", fireredlid="FireRedLIDConfig", funaudiochat="FunAudioChatConfig", + granite4_vision="Granite4VisionConfig", hunyuan_vl="HunYuanVLConfig", isaac="IsaacConfig", kimi_k2="DeepseekV3Config", # Kimi K2 uses same architecture as DeepSeek V3 diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py index 5f0d95d7b8..a45ea865db 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -32,6 +32,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "FlexOlmoConfig": "vllm.transformers_utils.configs.flex_olmo", "FunAudioChatConfig": "vllm.transformers_utils.configs.funaudiochat", "FunAudioChatAudioEncoderConfig": "vllm.transformers_utils.configs.funaudiochat", + "Granite4VisionConfig": "vllm.transformers_utils.configs.granite4_vision", "HunYuanVLConfig": "vllm.transformers_utils.configs.hunyuan_vl", "HunYuanVLTextConfig": "vllm.transformers_utils.configs.hunyuan_vl", "HunYuanVLVisionConfig": "vllm.transformers_utils.configs.hunyuan_vl", @@ -92,6 +93,7 @@ __all__ = [ "FireRedLIDConfig", "FunAudioChatConfig", "FunAudioChatAudioEncoderConfig", + "Granite4VisionConfig", "HunYuanVLConfig", "HunYuanVLTextConfig", "HunYuanVLVisionConfig", diff --git a/vllm/transformers_utils/configs/granite4_vision.py b/vllm/transformers_utils/configs/granite4_vision.py new file mode 100644 index 0000000000..a67e6ffeb8 --- /dev/null +++ b/vllm/transformers_utils/configs/granite4_vision.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from typing import Any + +import transformers + + +class Granite4VisionConfig(transformers.PretrainedConfig): + """Configuration for Granite 4 Vision model. + + This config is needed because the granite4_vision model type is not yet + in the transformers version pinned by vLLM. Once transformers adds native + support, this file can be removed and the _CONFIG_REGISTRY entry dropped. + """ + + model_type = "granite4_vision" + is_composition = False + + def __init__( + self, + vision_config: dict[str, Any] | None = None, + text_config: dict[str, Any] | None = None, + image_token_index: int = 100352, + image_seq_length: int = 576, + image_grid_pinpoints: list[list[int]] | None = None, + vision_feature_select_strategy: str = "full", + vision_feature_layer: int | list[int] = -2, + projector_hidden_act: str = "gelu", + projector_dropout: float = 0.1, + downsample_rate: str | None = None, + use_image_newline_parameter: bool = True, + deepstack_layer_map: list[list[int]] | None = None, + use_spatial_sampling: bool = False, + spatial_stride: int = 2, + spatial_vision_layer: int = -1, + spatial_target_layers: list[int] | None = None, + # Hub aliases — base model config uses different field names + vision_layer_to_llm_layer: list[list[int]] | None = None, + use_checkerboard_sampling: bool | None = None, + checkerboard_stride: int | None = None, + checkerboard_vision_layer: int | None = None, + checkerboard_llm_layers: list[int] | None = None, + **kwargs: Any, + ): + self.image_token_index = image_token_index + self.image_seq_length = image_seq_length + self.image_grid_pinpoints = image_grid_pinpoints or [] + self.vision_feature_select_strategy = vision_feature_select_strategy + self.vision_feature_layer = vision_feature_layer + self.projector_hidden_act = projector_hidden_act + self.projector_dropout = projector_dropout + self.downsample_rate = downsample_rate + self.use_image_newline_parameter = use_image_newline_parameter + self.deepstack_layer_map = deepstack_layer_map or vision_layer_to_llm_layer + self.use_spatial_sampling = ( + use_spatial_sampling + if use_checkerboard_sampling is None + else use_checkerboard_sampling + ) + self.spatial_stride = ( + spatial_stride if checkerboard_stride is None else checkerboard_stride + ) + self.spatial_vision_layer = ( + spatial_vision_layer + if checkerboard_vision_layer is None + else checkerboard_vision_layer + ) + self.spatial_target_layers = ( + spatial_target_layers or checkerboard_llm_layers or [0, 10, 20, 30] + ) + + if vision_config is None: + vision_config = {} + if text_config is None: + text_config = {} + + vision_model_type = vision_config.get("model_type", "siglip_vision_model") + if vision_model_type in transformers.CONFIG_MAPPING: + self.vision_config = transformers.CONFIG_MAPPING[vision_model_type]( + **vision_config + ) + else: + self.vision_config = transformers.PretrainedConfig(**vision_config) + + text_model_type = text_config.get("model_type", "granite") + if text_model_type in transformers.CONFIG_MAPPING: + self.text_config = transformers.CONFIG_MAPPING[text_model_type]( + **text_config + ) + else: + self.text_config = transformers.PretrainedConfig(**text_config) + + super().__init__(**kwargs) diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index 750d6291e0..c1fe9eaf93 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -19,6 +19,7 @@ __all__ = [ "FireRedLIDProcessor", "FunASRProcessor", "GLM4VProcessor", + "Granite4VisionProcessor", "H2OVLProcessor", "HunYuanVLProcessor", "HunYuanVLImageProcessor", @@ -48,6 +49,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "FireRedLIDProcessor": "vllm.transformers_utils.processors.fireredlid", "FunASRProcessor": "vllm.transformers_utils.processors.funasr", "GLM4VProcessor": "vllm.transformers_utils.processors.glm4v", + "Granite4VisionProcessor": "vllm.transformers_utils.processors.granite4_vision", "H2OVLProcessor": "vllm.transformers_utils.processors.h2ovl", "HunYuanVLProcessor": "vllm.transformers_utils.processors.hunyuan_vl", "HunYuanVLImageProcessor": "vllm.transformers_utils.processors.hunyuan_vl_image", diff --git a/vllm/transformers_utils/processors/granite4_vision.py b/vllm/transformers_utils/processors/granite4_vision.py new file mode 100644 index 0000000000..7972015adf --- /dev/null +++ b/vllm/transformers_utils/processors/granite4_vision.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from fractions import Fraction + +from transformers import LlavaNextProcessor +from transformers.image_processing_utils import select_best_resolution + + +class Granite4VisionProcessor(LlavaNextProcessor): + """Processor for Granite 4 Vision. + + Extends LlavaNextProcessor to account for the Window Q-Former + downsampling when computing the number of image features. + + This processor is needed because the granite4_vision processor type + is not yet in the transformers version pinned by vLLM. + """ + + model_type = "granite4_vision" + + def __init__( + self, + image_processor=None, + tokenizer=None, + patch_size=None, + vision_feature_select_strategy=None, + chat_template=None, + image_token="", + num_additional_image_tokens=0, + downsample_rate=None, + **kwargs, + ): + super().__init__( + image_processor=image_processor, + tokenizer=tokenizer, + patch_size=patch_size, + vision_feature_select_strategy=vision_feature_select_strategy, + chat_template=chat_template, + image_token=image_token, + num_additional_image_tokens=num_additional_image_tokens, + ) + self.downsample_rate = downsample_rate + + def _get_number_of_features( + self, + orig_height: int, + orig_width: int, + height: int, + width: int, + ) -> int: + image_grid_pinpoints = self.image_processor.image_grid_pinpoints + + height_best_resolution, width_best_resolution = select_best_resolution( + [orig_height, orig_width], image_grid_pinpoints + ) + scale_height = height_best_resolution // height + scale_width = width_best_resolution // width + + patches_height = height // self.patch_size + patches_width = width // self.patch_size + if self.downsample_rate is not None: + ds_rate = Fraction(self.downsample_rate) + patches_height = int(patches_height * ds_rate) + patches_width = int(patches_width * ds_rate) + + unpadded_features, newline_features = self._get_unpadded_features( + orig_height, + orig_width, + patches_height, + patches_width, + scale_height, + scale_width, + ) + base_features = ( + patches_height * patches_width + self.num_additional_image_tokens + ) + return unpadded_features + newline_features + base_features