[CUDA] Add LinearAttentionGate and GatedRMSNorm contrib operators

[tianleiwu]

Description

Adds two CUDA contrib operators that fuse the elementwise chains around the gated-delta
(linear attention) block:

com.microsoft::LinearAttentionGate

decay = decay_scale * Softplus(a + dt_bias)
beta  = Sigmoid(b)                            # optional second output

dt_bias and decay_scale are float32 per-head vectors of length H; decay_scale is the
already-negated -exp(A_log) factor.

com.microsoft::GatedRMSNorm

Y = X * rsqrt(mean(X^2) + epsilon) * scale * SiLU(gate)

Both are registered for float, float16 and bfloat16.

Motivation and Context

Reference implementations compute the decay in float32 because exp(decay) inside the
recurrence exponentially amplifies any precision loss. Exporters therefore emit
Cast -> Add -> Softplus -> Mul -> Cast — five kernel launches over a tensor with only
num_heads elements per token. At that size every launch is pure overhead: the work per kernel
is a few hundred elements while the launch/replay cost is fixed. The gated RMSNorm chain has the
same shape of problem.

These ops keep the intermediates in float32 registers, so a single launch replaces each chain.
The float32 intermediate is preserved exactly, so the fused output is bit-identical to the
unfused graph and no accuracy gate is needed.

Measured impact

1x H200 SXM (SM90, 132 SM, ~4.8 TB/s HBM), CUDA 13.0, Qwen3.6-35B-A3B-NVFP4 + MTP N=3
(verify batch M=4), CUDA graph on, 200-step decode, 4 interleaved reps with the first
discarded:

rep	before (ms/step)	after (ms/step)
2	10.295	9.824
3	10.331	9.817
4	10.340	9.832
mean	10.322	9.824

−0.50 ms/step (−4.8%). No overlap between the two sets.

Graph-level accounting confirms the mechanism is launch/replay overhead, not arithmetic:

• text.onnx node count: 1973 -> 1583 (390 nodes removed: 30 decay chains + 30 gated norms).
• In-graph time: 8.945 -> 8.604 ms/step (−0.341 ms).
• The per-node graph-replay gap model predicts 0.89 us x 390 = 0.347 ms, against a measured
  −0.341 ms.

Output is bit-identical, and the MTP acceptance rate is unchanged in practice
(495–537 vs 503–544 tokens per 200 steps).

Tests

onnxruntime/test/contrib_ops/linear_attention_gates_op_test.cc — fp32/fp16/bf16, with and
without the optional b / beta output, checked against the unfused chain.

Note

Stacked on #31157 (this PR's base branch) because both edit bert_defs.cc and
docs/ContribOperators.md. docs/ContribOperators.md still needs regenerating
(python tools/python/gen_contrib_doc.py after a build) — draft until then. Review only the
top commit.

[Copilot]

Pull request overview

This PR extends ONNX Runtime’s contrib/CUDA LLM stack with new fusions and several related correctness/perf improvements across attention, MoE, transpose optimization, and EPContext utilities. The primary functional addition is two new com.microsoft contrib ops intended to collapse small elementwise chains into single CUDA launches.

Changes:

• Add new CUDA contrib ops LinearAttentionGate and GatedRMSNorm, plus unit tests and schema registration.
• Add/extend state_window support and validation for recurrent-state contrib ops (LinearAttention, CausalConvWithState) across CUDA/CPU/WebGPU and tests.
• Improve MoE FP4 (MXFP4) memory/dispatch behavior and top-k kernels; add additional optimizer/tests and minor WebNN builder fix.

Reviewed changes

