something wrong with update_out_and_lse

[fy1214]

你好，我私下profile了一下ring attention，发现有一些奇怪的现象：
环境：
8 gpu
pytorch2.12 + cu12.2

现象，当我使用输入为[1,2048,32,128]的向量时，profile单个ring-attention过程，得到的结果如下：

我发现有三次update_out_and_lse计算时间比较长，我把它打印出来结果如下

Rank 0 第0次计算update_out_and_lse的时间为: 0.0026123523712158203
Rank 0 第1次计算update_out_and_lse的时间为: 0.04419445991516113
Rank 0 第2次计算update_out_and_lse的时间为: 0.2504105567932129
Rank 0 第3次计算update_out_and_lse的时间为: 0.20969200134277344
Rank 0 第4次计算update_out_and_lse的时间为: 0.00016427040100097656
Rank 0 第5次计算update_out_and_lse的时间为: 0.00013637542724609375
Rank 0 第6次计算update_out_and_lse的时间为: 0.00012493133544921875
Rank 0 第7次计算update_out_and_lse的时间为: 0.0001404285430908203

中间的update_out_and_lse耗时很长，不知道有没有头绪