CUDA: fix Nemotron-H relu2, RMSNorm, and activation precision

[Rui Ren (rui-ren)]

Summary

Address CUDA operator gaps encountered while onboarding Nemotron-H: squared-ReLU activation in MoE/QMoE, the standard ONNX RMSNormalization rounding contract, and cancellation in Softplus, Sigmoid and fused QuickGelu/SiLU. The CUDA MoE parser previously rejected relu2, internal Relu2 dispatch selected ordinary ReLU, the RMS kernel scaled an FP32 standardized value without first rounding it back to the input type, Softplus lost small positive tails by rounding exp(x)+1 before taking log, and both Sigmoid and fused QuickGelu subtracted a rounded reciprocal from one for negative inputs.

• Parse relu2 and dispatch a dedicated Cutlass squared-ReLU functor in standalone and gated activation paths.
• Preserve the legacy FC1 output-type rounding and apply bias exactly once before squaring.
• Explicitly reject a Relu2 epilogue where the direct GEMM path cannot implement it.
• Document the CUDA activation and add independent dyadic numeric tests for FP32, FP16, BF16, raw INT4 and raw INT8 weights with CPU fallback disabled.
• Restore the required rounding between ONNX RMSNormalization's two stages: round the standardized value back to X's type, then apply scale. Cover scalar and half2 paths while keeping historical LayerNormalization/SimplifiedLayerNormalization behavior on the existing template path.
• Honor the default FP32 stage-one stash even for double inputs, while retaining double-precision scaling. Preserve explicit double stash for double inputs/scales, and reject unsupported stash choices rather than silently ignoring the attribute.
• Add exact BF16, FP16 and FP16-to-FP32 CUDA rounding regressions at odd, even and Nemotron-sized widths, with CPU fallback disabled. Update the registered-type oracle to follow the specified stage boundary, not a one-stage FP32 product.
• Evaluate Softplus with log1p and FP32 intermediates for low-precision inputs, retaining FP64 computation for doubles. Add four CUDA-only precision regressions against an independent FP64 reference, with zero absolute tolerance and CPU fallback disabled.
• Compute Sigmoid's negative branch directly, preserving representable small probabilities. Use the same FP32/FP64 intermediate policy and remove the redundant half specialization. The pinned decoder uses 23 FP32 Sigmoid router nodes and 23 FP32 Softplus nodes.
• Apply the same stable sigmoid policy inside fused QuickGelu and round only after its final product. Preserve the FP32 alpha attribute for low-precision inputs, and validate positive, negative and zero alpha. Profiling the unchanged pinned decoder shows 23 FP32 Swish gates actually execute as CUDA QuickGelu, so fixing the separate Sigmoid kernel alone does not correct this fused path.

CUDA validation

Built matching ORT 1.32 development sources with an isolated verified CUDA 13.0.88 SDK and cuDNN 9.24 on H100. All five new activation tests passed. The complete MoETest.* suite passed 73 tests, with 13 existing architecture/provider-dependent skips.

All three exact RMS rounding regressions fail against the original kernel and pass after the fix. The combined RMSNormalization, LayerNormalization, SkipLayerNormalization and MoE suites pass 172 tests, with the same 13 existing MoE skips. A separate synthetic-input probe using a pinned 2688-channel checkpoint scale and an independent two-stage PyTorch CPU oracle changes BF16 native mismatches from 2142/8064 to zero, and FP16 mismatches from 2119/8064 to zero. This is component arithmetic evidence, not full-checkpoint HuggingFace parity.

All four Softplus, four Sigmoid and four QuickGelu precision regressions fail before their respective corrections and pass afterwards. They use independent FP64 references, zero absolute tolerance, and CPU fallback disabled. QuickGelu covers alpha 1 (SiLU), 1.702, -1.702 and 0 across FP32, FP64, FP16 and BF16. Including the existing activation suites, the combined activation, normalization and MoE suites pass 224 cases with 13 existing MoE skips. Run the test binary from its build directory so the existing BERT LayerNorm fixtures are found.

The latest review follow-up at 8c36c6c8e8a7ec43fe503491b82778c2bddce262 fixes the remaining default-stash conformance issue for double RMS inputs and corrects the registered-type oracle. An exact regression fails against the old double accumulator and passes with FP32 standardization plus double scaling. Separate tests cover explicit double stash and explicit rejection of unsupported choices. The complete coupled suites now pass 227 cases with 13 existing skips, and GenAI's 21 public, 34 engine and eight independent CUDA fixture cases pass against the exact final provider SHA256 068c67babdffb1e14a6ed0c5d32be352ac0bc1be8f3d1de8ca1e87e5dfb812f2. The full-model results below are tied to the earlier post-QuickGelu binary, not claimed as rerun against this later provider.

This focused build disables Flash Attention and FP4 QMoE; neither is claimed as validated. The raw integer tests require SM75+; FP32 and quantized legacy paths and H100 FP16/BF16 paths are covered.

Scope and readiness

