transformer_engine/jax/csrc/extensions/attention.cpp · 16f3f897c91329ddec9756a85e8ce5dd8f1a79fc · OpenDAS / TransformerEngine

[JAX] Splitting `csrc/modules.cpp` by category (#883) · 16f3f897

Phuong Nguyen authored Jun 08, 2024



* categorized `csrc/modules.cpp`
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

* adapted the build tool
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

16f3f897

attention.cpp 31.6 KB

Replace attention.cpp