onnxruntime
0e18025d - Add `activation` mode to `com.microsoft::GatedRMSNorm` for SiLU/Swish/Sigmoid across CPU, CUDA, and WebGPU (#32512)

Commit
6 days ago
Add `activation` mode to `com.microsoft::GatedRMSNorm` for SiLU/Swish/Sigmoid across CPU, CUDA, and WebGPU (#32512) ### Description `com.microsoft::GatedRMSNorm` previously hardcoded SiLU gating; this PR keeps backward-compatible default behavior (`silu`) and adds explicit activation selection across CPU, CUDA, and native WebGPU. - **Schema + docs** - Added string attribute `activation` to `GatedRMSNorm` schema with default `"silu"`. - Accepted values: `"silu"`, `"swish"` (alias), `"sigmoid"`. - Updated operator documentation to describe both formulas. - **CPU kernel** - Parse/validate `activation` during kernel construction. - Preserve existing FP32-intermediate RMSNorm flow and dtype coverage. - Apply stable gate function per mode: - SiLU/Swish: `z * sigmoid(z)` - Sigmoid: `sigmoid(z)` - **CUDA kernel** - Parse/validate `activation` in host wrapper. - Pass a compact activation selector to `LaunchGatedRMSNormKernel` (no device-side string handling). - Keep existing launch structure, dtype registrations (float/float16/bfloat16), and FP32 arithmetic. - **WebGPU kernel** - Parse/validate `activation`. - Specialize generated WGSL by activation mode and include mode in `CacheHint` to avoid shader-cache collisions. - Use stable sigmoid helper in both paths: - `normalized * stable_sigmoid(z)` (sigmoid) - `normalized * (z * stable_sigmoid(z))` (silu/swish) - **Cross-EP tests (`linear_attention_gates_op_test.cc`)** - Extended `GatedRMSNorm` test helper to take activation mode and activation attribute. - Added coverage for: - default (attribute omitted) SiLU behavior - explicit `"silu"` - `"swish"` alias parity with SiLU - `"sigmoid"` reference parity - float32 + float16 decode/prefill geometries - norm-size boundary cases - CUDA bfloat16 sigmoid path - invalid activation rejection - sigmoid stability with large-magnitude gate values - shape-validation behavior retention ```cpp OpTester tester("GatedRMSNorm", 1, onnxruntime::kMSDomain); tester.AddAttribute<std::string>("activation", "sigmoid"); // default is "silu" ``` ### Motivation and Context Models can require sigmoid output gating while existing ORT `GatedRMSNorm` was SiLU-only. This change enables both gating modes under the same operator ABI and keeps existing models behavior-identical by default. --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Author
Parents
Loading