These are CUDA operator corrections, not an expert-routing fusion, native NVFP4 checkpoint support, or a claim that Nemotron 3.5 Lightning is production-ready. The pinned export in onnxruntime/mobius#761 still unrolls experts rather than emitting fused MoE. Its pre-fused SkipSimplifiedLayerNormalization nodes retain their historical semantics; the standard RMS correction does not repair those serialized nodes or establish full-decoder parity. Published full-decoder numerical gates remain unresolved; a recipe and a loadable graph do not clear those gates.

A full native GenAI control after the RMS correction, before the Softplus correction, reaches 257 ordinary tokens but still diverges in both N1 and N2 MTP at index 64 (87/129 and 122/169 accepted drafts, respectively), with CUDA TF32 disabled consistently. The earlier supplied full-model reference scorecard was not rerun or replaced by the component checks.

A fresh unchanged-export control with both corrections still fails: N1 first diverges at index 70, accepting 86/131 drafts, and N2 at index 100, accepting 107/179. All three sequences contain 257 tokens. GenAI's 21 public tests, 34 engine tests and eight independent CUDA fixture cases pass against this exact provider; these do not substitute for full-model generation or reference parity.

The Sigmoid-stage provider also includes the Sigmoid correction, and the same 21 public tests, 34 engine tests and eight independent CUDA fixture cases pass against it. A fresh unchanged-export control against this provider still fails ordinary/MTP equivalence: N1 first diverges at index 90 (90/126 accepted drafts), and N2 at index 119 (119/167). All three sequences contain 257 tokens; indices are zero-based and include the 40-token prompt. CUDA TF32 and graph capture are disabled. The original decoder/head graph hashes remain unchanged. The 70/100 divergence above was measured before the Sigmoid correction.

A subsequent unchanged-gate rerun against the Sigmoid-corrected provider fails 139/240 comparisons under both TF32-on and TF32-off policies (maximum NRMSE 0.10676754207476785, minimum cosine 0.9942840098112892). This is the supplied independent FP32 source-checkpoint equation reference, not a direct HuggingFace model-class invocation. All floating-point model computation remains on CUDA; CPU placements are verified integer-only. A target-only replay also reproduces the N2 index-119 divergence without any head/controller call, while the index-90 replay retains the ordinary argmax; this does not establish the cause of the remaining N1 mismatch.

The later provider additionally fixes fused QuickGelu at 0acd2ee5ca0d738767d021be52d1e816ad70913b, and all 21 public, 34 engine and eight independent CUDA fixture cases remain green against its exact binary (CUDA provider SHA256 c54d757853ad295a933d20a55435341694fd1e7845695117f0371f27cc9315ea). The fresh unchanged-export 257-token control still fails: N1 first diverges at index 109, accepting 92/125 drafts; N2 at index 92, accepting 121/165. CUDA TF32 and graph capture are disabled. Under the original TF32-on policy, the unchanged FP32-equation reference gates fail 141/240 comparisons, with maximum NRMSE 0.1106184779356323 and minimum cosine 0.9938640906045041; floating-point computation remains CUDA-only. The 90/119 generation results and 139/240 scorecard failures above were measured before QuickGelu's correction. Historical target-only attribution does not establish the cause of these current failures, and no direct HuggingFace model-class parity is claimed.

A separate direct upstream transformers==5.19.0 NemotronHForCausalLM run now loads the full pinned checkpoint without missing/mismatched decoder weights, using BF16 parameters, eager CUDA attention, TF32-on and upstream's PyTorch Mamba reference path (mamba_ssm is unavailable). Four independently cached forwards compare logits and final normalized hidden against the saved post-QuickGelu native outputs. Only 1/8 comparisons passes unchanged NRMSE/cosine thresholds; prefill logits pass, but hidden and all cached outputs fail (maximum NRMSE 0.11662261092503201, minimum cosine 0.9941554758121095, maximum absolute hidden error 5.5). Sampled argmaxes all agree, which is not a numerical/generation pass. This is additional direct BF16-class evidence, not a replacement for the original 240-output FP32-equation scorecard or full upstream generation.

A completed private control changes only the initial standard RMS node to scale in FP32 before its final BF16 cast. Against the final provider, it passes 0/8 direct upstream comparisons and 168/240 original FP32-equation comparisons (72 failures). It is not adopted or published as a fix; the original decoder/head graphs remain unchanged.

The same private graph with TF32 disabled still fails 71/240 equation comparisons (169 pass): all 60 prefill comparisons pass, but cached-step failures remain. Neither the private node change nor this precision policy is adopted as a fix.

The GenAI review follow-up additionally validates symbolic prompt-hidden shapes using device-allocated output, instead of pre-binding the expected sequence length. The updated GenAI contribution passes 23 public, 34 engine and eight independent CUDA fixture cases against this final ORT provider; no full-model pass is inferred.

Related work: onnxruntime/mobius#761, microsoft/olive-recipes#665 and the experimental GenAI onboarding microsoft/onnxruntime-genai#2710.

