question: backward pass

[pengzhangzhi]

Hi,

In the current imple. e.g., the flash attention version, the backward pass isn't implemented. I wonder if tilelang will aotumatically implement the backward pass or some other reason?

Also a side question, I want to optimize some existing ops, the first thing is to profile the sub ops. I wonder if you have any ideas on profiling the operations.

[LeiWang1999]

CC @chengyupku

[LeiWang1999]

@pengzhangzhi Thanks for your attention! We do have bwd kernels with tilelang and will be uploading some kernels soon. Sorry but tilelang currently cannolt apply auto diff for a fwd kernel (though I believe that will be a great research topic).

[pengzhangzhi]

sounds good! I have keen interest in using tilelang instead of triton to optimize an op for training (like what i emailed u before). So the backward pass is important to me : ) Could you please let me know if the backward pass code is available for flashattention.

Not an expert in this, but since triton doesn't support auto diff, i wonder how feasible it is for tilelang to do that.

[chengyupku]

@pengzhangzhi Hi, I have added an example of mha backward, please refer to example_mha_bwd.py.