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 14, 23:45 UTC
Remove xformers requirement for Mistral-format Pixtral and Mistral3 (#21154)
Signed-off-by: Wenchen Lo <charles761013@gmail.com>
Wenchen Lo committed
Jul 26, 2025 at 16:20 UTC
6c66f28fa5dc88ce6f7ab30dfa733f9ddb927d3c
1 file changed
+18
-3
vllm/model_executor/models/pixtral.py
+18
-3
@@ -671,7 +671,19 @@ class Attention(nn.Module):
671
v = v.reshape(batch, patches, self.n_heads, self.head_dim)
672
673
q, k = apply_rotary_emb_vit(q, k, freqs_cis=freqs_cis)
674
- out = xops.memory_efficient_attention(q, k, v, attn_bias=mask)
674
+
675
+ if USE_XFORMERS_OPS:
676
+ out = xops.memory_efficient_attention(q, k, v, attn_bias=mask)
677
+ else:
678
+ q = q.transpose(1, 2)
679
+ k = k.transpose(1, 2)
680
+ v = v.transpose(1, 2)
681
+ out = nn.functional.scaled_dot_product_attention(q,
682
+ k,
683
+ v,
684
+ attn_mask=mask)
685
+ out = out.transpose(1, 2)
686
+
687
out = out.reshape(batch, patches, self.n_heads * self.head_dim)
688
return self.wo(out)
689
@@ -814,8 +826,11 @@ class VisionTransformer(nn.Module):
826
mask = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(
827
[p.shape[-2] * p.shape[-1] for p in patch_embeds_list], )
828
else:
817
- raise ImportError("Xformers is required for Pixtral inference "
818
- "with the Mistral format")
829
+ from transformers.models.pixtral.modeling_pixtral import (
830
+ generate_block_attention_mask)
831
+ mask = generate_block_attention_mask(
832
+ [p.shape[-2] * p.shape[-1] for p in patch_embeds_list],
833
+ patch_embeds)
834
out = self.transformer(patch_embeds, mask=mask, freqs_cis=freqs_cis)
835
836
# squeeze dim 0 and split into separate tensors for each image