[Akshay Sonawane (apsonawane)]

Description

Adds BFloat16 support to standard ai.onnx::ArgMax and ai.onnx::ArgMin (opset 13+) on the CUDA Execution Provider and extends com.microsoft::GatedAdd with an optional projected-gate weight (gate_weight) and 8-element vectorized (uint4) CUDA kernels. Together, these changes allow Qwen3.8-Flash-Next RC2/RC3 models to use standard ArgMax and GatedAdd without introducing standalone ArgMaxLastAxis or SigmoidGateMul contrib operators.

Summary of Changes

Standard ArgMax / ArgMin BFloat16 (CUDA)

File	Change
onnxruntime/core/providers/cuda/reduction/reduction_ops.cc	Registers BFloat16 for ArgMax and ArgMin (opset 13+) and casts BFloat16 inputs to FP32 in ReduceComputeCore when falling back to cuDNN for non-last-axis or select_last_index=1 reductions
onnxruntime/core/providers/cuda/reduction/reduction_functions.cu	Defines ArgReductionIdentity<BFloat16> (-inf / +inf bit patterns) and instantiates reduce_lastdim_max_min<BFloat16> for fast last-axis reductions
onnxruntime/core/providers/cuda/cuda_execution_provider.cc	Registers the opset 13 BFloat16 ArgMax and ArgMin kernel creation info in the CUDA provider table

com.microsoft::GatedAdd Projected Gate & Vectorized CUDA Kernel

File	Change
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Adds optional 4th input gate_weight ([K, 1]) to GatedAdd(1), updates type/shape inference to validate both plain mode (gate shape (..., 1)) and projected mode (gate shape (..., K), gate_weight shape (K, 1)), and documents both modes including FP32 accumulation/sigmoid and BFloat16 rounding points
onnxruntime/contrib_ops/cuda/bert/gated_add.cc	Validates gate_weight shape ([K, 1] with K > 0), restricts projected mode to BFloat16 on CUDA (returning NOT_IMPLEMENTED for other types), and dispatches LaunchProjectedSigmoidGateMul
onnxruntime/contrib_ops/cuda/bert/gated_add_impl.h, onnxruntime/contrib_ops/cuda/bert/gated_add_impl.cu	Adds ProjectedSigmoidGateMulKernel (block-reduced FP32 dot product + FP32 stable sigmoid with BFloat16 rounding after projection, sigmoid, multiplication, and addition, using uint4 8-element vectorized loads/stores when aligned) and GatedAddVec8Kernel for 16-byte aligned FP16/BF16 plain-mode GatedAdd
onnxruntime/contrib_ops/cpu/bert/gated_add.cc, onnxruntime/contrib_ops/webgpu/bert/gated_add.cc	Returns NOT_IMPLEMENTED when gate_weight is provided

Tests and Documentation

• onnxruntime/test/contrib_ops/gated_add_op_test.cc: Adds coverage for non-vectorized FP16/BF16 hidden sizes, BFloat16 projected-gate GatedAdd across multiple (width, K) shapes ({8, 16}, {2560, 16}, {2560, 2560}, {259, 300}), negative shape/rank validation (gate vs X rank/leading dims, invalid gate_weight rank/columns), and non-BFloat16 rejection.
• onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc, onnxruntime/test/providers/cuda/test_cases/reduction_functions_test.cc: Adds BFloat16 ArgMax and ArgMin tests covering keepdims=0/1, single-block and multi-block last-axis reductions (N=2048), non-last-axis reductions (axis=0), select_last_index=1, ties, and NaN propagation.
• docs/ContribOperators.md, docs/OperatorKernels.md: Regenerates operator schema and kernel registration documentation to match the updated definitions.

Testing

• Run provider unit tests:
  onnxruntime_provider_test --gtest_filter="ContribOpGatedAddTest.*:ReductionOpTest.*bfloat16*:ReductionFunctionsTest.*"
• End-to-end verification on qwen_flash_next_int4_int8_rc3_contrib_dflash2 ($S=8$ DFlash2 speculative decoding): 27/27 engine_behaviours.py pass, 14.96–15.30 ms/cycle across 2K–32K context lengths, and 665 / 800 (83.125%) on MMLU-Pro.

Checklist

• Tests added/updated
• No breaking changes (or documented in description)
• Documentation updated (if applicable)

[Copilot]

🟡 Changes recommended

Unresolved memory-safety, dependency-ordering and schema-integration defects block approval.

7 open issues

• Regenerate operator documentation and CUDA registration listings · New
• Complete or remove unsupported FP8/BF16 dequantization specialization · New
• Prevent projection extent overflow before bounds validation · New
• Return a valid index for all-NaN reduction rows · New
• Validate embedding IDs before device-memory table access · New
• Prevent dependent reads before programmatic launch synchronization · New
• Support packed EngramGate inputs during graph resolution · New

