Question about FP8 matmul coverage in FP8-LM about ms-amp

[wkcn]

Hi @stakahashy, thanks for your attention to our work!
Regarding your questions,

1. Are matmul in mult-head attention or flash attention modules executed in fp8?

The two linear projections (QKV=proj1(x), output=proj2(SV), where S is the attention score) in the multi-head attention are executed in FP8, but flash attention (SV = flash_attn(QKV)) is executed in BF16 or FP16. Besides, the two linears in MLP are also executed in FP8.

2. Does the input and positional embedding layer remain fp16 for matmul input?

Yes. It is executed in FP16 or BF16.

[stakahashy]

@wkcn Thank you very much for the precise answers!