transformer_engine/pytorch/attention.py · 49a4535d1addd2c5743a7e280e2f4f2640f0bedf · OpenDAS / TransformerEngine

Add NVTX ranges to categorize execution (#1447) · 49a4535d

Jaemin Choi authored Feb 11, 2025


Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Co-authored-by: Jaemin Choi <jaeminc@nvidia.com>
Co-authored-by: Tim Moon <tmoon@nvidia.com>

49a4535d

attention.py 378 KB

Replace attention.py