GitLab

[JAX] XLA Custom Calls with FFI for FusedAttnFwd, Quantize, Transpose, ActLuFP8, LayerNormForwardFP8FFI, and LayerNormBackwardFFI (#1263)

* Add TransposeFFI, test passed
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add ActLuFP8FFI; fix TransposeFFI
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add QuantizeFFI
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add FusedAttnForwardFFI and some unit tests
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Minor fix
Signed-off-by: Hua Huang <huah@nvidia.com>

* Add LayerNormForwardFP8FFI & LayerNormBackwardFFI
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Revise FusedAttnForwardFFI()
Signed-off-by: Hua Huang <huah@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Add FFI_CudaGraph_Traits

All tests passed, ready for merge
Signed-off-by: Hua Huang <huah@nvidia.com>

* Bug fix for FFI data type mismatch

Also add a safeguard on the entrance to FFI function
Signed-off-by: Hua Huang <huah@nvidia.com>

---------
Signed-off-by: Hua Huang <huah@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

* Specify python version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add classifiers for python
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add utils to build wheels
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* make wheel scripts
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add aarch
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix paddle wheel
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* PaddlePaddle only builds for x86
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add optional fwk deps
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Python3.8; catch install error
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [wip] cudnn9 compile with paddle support
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [wip] dont link cudnn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* dlopen cudnn
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* dynamically load nvrtc
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Lint
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove residual packages; exclude stub from nvrtc .so search
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Exclude builtins from nvrtc .so search
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* properly include files for sdist
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* paddle wheel tie to python version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix paddle build from src [wip]
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix workflow paddle build
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* Fix paddle
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix paddle
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix lint from pr986
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* fix
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add sanity wheel test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Add sanity import to wheel test
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* remove upper limit on paddlepaddle version
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove unused imports
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Remove pybind11 dependency
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix cpp tests
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Search .sos in cuda home
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci



* fixes
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* CLeanup, remove residual code
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

replaced plain C asserts with NVTE_CHECK to avoid unused-variable warnings
Signed-off-by: Alp Dener <adener@nvidia.com>

* Apply formatting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Apply formatting
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* 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>