@@ -452,2 +424,2 @@
-        if attn_mask:
-            _attn_mask = self._process_mask(attn_mask, _batch_size)
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"

@@ -454,0 +425,4 @@
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+        if self._qkv_same_embed_dim is False:
+            self.q_proj_weight = Parameter(ms.ops.zeros((embed_dim, embed_dim), dtype=dtype))

@@ -511,3 +511,3 @@
 
 #TODO: multiheadattention need reconstruct
-'''
+

@@ -510,8 +511,9 @@ def test_hardsigmoid():
     assert ms_out.asnumpy().dtype == torch_out.numpy().dtype
 
 #TODO: multiheadattention need reconstruct

@@ -514,0 +441,4 @@
+            self.in_proj_bias = Parameter(empty(3 * embed_dim, dtype=dtype))
+        else:
+            self.in_proj_bias = None
+        # TODO: NonDynamicallyQuantizableLinear

@@ -525,0 +521,4 @@
+            attn_output = attn_output.swapaxes(1, 0)
+        if need_weights:
+            return attn_output, attn_output_weights
+        return (attn_output,)