[BUG] triton attn not work for mpt-1b-redpajama-200b

[ustcwhy]

Thanks for your wonderful work. I try to use triton for mpt-1b-redpajama-200b. I use the default version of flash attention. However, it doesnot work. Do you meet the similar issues?

The error info:

Traceback (most recent call last):
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 937, in build_triton_ir
generator.visit(fn.parse())
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
return super().visit(node)
File "/opt/conda/lib/python3.8/ast.py", line 371, in visit
return visitor(node)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 183, in visit_Module
ast.NodeVisitor.generic_visit(self, node)
File "/opt/conda/lib/python3.8/ast.py", line 379, in generic_visit
self.visit(item)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
return super().visit(node)
File "/opt/conda/lib/python3.8/ast.py", line 371, in visit
return visitor(node)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 252, in visit_FunctionDef
has_ret = self.visit_compound_statement(node.body)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 177, in visit_compound_statement
self.last_ret_type = self.visit(stmt)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
return super().visit(node)
File "/opt/conda/lib/python3.8/ast.py", line 371, in visit
return visitor(node)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 678, in visit_For
self.visit_compound_statement(node.body)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 177, in visit_compound_statement
self.last_ret_type = self.visit(stmt)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
return super().visit(node)
File "/opt/conda/lib/python3.8/ast.py", line 371, in visit
return visitor(node)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 319, in visit_AugAssign
self.visit(assign)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
return super().visit(node)
File "/opt/conda/lib/python3.8/ast.py", line 371, in visit
return visitor(node)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 301, in visit_Assign
values = self.visit(node.value)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
return super().visit(node)
File "/opt/conda/lib/python3.8/ast.py", line 371, in visit
return visitor(node)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 339, in visit_BinOp
rhs = self.visit(node.right)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 855, in visit
return super().visit(node)
File "/opt/conda/lib/python3.8/ast.py", line 371, in visit
return visitor(node)
File "/opt/conda/lib/python3.8/site-packages/triton/compiler.py", line 797, in visit_Call
return fn(*args, _builder=self.builder, **kws)
File "/opt/conda/lib/python3.8/site-packages/triton/impl/base.py", line 22, in wrapper
return fn(*args, **kwargs)
TypeError: dot() got an unexpected keyword argument 'trans_b'