Fix f16 overflow in SkipLayerNormalization CUDA kernel - #28442
Open
justinchuby wants to merge 1 commit into
Open
Fix f16 overflow in SkipLayerNormalization CUDA kernel#28442justinchuby wants to merge 1 commit into
justinchuby wants to merge 1 commit into
Conversation
Comment on lines
+155
to
+156
| const auto sum_kv = BlockReduceKV(temp_storage_kv).Reduce( | ||
| AccKeyValue(acc_mean, acc_var), acc_pair_sum); |
Contributor
There was a problem hiding this comment.
Suggested change
| const auto sum_kv = BlockReduceKV(temp_storage_kv).Reduce( | |
| AccKeyValue(acc_mean, acc_var), acc_pair_sum); | |
| const auto sum_kv = BlockReduceKV(temp_storage_kv).Reduce(AccKeyValue(acc_mean, acc_var), acc_pair_sum); |
Comment on lines
+262
to
+263
| const auto sum_kv = BlockReduceKV(temp_storage_kv).Reduce( | ||
| AccKeyValue(acc_mean, acc_var), acc_pair_sum); |
Contributor
There was a problem hiding this comment.
Suggested change
| const auto sum_kv = BlockReduceKV(temp_storage_kv).Reduce( | |
| AccKeyValue(acc_mean, acc_var), acc_pair_sum); | |
| const auto sum_kv = BlockReduceKV(temp_storage_kv).Reduce(AccKeyValue(acc_mean, acc_var), acc_pair_sum); |
Add stash_type attribute (default=1, float32) to both SkipLayerNormalization and SkipSimplifiedLayerNormalization schemas, matching the existing attribute on standard LayerNormalization. When stash_type=1 and the input type is f16/bf16, the CUDA kernel routes to HostApplyLayerNorm which accumulates variance in float32. This prevents f16 overflow in deep networks where residual values grow large through skip connections (e.g. Qwen3-VL 28-layer with QK-norm produces residuals with absmax > 100 by layer 22). The fix addresses the root cause of f16 NaN in decoder models: SkipSimplifiedLayerNormalization accumulated x²/ld in f16, which overflows when x > ~180 for hidden_size=2048 (x²/2048 > 65504). Changes: - bert_defs.cc: Add stash_type attribute to both SkipLayerNorm and SkipSimplifiedLayerNorm schemas - skip_layer_norm.h: Add use_float_accumulation_ member - skip_layer_norm.cc: Read stash_type, use HostApplyLayerNorm (float accumulation) when stash_type=1 for f16/bf16 Tested: Qwen3-VL-2B f16 CUDA decoder produces correct output at all input magnitudes (std=0.01 to 0.5) with zero NaN. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
force-pushed
the
fix-skip-layer-norm-f16-overflow
branch
from
May 9, 2026 16:41
5ea2c18 to
ce40b70
Compare
Contributor
|
There already exists a strict mode flag that is supposed to handle this. Why is this PR needed? |
Contributor
Author
Is it preferred to handle model precision behavior via runtime flags? I wonder if this would make the model under specified. The similar ONNX operators all have the stash_type attributes to control this behavior. What do you think? |
tianleiwu
added a commit
that referenced
this pull request
May 27, 2026
…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 #28442 #15660
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Fixes f16/bf16 overflow in the
SkipSimplifiedLayerNormalizationCUDA kernel by switching variance accumulation from native type T to float.Problem
Deep transformer models with QK-norm (e.g. Qwen3-VL 28 layers, 56 standalone
RMSNormalizationops) produce growing residual values through skip connections. The fusedSkipSimplifiedLayerNormalizationkernel accumulatesx²/ldin the native half type, which overflows f16's max of 65504 when residual values grow large (absmax > ~180 for hidden_size=2048).Symptom: NaN output from f16 CUDA decoder starting at layer 23 for Qwen3-VL-2B with small input embeddings (std < 0.3). The standalone
RMSNormalizationkernel (used by QK-norms) already uses float accumulation and does not have this issue.Root cause analysis
Layer-by-layer profiling of Qwen3-VL-2B f16 CUDA decoder with input std=0.01:
The overflow occurs in the
SkipLayerNormKernelSmallvectorized kernel at line 171:rldvalsq_sum += rldval * sum_v[i]; // f16 overflow when sum_v[i] > ~180Fix
Use
floatfor:x²/ldaccumulationx/ldaccumulationcub::BlockReduce<float>)rsqrtf()computationBoth the non-vectorized (
SkipLayerNormKernel) and vectorized (SkipLayerNormKernelSmall) paths are fixed.Testing
Tested with Qwen3-VL-2B f16 CUDA decoder. This fix moves the first NaN from layer 23 → layer 26. The remaining NaN at L26 appears to come from upstream MatMul overflow, not from SkipLayerNorm.
Note
This is a partial fix for the complete f16 CUDA precision issue with QK-norm models. Full f16 support may require similar f32 accumulation fixes in other CUDA kernels (e.g., cuBLAS MatMul or attention kernels). The standalone
RMSNormalizationkernel already uses float accumulation correctly.Co-authored-by: Copilot 223556219+Copilot@users.noreply.github.com