The calculation of scores_max in the flash_attn example #948

You're right that the code you quoted never combines the new block max with the previous one. That has since been added: #1269 (merged 2025-11-17) put

for i in T.Parallel(block_M):
    scores_max[i] = T.max(scores_max[i], scores_max_prev[i])

right after the T.reduce_max in this example and in other attention examples. #1148, merged earlier the same day, made the same change in example_mha_fwd_bhsd.py.

Without that line, scores_max is the max of the current block only. As long as that value is finite and the rescale factor doesn't overflow, the result can still be right, because scores_scale rescales acc_o and logsum to the new reference value. It goes wrong in two ways that I could reproduce:

• Causal masking. With is_causal=True and block_M larger than block_N, the last key block that every full query tile visits is entirely in the future for some of the tile's rows. For those rows acc_s is all -inf, reduce_max gives -inf, and acc_s[i, j] * scale - scores_max[i] * scale is -inf minus -inf, which is NaN. With the running max, those rows keep the finite max from earlier blocks.
• Overflow. When a later key block has a much smaller row max than an earlier one, exp2(scores_max_prev[i] * scale - scores_max[i] * scale) overflows to inf, and the output becomes NaN.

I ran example_gqa_fwd_bshd.py from main (74369d9), and a copy with only those two lines removed, with tilelang 0.1.14 (the current PyPI release) on an RTX 5090, fp16, batch 1, 1 head, dim 128, num_stages=2, threads=128:

block_M=128, block_N=64 is one of the configs the example's get_configs() returns. In the two constant-input rows every Q entry is c (0.05 or 3.0), the first 64 keys are all c, the last 64 keys are all -c, and V is 1.0 in the first key block and 2.0 in the second.

On your second question: yes. The reduce_max docstring says that with clear=True "the output buffer will first be initialized to -inf", and replacing the T.fill and clear=False pair with T.reduce_max(acc_s, scores_max, dim=1, clear=True) gave the same max abs diff as main in all three rows. With clear=False and no T.fill, the reduction starts from the previous row max, so it produces the running max in one step: dropping both the T.fill and the new loop also gave the same max abs diff as main in all three rows.

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?

You're right that the code you quoted never combines the new block max with the previous one. That has since been added: #1269 (merged 2025-11-17) put

for i in T.Parallel(block_M):
    scores_max[i] = T.max(scores_max[i], scores_max_prev[i])

right after the T.reduce_max in this example and in other attention examples. #1148, merged earlier the same day, made the same change in example_mha_fwd_bhsd.py.

Without that line, scores_max is the max of the current block only. As long as that value is finite and the rescale factor doesn't overflow, the result can still be right, because scores_scale rescales acc_o and logsum to the new reference value. It goes wrong in two ways that I could reproduce:

• Causal masking. With is_causal=True and block_M larger than block_N, the last key block that every full query tile visits is entirely in the future for some of the tile's rows. For those rows acc_s is all -inf, reduce_max gives -inf, and acc_s[i, j] * scale - scores_max[i] * scale is -inf minus -inf, which is NaN. With the running max, those rows keep the finite max from earlier blocks.
• Overflow. When a later key block has a much smaller row max than an earlier one, exp2(scores_max_prev[i] * scale - scores_max[i] * scale) overflows to inf, and the output becomes NaN.

I ran example_gqa_fwd_bshd.py from main (74369d9), and a copy with only those two lines removed, with tilelang 0.1.14 (the current PyPI release) on an RTX 5090, fp16, batch 1, 1 head, dim 128, num_stages=2, threads=128:

case	example on main	without the max line
causal, seq_len 256, block_M=128, block_N=64, torch.randn inputs with torch.manual_seed(0)	max abs diff vs ref_program 0.001953125	NaN in 16384 of 32768 outputs (query rows 0-63 and 128-191)
non-causal, seq_len 128, 64x64 blocks, constant inputs, raw scores 0.32 and -0.32	max abs diff 0.0009765625	max abs diff 0.0
non-causal, seq_len 128, 64x64 blocks, constant inputs, raw scores 1152 and -1152	max abs diff 0.0	NaN in 16384 of 16384 outputs

block_M=128, block_N=64 is one of the configs the example's get_configs() returns. In the two constant-input rows every Q entry is c (0.05 or 3.0), the first 64 keys are all c, the last 64 keys are all -c, and V is 1.0 in the first key block and 2.0 in the second.

