Skip to content

Add weightless hyper-connection fusion operators - #32687

Open
kunal-vaishnavi with Copilot wants to merge 11 commits into
mainfrom
copilot/add-hyper-connection-operators
Open

kunal-vaishnavi with Copilot wants to merge 11 commits into
mainfrom
copilot/add-hyper-connection-operators

Conversation

Copilot AI commented Sep 18, 2026

Copy link
Copy Markdown
Contributor

Description

Adds four com.microsoft operators for weightless hyper-connections:

  • BranchwiseRMSNorm
  • ScaledSiLU
  • HyperConnectionPreMix
  • HyperConnectionPostMix

Key changes:

  • Adds schemas, CPU/CUDA/WebGPU registrations, and per-operator source files.
  • Reuses existing SimplifiedLayerNormalization implementations for compatible BranchwiseRMSNorm paths.
  • Retains targeted fallbacks for optional, mixed-type, and branch-specific scales.
  • Adds dedicated symbolic shape-inference handlers for each operator.
  • Adds operator coverage and generated documentation.

Motivation and Context

Provides portable fused primitives for weightless hyper-connection models while sharing established normalization infrastructure where semantics permit. Dedicated schemas and shape-inference handlers preserve each operator’s distinct contract without coupling them to unrelated operators.

Copilot AI and others added 6 commits September 18, 2026 05:43
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

Comment thread onnxruntime/python/tools/symbolic_shape_infer.py Fixed
Comment thread onnxruntime/python/tools/symbolic_shape_infer.py Fixed
@kunal-vaishnavi
kunal-vaishnavi marked this pull request as ready for review September 18, 2026 08:13
Copilot AI balanced review requested due to automatic review settings September 18, 2026 08:13

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

The flattened symbolic-shape test relies on an expression-like dimension name that the inference engine treats atomically, and important normalization paths remain untested.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Adds four fused hyper-connection contrib operators across CPU, CUDA, and WebGPU, including schemas, shape inference, tests, and generated documentation.

Changes:

  • Implements BranchwiseRMSNorm, ScaledSiLU, PreMix, and PostMix kernels.
  • Adds shared validation and symbolic shape inference.
  • Registers and documents the operators across supported providers.