What changed in this PR

Adds fused CUDA operators for transformer workloads and extends EngramGate to accept packed inputs on CPU and CUDA.

Changes:

• Adds HyperConnectionLowRankMix and five fused tensor operators.
• Adds numerical, shape-validation and CUDA graph regression tests.
• Adds schemas, kernel registrations and partial operator documentation.

File	Description
onnxruntime/​test/​python/​transformers/​test_hyper_connection_low_rank_mix.py	Adds reference comparisons and graph-capture tests.
onnxruntime/​test/​contrib_ops/​hyper_connection_low_rank_mix_op_test.cc	Tests fused and fallback paths.
onnxruntime/​test/​contrib_ops/​fused_tensor_ops_test.cc	Tests fused tensor operators.
onnxruntime/​test/​contrib_ops/​engram_ops_test.cc	Adds packed EngramGate cases.
onnxruntime/​core/​graph/​contrib_ops/​ms_opset.h	Registers new schemas.
onnxruntime/​core/​graph/​contrib_ops/​bert_defs.cc	Defines operator contracts and shape inference.
onnxruntime/​contrib_ops/​cuda/​cuda_contrib_kernels.cc	Registers new CUDA kernels.
onnxruntime/​contrib_ops/​cuda/​bert/​hyper_connection_low_rank_mix.h	Declares the operator class.
onnxruntime/​contrib_ops/​cuda/​bert/​hyper_connection_low_rank_mix.cc	Validates inputs and dispatches execution.
onnxruntime/​contrib_ops/​cuda/​bert/​hyper_connection_low_rank_mix_impl.h	Defines launch interfaces and configuration.
onnxruntime/​contrib_ops/​cuda/​bert/​hyper_connection_low_rank_mix_impl.cu	Implements fused kernels and cuBLAS fallback.
onnxruntime/​contrib_ops/​cuda/​bert/​fused_tensor_ops.h	Declares tensor operators and launchers.
onnxruntime/​contrib_ops/​cuda/​bert/​fused_tensor_ops.cu	Implements tensor CUDA kernels.
onnxruntime/​contrib_ops/​cuda/​bert/​fused_tensor_ops.cc	Adds validation and kernel wrappers.
onnxruntime/​contrib_ops/​cuda/​bert/​engram_gate.cc	Accepts packed CUDA inputs.
onnxruntime/​contrib_ops/​cuda/​bert/​engram_gate_impl.h	Updates the launcher interface.
onnxruntime/​contrib_ops/​cuda/​bert/​engram_gate_impl.cu	Uses flattened token counts.
onnxruntime/​contrib_ops/​cpu/​bert/​engram_gate.cc	Accepts packed CPU inputs.
docs/​ContribOperators.md	Documents HyperConnectionLowRankMix.

_{Review effort: Balanced}

[Copilot]

🟡 Changes recommended

Layout validation, provider dispatch, no-exceptions tests, and generated documentation still contain actionable defects.

8 open issues

• Regenerate documentation sections in renderer order · New
• Preserve literal T constraint keys in kernel documentation · New
• Checked narrowing for ONNX stride attributes · New
• Correct projection extent calculation for padded head strides · New
• Add WebGPU fallback for unsupported packed EngramGate layouts · New
• Skip fail-shape-inference test without exception support · New
• Skip aborting invalid modes when exceptions are disabled · New
• Regenerate operator documentation and CUDA registration listings

6 resolved since last review

• Support packed EngramGate inputs during graph resolution
• Prevent dependent reads before programmatic launch synchronization
• Validate embedding IDs before device-memory table access
• Return a valid index for all-NaN reduction rows
• Prevent projection extent overflow before bounds validation
• Complete or remove unsupported FP8/BF16 dequantization specialization

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

Pushed e587c56 addressing all eight outstanding review threads: checked narrowing of packed gate attributes, precise padded-stride projection bounds, WebGPU-to-CPU packed EngramGate fallback with real session regressions, safe no-exception regression guards, and generated schema ordering/type-constraint corrections.

Validation: full-PR lintrunner -a and check-only lint passed in qwen-flash; official renderer/source-metadata checks verify complete schema ordering and exact CPU/CUDA type sets; projection extent arithmetic was checked across 2,304 layouts. Source whitespace checks pass. Renderer-produced blank-line indentation was preserved in the moved generated documentation.

No local builds, compilation, native tests, or model execution were performed. New C++ regressions are added but unexecuted locally. At inspection, no head checks were failed; several were still running. Fresh post-push CI, including native --gen_doc validate, is pending and is not claimed green.

[Copilot]

🔵 Needs a closer look

WebGPU adapter builds still assign packed EngramGate inputs to an unsupported kernel instead of falling back to CPU.