On your second question: yes. The reduce_max docstring says that with clear=True "the output buffer will first be initialized to -inf", and replacing the T.fill and clear=False pair with T.reduce_max(acc_s, scores_max, dim=1, clear=True) gave the same max abs diff as main in all three rows. With clear=False and no T.fill, the reduction starts from the previous row max, so it produces the running max in one step: dropping both the T.fill and the new loop also gave the same max abs diff as main in all three rows.

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?

You're right that the code you quoted never combines the new block max with the previous one. That has since been added: #1269 (merged 2025-11-17) put

for i in T.Parallel(block_M):
    scores_max[i] = T.max(scores_max[i], scores_max_prev[i])

right after the T.reduce_max in this example and in other attention examples. #1148, merged earlier the same day, made the same change in example_mha_fwd_bhsd.py.

Without that line, scores_max is the max of the current block only. As long as that value is finite and the rescale factor doesn't overflow, the result can still be right, because scores_scale rescales acc_o and logsum to the new reference value. It goes wrong in two ways that I could reproduce:

• Causal masking. With is_causal=True and block_M larger than block_N, the last key block that every full query tile visits is entirely in the future for some of the tile's rows. For those rows acc_s is all -inf, reduce_max gives -inf, and acc_s[i, j] * scale - scores_max[i] * scale is -inf minus -inf, which is NaN. With the running max, those rows keep the finite max from earlier blocks.
• Overflow. When a later key block has a much smaller row max than an earlier one, exp2(scores_max_prev[i] * scale - scores_max[i] * scale) overflows to inf, and the output becomes NaN.

I ran example_gqa_fwd_bshd.py from main (74369d9), and a copy with only those two lines removed, with tilelang 0.1.14 (the current PyPI release) on an RTX 5090, fp16, batch 1, 1 head, dim 128, num_stages=2, threads=128:

case	example on main	without the max line
causal, seq_len 256, block_M=128, block_N=64, torch.randn inputs with torch.manual_seed(0)	max abs diff vs ref_program 0.001953125	NaN in 16384 of 32768 outputs (query rows 0-63 and 128-191)
non-causal, seq_len 128, 64x64 blocks, constant inputs, raw scores 0.32 and -0.32	max abs diff 0.0009765625	max abs diff 0.0
non-causal, seq_len 128, 64x64 blocks, constant inputs, raw scores 1152 and -1152	max abs diff 0.0	NaN in 16384 of 16384 outputs

block_M=128, block_N=64 is one of the configs the example's get_configs() returns. In the two constant-input rows every Q entry is c (0.05 or 3.0), the first 64 keys are all c, the last 64 keys are all -c, and V is 1.0 in the first key block and 2.0 in the second.

On your second question: yes. The reduce_max docstring says that with clear=True "the output buffer will first be initialized to -inf", and replacing the T.fill and clear=False pair with T.reduce_max(acc_s, scores_max, dim=1, clear=True) gave the same max abs diff as main in all three rows. With clear=False and no T.fill, the reduction starts from the previous row max, so it produces the running max in one step: dropping both the T.fill and the new loop also gave the same max abs diff as main in all three rows.

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?

You're right that the code you quoted never combines the new block max with the previous one. That has since been added: #1269 (merged 2025-11-17) put

for i in T.Parallel(block_M):
    scores_max[i] = T.max(scores_max[i], scores_max_prev[i])

right after the T.reduce_max in this example and in other attention examples. #1148, merged earlier the same day, made the same change in example_mha_fwd_bhsd.py.

Without that line, scores_max is the max of the current block only. As long as that value is finite and the rescale factor doesn't overflow, the result can still be right, because scores_scale rescales acc_o and logsum to the new reference value. It goes wrong in two ways that I could reproduce:

• Causal masking. With is_causal=True and block_M larger than block_N, the last key block that every full query tile visits is entirely in the future for some of the tile's rows. For those rows acc_s is all -inf, reduce_max gives -inf, and acc_s[i, j] * scale - scores_max[i] * scale is -inf minus -inf, which is NaN. With the running max, those rows keep the finite max from earlier blocks.
• Overflow. When a later key block has a much smaller row max than an earlier one, exp2(scores_max_prev[i] * scale - scores_max[i] * scale) overflows to inf, and the output becomes NaN.

