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 norm (LayerNorm / RMSNorm) and softmax kernels from ahead-of-time NVCC compilation to JIT compilation via NVRTC, reducing the libtransformer_engine.so binary size by ~36% and build wall time by ~26%. A new rtc_dispatch.cpp replaces the static registry macro with NVRTC-backed closures that compile kernels on first use, with an optional static fallback controlled by NVTE_BUILD_LEGACY_STATIC_NORM / NVTE_BUILD_LEGACY_STATIC_FUSED_SOFTMAX.

• Adds extra_options and extra_headers parameters to KernelManager::compile, enabling domain-specific headers (norm kernel params, traits) to be injected only where needed; upgrades KernelManager::lock_ from std::mutex to std::shared_mutex with a double-checked compile under unique_lock to fix the prior TOCTOU concern.
• Replaces std::exp / std::numeric_limits in softmax device code with expf / neg_infinity<T>() for NVRTC compatibility, and guards host-only includes behind #ifndef __CUDACC_RTC__ in all three softmax translation units.
• Introduces register_launcher and per-variant register_*_tuned/general functions in rtc_dispatch.cpp that replicate the smem, barrier, workspace, and grid-dimension formulae previously embedded in the static launch helpers, with static_assert(ADD_FLAG, \u2026) guards on the RMSNorm BackwardAdd macros to address a previously flagged silent parameter drop.

Confidence Score: 5/5

Safe to merge — the migration is well-contained, the thread-safety model is sound, and the static fallback escape hatches provide a reliable rollback path.

The KernelManager locking upgrade is correctly implemented and resolves the prior TOCTOU concern. All dispatch paths correctly handle the NVRTC-disabled case. The only finding is a dead needs_cooperative capture in register_launcher, which has no behavioral impact.

transformer_engine/common/normalization/rtc_dispatch.cpp — the register_launcher helper carries a dead needs_cooperative parameter/capture that can be cleaned up, but does not affect correctness.

Important Files Changed

Filename	Overview
transformer_engine/common/normalization/rtc_dispatch.cpp	New 742-line file implementing NVRTC-backed closures for all norm kernel variants; correctly double-checks compilation under unique_lock, but captures needs_cooperative in register_launcher without ever using it.
transformer_engine/common/util/rtc.cpp	Upgrades lock from std::mutex to std::shared_mutex with double-checked compile under unique_lock; adds extra_options/extra_headers support; moves is_compiled and launch to shared_lock.
transformer_engine/common/util/rtc.h	Clean additions: Header struct, launch_cooperative template, set_function_attribute, occupancy_max_active_blocks_per_sm, and mutable shared_mutex.
transformer_engine/common/fused_softmax/scaled_masked_softmax.cu	Properly guards host vs. RTC code paths with CUDACC_RTC, replaces std::exp/std::numeric_limits with expf/neg_infinity(), correct rtc_scale cast in forward path.
transformer_engine/common/normalization/rmsnorm/rmsnorm_bwd_semi_cuda_kernel.cu	Adds static_assert(ADD_FLAG) guards on BackwardAdd macros and passes ADD_FLAG through to register_rmsnorm_bwd_tuned/general, addressing the previously flagged silent drop.
transformer_engine/common/normalization/kernel_params.h	New file factoring out norm kernel parameter structs from common.h; enables NVRTC to access them as an injected header.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Kernel dispatch call] --> B{NVRTC enabled?}
    B -- No --> C{Static fallback registered?}
    C -- Yes --> D[Call static fallback fn ptr]
    C -- No --> E[NVTE_ERROR: rebuild with NVTE_BUILD_LEGACY_STATIC_NORM]
    B -- Yes --> F[KernelManager::is_compiled? shared_lock]
    F -- No --> G[KernelManager::compile unique_lock + double-check]
    G --> H[nvrtcCreateProgram with extra_headers]
    H --> I[nvrtcCompileProgram with extra_options]
    I --> J[Cache Kernel in kernel_cache_]
    F -- Yes --> K[configure_params?]
    J --> K
    K -- Yes --> L[occupancy_max_active_blocks_per_sm shared_lock]
    L --> M[Set ctas_per_col, barrier_bytes, workspace_bytes]
    K -- No --> N{ctas_per_row == 1?}
    N -- Yes --> O[launch shared_lock cuLaunchKernel]
    N -- No --> P[launch_cooperative shared_lock cuLaunchCooperativeKernel]
    P --> Q[For bwd: also launch finalize kernel]
    O --> Q

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"}}}%%
flowchart TD
    A[Kernel dispatch call] --> B{NVRTC enabled?}
    B -- No --> C{Static fallback registered?}
    C -- Yes --> D[Call static fallback fn ptr]
    C -- No --> E[NVTE_ERROR: rebuild with NVTE_BUILD_LEGACY_STATIC_NORM]
    B -- Yes --> F[KernelManager::is_compiled? shared_lock]
    F -- No --> G[KernelManager::compile unique_lock + double-check]
    G --> H[nvrtcCreateProgram with extra_headers]
    H --> I[nvrtcCompileProgram with extra_options]
    I --> J[Cache Kernel in kernel_cache_]
    F -- Yes --> K[configure_params?]
    J --> K
    K -- Yes --> L[occupancy_max_active_blocks_per_sm shared_lock]
    L --> M[Set ctas_per_col, barrier_bytes, workspace_bytes]
    K -- No --> N{ctas_per_row == 1?}
    N -- Yes --> O[launch shared_lock cuLaunchKernel]
    N -- No --> P[launch_cooperative shared_lock cuLaunchCooperativeKernel]
    P --> Q[For bwd: also launch finalize kernel]
    O --> Q

Loading

_{Reviews (9): Last reviewed commit: "Merge branch 'main' into cgomes/nvrtc-ph..." | 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

[timmoon10]

/te-ci

[greptile-apps]

Want your agent to iterate on Greptile's feedback? Try greploops.