Migrate norms and softmax kernels to NVRTC

[CarlosGomes98]

Description

Enables JIT compilation through NVRTC for Norm and Softmax kernels.
Reduces TE binary size by 36%, sequential build time by 5% (measured in cpu_user total time, hard to measure real impact due to parallelization, machine specs)
This is the first chunk of work related to #3054 .

The softmax kernels were chosen as they seemed like one of the simplest to migrate, for my understanding of the system.
The norm kernels include normalization/layernorm/ln_fwd_cuda_kernel.cu, which is one of the heaviest kernel compilations in the build.

It is still possible to enable nvcc static compilation through NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX and NVTE_BUILD_LEGACY_STATIC_NORM, which then allow for NVTE_DISABLE_NVRTC=1 to be used during runtime.

Build time results:

Measured on RTX 6000 Ada, CUDA 12.8, single arch sm_89, 32-core host. AOT = -DNVTE_BUILD_LEGACY_STATIC_{FUSED_SOFTMAX,NORM}=ON (old behavior); NVRTC = default.

Per TU build time

translation unit	compile AOT (s)	compile NVRTC (s)	Δ time (s)	obj AOT (KB)	obj NVRTC (KB)	Δ size (KB)
scaled_masked_softmax.cu	10.80	2.64	−8.2	2348	274	−2074
scaled_upper_triang_masked_softmax.cu	11.75	2.51	−9.2	1891	238	−1653
scaled_aligned_causal_masked_softmax.cu	10.24	2.52	−7.7	2097	233	−1864
ln_fwd_cuda_kernel.cu	63.90	28.47	−35.4	9092	131	−8961
ln_bwd_semi_cuda_kernel.cu	43.80	28.45	−15.4	5295	121	−5174
rmsnorm_fwd_cuda_kernel.cu	38.28	28.39	−9.9	2499	114	−2385
rmsnorm_bwd_semi_cuda_kernel.cu	37.34	28.55	−8.8	2566	115	−2451
total	225.0	130.5	−94.5 (−42%)	26060	1498	−24562 (−94%)

Binary size

target	AOT (MB)	NVRTC (MB)	Δ (MB)
libtransformer_engine.so	64.4	41.2	−23.2 (−36%)

Total build time

metric	AOT	NVRTC	Δ
wall (s)	283.6	210.6	−72.9 (−26%)
cpu_user (s)	2584.6	2455.0	−129.6 (−5%)
max_rss (MB)	4037	4037	0

JIT compilation cost

kernel	NVRTC cold (ms)	static (ms)
layernorm_fwd	97.1	3.5
layernorm_bwd	133.1	2.7
rmsnorm_fwd	83.6	1.0
rmsnorm_bwd	106.4	1.1
scaled_masked_softmax_fwd	52.5	0.7
scaled_masked_softmax_bwd	38.8	—
scaled_upper_triang_softmax_fwd	48.3	0.4
scaled_upper_triang_softmax_bwd	42.3	—
scaled_aligned_causal_softmax_fwd	48.5	0.3

Fixes # (issue)

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 refactoring

Changes

Please list the changes introduced in this PR:

• Functionality to pass build options to the NVRTC compile manager
• Softmax kernels through NVRTC
• rtc_dispatch.cpp to allow NVRTC to work with the registry used by norms. This is the largest chunk of new code.
• Norm kernels through NVRTC

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

[CarlosGomes98]

/te-ci pytorch

[greptile-apps]

Greptile Summary

This PR migrates the fused softmax (3 variants) and norm (LayerNorm/RMSNorm forward and backward) kernels from ahead-of-time NVCC compilation to JIT compilation via NVRTC, reducing the library binary size by 36% and total build time by ~26%. The legacy AOT path is preserved behind NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX / NVTE_BUILD_LEGACY_STATIC_NORM flags.

