transformer_engine/pytorch/module/linear.py · 2a95efd39128955081c60b67d49351d89f003324 · OpenDAS / TransformerEngine

Add NVTX ranges to FP8 amax AR and grad output preprocessing (#1530) · de06a34c

Jaemin Choi authored Mar 06, 2025



Add NVTX ranges
Signed-off-by: Jaemin Choi <jaeminc@nvidia.com>
Co-authored-by: Jaemin Choi <jaeminc@nvidia.com>
Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>

de06a34c

linear.py 46 KB

Replace linear.py