Precision used in multi head attention layer

[kimdaeun00]

Hi. Thanks for sharing your work!

I have a question for the precision used in multi-head attention layer.
In the current code(Flux), it seems that the activations are not casted to FP16 before the attention layer(mha_forward). Do other models, like PixArt-Sigma, use cast_fp16 in their mha layers?

[lmxyy]

The activations in FLUX are in BF16 precision. We do not use cast_fp16 in other models.