[Tianlei Wu (tianleiwu)]

Description

Adds a HyperConnectionMix CUDA contrib operator that fuses the post-mix, pre-mix and layer norm that sit between two consecutive blocks of a hyper-connection network. Expressed as ONNX primitives this is a ~20-node subgraph of small reductions, a Sinkhorn normalization and two RMS norms, all of which are launch-bound at decode batch sizes; the fused kernel replaces it with a small fixed number of launches.

Hyper-connections (multiple parallel residual streams mixed by learned, input-dependent weights) are used by DeepSeek-V4 and are a general architectural pattern, so the operator is defined in terms of the stream count hc and the mixing weights rather than any one model's constants.

Summary of Changes

New operator

File	Change
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	HyperConnectionMix schema, shape inference and the formula the kernel implements
onnxruntime/core/graph/contrib_ops/ms_opset.h	Schema registration
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	CUDA kernel registration for float, float16 and bfloat16
onnxruntime/contrib_ops/cuda/math/hyper_connection_mix.{h,cc}	OpKernel — attribute parsing, input validation, output shapes
onnxruntime/contrib_ops/cuda/math/hyper_connection_mix_impl.{h,cu}	Fused post-mix / pre-mix / norm kernels
onnxruntime/contrib_ops/cuda/math/sinkhorn_normalize_impl.cuh	Sinkhorn normalization used by the combination-mix branch

The computation, quoted from the schema:

residual_out[..,h,:] = post_mix[..,h] * x[..,:] + ReduceSum(comb_mix[..,g,h] * residual[..,g,:], g)
flat         = Reshape(residual_out, (.., hc * dim))
mixes        = MatMul(flat, fn) / Sqrt(ReduceMean(flat * flat, -1) + epsilon)
pre          = Sigmoid(mixes[..,0:hc] * scale[0] + base[0:hc]) + hc_epsilon
post_mix_out = Sigmoid(mixes[..,hc:2*hc] * scale[1] + base[hc:2*hc]) * post_alpha
comb         = Reshape(mixes[..,2*hc:] * scale[2] + base[2*hc:], (.., hc, hc))
comb_mix_out = SinkhornNormalize(Softmax(comb, -1) + hc_epsilon)
y            = ReduceSum(pre[..,h] * residual_out[..,h,:], h)
layer_input  = norm_weight * y / Sqrt(ReduceMean(y * y, -1) + epsilon)

Intermediates are computed in float; y and residual_out are rounded to T at exactly the points the unfused subgraph rounds, so the fused result matches the primitive subgraph rather than merely approximating it.

Documentation

• docs/ContribOperators.md, docs/OperatorKernels.md — regenerated entries for the new operator.

Testing

• onnxruntime/test/python/transformers/test_hyper_connection_mix.py — compares the operator against a PyTorch reference built from the primitive formula above, over float32/float16/bfloat16 and several hc/dim/token counts, including the rounding points.

python -m pytest onnxruntime/test/python/transformers/test_hyper_connection_mix.py
# 3 passed, 9 subtests passed

Additive only — no existing operator, schema or kernel is modified, so there is no behavior change for existing models.

Checklist

• Tests added
• No breaking changes
• Documentation updated

[Copilot]

Pull request overview

Adds the CUDA-only HyperConnectionMix contrib operator for fused hyper-connection mixing and RMS normalization.

Changes:

• Defines and registers the operator schema and CUDA kernels.
• Implements fused mixing, Sinkhorn normalization, and dtype rounding.
• Adds generated documentation and CUDA numeric tests.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Defines the operator schema.
onnxruntime/core/graph/contrib_ops/ms_opset.h	Registers the schema.
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registers CUDA kernels by dtype.
onnxruntime/contrib_ops/cuda/math/hyper_connection_mix.h	Declares the CUDA operator.
onnxruntime/contrib_ops/cuda/math/hyper_connection_mix.cc	Validates inputs and dispatches execution.
onnxruntime/contrib_ops/cuda/math/hyper_connection_mix_impl.h	Defines parameters and launch helpers.
onnxruntime/contrib_ops/cuda/math/hyper_connection_mix_impl.cu	Implements the fused CUDA kernels.
onnxruntime/contrib_ops/cuda/math/sinkhorn_normalize_impl.cuh	Implements warp-level Sinkhorn normalization.
onnxruntime/test/python/transformers/test_hyper_connection_mix.py	Adds reference-based numeric tests.
docs/ContribOperators.md	Documents the contrib operator.
docs/OperatorKernels.md	Lists CUDA kernel support.

Suppressed comments (1)

onnxruntime/contrib_ops/cuda/math/hyper_connection_mix.cc:103

• Checking only the element count accepts a transposed or otherwise mis-shaped fn. The kernel always indexes this buffer as row-major [hc * dim, mix_dim], so such an input silently computes a different matrix multiplication instead of returning INVALID_ARGUMENT; validate both rank and dimensions.

  if (fn->Shape().Size() != hc * dim * mix_dim) {
    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "fn must be shaped (hc * dim, ", mix_dim,
                           "), which is ", hc * dim * mix_dim, " elements, got ",
                           fn->Shape().Size());
  }

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.