I ran example_gqa_fwd_bshd.py from main (74369d9), and a copy with only those two lines removed, with tilelang 0.1.14 (the current PyPI release) on an RTX 5090, fp16, batch 1, 1 head, dim 128, num_stages=2, threads=128:

case	example on main	without the max line
causal, seq_len 256, block_M=128, block_N=64, torch.randn inputs with torch.manual_seed(0)	max abs diff vs ref_program 0.001953125	NaN in 16384 of 32768 outputs (query rows 0-63 and 128-191)
non-causal, seq_len 128, 64x64 blocks, constant inputs, raw scores 0.32 and -0.32	max abs diff 0.0009765625	max abs diff 0.0
non-causal, seq_len 128, 64x64 blocks, constant inputs, raw scores 1152 and -1152	max abs diff 0.0	NaN in 16384 of 16384 outputs

block_M=128, block_N=64 is one of the configs the example's get_configs() returns. In the two constant-input rows every Q entry is c (0.05 or 3.0), the first 64 keys are all c, the last 64 keys are all -c, and V is 1.0 in the first key block and 2.0 in the second.

On your second question: yes. The reduce_max docstring says that with clear=True "the output buffer will first be initialized to -inf", and replacing the T.fill and clear=False pair with T.reduce_max(acc_s, scores_max, dim=1, clear=True) gave the same max abs diff as main in all three rows. With clear=False and no T.fill, the reduction starts from the previous row max, so it produces the running max in one step: dropping both the T.fill and the new loop also gave the same max abs diff as main in all three rows.

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?

You're right that the code you quoted never combines the new block max with the previous one. That has since been added: #1269 (merged 2025-11-17) put

for i in T.Parallel(block_M):
    scores_max[i] = T.max(scores_max[i], scores_max_prev[i])

right after the T.reduce_max in this example and in other attention examples. #1148, merged earlier the same day, made the same change in example_mha_fwd_bhsd.py.

Without that line, scores_max is the max of the current block only. As long as that value is finite and the rescale factor doesn't overflow, the result can still be right, because scores_scale rescales acc_o and logsum to the new reference value. It goes wrong in two ways that I could reproduce:

• Causal masking. With is_causal=True and block_M larger than block_N, the last key block that every full query tile visits is entirely in the future for some of the tile's rows. For those rows acc_s is all -inf, reduce_max gives -inf, and acc_s[i, j] * scale - scores_max[i] * scale is -inf minus -inf, which is NaN. With the running max, those rows keep the finite max from earlier blocks.
• Overflow. When a later key block has a much smaller row max than an earlier one, exp2(scores_max_prev[i] * scale - scores_max[i] * scale) overflows to inf, and the output becomes NaN.

I ran example_gqa_fwd_bshd.py from main (74369d9), and a copy with only those two lines removed, with tilelang 0.1.14 (the current PyPI release) on an RTX 5090, fp16, batch 1, 1 head, dim 128, num_stages=2, threads=128:

case	example on main	without the max line
causal, seq_len 256, block_M=128, block_N=64, torch.randn inputs with torch.manual_seed(0)	max abs diff vs ref_program 0.001953125	NaN in 16384 of 32768 outputs (query rows 0-63 and 128-191)
non-causal, seq_len 128, 64x64 blocks, constant inputs, raw scores 0.32 and -0.32	max abs diff 0.0009765625	max abs diff 0.0
non-causal, seq_len 128, 64x64 blocks, constant inputs, raw scores 1152 and -1152	max abs diff 0.0	NaN in 16384 of 16384 outputs

block_M=128, block_N=64 is one of the configs the example's get_configs() returns. In the two constant-input rows every Q entry is c (0.05 or 3.0), the first 64 keys are all c, the last 64 keys are all -c, and V is 1.0 in the first key block and 2.0 in the second.

On your second question: yes. The reduce_max docstring says that with clear=True "the output buffer will first be initialized to -inf", and replacing the T.fill and clear=False pair with T.reduce_max(acc_s, scores_max, dim=1, clear=True) gave the same max abs diff as main in all three rows. With clear=False and no T.fill, the reduction starts from the previous row max, so it produces the running max in one step: dropping both the T.fill and the new loop also gave the same max abs diff as main in all three rows.

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?