Fix SDPA implementation in Qwen2-VL (issues with torch==2.6.0) (#36891)

* fix sdpa implementation

* ruff

* also modify 2_5 for consistency
This commit is contained in:
Manuel Faysse
2025-03-28 09:54:21 +01:00
committed by GitHub
parent 348f3285c5
commit 064cd7cdac
2 changed files with 8 additions and 4 deletions

View File

@@ -316,8 +316,10 @@ class Qwen2_5_VLVisionSdpaAttention(nn.Module):
q = q.transpose(0, 1)
k = k.transpose(0, 1)
v = v.transpose(0, 1)
attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
attn_output = attn_output.transpose(0, 1)
attn_output = F.scaled_dot_product_attention(
q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, dropout_p=0.0
)
attn_output = attn_output.squeeze(0).transpose(0, 1)
attn_output = attn_output.reshape(seq_length, -1)
attn_output = self.proj(attn_output)
return attn_output

View File

@@ -416,8 +416,10 @@ class VisionSdpaAttention(nn.Module):
q = q.transpose(0, 1)
k = k.transpose(0, 1)
v = v.transpose(0, 1)
attn_output = F.scaled_dot_product_attention(q, k, v, attention_mask, dropout_p=0.0)
attn_output = attn_output.transpose(0, 1)
attn_output = F.scaled_dot_product_attention(
q.unsqueeze(0), k.unsqueeze(0), v.unsqueeze(0), attention_mask, dropout_p=0.0
)
attn_output = attn_output.squeeze(0).transpose(0, 1)
attn_output = attn_output.reshape(seq_length, -1)
attn_output = self.proj(attn_output)
return attn_output