File summaries
File Description
onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py Adds symbolic shape tests.
onnxruntime/test/contrib_ops/hyper_connection_op_test.cc Adds operator tests.
onnxruntime/python/tools/symbolic_shape_infer.py Adds symbolic inference handlers.
onnxruntime/core/providers/cpu/nn/layer_norm_impl.h Exposes parameterized normalization construction.
onnxruntime/core/providers/cpu/nn/layer_norm_impl.cc Implements reusable normalization setup.
onnxruntime/core/graph/contrib_ops/ms_opset.h Registers operator schemas.
onnxruntime/core/graph/contrib_ops/bert_defs.cc Defines schemas and shape inference.
onnxruntime/contrib_ops/hyper_connection_helper.h Provides shared shape validation.
onnxruntime/contrib_ops/cpu/cpu_contrib_kernels.cc Registers CPU kernels.
onnxruntime/contrib_ops/cpu/bert/scaled_silu.h Declares CPU ScaledSiLU.
onnxruntime/contrib_ops/cpu/bert/scaled_silu.cc Implements CPU ScaledSiLU.
onnxruntime/contrib_ops/cpu/bert/hyper_connection_pre_mix.h Declares CPU PreMix.
onnxruntime/contrib_ops/cpu/bert/hyper_connection_pre_mix.cc Implements CPU PreMix.
onnxruntime/contrib_ops/cpu/bert/hyper_connection_post_mix.h Declares CPU PostMix.
onnxruntime/contrib_ops/cpu/bert/hyper_connection_post_mix.cc Implements CPU PostMix.
onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.h Declares CPU branchwise RMS normalization.
onnxruntime/contrib_ops/cpu/bert/branchwise_rms_norm.cc Implements CPU branchwise RMS normalization.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc Registers CUDA kernels.
onnxruntime/contrib_ops/cuda/bert/scaled_silu.h Declares CUDA ScaledSiLU.
onnxruntime/contrib_ops/cuda/bert/scaled_silu.cc Integrates CUDA ScaledSiLU kernel.
onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.h Declares CUDA launcher.
onnxruntime/contrib_ops/cuda/bert/scaled_silu_impl.cu Implements CUDA ScaledSiLU.
onnxruntime/contrib_ops/cuda/bert/hyper_connection_pre_mix.h Declares CUDA PreMix.
onnxruntime/contrib_ops/cuda/bert/hyper_connection_pre_mix.cc Integrates CUDA PreMix kernel.
onnxruntime/contrib_ops/cuda/bert/hyper_connection_pre_mix_impl.h Declares CUDA PreMix launcher.
onnxruntime/contrib_ops/cuda/bert/hyper_connection_pre_mix_impl.cu Implements CUDA PreMix.
onnxruntime/contrib_ops/cuda/bert/hyper_connection_post_mix.h Declares CUDA PostMix.
onnxruntime/contrib_ops/cuda/bert/hyper_connection_post_mix.cc Integrates CUDA PostMix kernel.
onnxruntime/contrib_ops/cuda/bert/hyper_connection_post_mix_impl.h Declares CUDA PostMix launcher.
onnxruntime/contrib_ops/cuda/bert/hyper_connection_post_mix_impl.cu Implements CUDA PostMix.
onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.h Declares CUDA branchwise RMS normalization.
onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm.cc Integrates CUDA normalization paths.
onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.h Declares mixed-scale CUDA launcher.
onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu Implements mixed-scale CUDA normalization.
onnxruntime/contrib_ops/webgpu/webgpu_contrib_kernels.cc Registers WebGPU kernels.
onnxruntime/contrib_ops/webgpu/bert/scaled_silu.h Declares WebGPU ScaledSiLU.
onnxruntime/contrib_ops/webgpu/bert/scaled_silu.cc Implements WebGPU ScaledSiLU.
onnxruntime/contrib_ops/webgpu/bert/hyper_connection_pre_mix.h Declares WebGPU PreMix.
onnxruntime/contrib_ops/webgpu/bert/hyper_connection_pre_mix.cc Implements WebGPU PreMix.
onnxruntime/contrib_ops/webgpu/bert/hyper_connection_post_mix.h Declares WebGPU PostMix.
onnxruntime/contrib_ops/webgpu/bert/hyper_connection_post_mix.cc Implements WebGPU PostMix.
onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.h Declares WebGPU branchwise RMS normalization.
onnxruntime/contrib_ops/webgpu/bert/branchwise_rms_norm.cc Implements WebGPU normalization paths.
docs/OperatorKernels.md Documents provider kernel support.
docs/ContribOperators.md Documents the new operator schemas.
Review details
  • Files reviewed: 45/45 changed files
  • Comments generated: 3
  • Review effort level: Balanced

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

Comment thread onnxruntime/test/python/onnxruntime_test_python_symbolic_shape_infer.py Outdated
Comment thread onnxruntime/contrib_ops/hyper_connection_helper.h Outdated
Comment thread onnxruntime/test/contrib_ops/hyper_connection_op_test.cc
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>

template <typename T>
Status BranchwiseRMSNorm<T>::ComputeInternal(OpKernelContext* context) const {
using CudaT = typename ToCudaType<T>::MappedType;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

For T = BFloat16, this maps CudaT to onnxruntime::BFloat16, but branchwise_rms_norm_impl.cu only explicitly instantiates LaunchMixedScaleBranchwiseRMSNorm<__nv_bfloat16>. Because these are different template specializations, the registered BF16 kernel can produce an unresolved linker symbol.

Could this use OrtToCudaType<T>::type, as the other new CUDA operators do, so the host call matches the __nv_bfloat16 explicit instantiation?

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟡 Changes recommended

Generated documentation is inconsistent with its schema source, and the CUDA RMS normalization kernel is severely under-parallelized.

Get a fresh assessment by requesting another Copilot review.

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

onnxruntime/contrib_ops/cuda/bert/branchwise_rms_norm_impl.cu:40

  • The launch assigns one thread to an entire branch vector, and that thread serially traverses hidden twice. For common decode shapes such as one row with four branches, this launches only four active threads while each performs thousands of operations, making the CUDA fused op severely underutilize the GPU. Use one block per group, distribute h across threads, and reduce the sum of squares cooperatively (as the existing per-head RMS normalization kernel does).
  • Files reviewed: 45/45 changed files
  • Comments generated: 1
  • Review effort level: Balanced

constexpr const char* HyperConnectionPostMix_ver1_doc = R"DOC(
Mixes existing streams and injects one branch output:
Y[..., k, h] = sum_c(stream_mix[..., c, k] * streams[..., c, h])
+ post_mix[..., k, h] * branch_output[..., h].
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.

5 participants