tests/pytorch/test_fusible_ops.py · c9e8e305f707d94cc1f829ea6899c4cfdd259548 · OpenDAS / TransformerEngine

[PyTorch] Add option in activation ops to cache input in FP8 (#1665) · 04642bf8

Tim Moon authored Apr 11, 2025



* Add option to cache activation input in FP8
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Avoid casting to FP8 transpose
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Skip input caching if device is not supported
Signed-off-by: Tim Moon <tmoon@nvidia.com>

* Add documentation that FP8 input caching is experimental
Signed-off-by: Tim Moon <tmoon@nvidia.com>

---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>

04642bf8

test_fusible_ops.py 70.5 KB

Replace test_fusible_ops.py