• NVRTC infra improvements: KernelManager::compile is now idempotent under a shared_mutex (the TOCTOU race from previous review is resolved), extra_options/extra_headers parameters allow per-family header injection, and new launch_cooperative/occupancy_max_active_blocks_per_sm/set_function_attribute wrappers are added.
• Softmax migration: All three .cu files are made dual-mode (__CUDACC_RTC__ guard), with device-safe neg_infinity<T>() replacing std::numeric_limits<>::infinity() and expf replacing std::exp, both appropriate since acc_t is always float.
• Norm migration: Kernel param structs are extracted to a standalone kernel_params.h (NVRTC-safe), type-trait helpers moved to transformer_engine::detail to avoid namespace std injection, the zero-length array in dx_add_t is fixed via a NeedsPadding template specialization, and per-family REGISTER_NORM_LAUNCHER_* macros wire closures into the existing registry with optional static fallbacks.

Confidence Score: 5/5

The change is safe to merge. The NVRTC dispatch path is additive, the legacy AOT path is fully preserved behind build flags, and all prior review concerns have been addressed.

Thread-safety is correctly handled via shared_mutex with a double-check inside the unique lock. The dx_add_t specialization fix for zero-length arrays is correct. The softmax kernel changes are semantically equivalent to the original. Only minor inefficiencies remain, neither of which affects correctness.

transformer_engine/common/normalization/rtc_dispatch.cpp — the set_function_attribute call inside the launch closure and the two-compilation pattern for backward kernels are worth a follow-up optimization.

Important Files Changed

Filename	Overview
transformer_engine/common/normalization/rtc_dispatch.cpp	New file implementing NVRTC-backed launchers for all four norm kernel families. Correctly handles thread-safe double-check-lock pattern via shared_mutex. Minor inefficiencies: set_function_attribute called on every dispatch and backward kernels trigger two separate NVRTC compilations for the same source.
transformer_engine/common/util/rtc.cpp	Upgraded mutex to shared_mutex; added is_enabled() as static const; added extra_options/extra_headers to compile(); added set_function_attribute, occupancy_max_active_blocks_per_sm, and launch_cooperative. TOCTOU race fixed with double-check inside unique lock.
transformer_engine/common/util/rtc.h	Added Header struct, compile overload with extra_options/extra_headers, set_function_attribute, occupancy_max_active_blocks_per_sm, and launch_cooperative. Switched to mutable shared_mutex.
transformer_engine/common/normalization/rtc_dispatch.h	New header declaring the six register_* functions with optional StaticFallback pointer. Clean interface matching the implementation.
transformer_engine/common/fused_softmax/scaled_masked_softmax.cu	Migrated to NVRTC-first dispatch. Device-safe neg_infinity() and expf replacements correct for float acc_t. Legacy static path guarded behind NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX.
transformer_engine/common/fused_softmax/scaled_upper_triang_masked_softmax.cu	Same NVRTC migration pattern as scaled_masked_softmax.cu. No label collisions with other softmax variants.
transformer_engine/common/fused_softmax/scaled_aligned_causal_masked_softmax.cu	Migrated to NVRTC. RTC label correctly encodes 'aligned_causal' variant to avoid cache collisions.
transformer_engine/common/utils.cuh	Added transformer_engine::detail namespace with is_same/conditional traits. Resolves the std namespace concern from prior review.
transformer_engine/common/normalization/kernel_params.h	New header extracting kernel param structs from common.h. Intentionally free of host-only includes so NVRTC can consume it.
transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_kernels.cuh	Uses detail::conditional_t. Added NeedsPadding specialization to dx_add_t to avoid zero-length array rejected by NVRTC. Correct fix.
transformer_engine/common/normalization/layernorm/ln_fwd_cuda_kernel.cu	Adds REGISTER_NORM_LAUNCHER_LN_FWD_tuned/general macros. Static fallback compiled only when NVTE_BUILD_LEGACY_STATIC_NORM=ON.
transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu	ADD_FLAG now forwarded correctly and guarded by static_assert, addressing the prior review concern.

Sequence Diagram

