They don't mention it explicitly except in one place: > GPT-NeoX achieves 150 TFLOP/s/A100 with normal attention and 180 FLOP/s/A100 with Flash Attention. Thi

They don't mention it explicitly except in one place: > GPT-NeoX achieves 150 TFLOP/s/A100 with normal attention and 180 FLOP/s/A100 with Flash Attention. This advice implies they are using flash attention.