A high-throughput and memory-efficient inference and serving engine for LLMs
Fix MiDashengLM TP>1 crash in audio encoder attention (#44408)
Signed-off-by: Michał Ganczarenko <michal.ganczarenko@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Michał Ganczarenko committed
Jun 9, 2026 at 19:29 UTC
766ce2bb6bb614075fa7b0998c06d257f9c49f22
1 file changed
+3
-3
vllm/model_executor/models/midashenglm.py
+3
-3
@@ -229,10 +229,10 @@ class DashengAttention(nn.Module):
229
)
230
231
def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None):
232
- B, N, C = x.shape
232
+ B, N, _ = x.shape
233
234
qkv, _ = self.qkv(x)
235
- qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads)
235
+ qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim)
236
qkv = qkv.permute(2, 0, 3, 1, 4)
237
q, k, v = qkv.unbind(0)
238
@@ -243,7 +243,7 @@ class DashengAttention(nn.Module):
243
attn_mask=mask[:, None, None, :] if mask is not None else None,
244
)
245
246
- x = x.transpose(1, 2).reshape(B, N, C)
246
+ x = x.transpose(1, 2).reshape(B, N, self.q_size)
247
x, _ = self.proj(x)
248
return x
249