export static llama with masked softmax

Differential Revision: D81248691

Pull Request resolved: pytorch#13832

Author: billmguo

examples/models/llama/static_attention.py

@@ -1044,5 +1044,5 @@ def transfer_weight(linear, conv2d):
 
 @register_attention("static_mha")
 class StaticAttentionMHA(StaticAttention):
-    def __init__(self, config: ModelArgs, layer_id: int, rope: Rope):
-        super().__init__(config, layer_id, rope, split_mha=False)
+    def __init__(self, config: ModelArgs, layer_id: int, rope: Rope, **kwargs: Any):
+        super().__init__(config, layer_id, rope, split_mha=False, **kwargs)