From 1fa7cb38a291f8c305c284a86ce95882944c500b Mon Sep 17 00:00:00 2001 From: artemspector Date: Thu, 16 Apr 2026 14:37:30 +0300 Subject: [PATCH 01/17] Add Granite 4 Vision as built-in multimodal model Adds granite4_vision (Granite4VisionForConditionalGeneration) with GraniteForCausalLM backbone, SigLIP vision encoder, and deepstack feature injection via WindowQFormer projectors. Includes config/processor for _CONFIG_REGISTRY bypass, model registry, docs, and test entry. Signed-off-by: Artem Spector Co-Authored-By: Claude Opus 4.6 Signed-off-by: artemspector --- docs/models/supported_models.md | 1 + tests/models/registry.py | 3 + vllm/model_executor/models/granite4_vision.py | 981 ++++++++++++++++++ vllm/model_executor/models/registry.py | 4 + vllm/transformers_utils/config.py | 1 + vllm/transformers_utils/configs/__init__.py | 2 + .../configs/granite4_vision.py | 76 ++ .../transformers_utils/processors/__init__.py | 2 + .../processors/granite4_vision.py | 75 ++ 9 files changed, 1145 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 746980b8fe16..f1ba46b707d6 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/tests/models/registry.py b/tests/models/registry.py index f5968438cbde..5e82e77390b0 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -899,6 +899,9 @@ def check_available_online( "zai-org/GLM-ASR-Nano-2512", min_transformers_version="5.0.0", ), + "Granite4VisionForConditionalGeneration": _HfExamplesInfo( + "granite-vision-dev/granite-4.1-3b-vision", + ), "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 000000000000..af96d85d3fc2 --- /dev/null +++ b/vllm/model_executor/models/granite4_vision.py @@ -0,0 +1,981 @@ +# 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: +- Full merge (--hf-overrides '{"adapter_path": "..."}') merges LM-only LoRA + deltas into base weights at load time. +- Native LoRA (--enable-lora --default-mm-loras) lets vLLM runtime serve + LM LoRA deltas per-request. +Both modes expect a LM-only adapter (no modules_to_save). +""" + +import json +import math +import os +from collections.abc import Iterable, Mapping +from fractions import Fraction + +import torch +import torch.nn as nn +from safetensors.torch import load_file +from transformers import BatchFeature +from transformers.models.blip_2.configuration_blip_2 import Blip2QFormerConfig +from transformers.models.blip_2.modeling_blip_2 import Blip2QFormerModel +from transformers.models.llava_next.modeling_llava_next import ( + get_anyres_image_grid_shape, + image_size_to_num_patches, + unpad_image, +) + +from vllm.config import VllmConfig +from vllm.distributed.parallel_state import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.logger import init_logger +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import MultiModalFieldConfig +from vllm.sequence import IntermediateTensors + +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.llava import ( + LlavaDummyInputsBuilder, + init_vision_tower_for_llava, +) +from vllm.model_executor.models.llava_next import ( + BaseLlavaNextMultiModalProcessor, + LlavaNextProcessingInfo, + LlavaNextImagePixelInputs, + LlavaNextImageEmbeddingInputs, + LlavaNextImageInputs, +) +from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + _merge_multimodal_embeddings, + init_vllm_registered_model, + maybe_prefix, +) + +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, spatial_offset=None): + 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) + + 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 HW == self.image_side * self.image_side + 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, + return_dict=True, + ).last_hidden_state + + out = self._unwin(out_w, n=n, win=self.query_side) + out = self.dropout(out) + return self.out_linear(out) + + +# --------------------------------------------------------------------------- +# 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"): + self.vision_tower = init_vision_tower_for_llava( + config, + quant_config=quant_config, + require_post_norm=False, + prefix=maybe_prefix(prefix, "vision_tower"), + ) + + # image_newline parameter + if getattr(config, "use_image_newline_parameter", False): + self.image_newline = nn.Parameter( + torch.empty(config.text_config.hidden_size) + ) + else: + self.image_newline = None + + # Deepstack projectors: one per (vision_layer, llm_layer) pair + self.layerwise_projectors = nn.ModuleList([ + WindowQFormerDownsampler(config) + for _ in range(len(config.deepstack_layer_map)) + ]) + + # Spatial projectors: 4 offset groups + self.spatial_projectors = None + if getattr(config, "use_spatial_sampling", False): + self.spatial_projectors = nn.ModuleList([ + WindowQFormerDownsampler(config, spatial_offset=i) + for i in range(4) + ]) + + # ----- Language model (marked as LM) ----- + with self._mark_language_model(vllm_config): + self.language_model = init_vllm_registered_model( + vllm_config=vllm_config, + 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) + + # Deepstack state — set during embed_input_ids, consumed during forward + # list of (llm_layer_idx, features_buffer) where buffer is (N, hidden_size) + self._ds_features: list[tuple[int, torch.Tensor]] = [] + self._ds_vision_mask: torch.Tensor | None = None + + # ----- 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, + ) -> list[tuple[int, list[torch.Tensor]]]: + """Extract deepstack + spatial features. + + Returns list of (llm_layer_idx, [per_image_features, ...]) tuples. + This is the vLLM equivalent of HF's get_image_features. + """ + select_strategy = self._vision_feature_select_strategy + + # Count patches per image for splitting + 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 + ] + + # Flatten 5D → 4D if needed + if pixel_values.dim() == 5: + _pv_list = [ + pv[:np_] + for pv, np_ in zip(pixel_values, image_num_patches) + ] + pixel_values = torch.cat(_pv_list, dim=0) + + # Run vision tower once, get all hidden states + all_hidden_states = self._get_vision_hidden_states(pixel_values) + + all_features = [] + + # ----- Deepstack features ----- + 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:] # remove CLS + + projected = self.layerwise_projectors[proj_idx](selected) + projected_split = torch.split(projected, image_num_patches, dim=0) + + packed = self._pack_and_unpad_image_features( + projected_split, image_sizes + ) + all_features.append((llm_layer, packed)) + + # ----- Spatial features ----- + 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) + projected_split = torch.split( + projected, image_num_patches, dim=0 + ) + packed = self._pack_and_unpad_image_features( + projected_split, image_sizes + ) + all_features.append((llm_layer, packed)) + + return all_features + + # ----- 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: + """Convert pixel values → per-image placeholder tensors. + + The actual vision features are stored in self._ds_level_features and + injected during the forward loop (like HF does). We return zero- + tensors with the right shape so that _merge_multimodal_embeddings + fills image positions with zeros (matching HF's masked_fill(mask, 0)). + """ + 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) + + # Get all (llm_layer, [per_image_features]) pairs + all_features = self._get_all_layer_features(pixel_values, image_sizes) + + # Store ALL level features for deepstack injection in forward() + self._ds_level_features = all_features + + # Return zero-tensors matching the shape of level-0 features. + # This makes _merge_multimodal_embeddings write zeros at image + # positions (equivalent to HF's inputs_embeds.masked_fill(mask, 0)). + # All real features are injected during the layer loop. + if all_features: + return [ + torch.zeros_like(feat) + for feat in all_features[0][1] + ] + return [] + + 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 with deepstack injection via masked_scatter + + vLLM's GraniteModel.forward: + - if inputs_embeds given: hidden_states = inputs_embeds (NO multiplier) + - if input_ids given: hidden_states = embed(input_ids) * multiplier + + So we apply embedding_multiplier here and pass inputs_embeds to forward. + """ + # Access the inner GraniteModel + lm_inner = self.language_model.model # GraniteModel + + 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: + # Text-only or decode: clear deepstack state + self._ds_features = [] + self._ds_vision_mask = None + self._ds_level_features = [] + # Apply embedding_multiplier here because forward() always receives + # inputs_embeds and skips the inner model's embed+multiply path. + embeds = lm_inner.embed_input_ids(input_ids) + return embeds * lm_inner.config.embedding_multiplier + + # --- Vision path --- + # HF flow: embed -> masked_fill(0) -> multiply + # by embedding_multiplier. Layer loop adds vision + # features via masked_scatter (never in inputs_embeds). + + # 1. Get text embeddings + text_embeds = lm_inner.embed_input_ids(input_ids) + + # 2. Zero out image positions (HF: inputs_embeds.masked_fill(vision_mask, 0.0)) + # _merge_multimodal_embeddings writes our zero-tensors here, same effect. + _merge_multimodal_embeddings( + inputs_embeds=text_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + # 3. Apply embedding_multiplier to ALL positions (text + vision=0) + embedding_multiplier = lm_inner.config.embedding_multiplier + inputs_embeds = text_embeds * embedding_multiplier + + # 5. Prepare deepstack feature buffers for the layer loop + N = inputs_embeds.size(0) + hidden_size = inputs_embeds.size(1) + + prepared = [] + for llm_layer, per_image_features in getattr(self, "_ds_level_features", []): + concat_features = torch.cat(per_image_features, dim=0) + buf = torch.zeros( + N, hidden_size, + dtype=inputs_embeds.dtype, + device=inputs_embeds.device, + ) + buf[is_multimodal] = concat_features.to(dtype=inputs_embeds.dtype) + prepared.append((llm_layer, buf)) + + self._ds_features = prepared + self._ds_vision_mask = is_multimodal + self._ds_level_features = [] # consumed + + 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: + """Forward pass with deepstack injection. + + Runs the LLM layer loop directly (like HF Granite4VisionModel.forward) + to inject vision features at target layers via masked addition. + """ + if intermediate_tensors is not None: + inputs_embeds = None + + # Access the inner GraniteModel + lm_inner = self.language_model.model + + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + # embed_input_ids already applied embedding_multiplier + hidden_states = inputs_embeds + else: + # Text-only path: embed + multiply + hidden_states = lm_inner.embed_input_ids(input_ids) + hidden_states = hidden_states * lm_inner.config.embedding_multiplier + else: + if intermediate_tensors is None: + raise RuntimeError("Intermediate tensors may not be None!") + hidden_states = intermediate_tensors["hidden_states"] + + num_tokens = hidden_states.size(0) + + # Build O(1) lookup for deepstack features + ds_map: dict[int, torch.Tensor] = {} + for llm_layer_idx, features in self._ds_features: + ds_map[llm_layer_idx] = features + + vision_mask = self._ds_vision_mask + if vision_mask is not None: + vision_mask = vision_mask[:num_tokens] + + # Run through decoder layers with deepstack injection + for i in range(lm_inner.start_layer, lm_inner.end_layer): + layer = lm_inner.layers[i] + + # Inject deepstack features at target layers (before layer forward) + if i in ds_map and vision_mask is not None and vision_mask.any(): + features = ds_map[i][:num_tokens] + hidden_states = hidden_states.clone() + hidden_states[vision_mask] = ( + hidden_states[vision_mask] + features[vision_mask] + ) + + hidden_states = layer(positions, hidden_states) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states}) + + hidden_states = lm_inner.norm(hidden_states) + 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) + + # ----- Full-merge LoRA support ----- + + # HF→vLLM key prefix mapping (same transforms as hf_to_vllm_mapper) + _ADAPTER_PREFIX_MAP = [ + ("model.language_model.", "language_model.model."), + ] + + # vLLM fuses q/k/v_proj into qkv_proj and gate/up_proj into gate_up_proj. + _STACKED_PARAMS_MAPPING = [ + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + ] + + @staticmethod + def _peft_to_vllm(peft_key: str) -> str: + """Strip 'base_model.model.' and apply HF→vLLM prefix mapping.""" + name = peft_key + if name.startswith("base_model.model."): + name = name[len("base_model.model."):] + for old_pfx, new_pfx in ( + Granite4VisionForConditionalGeneration._ADAPTER_PREFIX_MAP + ): + if name.startswith(old_pfx): + name = new_pfx + name[len(old_pfx):] + break + return name + + @staticmethod + def _load_adapter(adapter_path: str) -> tuple[dict, dict[str, torch.Tensor]]: + """Load adapter config and safetensors from a directory or HF hub ID.""" + # Resolve HF hub IDs to local cache path + if not os.path.isdir(adapter_path): + from huggingface_hub import snapshot_download + adapter_path = snapshot_download(adapter_path) + config_path = os.path.join(adapter_path, "adapter_config.json") + weights_path = os.path.join(adapter_path, "adapter_model.safetensors") + if not os.path.exists(config_path): + raise FileNotFoundError(f"No adapter_config.json in {adapter_path}") + if not os.path.exists(weights_path): + raise FileNotFoundError( + f"No adapter_model.safetensors in {adapter_path}") + with open(config_path) as f: + config = json.load(f) + weights = load_file(weights_path) + return config, weights + + def _merge_lora_deltas( + self, + adapter_config: dict, + adapter_weights: dict[str, torch.Tensor], + ) -> int: + """Merge LM-only LoRA deltas into model weights: W += scaling * B @ A. + + Uses _STACKED_PARAMS_MAPPING + module._get_shard_offset_mapping() + to handle packed QKV correctly (works with GQA automatically). + """ + lora_alpha = adapter_config.get("lora_alpha", 1) + lora_r = adapter_config.get("r", 1) + scaling = lora_alpha / lora_r + + # Collect lora_A / lora_B by vLLM module key + lora_a: dict[str, torch.Tensor] = {} + lora_b: dict[str, torch.Tensor] = {} + for peft_key, tensor in adapter_weights.items(): + if ".lora_A." in peft_key: + module_key = self._peft_to_vllm( + peft_key.replace(".lora_A.weight", "")) + lora_a[module_key] = tensor + elif ".lora_B." in peft_key: + module_key = self._peft_to_vllm( + peft_key.replace(".lora_B.weight", "")) + lora_b[module_key] = tensor + + params_dict = dict(self.named_parameters()) + modules_dict = dict(self.named_modules()) + + def _add_delta(name: str, delta: torch.Tensor) -> bool: + # Try stacked/fused params first (qkv_proj, gate_up_proj) + for fused_name, orig_name, shard_id in self._STACKED_PARAMS_MAPPING: + if orig_name not in name: + continue + fused_param_name = name.replace(orig_name, fused_name) + if fused_param_name not in params_dict: + continue + param = params_dict[fused_param_name] + module_path = fused_param_name.rsplit(".weight", 1)[0] + module = modules_dict.get(module_path) + if module is None: + continue + + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + if hasattr(module, "_get_shard_offset_mapping"): + # QKVParallelLinear: string shard_id ("q", "k", "v") + shard_offset = module._get_shard_offset_mapping(shard_id) + if shard_offset is not None: + shard_size = delta.shape[0] // tp_size + tp_delta = delta.narrow( + 0, tp_rank * shard_size, shard_size) + shard = param.data[shard_offset:shard_offset + shard_size] + param.data[shard_offset:shard_offset + shard_size] = ( + shard.float() + tp_delta.to(shard.device) + ).to(shard.dtype) + return True + elif hasattr(module, "output_sizes") and isinstance(shard_id, int): + # MergedColumnParallelLinear: integer shard_id (0, 1) + shard_size = module.output_sizes[shard_id] // tp_size + shard_offset = sum( + s // tp_size for s in module.output_sizes[:shard_id] + ) + tp_delta = delta.narrow( + 0, tp_rank * (delta.shape[0] // tp_size), + delta.shape[0] // tp_size) + shard = param.data[shard_offset:shard_offset + shard_size] + param.data[shard_offset:shard_offset + shard_size] = ( + shard.float() + tp_delta.to(shard.device) + ).to(shard.dtype) + return True + # Direct param (o_proj, down_proj) + if name in params_dict: + param = params_dict[name] + # Under TP, param is already sharded but delta is full-size. + # Slice delta to match: dim 0 for column-parallel, dim 1 for + # row-parallel. + if delta.shape != param.data.shape: + tp_rank = get_tensor_model_parallel_rank() + for dim in range(delta.dim()): + if delta.shape[dim] != param.data.shape[dim]: + shard_size = param.data.shape[dim] + offset = tp_rank * shard_size + delta = delta.narrow(dim, offset, shard_size) + break + merged = param.data.float() + delta.to(param.device) + param.data = merged.to(param.dtype) + return True + return False + + merge_device = next(self.parameters()).device + merged = 0 + for module_key in sorted(lora_a): + if module_key not in lora_b: + logger.warning("LoRA B missing for %s, skipping", module_key) + continue + A = lora_a[module_key].to(merge_device).float() + B = lora_b[module_key].to(merge_device).float() + delta = scaling * (B @ A) + if _add_delta(module_key + ".weight", delta): + merged += 1 + else: + logger.warning("LoRA target not found: %s", module_key) + + return merged + + def _apply_adapter(self) -> None: + """Full-merge entry point: called when config.adapter_path is set.""" + adapter_path = getattr(self.config, "adapter_path", None) + if not adapter_path: + return + logger.info("Full-merge LoRA from %s", adapter_path) + adapter_config, adapter_weights = self._load_adapter(adapter_path) + + if adapter_config.get("modules_to_save"): + raise ValueError( + "Adapter has modules_to_save — only LM-only adapters " + "(no modules_to_save) are supported." + ) + + n = self._merge_lora_deltas(adapter_config, adapter_weights) + logger.info("Merged %d LoRA pairs into base weights", n) + + def load_weights( + self, weights: Iterable[tuple[str, torch.Tensor]] + ) -> set[str]: + loader = AutoWeightsLoader(self) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + self._apply_adapter() + return loaded diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index b3cebd5100bb..4ba774a9fe80 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -407,6 +407,10 @@ "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 5f4b5a3b2a48..ff113e1a5f2c 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -92,6 +92,7 @@ def __getitem__(self, key): 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 5f0d95d7b8f6..a45ea865db81 100644 --- a/vllm/transformers_utils/configs/__init__.py +++ b/vllm/transformers_utils/configs/__init__.py @@ -32,6 +32,7 @@ "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 @@ "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 000000000000..b79702514d9a --- /dev/null +++ b/vllm/transformers_utils/configs/granite4_vision.py @@ -0,0 +1,76 @@ +# 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, + **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 + self.use_spatial_sampling = use_spatial_sampling + self.spatial_stride = spatial_stride + self.spatial_vision_layer = spatial_vision_layer + self.spatial_target_layers = spatial_target_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 750d6291e0f1..c1fe9eaf9340 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -19,6 +19,7 @@ "FireRedLIDProcessor", "FunASRProcessor", "GLM4VProcessor", + "Granite4VisionProcessor", "H2OVLProcessor", "HunYuanVLProcessor", "HunYuanVLImageProcessor", @@ -48,6 +49,7 @@ "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 000000000000..b1f0cdc4ab03 --- /dev/null +++ b/vllm/transformers_utils/processors/granite4_vision.py @@ -0,0 +1,75 @@ +# 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 From ce27f978f6194996c80c1c470784ece03d46c858 Mon Sep 17 00:00:00 2001 From: artemspector Date: Sun, 19 Apr 2026 11:18:12 +0300 Subject: [PATCH 02/17] Fix vision tower: use SiglipVisionModel directly for deepstack init_vision_tower_for_llava truncates the encoder to vision_feature_layer depth, but deepstack needs ALL hidden states (deepstack_layer_map uses indices into the full encoder output list). Use SiglipVisionModel directly and update the weight mapping prefix accordingly. Also removes debug dump instrumentation. Co-Authored-By: Claude Opus 4.6 Signed-off-by: artemspector --- vllm/model_executor/models/granite4_vision.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index af96d85d3fc2..646c71e83ecd 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -52,6 +52,7 @@ LlavaDummyInputsBuilder, init_vision_tower_for_llava, ) +from vllm.model_executor.models.siglip import SiglipVisionModel from vllm.model_executor.models.llava_next import ( BaseLlavaNextMultiModalProcessor, LlavaNextProcessingInfo, @@ -328,7 +329,7 @@ class Granite4VisionForConditionalGeneration( "model.layerwise_projectors.": "layerwise_projectors.", "model.spatial_projectors.": "spatial_projectors.", "model.image_newline": "image_newline", - "model.vision_tower.": "vision_tower.", + "model.vision_tower.": "vision_tower.vision_model.", "lm_head.": "language_model.lm_head.", } ) @@ -356,8 +357,12 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # ----- Vision tower + projectors (marked as tower) ----- with self._mark_tower_model(vllm_config, "image"): - self.vision_tower = init_vision_tower_for_llava( - config, + # 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"), @@ -549,11 +554,13 @@ def _get_all_layer_features( selected = selected[:, 1:] # remove CLS projected = self.layerwise_projectors[proj_idx](selected) + projected_split = torch.split(projected, image_num_patches, dim=0) packed = self._pack_and_unpad_image_features( projected_split, image_sizes ) + all_features.append((llm_layer, packed)) # ----- Spatial features ----- From 7270605a3cceed356fb1ffda7229e4dcbb2375d0 Mon Sep 17 00:00:00 2001 From: artemspector Date: Sun, 19 Apr 2026 14:49:26 +0300 Subject: [PATCH 03/17] Add granite4_vision to VLM test suite Add VLMTestInfo entry for Granite 4.1 Vision in test_common.py: - Single image correctness test (HF vs vLLM output comparison) - LoRA adapter support via default_mm_loras (same-repo adapter) - Self-contained post-processor to avoid trust_remote_code issues with AutoConfig/AutoTokenizer for models not yet in upstream HF Co-Authored-By: Claude Opus 4.6 Signed-off-by: artemspector --- .../multimodal/generation/test_common.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 1147ccef35b4..44cd5cb78bda 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -86,6 +86,26 @@ # 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 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( @@ -425,6 +445,20 @@ auto_cls=AutoModelForImageTextToText, vllm_output_post_proc=model_utils.llava_image_vllm_to_hf_output, ), + "granite4_vision": VLMTestInfo( + models=["granite-vision-dev/granite-4.1-3b-vision"], + 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": "granite-vision-dev/granite-4.1-3b-vision"}, + }, + ), "glm4v": VLMTestInfo( models=["zai-org/glm-4v-9b"], test_type=VLMTestType.IMAGE, From c71c2ef24bb625584e83fb8ad2bf06cd02c40838 Mon Sep 17 00:00:00 2001 From: artemspector Date: Sun, 19 Apr 2026 15:21:31 +0300 Subject: [PATCH 04/17] Use in-place addition for deepstack feature injection Replace hidden_states.clone() + indexed assignment with in-place +=. No autograd in vLLM inference, so the defensive copy is unnecessary. Eliminates up to 8 full tensor clones per forward pass. Co-Authored-By: Claude Opus 4.6 Signed-off-by: artemspector --- vllm/model_executor/models/granite4_vision.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index 646c71e83ecd..a1908761ce63 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -784,10 +784,7 @@ def forward( # Inject deepstack features at target layers (before layer forward) if i in ds_map and vision_mask is not None and vision_mask.any(): features = ds_map[i][:num_tokens] - hidden_states = hidden_states.clone() - hidden_states[vision_mask] = ( - hidden_states[vision_mask] + features[vision_mask] - ) + hidden_states[vision_mask] += features[vision_mask] hidden_states = layer(positions, hidden_states) From 41ec9886b457b78e07211dbf329c360826edcec8 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 09:05:54 +0300 Subject: [PATCH 05/17] Address review comments: alphabetical order, remove getattr - Move granite4_vision test entry to alphabetical position - Replace getattr(config, ...) with direct config attribute access Co-Authored-By: Claude Opus 4.6 Signed-off-by: artemspector --- .../multimodal/generation/test_common.py | 28 +++++++++---------- vllm/model_executor/models/granite4_vision.py | 4 +-- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 44cd5cb78bda..bb2c7d33ff8e 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -445,20 +445,6 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): auto_cls=AutoModelForImageTextToText, vllm_output_post_proc=model_utils.llava_image_vllm_to_hf_output, ), - "granite4_vision": VLMTestInfo( - models=["granite-vision-dev/granite-4.1-3b-vision"], - 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": "granite-vision-dev/granite-4.1-3b-vision"}, - }, - ), "glm4v": VLMTestInfo( models=["zai-org/glm-4v-9b"], test_type=VLMTestType.IMAGE, @@ -526,6 +512,20 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): auto_cls=AutoModelForImageTextToText, marks=[large_gpu_mark(min_gb=32)], ), + "granite4_vision": VLMTestInfo( + models=["granite-vision-dev/granite-4.1-3b-vision"], + 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": "granite-vision-dev/granite-4.1-3b-vision"}, + }, + ), "h2ovl": VLMTestInfo( models=[ "h2oai/h2ovl-mississippi-800m", diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index a1908761ce63..64a636275de0 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -369,7 +369,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) # image_newline parameter - if getattr(config, "use_image_newline_parameter", False): + if config.use_image_newline_parameter: self.image_newline = nn.Parameter( torch.empty(config.text_config.hidden_size) ) @@ -384,7 +384,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # Spatial projectors: 4 offset groups self.spatial_projectors = None - if getattr(config, "use_spatial_sampling", False): + if config.use_spatial_sampling: self.spatial_projectors = nn.ModuleList([ WindowQFormerDownsampler(config, spatial_offset=i) for i in range(4) From 4c7c5d327e5bf7c867423cf6db140a276e7b5c6d Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 09:06:34 +0300 Subject: [PATCH 06/17] Accept Hub config field name aliases in Granite4VisionConfig The base model repo uses different field names (vision_layer_to_llm_layer, checkerboard_*) than our config class (deepstack_layer_map, spatial_*). Accept both naming conventions so the model loads from either source. Co-Authored-By: Claude Opus 4.6 Signed-off-by: artemspector --- .../configs/granite4_vision.py | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/vllm/transformers_utils/configs/granite4_vision.py b/vllm/transformers_utils/configs/granite4_vision.py index b79702514d9a..1e9aee622119 100644 --- a/vllm/transformers_utils/configs/granite4_vision.py +++ b/vllm/transformers_utils/configs/granite4_vision.py @@ -35,6 +35,12 @@ def __init__( 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 @@ -46,11 +52,22 @@ def __init__( 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 - self.use_spatial_sampling = use_spatial_sampling - self.spatial_stride = spatial_stride - self.spatial_vision_layer = spatial_vision_layer - self.spatial_target_layers = spatial_target_layers or [0, 10, 20, 30] + 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 = {} From 0a70fa34ef7370d5228ac8ea7cd00bd9602b8161 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 10:03:56 +0300 Subject: [PATCH 07/17] Switch WindowQFormerDownsampler to use vLLM's Blip2QFormerModel Replaces the transformers import with vLLM's built-in Blip2QFormerModel from blip2.py. Passes quant_config, cache_config, and prefix through WindowQFormerDownsampler to the QFormer, matching the pattern used by GraniteSpeech. Removes return_dict=True (vLLM returns raw tensor). Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- vllm/model_executor/models/granite4_vision.py | 44 +++++++++++++++---- 1 file changed, 35 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index 64a636275de0..a7646c7f7510 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -24,14 +24,16 @@ from safetensors.torch import load_file from transformers import BatchFeature from transformers.models.blip_2.configuration_blip_2 import Blip2QFormerConfig -from transformers.models.blip_2.modeling_blip_2 import Blip2QFormerModel + +from .blip2 import Blip2QFormerModel from transformers.models.llava_next.modeling_llava_next import ( get_anyres_image_grid_shape, image_size_to_num_patches, unpad_image, ) -from vllm.config import VllmConfig +from vllm.config import CacheConfig, VllmConfig +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.distributed.parallel_state import ( get_pp_group, get_tensor_model_parallel_rank, @@ -125,7 +127,14 @@ def __call__(self, image_features: torch.Tensor) -> torch.Tensor: class WindowQFormerDownsampler(nn.Module): """Window-based QFormer downsampler (matches HF downsampling.py exactly).""" - def __init__(self, config, spatial_offset=None): + 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 @@ -147,7 +156,12 @@ def __init__(self, config, spatial_offset=None): max_position_embeddings=2048, use_qformer_text_input=False, ) - self.qformer = Blip2QFormerModel(qformer_config) + 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 @@ -208,8 +222,7 @@ def forward(self, image_features: torch.Tensor) -> torch.Tensor: out_w = self.qformer( query_embeds=query_embeds, encoder_hidden_states=encoder_embeds, - return_dict=True, - ).last_hidden_state + ) out = self._unwin(out_w, n=n, win=self.query_side) out = self.dropout(out) @@ -376,17 +389,30 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: 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) - for _ in range(len(config.deepstack_layer_map)) + 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, spatial_offset=i) + 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) ]) From 77b778e10f1c1575953b026531d3a464dc24e8de Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 12:58:01 +0300 Subject: [PATCH 08/17] Fix deepstack injection for CUDA graph / compiled mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three bugs prevented vision from working with enforce_eager=False: 1. _ds_layer_indices was populated lazily (first embed_multimodal call), so it was empty during CUDA graph capture. Forward passed ds=None, capturing the graph without any injection code path. Fix: pre-populate _ds_layer_indices from config in __init__. 2. forward() only passed deepstack when _ds_num_tokens > 0, so CUDA graph capture (which has no real images) captured without injection. Fix: always pass deepstack buffers (zero-filled = no-op) when inputs_embeds is non-None, so the graph captures the injection path. 3. pbuf[:N][is_multimodal] = feat is a PyTorch no-op — boolean indexing on a slice returns a copy, not a view. Buffers stayed all zeros. Fix: build a full (N, lm_h) buffer tensor first, then copy_ into the persistent pre-allocated buffer (matches Qwen3-VL pattern). Also fixes stale buffer leak: after each prefill, zero the buffers and reset _ds_num_tokens so the next request (text-only or new vision) does not inherit features from the previous request. Also adds Granite4VisionLLMModel / Granite4VisionLLMForCausalLM classes (DarkLight1337's request) so the deepstack layer loop lives in a proper LLM subclass rather than the outer model's forward(). Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- vllm/model_executor/models/granite4_vision.py | 388 ++++++++++-------- 1 file changed, 228 insertions(+), 160 deletions(-) diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index a7646c7f7510..a3bfb5fa196d 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -18,6 +18,7 @@ import os from collections.abc import Iterable, Mapping from fractions import Fraction +from itertools import islice import torch import torch.nn as nn @@ -50,10 +51,7 @@ SupportsMultiModal, SupportsPP, ) -from vllm.model_executor.models.llava import ( - LlavaDummyInputsBuilder, - init_vision_tower_for_llava, -) +from vllm.model_executor.models.llava import LlavaDummyInputsBuilder from vllm.model_executor.models.siglip import SiglipVisionModel from vllm.model_executor.models.llava_next import ( BaseLlavaNextMultiModalProcessor, @@ -63,11 +61,14 @@ LlavaNextImageInputs, ) from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.model_executor.models.granite import GraniteForCausalLM, GraniteModel +from vllm.compilation.decorators import support_torch_compile +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.utils import ( AutoWeightsLoader, + PPMissingLayer, WeightsMapper, - _merge_multimodal_embeddings, - init_vllm_registered_model, maybe_prefix, ) @@ -229,6 +230,95 @@ def forward(self, image_features: torch.Tensor) -> torch.Tensor: 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"] + + 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: + return IntermediateTensors({"hidden_states": hidden_states}) + + 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() + + # --------------------------------------------------------------------------- # Processing info / processor (reuses LlavaNext patterns) # --------------------------------------------------------------------------- @@ -418,9 +508,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # ----- Language model (marked as LM) ----- with self._mark_language_model(vllm_config): - self.language_model = init_vllm_registered_model( - vllm_config=vllm_config, - hf_config=config.text_config, + self.language_model = Granite4VisionLLMForCausalLM( + vllm_config=vllm_config.with_hf_config(config.text_config), prefix=maybe_prefix(prefix, "language_model"), ) @@ -438,10 +527,30 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: ) self._downsample_rate = Fraction(config.downsample_rate) - # Deepstack state — set during embed_input_ids, consumed during forward - # list of (llm_layer_idx, features_buffer) where buffer is (N, hidden_size) - self._ds_features: list[tuple[int, torch.Tensor]] = [] - self._ds_vision_mask: torch.Tensor | None = None + # 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", [])) + ) + + # 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 ----- @@ -539,15 +648,20 @@ def _get_all_layer_features( self, pixel_values: torch.Tensor, image_sizes: torch.Tensor, - ) -> list[tuple[int, list[torch.Tensor]]]: - """Extract deepstack + spatial features. + ) -> tuple[list[int], list[torch.Tensor]]: + """Extract deepstack + spatial features for all levels. - Returns list of (llm_layer_idx, [per_image_features, ...]) tuples. - This is the vLLM equivalent of HF's get_image_features. + 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 - # Count patches per image for splitting image_num_patches = [ image_size_to_num_patches( image_size=imsize, @@ -557,56 +671,50 @@ def _get_all_layer_features( for imsize in image_sizes ] - # Flatten 5D → 4D if needed if pixel_values.dim() == 5: - _pv_list = [ - pv[:np_] - for pv, np_ in zip(pixel_values, image_num_patches) - ] - pixel_values = torch.cat(_pv_list, dim=0) + pixel_values = torch.cat( + [pv[:np_] for pv, np_ in zip(pixel_values, image_num_patches)], + dim=0, + ) - # Run vision tower once, get all hidden states all_hidden_states = self._get_vision_hidden_states(pixel_values) - all_features = [] + # Collect per-level: (llm_layer, [per_image_tensor, ...]) + levels: list[tuple[int, list[torch.Tensor]]] = [] - # ----- Deepstack features ----- 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:] # remove CLS - + selected = selected[:, 1:] projected = self.layerwise_projectors[proj_idx](selected) - - projected_split = torch.split(projected, image_num_patches, dim=0) - - packed = self._pack_and_unpad_image_features( - projected_split, image_sizes + per_image = self._pack_and_unpad_image_features( + torch.split(projected, image_num_patches, dim=0), image_sizes ) + levels.append((llm_layer, per_image)) - all_features.append((llm_layer, packed)) - - # ----- Spatial features ----- 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) - projected_split = torch.split( - projected, image_num_patches, dim=0 - ) - packed = self._pack_and_unpad_image_features( - projected_split, image_sizes + per_image = self._pack_and_unpad_image_features( + torch.split(projected, image_num_patches, dim=0), image_sizes ) - all_features.append((llm_layer, packed)) + 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 all_features + return llm_layer_indices, per_image_packed # ----- Multimodal interface ----- @@ -638,12 +746,14 @@ def _parse_and_validate_image_input( raise AssertionError("Unreachable") def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: - """Convert pixel values → per-image placeholder tensors. + """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. - The actual vision features are stored in self._ds_level_features and - injected during the forward loop (like HF does). We return zero- - tensors with the right shape so that _merge_multimodal_embeddings - fills image positions with zeros (matching HF's masked_fill(mask, 0)). + 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: @@ -658,22 +768,11 @@ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: if isinstance(pixel_values, list): pixel_values = torch.cat(pixel_values, dim=0) - # Get all (llm_layer, [per_image_features]) pairs - all_features = self._get_all_layer_features(pixel_values, image_sizes) - - # Store ALL level features for deepstack injection in forward() - self._ds_level_features = all_features - - # Return zero-tensors matching the shape of level-0 features. - # This makes _merge_multimodal_embeddings write zeros at image - # positions (equivalent to HF's inputs_embeds.masked_fill(mask, 0)). - # All real features are injected during the layer loop. - if all_features: - return [ - torch.zeros_like(feat) - for feat in all_features[0][1] - ] - return [] + 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, @@ -689,16 +788,14 @@ def embed_input_ids( 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 with deepstack injection via masked_scatter + 4. layer loop injects deepstack features at target layers - vLLM's GraniteModel.forward: - - if inputs_embeds given: hidden_states = inputs_embeds (NO multiplier) - - if input_ids given: hidden_states = embed(input_ids) * multiplier - - So we apply embedding_multiplier here and pass inputs_embeds to forward. + 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(). """ - # Access the inner GraniteModel - lm_inner = self.language_model.model # GraniteModel + lm_inner = self.language_model.model has_vision = ( multimodal_embeddings is not None @@ -708,54 +805,44 @@ def embed_input_ids( ) if not has_vision: - # Text-only or decode: clear deepstack state - self._ds_features = [] - self._ds_vision_mask = None - self._ds_level_features = [] - # Apply embedding_multiplier here because forward() always receives - # inputs_embeds and skips the inner model's embed+multiply path. + self._ds_num_tokens = 0 embeds = lm_inner.embed_input_ids(input_ids) return embeds * lm_inner.config.embedding_multiplier - # --- Vision path --- - # HF flow: embed -> masked_fill(0) -> multiply - # by embedding_multiplier. Layer loop adds vision - # features via masked_scatter (never in inputs_embeds). - - # 1. Get text embeddings + # 1. Text embeddings text_embeds = lm_inner.embed_input_ids(input_ids) - # 2. Zero out image positions (HF: inputs_embeds.masked_fill(vision_mask, 0.0)) - # _merge_multimodal_embeddings writes our zero-tensors here, same effect. - _merge_multimodal_embeddings( - inputs_embeds=text_embeds, - multimodal_embeddings=multimodal_embeddings, - is_multimodal=is_multimodal, - ) + # 2. Zero image positions (matches HF masked_fill(vision_mask, 0.0)) + text_embeds[is_multimodal] = 0.0 - # 3. Apply embedding_multiplier to ALL positions (text + vision=0) - embedding_multiplier = lm_inner.config.embedding_multiplier - inputs_embeds = text_embeds * embedding_multiplier - - # 5. Prepare deepstack feature buffers for the layer loop - N = inputs_embeds.size(0) - hidden_size = inputs_embeds.size(1) - - prepared = [] - for llm_layer, per_image_features in getattr(self, "_ds_level_features", []): - concat_features = torch.cat(per_image_features, dim=0) - buf = torch.zeros( - N, hidden_size, - dtype=inputs_embeds.dtype, - device=inputs_embeds.device, - ) - buf[is_multimodal] = concat_features.to(dtype=inputs_embeds.dtype) - prepared.append((llm_layer, buf)) + # 3. Apply embedding_multiplier + inputs_embeds = text_embeds * lm_inner.config.embedding_multiplier - self._ds_features = prepared - self._ds_vision_mask = is_multimodal - self._ds_level_features = [] # consumed + # 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)): + buf_data = torch.zeros(N, lm_h, dtype=inputs_embeds.dtype, device=inputs_embeds.device) + buf_data[is_multimodal] = level_features[level_idx] + self._ds_buffers[level_idx][:N].copy_(buf_data) + self._ds_num_tokens = N return inputs_embeds # ----- Forward ----- @@ -768,56 +855,37 @@ def forward( inputs_embeds: torch.Tensor | None = None, **kwargs: object, ) -> torch.Tensor | IntermediateTensors: - """Forward pass with deepstack injection. - - Runs the LLM layer loop directly (like HF Granite4VisionModel.forward) - to inject vision features at target layers via masked addition. - """ if intermediate_tensors is not None: inputs_embeds = None - # Access the inner GraniteModel - lm_inner = self.language_model.model - - if get_pp_group().is_first_rank: - if inputs_embeds is not None: - # embed_input_ids already applied embedding_multiplier - hidden_states = inputs_embeds - else: - # Text-only path: embed + multiply - hidden_states = lm_inner.embed_input_ids(input_ids) - hidden_states = hidden_states * lm_inner.config.embedding_multiplier + # 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: + n = self._ds_num_tokens if self._ds_num_tokens > 0 else inputs_embeds.size(0) + ds: IntermediateTensors | None = IntermediateTensors({ + f"ds_{llm_layer}": self._ds_buffers[lvl][:n] + for lvl, llm_layer in enumerate(self._ds_layer_indices) + }) else: - if intermediate_tensors is None: - raise RuntimeError("Intermediate tensors may not be None!") - hidden_states = intermediate_tensors["hidden_states"] - - num_tokens = hidden_states.size(0) - - # Build O(1) lookup for deepstack features - ds_map: dict[int, torch.Tensor] = {} - for llm_layer_idx, features in self._ds_features: - ds_map[llm_layer_idx] = features - - vision_mask = self._ds_vision_mask - if vision_mask is not None: - vision_mask = vision_mask[:num_tokens] - - # Run through decoder layers with deepstack injection - for i in range(lm_inner.start_layer, lm_inner.end_layer): - layer = lm_inner.layers[i] - - # Inject deepstack features at target layers (before layer forward) - if i in ds_map and vision_mask is not None and vision_mask.any(): - features = ds_map[i][:num_tokens] - hidden_states[vision_mask] += features[vision_mask] - - hidden_states = layer(positions, hidden_states) + 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, + ) - if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) + # 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 - hidden_states = lm_inner.norm(hidden_states) return hidden_states def compute_logits( From 399ce5743db1fecc817783d9027b93b819e033c6 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 14:13:26 +0300 Subject: [PATCH 09/17] Switch to ibm-granite/granite-vision-4.1-4b as the official model - Fix weight mapper: model.vision_tower. -> vision_tower. (the new checkpoint already includes vision_model. in the key path, so the old mapping was producing a double vision_model.vision_model. prefix) - Update test_common.py and registry.py to use the official model ID Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- tests/models/multimodal/generation/test_common.py | 4 ++-- tests/models/registry.py | 2 +- vllm/model_executor/models/granite4_vision.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index bb2c7d33ff8e..8a8a2241f4bb 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -513,7 +513,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): marks=[large_gpu_mark(min_gb=32)], ), "granite4_vision": VLMTestInfo( - models=["granite-vision-dev/granite-4.1-3b-vision"], + 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, @@ -523,7 +523,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): vllm_runner_kwargs={ "enable_lora": True, "max_lora_rank": 256, - "default_mm_loras": {"image": "granite-vision-dev/granite-4.1-3b-vision"}, + "default_mm_loras": {"image": "ibm-granite/granite-vision-4.1-4b"}, }, ), "h2ovl": VLMTestInfo( diff --git a/tests/models/registry.py b/tests/models/registry.py index 5e82e77390b0..4a7ab853c05e 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -900,7 +900,7 @@ def check_available_online( min_transformers_version="5.0.0", ), "Granite4VisionForConditionalGeneration": _HfExamplesInfo( - "granite-vision-dev/granite-4.1-3b-vision", + "ibm-granite/granite-vision-4.1-4b", ), "GraniteVision": _HfExamplesInfo("ibm-granite/granite-vision-3.3-2b"), "GraniteSpeechForConditionalGeneration": _HfExamplesInfo( diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index a3bfb5fa196d..63080f4a4de6 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -432,7 +432,7 @@ class Granite4VisionForConditionalGeneration( "model.layerwise_projectors.": "layerwise_projectors.", "model.spatial_projectors.": "spatial_projectors.", "model.image_newline": "image_newline", - "model.vision_tower.": "vision_tower.vision_model.", + "model.vision_tower.": "vision_tower.", "lm_head.": "language_model.lm_head.", } ) From 453479aafdfa3702b6728f60d0ba2953090d1df2 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 16:20:37 +0300 Subject: [PATCH 10/17] =?UTF-8?q?Fix=20pipeline=20parallelism=20(PP?= =?UTF-8?q?=E2=89=A52)=20for=20deepstack=20injection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With PP=2, rank 0 runs layers 0–19 and rank 1 runs layers 20–39. Deepstack feature buffers for rank 1's target layers (e.g. ds_21) must be included in the IntermediateTensors handed off between ranks. Three co-dependent fixes: 1. Override make_empty_intermediate_tensors on Granite4VisionLLMForCausalLM (not on Granite4VisionLLMModel): GraniteForCausalLM.make_empty_intermediate_tensors does not delegate to self.model, so the override must live on the causal wrapper to be reachable by vLLM's PP machinery. 2. Set self.language_model._ds_layer_indices from the outer model after construction so make_empty_intermediate_tensors can enumerate the ds keys (text_config alone has no deepstack_layer_map). 3. Send full-size buffers (shape max_tokens × H) rather than a sliced view [:n] when forwarding ds tensors to PP rank 1. The framework's sync_and_slice_intermediate_tensors copies with copy_len = padded token count (CUDA graph bucket size), which may exceed the actual token count n; sliced tensors caused a RuntimeError on the shape mismatch. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- vllm/model_executor/models/granite4_vision.py | 38 +++++++++++++++++-- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index 63080f4a4de6..673148b6e0d4 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -263,6 +263,13 @@ def forward( 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 @@ -283,7 +290,15 @@ def forward( hidden_states = layer(positions, hidden_states) if not get_pp_group().is_last_rank: - return IntermediateTensors({"hidden_states": hidden_states}) + # 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 @@ -318,6 +333,18 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: 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) @@ -535,6 +562,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + 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 @@ -863,9 +894,8 @@ def forward( # 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: - n = self._ds_num_tokens if self._ds_num_tokens > 0 else inputs_embeds.size(0) ds: IntermediateTensors | None = IntermediateTensors({ - f"ds_{llm_layer}": self._ds_buffers[lvl][:n] + f"ds_{llm_layer}": self._ds_buffers[lvl] for lvl, llm_layer in enumerate(self._ds_layer_indices) }) else: @@ -1049,7 +1079,7 @@ def _add_delta(name: str, delta: torch.Tensor) -> bool: if _add_delta(module_key + ".weight", delta): merged += 1 else: - logger.warning("LoRA target not found: %s", module_key) + logger.debug("LoRA target not found on this PP rank: %s", module_key) return merged From 53c0b8c9764ccf6aa77a56d88ae5633230c03364 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 16:30:11 +0300 Subject: [PATCH 11/17] Address Gemini review: idx==0 guard and eliminate temp buffer allocation Fix logic error in granite4_vision post-processor: when idx==0, output_ids[idx-1] wraps to the last element, incorrectly skipping the first token if it matches mm_token_id. Add an explicit idx==0 check to always keep the first token. In embed_input_ids, eliminate the per-call torch.zeros allocation inside the deepstack level loop. Instead zero the persistent buffer slice directly and scatter features into it, removing the intermediate allocation and copy. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- tests/models/multimodal/generation/test_common.py | 2 +- vllm/model_executor/models/granite4_vision.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 8a8a2241f4bb..2da2cbe06003 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -97,7 +97,7 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): hf_output_ids = [ token_id for idx, token_id in enumerate(output_ids) - if token_id != mm_token_id or output_ids[idx - 1] != mm_token_id + 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 diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index 673148b6e0d4..da727ddb1af7 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -869,9 +869,9 @@ def embed_input_ids( ] for level_idx in range(len(self._ds_layer_indices)): - buf_data = torch.zeros(N, lm_h, dtype=inputs_embeds.dtype, device=inputs_embeds.device) - buf_data[is_multimodal] = level_features[level_idx] - self._ds_buffers[level_idx][:N].copy_(buf_data) + target = self._ds_buffers[level_idx][:N] + target.zero_() + target[is_multimodal] = level_features[level_idx] self._ds_num_tokens = N return inputs_embeds From 2de1557fda4e7c910fe1f80924c4643e37114745 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 16:53:50 +0300 Subject: [PATCH 12/17] Fix ruff E501 line-length violations Wrap long lines flagged by pre-commit ruff-check to stay within the 88-character limit. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- .../multimodal/generation/test_common.py | 4 +++- vllm/model_executor/models/granite4_vision.py | 23 +++++++++++++++---- 2 files changed, 21 insertions(+), 6 deletions(-) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 2da2cbe06003..23ffe9a40a61 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -99,7 +99,9 @@ def _granite4_vision_vllm_to_hf_output(vllm_output, model): 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 + 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|>" diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index da727ddb1af7..1fde973b1065 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -265,7 +265,11 @@ def forward( 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_")] + 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} @@ -562,8 +566,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + 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). + # 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. @@ -893,7 +898,11 @@ def forward( # 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: + 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) @@ -910,7 +919,11 @@ def forward( ) # 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: + 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_() From 119eb73df2805d1ef0c5c225cd10df14c7b5bff2 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 17:00:04 +0300 Subject: [PATCH 13/17] Apply ruff auto-formatting (pre-commit) Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- .../multimodal/generation/test_common.py | 1 + vllm/model_executor/models/granite4_vision.py | 170 +++++++++--------- 2 files changed, 84 insertions(+), 87 deletions(-) diff --git a/tests/models/multimodal/generation/test_common.py b/tests/models/multimodal/generation/test_common.py index 23ffe9a40a61..29680bafd506 100644 --- a/tests/models/multimodal/generation/test_common.py +++ b/tests/models/multimodal/generation/test_common.py @@ -86,6 +86,7 @@ # 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. diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index 1fde973b1065..a281f8b1fdaa 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -25,26 +25,24 @@ from safetensors.torch import load_file from transformers import BatchFeature from transformers.models.blip_2.configuration_blip_2 import Blip2QFormerConfig - -from .blip2 import Blip2QFormerModel 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.model_executor.layers.quantization import QuantizationConfig from vllm.distributed.parallel_state import ( get_pp_group, get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) from vllm.logger import init_logger -from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import MultiModalFieldConfig -from vllm.sequence import IntermediateTensors - +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, @@ -52,25 +50,26 @@ SupportsPP, ) from vllm.model_executor.models.llava import LlavaDummyInputsBuilder -from vllm.model_executor.models.siglip import SiglipVisionModel from vllm.model_executor.models.llava_next import ( BaseLlavaNextMultiModalProcessor, - LlavaNextProcessingInfo, - LlavaNextImagePixelInputs, LlavaNextImageEmbeddingInputs, LlavaNextImageInputs, + LlavaNextImagePixelInputs, + LlavaNextProcessingInfo, ) from vllm.model_executor.models.module_mapping import MultiModelKeys -from vllm.model_executor.models.granite import GraniteForCausalLM, GraniteModel -from vllm.compilation.decorators import support_torch_compile -from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +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__) @@ -79,6 +78,7 @@ # Downsampler modules (translated from HF downsampling.py) # --------------------------------------------------------------------------- + class InterpolateDownsampler: """Spatial downsampling via area interpolation.""" @@ -208,7 +208,7 @@ def _unwin(self, xw: torch.Tensor, n: int, win: int) -> torch.Tensor: def forward(self, image_features: torch.Tensor) -> torch.Tensor: B, HW, C = image_features.shape - assert HW == self.image_side * self.image_side + assert self.image_side * self.image_side == HW n = self.image_side // self.window_side image_features = self.norm(image_features) @@ -234,6 +234,7 @@ def forward(self, image_features: torch.Tensor) -> torch.Tensor: # LLM subclasses with deepstack injection in the layer loop # --------------------------------------------------------------------------- + @support_torch_compile( dynamic_arg_dims={ "input_ids": 0, @@ -266,9 +267,7 @@ def forward( # 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_") + k for k in intermediate_tensors.tensors if k.startswith("ds_") ] if ds_keys: deepstack_input_embeds = IntermediateTensors( @@ -298,7 +297,8 @@ def forward( it = {"hidden_states": hidden_states} if deepstack_input_embeds is not None: remaining = { - k: v for k, v in deepstack_input_embeds.tensors.items() + k: v + for k, v in deepstack_input_embeds.tensors.items() if int(k.split("_")[1]) >= self.end_layer } it.update(remaining) @@ -332,8 +332,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: 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) + self.logits_processor = LogitsProcessor( + config.vocab_size, scale=logit_scale + ) else: self.lm_head = PPMissingLayer() @@ -354,8 +355,8 @@ def make_empty_intermediate_tensors( # Processing info / processor (reuses LlavaNext patterns) # --------------------------------------------------------------------------- -class Granite4VisionProcessingInfo(LlavaNextProcessingInfo): +class Granite4VisionProcessingInfo(LlavaNextProcessingInfo): def get_hf_config(self): return self.ctx.get_hf_config() @@ -417,6 +418,7 @@ def _get_mm_fields_config( # Top-level model # --------------------------------------------------------------------------- + @MULTIMODAL_REGISTRY.register_processor( Granite4VisionMultiModalProcessor, info=Granite4VisionProcessingInfo, @@ -513,29 +515,33 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> 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([ + self.layerwise_projectors = nn.ModuleList( + [ WindowQFormerDownsampler( config, quant_config=quant_config, cache_config=cache_config, - spatial_offset=i, - prefix=maybe_prefix(prefix, f"spatial_projectors.{i}"), + prefix=maybe_prefix(prefix, f"layerwise_projectors.{i}"), ) - for i in range(4) - ]) + 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): @@ -561,10 +567,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # 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", [])) - ) + 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 @@ -583,8 +588,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: 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) + 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 @@ -628,8 +632,7 @@ def _pack_and_unpad_image_features( image_feature = image_feature[1:] height = width = ( - config.vision_config.image_size - // config.vision_config.patch_size + config.vision_config.image_size // config.vision_config.patch_size ) # After QFormer downsampling height = int(height * ds_rate) @@ -645,13 +648,12 @@ def _pack_and_unpad_image_features( num_patch_height, num_patch_width, height, width, -1 ) image_feature = ( - image_feature.permute(4, 0, 2, 1, 3).contiguous() + 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] - ) + image_feature = unpad_image(image_feature, image_sizes[image_idx]) if self.image_newline is not None: image_feature = torch.cat( @@ -665,9 +667,7 @@ def _pack_and_unpad_image_features( ) image_feature = image_feature.flatten(1, 2).transpose(0, 1) - image_feature = torch.cat( - (base_image_feature, image_feature), dim=0 - ) + image_feature = torch.cat((base_image_feature, image_feature), dim=0) else: image_feature = image_feature[0] if self.image_newline is not None: @@ -718,9 +718,7 @@ def _get_all_layer_features( # 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 - ): + 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:] @@ -744,9 +742,7 @@ def _get_all_layer_features( 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 - ) + torch.cat([levels[lvl][1][img] for lvl in range(len(levels))], dim=-1) for img in range(num_images) ] @@ -903,10 +899,12 @@ def forward( 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) - }) + 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 @@ -960,12 +958,13 @@ def _peft_to_vllm(peft_key: str) -> str: """Strip 'base_model.model.' and apply HF→vLLM prefix mapping.""" name = peft_key if name.startswith("base_model.model."): - name = name[len("base_model.model."):] - for old_pfx, new_pfx in ( - Granite4VisionForConditionalGeneration._ADAPTER_PREFIX_MAP - ): + name = name[len("base_model.model.") :] + for ( + old_pfx, + new_pfx, + ) in Granite4VisionForConditionalGeneration._ADAPTER_PREFIX_MAP: if name.startswith(old_pfx): - name = new_pfx + name[len(old_pfx):] + name = new_pfx + name[len(old_pfx) :] break return name @@ -975,14 +974,14 @@ def _load_adapter(adapter_path: str) -> tuple[dict, dict[str, torch.Tensor]]: # Resolve HF hub IDs to local cache path if not os.path.isdir(adapter_path): from huggingface_hub import snapshot_download + adapter_path = snapshot_download(adapter_path) config_path = os.path.join(adapter_path, "adapter_config.json") weights_path = os.path.join(adapter_path, "adapter_model.safetensors") if not os.path.exists(config_path): raise FileNotFoundError(f"No adapter_config.json in {adapter_path}") if not os.path.exists(weights_path): - raise FileNotFoundError( - f"No adapter_model.safetensors in {adapter_path}") + raise FileNotFoundError(f"No adapter_model.safetensors in {adapter_path}") with open(config_path) as f: config = json.load(f) weights = load_file(weights_path) @@ -1007,12 +1006,10 @@ def _merge_lora_deltas( lora_b: dict[str, torch.Tensor] = {} for peft_key, tensor in adapter_weights.items(): if ".lora_A." in peft_key: - module_key = self._peft_to_vllm( - peft_key.replace(".lora_A.weight", "")) + module_key = self._peft_to_vllm(peft_key.replace(".lora_A.weight", "")) lora_a[module_key] = tensor elif ".lora_B." in peft_key: - module_key = self._peft_to_vllm( - peft_key.replace(".lora_B.weight", "")) + module_key = self._peft_to_vllm(peft_key.replace(".lora_B.weight", "")) lora_b[module_key] = tensor params_dict = dict(self.named_parameters()) @@ -1040,10 +1037,9 @@ def _add_delta(name: str, delta: torch.Tensor) -> bool: shard_offset = module._get_shard_offset_mapping(shard_id) if shard_offset is not None: shard_size = delta.shape[0] // tp_size - tp_delta = delta.narrow( - 0, tp_rank * shard_size, shard_size) - shard = param.data[shard_offset:shard_offset + shard_size] - param.data[shard_offset:shard_offset + shard_size] = ( + tp_delta = delta.narrow(0, tp_rank * shard_size, shard_size) + shard = param.data[shard_offset : shard_offset + shard_size] + param.data[shard_offset : shard_offset + shard_size] = ( shard.float() + tp_delta.to(shard.device) ).to(shard.dtype) return True @@ -1054,10 +1050,12 @@ def _add_delta(name: str, delta: torch.Tensor) -> bool: s // tp_size for s in module.output_sizes[:shard_id] ) tp_delta = delta.narrow( - 0, tp_rank * (delta.shape[0] // tp_size), - delta.shape[0] // tp_size) - shard = param.data[shard_offset:shard_offset + shard_size] - param.data[shard_offset:shard_offset + shard_size] = ( + 0, + tp_rank * (delta.shape[0] // tp_size), + delta.shape[0] // tp_size, + ) + shard = param.data[shard_offset : shard_offset + shard_size] + param.data[shard_offset : shard_offset + shard_size] = ( shard.float() + tp_delta.to(shard.device) ).to(shard.dtype) return True @@ -1113,9 +1111,7 @@ def _apply_adapter(self) -> None: n = self._merge_lora_deltas(adapter_config, adapter_weights) logger.info("Merged %d LoRA pairs into base weights", n) - def load_weights( - self, weights: Iterable[tuple[str, torch.Tensor]] - ) -> set[str]: + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) self._apply_adapter() From bb288aebceb428298abec52fcf1c00c5bafea968 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 17:09:23 +0300 Subject: [PATCH 14/17] Apply ruff formatting to configs and processors files Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- vllm/transformers_utils/configs/granite4_vision.py | 9 +++++---- vllm/transformers_utils/processors/granite4_vision.py | 9 ++++++--- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/vllm/transformers_utils/configs/granite4_vision.py b/vllm/transformers_utils/configs/granite4_vision.py index 1e9aee622119..a67e6ffeb8ee 100644 --- a/vllm/transformers_utils/configs/granite4_vision.py +++ b/vllm/transformers_utils/configs/granite4_vision.py @@ -54,15 +54,16 @@ def __init__( 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 + 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 + spatial_stride if checkerboard_stride is None else checkerboard_stride ) self.spatial_vision_layer = ( - spatial_vision_layer if checkerboard_vision_layer is None + spatial_vision_layer + if checkerboard_vision_layer is None else checkerboard_vision_layer ) self.spatial_target_layers = ( diff --git a/vllm/transformers_utils/processors/granite4_vision.py b/vllm/transformers_utils/processors/granite4_vision.py index b1f0cdc4ab03..7972015adf9d 100644 --- a/vllm/transformers_utils/processors/granite4_vision.py +++ b/vllm/transformers_utils/processors/granite4_vision.py @@ -65,9 +65,12 @@ def _get_number_of_features( 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, + orig_height, + orig_width, + patches_height, + patches_width, + scale_height, + scale_width, ) base_features = ( patches_height * patches_width + self.num_additional_image_tokens From 819e4ea77f80ba75c0a91e15b7989df66b81ae20 Mon Sep 17 00:00:00 2001 From: artemspector Date: Mon, 20 Apr 2026 17:30:41 +0300 Subject: [PATCH 15/17] Add granite4_vision to offline inference multi-image example Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- .../vision_language_multi_image.py | 33 +++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/examples/offline_inference/vision_language_multi_image.py b/examples/offline_inference/vision_language_multi_image.py index 1a3b7fd954cf..1b68a23b3bd0 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 @@ def load_molmo2(question: str, image_urls: list[str]) -> ModelRequestData: "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, From bb6a9636f9f47f55f0056bf0bfd32bdf2b960acb Mon Sep 17 00:00:00 2001 From: artemspector Date: Tue, 21 Apr 2026 12:24:00 +0300 Subject: [PATCH 16/17] Drop merge-on-load LoRA flow from Granite4Vision Remove _merge_lora_deltas, _apply_adapter, _load_adapter, _peft_to_vllm, and _STACKED_PARAMS_MAPPING. Native vLLM LoRA serving (--enable-lora --default-mm-loras) is the supported path; the manual merge-on-load path was not quantization-aware and fragile under TP. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- vllm/model_executor/models/granite4_vision.py | 194 +----------------- 1 file changed, 3 insertions(+), 191 deletions(-) diff --git a/vllm/model_executor/models/granite4_vision.py b/vllm/model_executor/models/granite4_vision.py index a281f8b1fdaa..147f02eced97 100644 --- a/vllm/model_executor/models/granite4_vision.py +++ b/vllm/model_executor/models/granite4_vision.py @@ -5,24 +5,16 @@ Uses GraniteForCausalLM as the language backbone with SigLIP vision encoder and deepstack feature injection via WindowQFormer projectors. -LoRA support: -- Full merge (--hf-overrides '{"adapter_path": "..."}') merges LM-only LoRA - deltas into base weights at load time. -- Native LoRA (--enable-lora --default-mm-loras) lets vLLM runtime serve - LM LoRA deltas per-request. -Both modes expect a LM-only adapter (no modules_to_save). +LoRA support: use --enable-lora --default-mm-loras for LM-only LoRA adapters. """ -import json import math -import os from collections.abc import Iterable, Mapping from fractions import Fraction from itertools import islice import torch import torch.nn as nn -from safetensors.torch import load_file from transformers import BatchFeature from transformers.models.blip_2.configuration_blip_2 import Blip2QFormerConfig from transformers.models.llava_next.modeling_llava_next import ( @@ -33,11 +25,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.distributed.parallel_state import ( - get_pp_group, - get_tensor_model_parallel_rank, - get_tensor_model_parallel_world_size, -) +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 @@ -937,182 +925,6 @@ def compute_logits( # LogitsProcessor(scale=1/logits_scaling) return self.language_model.compute_logits(hidden_states) - # ----- Full-merge LoRA support ----- - - # HF→vLLM key prefix mapping (same transforms as hf_to_vllm_mapper) - _ADAPTER_PREFIX_MAP = [ - ("model.language_model.", "language_model.model."), - ] - - # vLLM fuses q/k/v_proj into qkv_proj and gate/up_proj into gate_up_proj. - _STACKED_PARAMS_MAPPING = [ - (".qkv_proj", ".q_proj", "q"), - (".qkv_proj", ".k_proj", "k"), - (".qkv_proj", ".v_proj", "v"), - (".gate_up_proj", ".gate_proj", 0), - (".gate_up_proj", ".up_proj", 1), - ] - - @staticmethod - def _peft_to_vllm(peft_key: str) -> str: - """Strip 'base_model.model.' and apply HF→vLLM prefix mapping.""" - name = peft_key - if name.startswith("base_model.model."): - name = name[len("base_model.model.") :] - for ( - old_pfx, - new_pfx, - ) in Granite4VisionForConditionalGeneration._ADAPTER_PREFIX_MAP: - if name.startswith(old_pfx): - name = new_pfx + name[len(old_pfx) :] - break - return name - - @staticmethod - def _load_adapter(adapter_path: str) -> tuple[dict, dict[str, torch.Tensor]]: - """Load adapter config and safetensors from a directory or HF hub ID.""" - # Resolve HF hub IDs to local cache path - if not os.path.isdir(adapter_path): - from huggingface_hub import snapshot_download - - adapter_path = snapshot_download(adapter_path) - config_path = os.path.join(adapter_path, "adapter_config.json") - weights_path = os.path.join(adapter_path, "adapter_model.safetensors") - if not os.path.exists(config_path): - raise FileNotFoundError(f"No adapter_config.json in {adapter_path}") - if not os.path.exists(weights_path): - raise FileNotFoundError(f"No adapter_model.safetensors in {adapter_path}") - with open(config_path) as f: - config = json.load(f) - weights = load_file(weights_path) - return config, weights - - def _merge_lora_deltas( - self, - adapter_config: dict, - adapter_weights: dict[str, torch.Tensor], - ) -> int: - """Merge LM-only LoRA deltas into model weights: W += scaling * B @ A. - - Uses _STACKED_PARAMS_MAPPING + module._get_shard_offset_mapping() - to handle packed QKV correctly (works with GQA automatically). - """ - lora_alpha = adapter_config.get("lora_alpha", 1) - lora_r = adapter_config.get("r", 1) - scaling = lora_alpha / lora_r - - # Collect lora_A / lora_B by vLLM module key - lora_a: dict[str, torch.Tensor] = {} - lora_b: dict[str, torch.Tensor] = {} - for peft_key, tensor in adapter_weights.items(): - if ".lora_A." in peft_key: - module_key = self._peft_to_vllm(peft_key.replace(".lora_A.weight", "")) - lora_a[module_key] = tensor - elif ".lora_B." in peft_key: - module_key = self._peft_to_vllm(peft_key.replace(".lora_B.weight", "")) - lora_b[module_key] = tensor - - params_dict = dict(self.named_parameters()) - modules_dict = dict(self.named_modules()) - - def _add_delta(name: str, delta: torch.Tensor) -> bool: - # Try stacked/fused params first (qkv_proj, gate_up_proj) - for fused_name, orig_name, shard_id in self._STACKED_PARAMS_MAPPING: - if orig_name not in name: - continue - fused_param_name = name.replace(orig_name, fused_name) - if fused_param_name not in params_dict: - continue - param = params_dict[fused_param_name] - module_path = fused_param_name.rsplit(".weight", 1)[0] - module = modules_dict.get(module_path) - if module is None: - continue - - tp_rank = get_tensor_model_parallel_rank() - tp_size = get_tensor_model_parallel_world_size() - - if hasattr(module, "_get_shard_offset_mapping"): - # QKVParallelLinear: string shard_id ("q", "k", "v") - shard_offset = module._get_shard_offset_mapping(shard_id) - if shard_offset is not None: - shard_size = delta.shape[0] // tp_size - tp_delta = delta.narrow(0, tp_rank * shard_size, shard_size) - shard = param.data[shard_offset : shard_offset + shard_size] - param.data[shard_offset : shard_offset + shard_size] = ( - shard.float() + tp_delta.to(shard.device) - ).to(shard.dtype) - return True - elif hasattr(module, "output_sizes") and isinstance(shard_id, int): - # MergedColumnParallelLinear: integer shard_id (0, 1) - shard_size = module.output_sizes[shard_id] // tp_size - shard_offset = sum( - s // tp_size for s in module.output_sizes[:shard_id] - ) - tp_delta = delta.narrow( - 0, - tp_rank * (delta.shape[0] // tp_size), - delta.shape[0] // tp_size, - ) - shard = param.data[shard_offset : shard_offset + shard_size] - param.data[shard_offset : shard_offset + shard_size] = ( - shard.float() + tp_delta.to(shard.device) - ).to(shard.dtype) - return True - # Direct param (o_proj, down_proj) - if name in params_dict: - param = params_dict[name] - # Under TP, param is already sharded but delta is full-size. - # Slice delta to match: dim 0 for column-parallel, dim 1 for - # row-parallel. - if delta.shape != param.data.shape: - tp_rank = get_tensor_model_parallel_rank() - for dim in range(delta.dim()): - if delta.shape[dim] != param.data.shape[dim]: - shard_size = param.data.shape[dim] - offset = tp_rank * shard_size - delta = delta.narrow(dim, offset, shard_size) - break - merged = param.data.float() + delta.to(param.device) - param.data = merged.to(param.dtype) - return True - return False - - merge_device = next(self.parameters()).device - merged = 0 - for module_key in sorted(lora_a): - if module_key not in lora_b: - logger.warning("LoRA B missing for %s, skipping", module_key) - continue - A = lora_a[module_key].to(merge_device).float() - B = lora_b[module_key].to(merge_device).float() - delta = scaling * (B @ A) - if _add_delta(module_key + ".weight", delta): - merged += 1 - else: - logger.debug("LoRA target not found on this PP rank: %s", module_key) - - return merged - - def _apply_adapter(self) -> None: - """Full-merge entry point: called when config.adapter_path is set.""" - adapter_path = getattr(self.config, "adapter_path", None) - if not adapter_path: - return - logger.info("Full-merge LoRA from %s", adapter_path) - adapter_config, adapter_weights = self._load_adapter(adapter_path) - - if adapter_config.get("modules_to_save"): - raise ValueError( - "Adapter has modules_to_save — only LM-only adapters " - "(no modules_to_save) are supported." - ) - - n = self._merge_lora_deltas(adapter_config, adapter_weights) - logger.info("Merged %d LoRA pairs into base weights", n) - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) - loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - self._apply_adapter() - return loaded + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) From 0ba3f1772fa9c61d813804cb674458563e85d764 Mon Sep 17 00:00:00 2001 From: artemspector Date: Tue, 21 Apr 2026 13:34:35 +0300 Subject: [PATCH 17/17] Set is_available_online=False for Granite4Vision registry entry Model ibm-granite/granite-vision-4.1-4b is not yet public; setting is_available_online=False prevents CI from attempting to download it. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: artemspector --- tests/models/registry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/models/registry.py b/tests/models/registry.py index 4a7ab853c05e..ea1f0190562e 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -901,6 +901,7 @@ def check_available_online( ), "Granite4VisionForConditionalGeneration": _HfExamplesInfo( "ibm-granite/granite-vision-4.1-4b", + is_available_online=False, ), "GraniteVision": _HfExamplesInfo("ibm-granite/granite-vision-3.3-2b"), "GraniteSpeechForConditionalGeneration": _HfExamplesInfo(