From 766ce2bb6bb614075fa7b0998c06d257f9c49f22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Ganczarenko?= Date: Tue, 9 Jun 2026 19:29:14 +0300 Subject: [PATCH] Fix MiDashengLM TP>1 crash in audio encoder attention (#44408) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Michał Ganczarenko Co-authored-by: Kunshang Ji --- vllm/model_executor/models/midashenglm.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/models/midashenglm.py b/vllm/model_executor/models/midashenglm.py index 62aaed46f6..5ecc92e4d0 100644 --- a/vllm/model_executor/models/midashenglm.py +++ b/vllm/model_executor/models/midashenglm.py @@ -229,10 +229,10 @@ class DashengAttention(nn.Module): ) def forward(self, x: torch.Tensor, mask: torch.Tensor | None = None): - B, N, C = x.shape + B, N, _ = x.shape qkv, _ = self.qkv(x) - qkv = qkv.reshape(B, N, 3, self.num_heads, C // self.num_heads) + qkv = qkv.reshape(B, N, 3, self.num_heads, self.head_dim) qkv = qkv.permute(2, 0, 3, 1, 4) q, k, v = qkv.unbind(0) @@ -243,7 +243,7 @@ class DashengAttention(nn.Module): attn_mask=mask[:, None, None, :] if mask is not None else None, ) - x = x.transpose(1, 2).reshape(B, N, C) + x = x.transpose(1, 2).reshape(B, N, self.q_size) x, _ = self.proj(x) return x