The calculation of scores_max in the flash_attn example

[chenghuaWang]

In the FA2 algorithm, there is a step that updates $max_{new}$, which requires taking the element-wise maximum between the current block's max and the previous $max_{new}$.

However, in the code at

tilelang/examples/flash_attention/example_gqa_fwd_bshd.py

Line 116 in 7fb0677

def Softmax(

, I noticed the following:

T.copy(scores_max, scores_max_prev)
T.fill(scores_max, -T.infinity(accum_dtype))
T.reduce_max(acc_s, scores_max, dim=1, clear=False)

It seems that the operation scores_max = max(scores_max_prev, scores_max) is not actually performed due to T.fill. Could you clarify where my understanding might be incorrect?

Additionally, would T.fill(scores_max, -T.infinity(accum_dtype)) be equivalent to setting clear=True in T.reduce_max?