Skip to content

Add activation mode to com.microsoft::GatedRMSNorm for SiLU/Swish/Sigmoid across CPU, CUDA, and WebGPU - #32512

Merged
kunal-vaishnavi merged 9 commits into
mainfrom
copilot/extend-gatedrmsnorm-contrib-operator
Sep 18, 2026
Merged

kunal-vaishnavi merged 9 commits into
mainfrom
copilot/extend-gatedrmsnorm-contrib-operator

Conversation

Copilot AI commented Sep 9, 2026 •

Copy link
Copy Markdown
Contributor

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
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.

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

@kunal-vaishnavi
kunal-vaishnavi marked this pull request as ready for review September 9, 2026 18:29
Copilot AI balanced review requested due to automatic review settings September 9, 2026 18:29

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot wasn't able to review any files in this pull request.


💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Copilot AI changed the title [WIP] Extend GatedRMSNorm operator to support sigmoid-gated behavior Add activation mode to com.microsoft::GatedRMSNorm for SiLU/Swish/Sigmoid across CPU, CUDA, and WebGPU Sep 9, 2026
@titaiwangms
Ti-Tai Wang (titaiwangms) requested a balanced review from Copilot September 9, 2026 18:44

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

The CUDA negative sigmoid path loses precision, its stability test cannot detect that regression, and generated operator documentation is stale.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Review details
  • Files reviewed: 10/10 changed files
  • Comments generated: 3
  • Review effort level: Balanced

Comment thread onnxruntime/core/graph/contrib_ops/bert_defs.cc
Comment thread onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.cu
Comment thread onnxruntime/test/contrib_ops/linear_attention_gates_op_test.cc Outdated
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Comment thread onnxruntime/contrib_ops/cpu/bert/linear_attention_gates.cc Outdated
Comment thread onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc Outdated
This was referenced Oct 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

6 participants