GitLab

[C/JAX] Support more mask types for the arbitrary seqlen kernels and minor changes of JAX bias (#469)

* Move bias to float32
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enable varlen
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Increase neg infinity abs values
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enable varlen tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove unnecessary code
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix lint
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Support variable sequence length after cuDNN 8.9.6
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Use unique_ptr instead of shared_ptr
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Add a new mask type: PADDING_CAUSAL_MASK
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Support flash padding mask after 8.9.6
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Enhance the Max512 handling for causal masking and add the related tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Update the fused attn support lists
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Remove padding_aware from the caching
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix libtransformer.so issue
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Reduce the pad ratio tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix a bug with cuDNN 8.9.5
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Release backend resource after the module level unit test
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Clean the jax live arrays before running the unit tests
Signed-off-by: Reese Wang <rewang@nvidia.com>

* Fix too-few-public-methods lint
Signed-off-by: Reese Wang <rewang@nvidia.com>

---------
Signed-off-by: Reese Wang <rewang@nvidia.com>