transformer_engine/pytorch/attention.py · c42beef4e1d13f8d99e2892e875cb9e96c2d67f7 · OpenDAS / TransformerEngine

[PyTorch] Make FP8 MHA work with RoPE when CP is on (#1297) · c42beef4

Xin Yao authored Nov 04, 2024



* Let fp8 mha work with rope when cp is on
Signed-off-by: Xin Yao <xiny@nvidia.com>

* fix and update ut
Signed-off-by: Xin Yao <xiny@nvidia.com>

---------
Signed-off-by: Xin Yao <xiny@nvidia.com>

c42beef4

attention.py 401 KB

Replace attention.py