Copilot reviewed 49 out of 50 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
onnxruntime/test/python/transformers/test_qmoe_fp4_cuda.py	Adds SM80-regime FP4/QMoE tests for single-copy weights + initializer release behavior.
onnxruntime/test/python/transformers/test_gqa.py	Adds opset imports and a FlashDecode multi-token parity test.
onnxruntime/test/providers/cpu/llm/attention_op_test.cc	Adds negative tests for 3D Attention shape validation failures.
onnxruntime/test/optimizer/transpose_optimizer_test.cc	Adds tests for new Transpose→Reshape(split) rewrite behavior.
onnxruntime/test/contrib_ops/linear_attention_op_test.cc	Adds CUDA-only state_window coverage for LinearAttention across kernel families.
onnxruntime/test/contrib_ops/linear_attention_gates_op_test.cc	New tests for LinearAttentionGate and GatedRMSNorm fused CUDA ops.
onnxruntime/test/contrib_ops/causal_conv_with_state_op_test.cc	Adds CUDA-only state_window tests for CausalConvWithState.
onnxruntime/test/autoep/library/example_plugin_ep/ep.cc	Updates example EP to use zero-copy EpContextData read path.
onnxruntime/test/autoep/ep_context_data_utils_test.cc	Expands tests for path validation, file-type gates, symlink handling, and zero-copy read semantics.
onnxruntime/test/autoep/ep_context_data_callbacks.h	Extends callback state to capture allocated buffer for adoption assertions.
onnxruntime/core/providers/webnn/builders/impl/matMulNBits_op_builder.cc	Fixes dequant reshape to padded-K then slices to true K.
onnxruntime/core/providers/cuda/cu_inc/topk_warp_sort.cuh	Adds in-register bitonic sort + warp Top-N helper.
onnxruntime/core/providers/cpu/llm/attention_helper.h	Adds missing 3D Attention validation (K hidden width, K/V seq len consistency).
onnxruntime/core/optimizer/transpose_optimization/onnx_transpose_optimization.cc	Adds HandleReshapeSplit rewrite (Transpose→Reshape(split) → Reshape→Transpose).
onnxruntime/core/graph/contrib_ops/ms_opset.h	Adds schema declarations/registration for new contrib ops.
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Adds schemas/docs for LinearAttentionGate and GatedRMSNorm; extends state_window docs/inference.
onnxruntime/contrib_ops/webgpu/bert/linear_attention.cc	Enforces state_window == 0 for WebGPU LinearAttention.
onnxruntime/contrib_ops/webgpu/bert/causal_conv_with_state.cc	Enforces state_window == 0 for WebGPU CausalConvWithState.
onnxruntime/contrib_ops/cuda/moe/qmoe_kernels.cu	Adds warp-register TopK path for certain expert-count regimes.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.h	Adds flags/docs for SM80 single-copy FP4 weights and raw-initializer release.
onnxruntime/contrib_ops/cuda/moe/moe_quantization.cc	Implements SM80 single-copy logic, release gating, and additional safety checks/logging.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_kernels.cu	Switches expert ranker to BlockRadixRankMatch; stages routing metadata for finalize kernel.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.h	Extends GEMV FP4 support/launch API for SM80 pair-interleaved decode.
onnxruntime/contrib_ops/cuda/llm/moe_gemm/moe_gemv_fp4.cu	Adds SM80 pair-interleaved details/dispatch and layout-support gate.
onnxruntime/contrib_ops/cuda/llm/fpA_intB_gemv/details.h	Extends FP4 converter to decode pair-interleaved nibble layout.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers new CUDA contrib kernels (LinearAttentionGate, GatedRMSNorm).
onnxruntime/contrib_ops/cuda/bert/linear_attention.h	Adds state_window_ member.
onnxruntime/contrib_ops/cuda/bert/linear_attention.cc	Parses state_window and uses shared shape helper + updated kernel signature.
onnxruntime/contrib_ops/cuda/bert/linear_attention_impl.h	Extends kernel launch signature to include windowed state.
onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.h	New op kernel declarations for LinearAttentionGate/GatedRMSNorm.
onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc	New CUDA kernel registrations and ComputeInternal implementations.
onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.h	New CUDA launcher declarations for fused gate kernels.
onnxruntime/contrib_ops/cuda/bert/linear_attention_gates_impl.cu	New CUDA implementations for fused gate kernels.
onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc	Documents FlashDecoding multi-token decode parity assumptions.
onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.h	Adds state_window_ member.
onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state.cc	Parses state_window and uses shared shape helper + updated kernel signature.
onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.h	Extends kernel launch signature to include windowed state.
onnxruntime/contrib_ops/cuda/bert/causal_conv_with_state_impl.cu	Implements windowed-state indexing and write semantics.
onnxruntime/contrib_ops/cpu/bert/linear_attention.h	Adds state_window_ placeholder for shared helper parity.
onnxruntime/contrib_ops/cpu/bert/linear_attention.cc	Enforces CPU state_window == 0 and uses shared CheckInputs helper.
onnxruntime/contrib_ops/cpu/bert/linear_attention_helper.h	New shared helper for parsing/validating LinearAttention state_window shapes.
onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.h	Adds state_window_ placeholder for shared helper parity.
onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state.cc	Enforces CPU state_window == 0 and uses shared CheckInputs helper.
onnxruntime/contrib_ops/cpu/bert/causal_conv_with_state_helper.h	New shared helper for parsing/validating CausalConvWithState state_window shapes.
js/node/package.json	Bumps adm-zip dependency.
js/node/package-lock.json	Updates adm-zip lock entry (resolved URL + integrity + engines).
docs/contrib_ops/cuda/moe_qmoe.md	Documents SM80 single-copy MXFP4 weights behavior and related env vars.

Files not reviewed (1)

• js/node/package-lock.json: Generated file

Suppressed comments (2)

onnxruntime/contrib_ops/cuda/bert/linear_attention_gates.cc:44

• GatedRMSNorm's schema/type constraints include bfloat16, but only float and MLFloat16 kernels are registered here. This will make bf16 models fail kernel selection at runtime.

Consider registering the BFloat16 specialization as well (and ensure the .cu instantiations exist).

REGISTER_KERNEL_TYPED(GatedRMSNorm, float)
REGISTER_KERNEL_TYPED(GatedRMSNorm, MLFloat16)

onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc:424

• The central kernel-registration list also lacks BFloat16 entries for the new ops. Without these, even if the bf16 kernels are implemented/registered, they won't be discoverable via the CUDA contrib kernel registry.

      BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttention)>,
      BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttention)>,
      BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, float, LinearAttentionGate)>,
      BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, LinearAttentionGate)>,
      BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, float, GatedRMSNorm)>,
      BuildKernelCreateInfo<CUDA_MS_OP_TYPED_CLASS_NAME(1, MLFloat16, GatedRMSNorm)>,