Skip to content

[CUDA] Add GatedAdd contrib operator - #31835

Open
Tianlei Wu (tianleiwu) wants to merge 1 commit into
mainfrom
tlwu/20260810/gated_add
Open

[CUDA] Add GatedAdd contrib operator#31835
Tianlei Wu (tianleiwu) wants to merge 1 commit into
mainfrom
tlwu/20260810/gated_add

Conversation

@tianleiwu

Copy link
Copy Markdown
Contributor

Description

Add com.microsoft::GatedAdd, a CUDA contrib operator that computes X + round_to_T(Y * gate) with a per-row gate broadcast across the hidden dimension. The FP16 and BF16 implementations use separate PTX multiply and add instructions to preserve the rounding behavior of separate ONNX Mul and Add nodes rather than contracting to FMA semantics.

The operator is intended to remove one launch per gated residual/add pair in MoE shared-expert paths. On Qwen3.6-35B-A3B-NVFP4 with N=3 MTP, fusing 40 main-model pairs plus one MTP pair reduced median decode latency from 7.311 to 7.225 ms/round (-1.18%). A graph-off Nsight capture measured 40.35 fewer launches/round and 1.30% lower GPU kernel time.

Changes

  • Add the GatedAdd schema for float, float16, and bfloat16.
  • Add CUDA kernels with strict reduced-precision rounding semantics.
  • Register all three CUDA kernel types.
  • Add float, FP16, and BF16 contrib-op tests.

Validation

  • Runtime micro-model outputs are bit-exact with separate Mul + Add for float, FP16, and BF16.
  • onnxruntime and onnxruntime_providers_cuda targets build successfully.
  • clang-format --dry-run --Werror passes for all touched C++/CUDA files.
  • Real Qwen graph validation replaced exactly 40 main and one MTP Mul + Add pairs; all other nodes, initializers, inputs, and outputs remained unchanged after normalizing fused edge names.

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.

1 participant