Repository navigation
Use fp32 accumulation in SkipLayerNorm/EmbedLayerNorm CUDA kernels - #28682
Conversation
There was a problem hiding this comment.
Pull request overview
This PR improves numerical stability of the CUDA fused normalization kernels used in transformer models by switching mean/variance (and RMS) accumulation to FP32 when inputs are FP16/BF16, reducing overflow/NaN risk while keeping outputs in the original type. It also adds local profiling utilities to measure kernel performance with Nsight Systems.
Changes:
- Promote SkipLayerNorm and EmbedLayerNorm CUDA kernel intermediate accumulation (mean/variance/RMS) from FP16/BF16 to FP32.
- Refactor
layer_norm.cuhhelpers to operate onfloatthread/block reduction values (and remove half/bfloat16 reduction overloads). - Add
nsys-based profiling + parsing scripts for reproducible kernel timing analysis.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cuda/bert/layer_norm.cuh | Switch block-reduce inputs/outputs and normalization math to FP32 intermediates for improved stability. |
| onnxruntime/contrib_ops/cuda/bert/skip_layer_norm_impl.cu | Accumulate skip layer norm statistics in FP32; remove now-unneeded maybe2half epsilon casting. |
| onnxruntime/contrib_ops/cuda/bert/embed_layer_norm_impl.cu | Accumulate embedding sum statistics in FP32 and pass epsilon as float. |
| onnxruntime/test/python/transformers/profile_skip_layer_norm.py | New Python profiler that builds a minimal SLN model and benchmarks it (optionally with NVTX ranges). |
| onnxruntime/test/python/transformers/profile_skip_layer_norm.sh | New wrapper script to run nsys profile and parse results. |
| onnxruntime/test/python/transformers/parse_nsys.py | New SQLite parser for nsys --export=sqlite kernel timing summaries. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- Fix SQL injection in parse_nsys.py: use parameterized queries instead of string interpolation for kernel pattern matching - Add --nvtx-range option to parse_nsys.py to filter kernels by NVTX range (e.g., 'benchmark'), eliminating the need for --skip-first to exclude warmup - Update parse_nsys.py description/epilog to reflect current purpose - Remove pip install nvtx from shell script; just check availability and warn - Fix CodeQL import warning: use 'from onnx import save_model' instead of 'import onnx' + 'onnx.save'
kunal-vaishnavi
left a comment
There was a problem hiding this comment.
Shall we remove the strict mode flag then if FP32 accumulation is now always used?
Maybe later. That flag will be slower, so it can be deprecated later. |
### Description Removes the deprecated `enable_skip_layer_norm_strict_mode` CUDA provider option and its associated "strict mode" code path from the CUDA `SkipLayerNorm` kernel. Strict mode previously routed `SkipLayerNorm` through the `LayerNormalization` kernel to gain fp32-accumulation accuracy at the cost of performance. After #28682 made `SkipLayerNorm`/`EmbedLayerNorm` CUDA kernels always accumulate in fp32, the strict-mode path is redundant: the default kernel already provides the same accuracy with better performance. This change deletes the now-dead branch and the plumbing that fed it. ### Key Changes - **`skip_layer_norm.cc`**: Remove the `strict_` branch that called `HostApplyLayerNorm`; always launch `LaunchSkipLayerNormKernel`. Drop the now-unused `layer_norm_impl.h` include and the input/skip same-shape strict-mode shape check. - **`skip_layer_norm.h`**: Remove the `strict_` member. - **`cuda_execution_provider.h`**: Remove `IsSkipLayerNormInStrictMode()`. - **`cuda_kernel_adapter.h`**: Remove `GetCudaKernelAdapterSkipLayerNormStrictMode()` shim. - **`cuda_provider_options.h`**: Keep `enable_skip_layer_norm_strict_mode` field for ABI/back-compat but mark it deprecated and ignored. - **`skiplayernorm_op_test.cc`**: Drop the redundant strict-mode test passes; tests now run a single default path. The provider option is retained (ignored) to preserve backward compatibility — existing configs that set it continue to work without error. ### Motivation Follow-up cleanup to #28682, which switched the CUDA kernels to fp32 accumulation, making strict mode obsolete. ### Testing - `skiplayernorm_op_test.cc` covers fp16/fp32/bf16 default path; strict-mode passes removed.
- Graphs fused for "cuda" need onnxruntime>=1.27 (fp32 accumulation in the CUDA skip norms, microsoft/onnxruntime#28682): the fp16 skip-norm unfusing is dropped, with a warning when an older version is installed. - GroupQueryAttention: a known total length goes in as an int32 constant (1.29's shape inference reads one computed from shapes as int64), and a decode taking any number of tokens is fused only for a batch of one (1.29 rejects several new tokens after a past for a larger batch). It's a CPU-only step: CUDA's rejects the bias every masked attention has. - MoE: SiLU-gated experts only; 1.29's CUDA MoE miscomputes gated GELU. - The pipeline is a `(name, step, providers)` table. - The export tests fuse for the provider they run on (`torch_device`). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ns=...)`) (huggingface#50309) * Exporters: feed SDPA's boolean mask to ONNX as an additive one torchlib guards its boolean-mask softmax with an IsNaN/Where over the whole (query, key) weight matrix; the patch already zeroes fully masked rows, so the guard only costs time (2x on a 2520-patch vision attention on ORT CPU). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: fuse ONNX exports into ONNX Runtime's ops at export time `OnnxConfig(ort_fusions="cpu" | "cuda")` fuses the exported graph into ONNX Runtime's transformer ops before it is quantized and saved: RMS / layer / skip norms, rotary embeddings, MultiHeadAttention, GroupQueryAttention over the key-value cache (rotating itself), MoE (CUDA), and packed sibling projections. onnxscript's `optimize_for_ort` does the core fusions; `onnxruntime_fusions.py` prepares the graph for them, repairs the ones ONNX Runtime rejects or miscomputes, and adds the rest. Off by default (`None`): the fused ops aren't standard ONNX. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: explain the ORT fusion tests, plainer docs Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: grouped_mm without Slices; fuse float16 MoE blocks The grouped_mm translation takes each group's end and weights from one Split and gathers its rows from a Range: onnxscript's SlicesSplit rule tries every pair of Slices in the graph, and a Slice per expert and layer made it most of an MoE export's time (granite 1B-A400M: export 23 min -> 5 min, an ORT fusion pass 13 min -> 2 min). The MoE fusion steps over the casts a float16 model's router and expert sum compute through in float32, so real checkpoints fuse, not only float32 ones. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: on CPU, unfuse the attentions GroupQueryAttention didn't take ONNX Runtime's CPU MultiHeadAttention with an attention bias runs slower than the MatMuls and Softmax it fuses (2.5x on Qwen3-0.6B's decode step), so for `ort_fusions="cpu"` the ones left after the GroupQueryAttention fusion (encoders, attention without a cache) go back to plain ops. EmbeddingGemma2 on CPU: 0.98x -> 1.17x. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: ORT fusions on ONNX Runtime 1.29 - Graphs fused for "cuda" need onnxruntime>=1.27 (fp32 accumulation in the CUDA skip norms, microsoft/onnxruntime#28682): the fp16 skip-norm unfusing is dropped, with a warning when an older version is installed. - GroupQueryAttention: a known total length goes in as an int32 constant (1.29's shape inference reads one computed from shapes as int64), and a decode taking any number of tokens is fused only for a batch of one (1.29 rejects several new tokens after a past for a larger batch). It's a CPU-only step: CUDA's rejects the bias every masked attention has. - MoE: SiLU-gated experts only; 1.29's CUDA MoE miscomputes gated GELU. - The pipeline is a `(name, step, providers)` table. - The export tests fuse for the provider they run on (`torch_device`). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: trace MoE experts as batched_mm for CUDA fusions With `ort_fusions="cuda"`, the ONNX exporter traces mixture-of-experts layers with `batched_mm` (restoring the model's own implementation after), and the `MoE` fusion matches that form: each routed row's expert weight picked by `GatherND`, a batched `MatMul`, the SiLU gate, the weighted sum per token. Its graph is a fraction of `grouped_mm`'s per-expert one, so Granite 3.0 1B-A400M exports in 69 s instead of 302 s and fuses in 15 s instead of 130 s, to the same `MoE` graph. `"cpu"` and unfused exports keep `grouped_mm`, whose prefill stays fast unfused. A warning says when a block doesn't fuse. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: GroupQueryAttention on CUDA for decode with a dynamic cache On ONNX Runtime 1.29, CUDA's `GroupQueryAttention` takes the attention bias and is exact. It's a large win for decode (Qwen3-0.6B: 7.48 -> 4.71 ms, 3.71x with a CUDA graph) but slower than `MultiHeadAttention` for a prefill, and a static cache's total length, computed on the GPU, is read as 0 (microsoft/onnxruntime#33867). So "cuda" fuses decode steps over a dynamic cache only; "cpu" fuses every case as before. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: CUDA GroupQueryAttention for static caches too A static cache attends over its whole buffer: `total_sequence_length` is the buffer's length as an int32 constant, the bias spans the buffer (the cache's mask already masks the slots not written), and `seqlens_k` carries each sequence's written length. A constant never lives on the GPU, which CUDA's kernel reads as 0 (microsoft/onnxruntime#33867), so static caches fuse on CUDA too. CUDA's step stays decode-only: with a bias, its prefill runs an unfused kernel, slower than `MultiHeadAttention`'s. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: GroupQueryAttention over transformers' folded decode attention Since huggingface#49743, a single-token SDPA step without `enable_gqa` folds the query heads onto their key-value heads (`[batch, kv_heads, groups, head_dim]`) instead of repeating the keys and values, which onnxscript fuses into a `MultiHeadAttention` over the key-value heads. The `GroupQueryAttention` fusion reads that form too: the query heads behind the fold, the bias behind its expand over the group rows, `num_heads` the key-value heads times the groups, its output reshaped back to the folded layout. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: pad GroupQueryAttention's rotary caches with a Concat ONNX Runtime's CUDA provider has no bfloat16 `Pad`, so a bfloat16 graph whose `GroupQueryAttention` rotates with a per-token cache failed to load with the CUDA provider alone. The zero rows are appended with a `Concat` instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: fuse GroupQueryAttention into dynamic decode graphs - The SDPA patch checks that the key, value and mask lengths are equal, which SDPA requires but doesn't check under export. A dynamic export then gives all caches one length symbol, so onnxscript's attention fusion matches every layer. - A rule turns torchlib's dynamic key transpose (Reshape, Transpose, Reshape with Shape-derived targets) into one Transpose. - Shape nodes reading the same dims are merged before CSE, so each layer's GroupQueryAttention lengths are computed once instead of 28 times with a host-to-device copy each. - The integer Equal rewrite is dropped: the CUDA provider has Equal kernels at every opset. - Comments cite microsoft/onnxruntime#33873-huggingface#33876. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: drop layout transposes around attention, fold q/k norms into GQA - A new step moves the transposes between [batch, length, heads, head_dim] and [batch, heads, length, head_dim] past rotary embeddings (run on the 3D layout), the key-value head repeat and the zeroing of unattended rows, until they cancel: MultiHeadAttention reads and writes the first layout. Qwen3's prefill keeps 2 transposes per layer instead of 8, about 1 ms faster on an A100. - On CUDA, decode GroupQueryAttention takes Qwen3's per-head query and key RMS norms as q_norm_weight / k_norm_weight, in half precision (CUDA's kernel is the one with them). The "cuda" minimum becomes onnxruntime 1.28, the first release with these inputs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: hand GroupQueryAttention its packed q/k/v projection When a GroupQueryAttention's query, key and value are the three outputs of the split after a packed projection, it now reads that projection packed, with key and value left empty: no Split, no reshapes. Outputs are bitwise identical; Qwen3 decode on an A100 goes from 4.71 to 4.53 ms (3.20 to 3.14 ms with a CUDA graph). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: simplify the ONNX Runtime fusions; tests in test_fusions.py - One `_add_node` helper builds and inserts the fusions' new nodes, and one `_merged_heads` helper matches the attention's merged heads. - The repeated key-value heads broadcast to a constant shape, instead of one built from the states' shape at run time. - GroupQueryAttention looks for per-head norms only on a key in its own layout. - The SDPA patch returns early without a mask; the grouped_mm translation splits its per-group inputs with one helper. - The module docstring, docs and version warning describe the current steps. - The ONNX Runtime fusion tests move to tests/exporters/test_fusions.py, one class per backend, and check that GroupQueryAttention reads the packed projection. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Docs: drop the onnxruntime-gpu installation details from the fusions section Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: inline three single-use fusion helpers `_unshare_key_transposes`, `_same_dims` and `_merge_equal_shapes` each had one caller, a few lines long; their bodies now sit there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: test the ORT fusions per layer on each family's tiny model; fuse float16 MoE - tests/exporters/test_fusions.py exports the tiny model of each decoder family's own tests (Llama, Qwen2, Qwen3, Phi-3, OLMo2, Granite, Mistral, Gemma 3, Mixtral, Qwen3-MoE), the same table as the OpenVINO fusion tests, and checks that every layer of the prefill and the decode runs ONNX Runtime's attention and, on CUDA, its `MoE`, with no rotary, attention or RMS norm left decomposed, for both providers. - Those tests found float16 mixture-of-experts blocks that didn't fuse: the experts' outputs weighted in float32 (Mixtral) and a router taking its top-k in float32 (Qwen3-MoE). The `MoE` fusion now steps over both casts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: give sliding-window cache layers a full cache during export The same patch as the OpenVINO exporter's (huggingface#50311): a `DynamicSlidingWindowLayer` gets a `DynamicLayer`'s update and mask sizes during the trace, and the sliding mask still applies the window. Temporary, until huggingface#49542 keeps sliding-window decode on one graph. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: name fuse_for_onnxruntime's target `provider`, not `device` It's the ONNX Runtime execution provider the graph is fused for (each step's `providers`), not the device the export runs on. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: call fuse_for_onnxruntime's target `ort_fusions`, as in the config Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: prepare the model for ORT fusions like ExecuTorch's backends, and warn about unfused MoE in the fusion `prepare_for_ort_fusions` sets the experts to `batched_mm` for "cuda", under the same condition as fusing (which also skips FX-graph quantizers). `fuse_for_onnxruntime` takes the original model to warn about mixture-of-experts blocks left unfused. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: add ort_fusions="webgpu" for ONNX Runtime's WebGPU provider WebGPU's GroupQueryAttention ignores position_ids when it rotates, so rotary stays a separate RotaryEmbedding there, and it has no FusedMatMul, so those go back to MatMul. MoE stays unfused (not in onnxruntime-webgpu 1.27). The docs list what each provider fuses in a table. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: say the ORT fusion table is what each provider supports, with the two slowdowns as temporary Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: cite the ORT issues behind the two fusions skipped for speed Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: shorter sliding-window patch docstring; build the fusion tests' models in their dtype Matches huggingface#50311 after its review: `from_config(dtype=...)` keeps the float32 buffers a model builds (rotary `inv_freq`) float32, which casting after didn't. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Tests: free the ORT fusion tests' memory with MemoryCleanupMixin; docs: the full sliding-window cache Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: one `ort_fusions` for the ONNX export's fusion target, and apply the backend patches first Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: simplify the ONNX Runtime fusions; restore the experts implementation after export - fuse_for_onnxruntime(model, ort_fusions): drop the steps / original_model params and the MoE-count warning - use_experts_implementation restores the model's experts implementation after a "cuda" export - OnnxConfig raises on ort_fusions without optimize; is_onnxruntime_available(min_version) - GroupQueryAttention fusion split into helpers; heads-second flag instead of a wrapper class - export tests fuse only for cpu / cuda with optimize on - ORT fusion tests laid out like the OpenVINO ones (expected-count tables, export helper, assert_fused) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: repair onnxscript's ORT fusions where ONNX Runtime rejects them; simplify the fusion steps - undo or fix what optimize_for_ort gets wrong: group norms (onnxscript#3182), position embeddings packed as attention biases (huggingface#3181), leftover SDPA ops (huggingface#3184), partial rotary without num_heads (huggingface#3185), loop-carried sequences typed as tensors (onnx/ir-py#539) - skip GroupQueryAttention when the bias reads the updated cache (Doge) - _patch_sdpa: leave a data-dependent mask width unchecked (SAM 3) - steps no longer return counts; _STEPS holds (step, providers); shared helpers for node loops, node replacement and the attention decomposition Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: mark the ORT fusion tests slow; small fusion-step cleanups Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: lay out the ORT fusion tests like the OpenVINO ones Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: import numpy and torch directly in the fusion tests Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: cite the onnxscript issues behind two ORT fusion repairs Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> * Exporters: ort_fusions="dml" for DirectML DirectML gets the shared fusions with MultiHeadAttention kept: its GroupQueryAttention ignores rotary and attention_bias and needs unpacked q/k/v. No MoE kernel. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
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
x²/ldacross thousands of elements in fp16 easily overflows or loses precision.Key Changes
layer_norm.cuhLayerNorm,SimplifiedLayerNorm,LayerNormSmall,SimplifiedLayerNormSmallto accept and operate onfloatfor thread_data, epsilon, mu, rsigma. Removed unusedKeyValuePairSumoverloads for half/bfloat16.skip_layer_norm_impl.cuSkipLayerNormKernelandSkipLayerNormKernelSmallto accumulate in fp32 (cub::KeyValuePair<float, float>). Removedmaybe2halfhelper (no longer needed).embed_layer_norm_impl.cuTtofloat, accumulation to usefloatthread_data.profile_skip_layer_norm.pyprofile_skip_layer_norm.shparse_nsys.pyPerformance Results
Profiled on NVIDIA GPU with nsys (B=1, seq_len=2048, fp16 data, 200 iterations, skip first 5 warmup):
No measurable performance regression. The kernel is memory-bandwidth-bound, so fp32 arithmetic is completely hidden behind memory latency.
Testing
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 5Related PRs
#28442
#15660