use triton kernel for training need warm up？

[mdy666]

Describe the issue

• I write a custom attention kernel fully by triton, it contain many triton kernels(>15).
• In training, it need nearly 200 steps(50 minutes) to achieve the best performance. If I use official flash-attn kernel, It will not appear this phenomenon .
• How can I solve this issue？
• PS: I have tuned the triton config and don't use autotune in training.

Environment details

Triton:3.2