onnxruntime

26871ab4 - Use fp32 accumulation in SkipLayerNorm/EmbedLayerNorm CUDA kernels (#28682)

Use fp32 accumulation in SkipLayerNorm/EmbedLayerNorm CUDA kernels (#28682) ## Description Use fp32 accumulation in SkipLayerNormalization, SkipSimplifiedLayerNormalization, and EmbedLayerNormalization CUDA kernels to avoid overflow and improve numerical accuracy when processing fp16/bf16 data. The original implementation accumulated mean and variance statistics in the input data type (fp16/bf16), which can overflow for large hidden sizes or when input values have large magnitude. This change promotes all intermediate accumulation (mean, variance, normalization math) to fp32, matching the approach used by TensorRT-LLM's LayerNorm kernels. ## Motivation - fp16 has limited range (max ~65504) and precision (10-bit mantissa). Accumulating `x²/ld` across thousands of elements in fp16 easily overflows or loses precision. - bf16 has even less precision (7-bit mantissa), making accumulation errors more severe. - The fix is straightforward: cast to float before accumulating, compute normalization in float, cast back to the output type. ## Key Changes | File | Change | |------|--------| | `layer_norm.cuh` | Changed `LayerNorm`, `SimplifiedLayerNorm`, `LayerNormSmall`, `SimplifiedLayerNormSmall` to accept and operate on `float` for thread_data, epsilon, mu, rsigma. Removed unused `KeyValuePairSum` overloads for half/bfloat16. | | `skip_layer_norm_impl.cu` | Changed `SkipLayerNormKernel` and `SkipLayerNormKernelSmall` to accumulate in fp32 (`cub::KeyValuePair<float, float>`). Removed `maybe2half` helper (no longer needed). | | `embed_layer_norm_impl.cu` | Changed epsilon from `T` to `float`, accumulation to use `float` thread_data. | | `profile_skip_layer_norm.py` | New profiling script for nsys-based kernel timing analysis. | | `profile_skip_layer_norm.sh` | Shell wrapper for running nsys profiling. | | `parse_nsys.py` | Utility to parse nsys SQLite output and extract CUDA kernel timings. | ## Performance Results Profiled on NVIDIA GPU with nsys (B=1, seq_len=2048, fp16 data, 200 iterations, skip first 5 warmup): | Hidden Size | fp16 accum (μs) | fp32 accum (μs) | Regression | |---|---|---|---| | 768 | 3.81 | 3.81 | **0.0%** | | 1024 | 4.22 | 4.22 | **0.0%** | | 4096 | 13.01 | 13.03 | **+0.15%** (noise) | | 8192 | 28.94 | 28.94 | **0.0%** | **No measurable performance regression.** The kernel is memory-bandwidth-bound, so fp32 arithmetic is completely hidden behind memory latency. ## Testing - Existing unit tests pass (SkipLayerNorm, EmbedLayerNorm ops). - Profiling scripts added for reproducible performance measurement: ```bash cd onnxruntime/test/python/transformers nsys profile -o sln_fp16 --export=sqlite python profile_skip_layer_norm.py --mode fp16 --warmup 5 --repeat 100 python parse_nsys.py sln_fp16.sqlite --skip-first 5 ``` ## Related PRs https://github.com/microsoft/onnxruntime/pull/28442 https://github.com/microsoft/onnxruntime/pull/15660