No open issues 🎉

2 previously missed issues

In code that hasn't changed since last review

Avoid unused scratch allocation for single-partition argmax

onnxruntime/​contrib_ops/​cuda/​bert/​fused_tensor_ops.cc:74

For last-axis widths up to 4096, LaunchTypedArgmax writes directly to the output and never uses either scratch buffer. These allocations still reserve 12 bytes per row: 12 MiB of unused scratch for a [1048576, 1] input. Set the scratch count to zero for a single partition; GetScratchBuffer returns nullptr for zero counts, avoiding both allocations.

Apply rank filtering in adapter capability path

onnxruntime/​core/​providers/​webgpu/​webgpu_execution_provider.cc:784

This rank filter is compiled out under ORT_USE_EP_API_ADAPTERS (the enclosing condition starts at line 706). Adapter builds instead use Ep::GetCapabilityImpl in core/providers/webgpu/ep/ep.cc:59–168, which accepts registered EngramGate nodes without checking rank. Packed FP32/FP16 inputs therefore still reach the WebGPU kernel's rank-4 check and fail instead of falling back to CPU. Add the equivalent filter to the plugin capability path and run the new fallback regressions with adapters enabled.

8 resolved since last review

• Skip aborting invalid modes when exceptions are disabled
• Skip fail-shape-inference test without exception support
• Add WebGPU fallback for unsupported packed EngramGate layouts
• Correct projection extent calculation for padded head strides
• Checked narrowing for ONNX stride attributes
• Preserve literal T constraint keys in kernel documentation
• Regenerate documentation sections in renderer order
• Regenerate operator documentation and CUDA registration listings

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

Pushed b679f3b: optional projected SigmoidGateMul, explicit prefill arithmetic control, and shared-expert activation/row-gate preparation. Existing layout/overflow validation and unrelated reviewer changes are preserved. Scoped CUDA host/device and schema translation units compile. Shared-expert tests in companion #33473 compare against these unfused operators; the combined local runtime passed its 38 focused shared-expert native tests. This does not establish independent PR-branch build/link qualification. Companion #33473 depends on these new helper declarations and implementations.

[Akshay Sonawane (apsonawane)]

Resolved conflicts against main (1f4b965) and pushed through dbd8b83. Both schema-registration conflicts retain the PR tensor operators and main’s ConvStateReplay declaration/registration.

Validation before pushing:

• Ran lintrunner -a, reviewed/applied formatting-only changes, then reran lintrunner -a and lintrunner: no lint issues.
• git diff --check against main passes; removed trailing whitespace in the PR operator documentation.
• Compiled merged environment.cc, bert_defs.cc, contrib_defs.cc, fused_tensor_ops.cc and fused_tensor_ops.cu using the exact isolated PR headers and main-compatible ONNX protobuf headers. All five compilation checks passed.

This is targeted merged-source compilation, not a full clean standalone build, link, or runtime-test qualification. Existing integration worktree, active model package, and LiveCodeBench snapshot were untouched.

[Copilot]

🟡 Changes recommended

Existing CPU and WebGPU kernels can silently ignore the newly accepted projection weight and return incorrect results.

4 open issues

• Legacy GatedAdd kernels ignore projected gate weights · New
• Expand BF16 ArgMax and non-last-axis reduction coverage · New
• Test odd output widths and large-K projection accumulation · New
• Document sigmoid_in_fp32 default annotation · New

_{Review effort: Balanced}

[Copilot]

🔵 Needs a closer look

Targeted documentation checks do not establish sufficient confidence to automatically approve the new CUDA kernels and numerical behavior.

1 open issue

• Document gate_weight mode and regenerate GatedAdd operator docs · New

4 resolved since last review

• Legacy GatedAdd kernels ignore projected gate weights
• Test odd output widths and large-K projection accumulation
• Expand BF16 ArgMax and non-last-axis reduction coverage
• Document sigmoid_in_fp32 default annotation

_{Review effort: Balanced}

[kunal-vaishnavi]

From Copilot:

Completed an implementation pass and a second edge-case/regression pass at edf701c. This PR merged after the review was completed, so I am recording the outcome here. No actionable finding.

Checked BF16 ArgMin/ArgMax registration, reduction identities, non-last-axis conversion, special-value/tie test coverage, projected-gate shape constraints and explicit BF16 rounding stages, vector-load alignment guards, and CPU/WebGPU rejection of unsupported projected mode. The projected gate preserves the documented rounding after projection, sigmoid, multiplication, and addition.

Validation: source and existing-test review only for this PR. No PR-specific ORT build, provider-suite execution, WebGPU execution, cross-architecture qualification, or performance benchmark was performed. This is a source-review outcome, not a claim that all tests and CI passed.