Verified commits: activation a275c830ad3425c59e6d821aed2f7a5ee616f49e; standard RMS rounding 6a0e7707f67838adf80215bf3770797b312a776b; stable Softplus 0f86545f17f92676a5cb5c2efca4076dc24eb39f; stable Sigmoid 19d99fa31eaf2623e8e20e2a77cf2e516b8be28f; stable fused QuickGelu 0acd2ee5ca0d738767d021be52d1e816ad70913b; default double-input RMS stash 8c36c6c8e8a7ec43fe503491b82778c2bddce262 (current HEAD).

[Rui Ren (rui-ren)]

Copilot Please review this PR and report actionable findings. Do not modify code.

[Copilot]

🔵 Needs a closer look

Numerically sensitive CUDA changes span multiple low-precision kernel paths and require final GPU-focused human validation.

0 open findings

What changed in this PR

Corrects CUDA numerical precision for Nemotron-H-related MoE activations, RMSNormalization, Softplus, and Sigmoid. No actionable defects were identified.

Changes:

• Adds squared-ReLU parsing and CUDA dispatch for MoE/QMoE.
• Restores RMSNormalization’s intermediate rounding contract.
• Improves Softplus/Sigmoid stability with CUDA-only regression coverage.

File	Description
onnxruntime/​test/​providers/​cpu/​nn/​rms_norm_op_test.cc	Adds exact CUDA rounding tests.
onnxruntime/​test/​providers/​cpu/​activation/​activation_op_test.cc	Adds activation precision regressions.
onnxruntime/​test/​contrib_ops/​moe_test.cc	Tests Relu2 across CUDA MoE variants.
onnxruntime/​core/​providers/​cuda/​nn/​rms_norm.cc	Enables pre-scale rounding.
onnxruntime/​core/​providers/​cuda/​nn/​layer_norm_impl.h	Adds rounding-mode template parameter.
onnxruntime/​core/​providers/​cuda/​nn/​layer_norm_impl.cu	Implements scalar and half2 rounding.
onnxruntime/​core/​providers/​cuda/​activation/​activations_impl.cu	Stabilizes Softplus and Sigmoid.
onnxruntime/​core/​graph/​contrib_ops/​contrib_defs.cc	Documents CUDA Relu2 support.
onnxruntime/​contrib_ops/​cuda/​moe/​moe_base.h	Parses the relu2 attribute.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_kernels.cu	Applies standalone squared ReLU correctly.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemm_template_dispatch.h	Rejects unsupported fused Relu2 epilogues.
onnxruntime/​contrib_ops/​cuda/​llm/​moe_gemm/​moe_gemm_activation_kernels.cuh	Adds squared-ReLU kernel dispatch.
docs/​ContribOperators.md	Regenerates MoE/QMoE documentation.

🧠 Review effort: Balanced

---

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

[Copilot]

🔵 Needs a closer look

The RMS test oracle incorrectly uses double precision for the default FP32 stash stage, masking nonconformant double-input behavior.

0 open findings

Previously missed (1)

In code that hasn't changed since last review

Double RMSNorm oracle masks incorrect FP32 stash conformance

onnxruntime/​test/​providers/​cpu/​nn/​rms_norm_op_test.cc:16

stash_type defaults to TensorProto::FLOAT even when InputType is double; the ONNX function casts X to that stage-one type, performs the RMS computation, and only then casts Normalized back to X's type. Selecting double here makes CheckCudaRmsNormTypes<double, double>() validate the CUDA kernel's existing double accumulation instead of the required FP32 stage-one result, so the updated oracle masks a remaining conformance bug. Use a float stash oracle for the default attribute and update the CUDA double specialization/dispatch accordingly (and handle explicit non-default stash types separately).

🧠 Review effort: Balanced

[Rui Ren (rui-ren)]

Addressed the actionable finding in Copilot review 5494829631 with commit 8c36c6c8e8a7ec43fe503491b82778c2bddce262.

The old double accumulator fails an exact CUDA-only regression using input values that round to one in the default FP32 stash and scale values that retain double precision. The fix casts stage-one inputs to FP32, rounds normalized values back to the input type, and scales double outputs in double precision. Explicit DOUBLE stash remains supported for double inputs/scales; unsupported choices fail explicitly instead of being silently ignored. The registered-type oracle now follows the default FP32 stash for every input type.

The complete coupled activation/normalization/MoE suites pass 227 tests with 13 existing skips. GenAI's 21 public, 34 engine and eight independent CUDA fixture cases also pass against the exact final provider (SHA256 068c67babdffb1e14a6ed0c5d32be352ac0bc1be8f3d1de8ca1e87e5dfb812f2). Full-checkpoint numerical and MTP gates remain unresolved and are documented separately; this operator conformance fix is not a readiness claim.

[Copilot]

🔵 Needs a closer look

Architecture-specific CUDA dispatch and exact mixed-precision semantics warrant final human review despite comprehensive tests.

0 open findings

🧠 Review effort: Balanced

[Rui Ren (rui-ren)]

Resolved the three documentation threads containing only - [ ], Rab, and dd as non-actionable, not as code fixes. Each already has a clarification request; no requested edit or clarification has been provided. No source or documentation change was invented for these comments. A concrete follow-up can be reopened and addressed independently of the substantive RMSNormalization finding, which was fixed and verified in 8c36c6c.