Assertion Failed in mamba2 Module During Training: "Unexpected mma -> mma layout conversion"

[xiaosa269]

Thank you very much for your work. We encountered the following issue while training with the mamba2 module:

At the beginning, the training works fine, but after some time, the program crashes with the following error:

python: /project/lib/Analysis/Allocation.cpp:40: std::pair<llvm::SmallVector, llvm::SmallVector > mlir::triton::getCvtOrder(mlir::Attribute, mlir::Attribute): Assertion `!(srcMmaLayout && dstMmaLayout) && "Unexpected mma -> mma layout conversion"' failed. Aborted (core dumped)

When we switch from the mamba2 module to the mamba module, the issue no longer occurs.

Do you have any suggestions on how to resolve this? We would greatly appreciate your help!