Enable FP32 Accumulate in Flash Attention and Flash Decode

[caixunshiren]

Description

We do not have support for fp32 accumulate in sdpa family kernels. This becomes a problem when number of chunks gets large and we see diverging pcc from ground truth. For models that requires 128K sequel, this is problematic.

This issue tracks the enabling of fp32 accumulate in the following kernels:

round 1:

• sdpa (bf16 cbs, fp32 accum)
• sdpa decode (bf16 cbs, fp32 accum)

round 2:

• sdpa (fp32 cbs, fp32 accum)
• sdpa decode (fp32 cbs, fp32 accum)

FYI @cglagovichTT

[caixunshiren]

Update:

• It appears that the largest amount of pcc drop are attributes to math approximation, not fp32 accumulate. The issue is tracked here: Fix Diverging PCC issue in SDPA Kernels #13866
• Based on my experiments, there is no issue of fp32 accumulate with bf16 cbs. With fp32 cbs, we need some inputs to mul_block_inplace and add_block_inplace to be bf16 cbs during im and stat accumulation/updates, otherwise we get pcc degradation compared to bf16. My WIP work is on branch sdpa-fp32-investigations
• I also found out that reconfig_dataformat doesn't work as expected and could hang/wrong result for in/out_cb. This explains the flash decode fp32 accumulate hang in reducer cores that I saw earlier: Deterministic Hang with unpack_reconfig_data_format at a specific spot in Flash Decode Kernel #9608
• Actions for now: Support fp32 accumulate with bf16 cbs on sdpa/sdpa-decode kernels. Leave fp32 cbs for future.