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