[Akshay Sonawane (apsonawane)]

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 sigmoid_in_fp32 attribute, eliminating the need for any new standalone contrib operators (ArgMaxLastAxis, SigmoidGateMul, HyperConnectionPostMix) in Qwen3.8-Flash-Next RC2.

Key Changes

Component	Files	Description
Standard ArgMax / ArgMin BFloat16 (CUDA)	onnxruntime/core/providers/cuda/reduction/reduction_ops.cc, reduction_functions.cu, cuda_execution_provider.cc	Registers BFloat16 for ArgMax and ArgMin (opset 13+) and instantiates reduce_lastdim_max_min<BFloat16> for fast last-axis vocabulary argmax without requiring com.microsoft::ArgMaxLastAxis.
com.microsoft::GatedAdd Projected Gate & Vectorized Kernel	onnxruntime/core/graph/contrib_ops/bert_defs.cc, onnxruntime/contrib_ops/cuda/bert/gated_add.{h,cc}, gated_add_impl.{h,cu}	Adds optional 4th input gate_weight ([K, 1]) and sigmoid_in_fp32 attribute to GatedAdd (output = X + Y * sigmoid(gate @ gate_weight)), plus GatedAddVec8Kernel (uint4 8-element vectorized loads/stores) and ProjectedSigmoidGateMulKernel.
Tests & Docs	onnxruntime/test/contrib_ops/gated_add_op_test.cc, docs/ContribOperators.md, docs/OperatorKernels.md	Adds unit tests for BFloat16 ArgMax (keepdims=0/1) and GatedAdd projected gate (sigmoid_in_fp32=0/1), and syncs operator schema/kernel tables.

Testing Notes

• Run unit tests: onnxruntime_provider_test --gtest_filter="ContribOpGatedAddTest.*"
• Verified end-to-end on qwen_flash_next_int4_int8_rc2_contrib_dflash2 ($S=8$ DFlash2 verification): 27/27 engine_behaviours.py pass, 15.54–15.88 ms/cycle at 2K–32K context, and 673 / 800 (84.125%) on MMLU-Pro (391.26 tok/s median, 444.08 tok/s P90, 474.01 tok/s max).

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

[github-actions]

You can commit the suggested changes from lintrunner.