A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 06:30 UTC
fix minicpmo4.5: fix attn_mask in vit attn && fix resampler pos_emb i… (#34127)
Signed-off-by: tc-mb <caitianchi@modelbest.cn> Co-authored-by: hezhihui <hezhihui@modelbest.cn>
tc-mb committed
Mar 5, 2026 at 01:46 UTC
bfdb512f111156a8f455dd9f396c1d15ba5bf655
2 files changed
+69
-7
vllm/model_executor/models/idefics2_vision_model.py
+67
-5
@@ -22,6 +22,7 @@ from collections.abc import Iterable
22
23
import torch
24
from torch import nn
25
+from torch.nn import functional as F
26
from transformers.models.idefics2.configuration_idefics2 import (
27
Idefics2Config,
28
Idefics2VisionConfig,
@@ -172,14 +173,41 @@ class Idefics2VisionAttention(nn.Module):
173
def forward(
174
self,
175
hidden_states: torch.Tensor,
176
+ attention_mask: torch.Tensor | None = None,
177
) -> torch.Tensor:
178
qkv, _ = self.qkv_proj(
179
hidden_states
180
) # batch_size, q_len, 3 * num_heads_per_partition * head_dim
181
query_states, key_states, value_states = qkv.chunk(3, dim=-1)
182
181
- # Use unified MMEncoderAttention implementation
182
- out = self.attn(query_states, key_states, value_states)
183
+ # If attention_mask is provided, prefer Torch SDPA so the mask is
184
+ # correctly applied (aligns with HuggingFace NaViT SigLIP behavior).
185
+ if attention_mask is None:
186
+ # Use unified MMEncoderAttention implementation
187
+ out = self.attn(query_states, key_states, value_states)
188
+ else:
189
+ bsz, q_len = query_states.size()[:2]
190
+ kv_len = key_states.size(1)
191
+
192
+ query = query_states.view(
193
+ bsz, q_len, self.num_heads_per_partition, self.head_dim
194
+ ).transpose(1, 2)
195
+ key = key_states.view(
196
+ bsz, kv_len, self.num_heads_per_partition, self.head_dim
197
+ ).transpose(1, 2)
198
+ value = value_states.view(
199
+ bsz, kv_len, self.num_heads_per_partition, self.head_dim
200
+ ).transpose(1, 2)
201
+
202
+ out = F.scaled_dot_product_attention(
203
+ query,
204
+ key,
205
+ value,
206
+ attn_mask=attention_mask,
207
+ dropout_p=0.0,
208
+ scale=self.scale,
209
+ )
210
+ out = out.transpose(1, 2).reshape(bsz, q_len, -1)
211
attn_output, _ = self.out_proj(out)
212
return attn_output
213
@@ -245,6 +273,7 @@ class Idefics2EncoderLayer(nn.Module):
273
def forward(
274
self,
275
hidden_states: torch.Tensor,
276
+ attention_mask: torch.Tensor | None = None,
277
) -> torch.Tensor:
278
"""
279
Args:
@@ -254,7 +283,7 @@ class Idefics2EncoderLayer(nn.Module):
283
"""
284
residual = hidden_states
285
hidden_states = self.layer_norm1(hidden_states)
257
- hidden_states = self.self_attn(hidden_states)
286
+ hidden_states = self.self_attn(hidden_states, attention_mask=attention_mask)
287
hidden_states += residual
288
residual = hidden_states
289
hidden_states = self.layer_norm2(hidden_states)
@@ -304,6 +333,7 @@ class Idefics2Encoder(nn.Module):
333
def forward(
334
self,
335
inputs_embeds: torch.Tensor,
336
+ attention_mask: torch.Tensor | None = None,
337
) -> torch.Tensor:
338
r"""
339
Args:
@@ -316,7 +346,7 @@ class Idefics2Encoder(nn.Module):
346
"""
347
hidden_states = inputs_embeds
348
for encoder_layer in self.layers:
319
- layer_outputs = encoder_layer(hidden_states)
349
+ layer_outputs = encoder_layer(hidden_states, attention_mask=attention_mask)
350
hidden_states = layer_outputs
351
return hidden_states
352
@@ -370,15 +400,47 @@ class Idefics2VisionTransformer(nn.Module):
400
patch_attention_mask: torch.BoolTensor | None = None,
401
tgt_sizes: torch.IntTensor | None = None,
402
) -> torch.Tensor:
403
+ batch_size = pixel_values.size(0)
404
+
405
+ if patch_attention_mask is None:
406
+ # No mask provided - create default all-ones mask for embeddings
407
+ # and skip attention masking (no padding to mask)
408
+ patch_attention_mask = torch.ones(
409
+ size=(
410
+ batch_size,
411
+ pixel_values.size(2) // self.config.patch_size,
412
+ pixel_values.size(3) // self.config.patch_size,
413
+ ),
414
+ dtype=torch.bool,
415
+ device=pixel_values.device,
416
+ )
417
+ flat_patch_mask = None
418
+ else:
419
+ flat_patch_mask = patch_attention_mask.view(batch_size, -1)
420
+
421
hidden_states = self.embeddings(
422
pixel_values=pixel_values,
423
patch_attention_mask=patch_attention_mask,
424
tgt_sizes=tgt_sizes,
425
)
426
+
427
+ # Align with HuggingFace NaViT SigLIP in MiniCPMV/O:
428
+ # - if patch_attention_mask was None, skip attention masking
429
+ # - if any padding exists, create an additive 4D mask and pass it
430
+ # to attention; else skip mask for performance.
431
+ if flat_patch_mask is None or not torch.any(~flat_patch_mask):
432
+ attention_mask = None
433
+ else:
434
+ # Additive mask: masked positions receive a large negative value.
435
+ # Shape: (B, 1, 1, L) broadcastable to (B, H, Q, K).
436
+ min_val = torch.finfo(hidden_states.dtype).min
437
+ attention_mask = (~flat_patch_mask).to(dtype=hidden_states.dtype) * min_val
438
+ attention_mask = attention_mask[:, None, None, :]
439
+
440
if self.use_data_parallel:
441
encoder_outputs = run_dp_sharded_vision_model(hidden_states, self.encoder)
442
else:
381
- encoder_outputs = self.encoder(hidden_states)
443
+ encoder_outputs = self.encoder(hidden_states, attention_mask=attention_mask)
444
last_hidden_state = self.post_layernorm(encoder_outputs)
445
return last_hidden_state
446
vllm/model_executor/models/minicpmv.py
+2
-2
@@ -387,8 +387,8 @@ class Resampler4_5(Resampler2_5):
387
pos_embed_2d, batch_first=True, padding_value=0.0
388
).permute(1, 0, 2) # BLD => L * B * D
389
390
- k = x
391
- v = x + pos_embed_2d
390
+ k = x + pos_embed_2d
391
+ v = x
392
if pos_embed_temporal:
393
k += torch.stack(pos_embed_temporal, dim=0)
394
bs = len(temporal_ids)