[C] Normalization Refactor + Adding CUDNN backend

[phu0ngng]

Description

1. cuDNN normalization was integrated into TE. Users can use this backend via env vars as:

export NVTE_NORM_FWD_USE_CUDNN=1
export NVTE_NORM_BWD_USE_CUDNN=1

CuDNN 9.6 is needed.

By default, TE Kernels are still used for normalization. Further performance evaluation on different supported GPUs is needed to make cuDNN backend as default.

2. Major changes were introduced in the TE normalization C APIs:

• Merging barrier, dgamma_part, and dbeta_part into the workspace. Thus those args are no longer needed in the APIs.
• Unifying APIs with and without zero centering.

These are the updated C APIs:

void nvte_layernorm_fwd(const NVTETensor x, const NVTETensor gamma, const NVTETensor beta, 
                        const float epsilon, NVTETensor z, NVTETensor mu, NVTETensor rsigma,
                        NVTETensor workspace, const int multiprocessorCount,
                        const bool zero_centered_gamma, cudaStream_t stream); 

void nvte_layernorm_bwd(const NVTETensor dz, const NVTETensor x,  const NVTETensor mu, 
                        const NVTETensor rsigma, const NVTETensor gamma,  NVTETensor dx, 
                        NVTETensor dgamma, NVTETensor dbeta, NVTETensor workspace,
                        const int multiprocessorCount, const bool zero_centered_gamma,
                        cudaStream_t stream);

void nvte_rmsnorm_fwd(const NVTETensor x, const NVTETensor gamma, const float epsilon, NVTETensor z,
                      NVTETensor rsigma, NVTETensor workspace, const int multiprocessorCount,
                      const bool zero_centered_gamma, cudaStream_t stream); 

void nvte_rmsnorm_bwd(const NVTETensor dz, const NVTETensor x, const NVTETensor rsigma,
                      const NVTETensor gamma, NVTETensor dx, NVTETensor dgamma,
                      NVTETensor workspace, const int multiprocessorCount,
                      const bool zero_centered_gamma, cudaStream_t stream);

The following APIs are deprecated:

nvte_layernorm1p_fwd
nvte_layernorm1p_bwd
nvte_rmsnorm1p_fwd
nvte_rmsnorm1p_bwd

3. TE normalization implementations were refactored to be more friendly for adding cuDNN backend.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactor

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[phu0ngng]

/te-ci

[timmoon10]

I have minor comments and we should make sure this passes CI, but overall LGTM.

[phu0ngng]

/te-ci