Why mamba2 is much slower than transformer (flash attn)?

[TimothyChen225]

No description provided.

[Dexterp37]

See #355 and #367

[AlwaysFHao]

I also encountered the same problem. Referring to the author's explanation in the #355 it is necessary to pre compile the model, which can be referred to https://discuss.pytorch.org/t/how-to-use-torch-compile-with-cuda-graphs-when-using-gradient-activation-checkpointing/179466.
You can wrap your model with the following code, although the first round of forward propagation will be very slow, the subsequent propagation will be very fast:
model_compile = torch.compile(model, mode="reduce-overhead")

[TimothyChen225]

I also encountered the same problem. Referring to the author's explanation in the #355 it is necessary to pre compile the model, which can be referred to https://discuss.pytorch.org/t/how-to-use-torch-compile-with-cuda-graphs-when-using-gradient-activation-checkpointing/179466. You can wrap your model with the following code, although the first round of forward propagation will be very slow, the subsequent propagation will be very fast: model_compile = torch.compile(model, mode="reduce-overhead")

thks for help, but it is still slow