%%{init: {'theme': 'neutral'}}%%
sequenceDiagram
    participant Caller as Caller
    participant Dispatch as Dispatch
    participant KMgr as KernelManager
    participant NVRTC as NVRTC
    participant Driver as CUDA Driver

    Caller->>Dispatch: forward/backward call
    Dispatch->>KMgr: is_compiled(label) [shared_lock]
    alt Not compiled - first call
        KMgr-->>Dispatch: false
        Dispatch->>KMgr: compile(label, kernel_expr, src, opts, headers)
        KMgr->>KMgr: acquire unique_lock and re-check cache
        KMgr->>NVRTC: nvrtcCreateProgram + nvrtcCompileProgram
        NVRTC-->>KMgr: PTX or cubin plus mangled name
        KMgr->>Driver: cuModuleLoadData
        KMgr->>KMgr: insert into cache and release lock
    else Already compiled
        KMgr-->>Dispatch: true returns early
    end
    Dispatch->>KMgr: "set_function_attribute if smem >= 48KB"
    Dispatch->>KMgr: occupancy_max_active_blocks_per_sm [shared_lock]
    Dispatch->>KMgr: launch or launch_cooperative [shared_lock]
    KMgr->>Driver: cuLaunchKernel or cuLaunchCooperativeKernel
    Driver-->>Caller: kernel running on GPU

Loading

%%{init: {'theme': 'base', 'themeVariables': {"darkMode": true, "background": "#0d1117", "primaryColor": "#21262d", "primaryTextColor": "#e6edf3", "primaryBorderColor": "#8b949e", "lineColor": "#8b949e", "textColor": "#e6edf3", "edgeLabelBackground": "#161b22", "actorBkg": "#21262d", "actorBorder": "#8b949e", "actorTextColor": "#e6edf3", "actorLineColor": "#8b949e", "signalColor": "#8b949e", "signalTextColor": "#e6edf3", "noteBkgColor": "#373320", "noteBorderColor": "#d4a72c", "noteTextColor": "#f0e6c0", "labelBoxBkgColor": "#21262d", "labelBoxBorderColor": "#8b949e", "labelTextColor": "#e6edf3", "loopTextColor": "#e6edf3", "activationBkgColor": "#30363d", "activationBorderColor": "#8b949e"}}}%%
sequenceDiagram
    participant Caller as Caller
    participant Dispatch as Dispatch
    participant KMgr as KernelManager
    participant NVRTC as NVRTC
    participant Driver as CUDA Driver

    Caller->>Dispatch: forward/backward call
    Dispatch->>KMgr: is_compiled(label) [shared_lock]
    alt Not compiled - first call
        KMgr-->>Dispatch: false
        Dispatch->>KMgr: compile(label, kernel_expr, src, opts, headers)
        KMgr->>KMgr: acquire unique_lock and re-check cache
        KMgr->>NVRTC: nvrtcCreateProgram + nvrtcCompileProgram
        NVRTC-->>KMgr: PTX or cubin plus mangled name
        KMgr->>Driver: cuModuleLoadData
        KMgr->>KMgr: insert into cache and release lock
    else Already compiled
        KMgr-->>Dispatch: true returns early
    end
    Dispatch->>KMgr: "set_function_attribute if smem >= 48KB"
    Dispatch->>KMgr: occupancy_max_active_blocks_per_sm [shared_lock]
    Dispatch->>KMgr: launch or launch_cooperative [shared_lock]
    KMgr->>Driver: cuLaunchKernel or cuLaunchCooperativeKernel
    Driver-->>Caller: kernel running on GPU

Loading

_{Reviews (8): Last reviewed commit: "Fix tests" | Re-trigger Greptile}

[ptrendx]

/te-ci pytorch

[ptrendx]

Generally OK. One general comment is that with this move to NVRTC we could actually expand the tuned kernel coverage not just to those predetermined cases, but also to other row lengths (previously we did not want to do that just because of the binary size and the compilation time). This should give a good benefit, since the general kernel is not very efficient compared with the tuned one. I'm not sure even if there is even a need for the general kernel (apart from the very long row lengths that would not fit) in that case.

[CarlosGomes98]

One general comment is that with this move to NVRTC we could actually expand the tuned kernel coverage not just to those predetermined cases, but also to other row lengths

I did consider this, but we also rely on the registration for the optimal launch parameters, for different shapes / archs right? I think its a cool idea but would probably defer it to a separate PR

[ptrendx]

/te-ci

[ptrendx]

Generally speaking the tuning should not be difficult, but I agree that this can be done in a subsequent PR.

[pggPL]

/te-ci

[timmoon10]

LGTM

[timmoon10]

/te-ci

[phu0ngng]

/te-ci