[Tianlei Wu (tianleiwu)]

Description

Adds com.microsoft.HyperConnectionLowRankMix, a CUDA contrib operator that fuses one hyper-connection block with a low-rank gate network (the Qwen3.8-Flash-Next block: 4 residual streams x 2560, rank 320, 96 blocks per model) into a single node: optional post-mix of the previous block, branchwise RMS norm, down projection, ScaledSiLU, up projection, sigmoid, pre-mix, and the post-mix inject gate. A chain of blocks needs one node per block; the unfused graph today needs 9-11 nodes (about 10 kernels) per block.

Motivation and Context

On a Qwen3.8-Flash-Next decode graph (batch 1, H200) the hyper-connection chain is ~16% of a decode cycle. In our research stack (a custom-op library, not part of this PR) fusing it removed 2.1 ms/token at S=1 (13.4 -> ~11 ms graph replay, bf16). This PR upstreams the part that is generic: standard row-major [K, N] weights used in place, no software grid barrier, no thread-block clusters, no process-wide state. The research one-kernel variant (pre-tiled weights + software grid barrier) is intentionally not proposed here.

HyperConnectionPreMix / HyperConnectionPostMix (already in main) are unchanged and stay the building blocks of the unfused graph. A graph fusion that matches the unfused pattern and emits this op is a follow-up (out of scope here).

Why one new op (and not extending PreMix/PostMix)

The fused block needs the norm weight, two weight matrices with different shapes, four activation constants and two consumers of the normalized tensor. Folding that into PreMix/PostMix attributes would create a second op under their names. One op with an optional prologue (prev_block_output, prev_post_gate) gives one node per block in a chain and one place to state the rounding contract.

Relation to #32123

#32123 (HyperConnectionMix, DeepSeek-V4 mHC) is a different op: a single linear gate on the flattened residual with Sinkhorn normalization and scalar-per-stream pre-mix. Qwen's block has a low-rank MLP gate (down -> SiLU -> up -> sigmoid), per-stream-per-feature pre-mix gates, no stream mixing and a per-branch weighted RMS norm, so no attribute setting of #32123 reproduces it. No code is shared; I followed its documentation and testing conventions (rounding points in the schema, torch reference test).

Schema summary

HyperConnectionLowRankMix(streams, norm_scale, down_weight, up_weight, inject_weight
                          [, prev_block_output, prev_post_gate])
    -> (pre_output, post_gate [, streams_out])

T = float / float16 / bfloat16 for every tensor. streams is flattened (..., C*H); down_weight (C*H, R), up_weight (R, C*H), inject_weight (C*H, C) row-major [K, N]; prev_block_output (..., H), prev_post_gate (..., C); pre_output (..., H), post_gate (..., C), streams_out = streams after the prologue (a copy of streams without prologue).
Attributes: num_branches (required), epsilon (1e-5), reduction_scale (1), silu_alpha (1), gate_div (1), gate_mul (1). Qwen3.8 uses num_branches=4, epsilon=1e-6, reduction_scale=0.25, silu_alpha=0.25, gate_div=4, gate_mul=2.

S  = streams + prev_post_gate * prev_block_output           (prologue; otherwise S = streams)
hn = S / sqrt(mean_h(S^2) + epsilon) * norm_scale           per branch
d  = ScaledSiLU(hn @ down_weight, silu_alpha)               (..., R)
g  = sigmoid(d @ up_weight)                                 per stream and feature
pre_output = reduction_scale * sum_c(hn * g)
post_gate  = sigmoid((hn @ inject_weight) / gate_div) * gate_mul

Numerics contract

Matrix products and reductions accumulate in fp32; the result of every step is rounded once to T where the equivalent standard operator would round (listed in the schema doc: S, hn, down, Z, SG, hidden, up, g, pre_output, inject, post_gate, as in the ScaledSiLU contract). The sigmoid is an accurate fp32 expf-based sigmoid rounded once to T: it does not reproduce the current CUDA Sigmoid<bf16> (round after every elementary operation; #33312 is the open fix), so on today's main the unfused chain and this op differ at the sigmoid sites, and after #33312 they agree up to accumulation order. Accumulation order is fixed: deterministic, and a row gives the same bits at any batch size.

Summary of Changes

File	Change
onnxruntime/core/graph/contrib_ops/bert_defs.cc, ms_opset.h	Schema + shape/type inference (next to the HyperConnectionPreMix/PostMix schemas)
onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.{h,cc}	Kernel class, input validation, path selection, scratch allocation via the allocator
onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.{h,cu}	Fused two-kernel path and the general cuBLAS path
onnxruntime/contrib_ops/cuda/cuda_contrib_kernels.cc	Registration (float, MLFloat16, BFloat16)
docs/ContribOperators.md, docs/OperatorKernels.md	New op entries
onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py, onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc	Tests

CUDA implementation

• Fused path (rows <= 8, fp16/bf16, hidden % 8 == 0 && <= 4096, rank % 8 == 0 && <= 512, num_branches in {1,2,4,8}, 16-byte aligned tensors, sm_80+): two kernels. HcDownKernel: optional post-mix + RMS norm + down GEMV + inject GEMV; K is split in slices inside one branch and each CTA writes fp32 partial sums. HcUpPreKernel: fixed-order partial reduction + ScaledSiLU + up GEMV + sigmoid + pre-mix + inject gate. Weights are streamed through shared memory (cp.async) before the kernel waits for its producer. One or two rows share a CTA (and its weight reads), more rows use more CTAs. No atomics, clusters, cooperative launch or spin waits.
• PDL: on sm_90+ (CUDA >= 12.3) both kernels are launched with programmatic stream serialization (cudaLaunchKernelEx); everything that consumes another kernel's output is after griddepcontrol.wait, so correctness does not depend on trigger placement. Older GPUs use plain launches (and plain copies instead of cp.async below sm_80 are compiled in, though the fused path is only selected on sm_80+).
• General path (any shape, any row count, fp32): norm kernel + 3 cuBLAS GEMMs (CUBLAS_COMPUTE_32F) + elementwise kernels with the same rounding points.
• CUDA-graph capturable: no host synchronization, scratch only through the allocator (verified by capture/replay in the tests and the benchmark).

Testing

All run for real on one H200 (CUDA_VISIBLE_DEVICES=2, sm_90, CUDA 13.0, Release build of this branch rebased on main dc89460):

• python -m pytest onnxruntime/test/python/transformers/test_hyper_connection_low_rank_mix.py: 9 tests, 192 subtests passed. Torch reference built from the ONNX primitives (BranchwiseRMSNorm / MatMul / ScaledSiLU / Sigmoid / PreMix / Div-Sigmoid-Mul / PostMix) in fp32 with the contract's rounding points: real shape (4 x 2560, rank 320) x {fp32, fp16, bf16} x rows {1, 2, 4, 16, 128} x {with, without prologue}; 8 small and irregular shapes (1/2/3/4/8 branches, hidden 32..4096 incl. not a multiple of 8, rank 8..512, multi-slice and multi-chunk cases); default and non-default attributes; two chained blocks; bitwise row independence; streams_out passthrough; zero rows; CUDA graph capture and replay (bitwise equal to eager). On the real bf16 shape the fused path is bit-identical to the reference (fp16: 99% of elements, rest 1 ulp; fp32 max abs error 4e-7).
• onnxruntime_provider_test --gtest_filter='HyperConnection*': 25/25 pass (10 new HyperConnectionLowRankMixTest.* cases: fused/unfused paths, optional output, K slices, weight chunks, fp32, error cases).
• compute-sanitizer memcheck, racecheck, synccheck on the fused and unfused gtest cases: 0 errors / 0 hazards.
• Cross-check against our research kernels (two-block real-shaped random case, bf16): bit-identical (all five outputs, rows 1 and 2) to the research two-kernel ops after swapping their bf16-Sigmoid emulation for the accurate sigmoid; with the original research kernels pre differs by exactly 1 bf16 ulp at the sigmoid sites (the contract change above).
• lintrunner -a on the changed files: no issues.

Benchmark

96-block chain, bf16, CUDA graph replay, 8 distinct weight sets (208 MB, larger than L2); us/block = slope between the 96- and 48-block graphs (min of 15 x 100 replays); H200. The unfused chain is the stock graph per block: PostMix, BranchwiseRMSNorm, MatMul, ScaledSiLU, MatMul, Sigmoid, PreMix, MatMul, Div, Sigmoid, Mul (11 nodes).

rows (S)	fused op	stock unfused chain	research two-kernel ops (clusters), same harness
1	10.2 us	33.2 us	10.5 us
2	12.9 us	36.8 us	17.5 us
4	21.1 us	36.3 us
8 (fused limit)	33.2 us	37.6 us
9 / 16 (general path)	33.9 / 32.8 us	38.7 / 37.9 us

fp16: 10.2 / 12.7 us (S=1 / 2) vs 32.2 / 36.0 us unfused. For 96 blocks at S=1 that is 3.2 ms -> 1.0 ms of replay.
For reference, our research one-kernel variant (pre-tiled weights, software grid barrier, not in this PR) measured 8.2 us/block (S=1) and 10.1 us (S=2) in the research stack with a different harness, so this two-kernel upstreamable variant is about 2 us (25%) per block slower at S=1 and ~2.8 us at S=2. In a real model the producer kernel does not trigger early, so expect +0.5..1 us per block over the chain benchmark.

Limitations

• CUDA only (no CPU / WebGPU kernel). fp32 uses the general path. Fused kernels need sm_80+; PDL needs sm_90+ (and CUDA >= 12.3 at build time).
• Fused path limits: rows <= 8, hidden <= 4096, rank <= 512, branches in {1, 2, 4, 8}, 16-byte aligned tensors; everything else (including prefill) takes the general path, which is correct but about as fast as the unfused chain.
• Exact agreement with the stock bf16 Sigmoid needs [CUDA] fix Nemotron-H relu2, RMSNorm, and activation precision #33312.
• No graph fusion transformer yet (follow-up).

Checklist

• Tests added/updated
• No breaking changes (new op, no existing behavior touched)
• Documentation updated (schema doc, docs/ContribOperators.md, docs/OperatorKernels.md)

[Copilot]

🟡 Changes recommended

Early weight reads under programmatic dependent launch can race with upstream producers and produce incorrect results.

1 open finding

• PDL enables racy prefetches for computed weight inputs · New

What changed in this PR

Adds a CUDA-only HyperConnectionLowRankMix contrib operator to fuse low-rank hyper-connection blocks. Existing PreMix/PostMix operators remain unchanged; graph fusion is deferred.

Changes:

• Adds the schema, type/shape inference, and CUDA registration.
• Implements a two-kernel fused path and a general cuBLAS fallback.
• Adds numerical, validation, and CUDA graph tests plus operator documentation.

File	Description
onnxruntime/​test/​python/​transformers/​test_hyper_connection_low_rank_mix.py	Reference comparisons and CUDA graph tests.
onnxruntime/​test/​contrib_ops/​hyper_connection_low_rank_mix_op_test.cc	Kernel correctness and validation tests.
onnxruntime/​core/​graph/​contrib_ops/​ms_opset.h	Registers the schema.
onnxruntime/​core/​graph/​contrib_ops/​bert_defs.cc	Defines schema, numerics, and shape inference.
onnxruntime/​contrib_ops/​cuda/​cuda_contrib_kernels.cc	Registers supported CUDA types.
onnxruntime/​contrib_ops/​cuda/​bert/​hyper_connection_low_rank_mix.h	Declares the kernel class.
onnxruntime/​contrib_ops/​cuda/​bert/​hyper_connection_low_rank_mix.cc	Validates inputs and selects execution paths.
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 fallback.
docs/​OperatorKernels.md	Lists CUDA support.
docs/​ContribOperators.md	Documents the operator contract.

🧠 Review effort: Balanced

---

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

[Copilot]

🔵 Needs a closer look

CUDA synchronization and numerical guarantees warrant human review, with unresolved precision and scratch-memory concerns.

2 open issues

• Avoid unused scratch allocation and writes on general norm path · New
• Prevent reduced-precision cuBLAS reductions on FP16 general path · New

1 resolved since last review

• PDL enables racy prefetches for computed weight inputs

_{Review effort: Balanced}

[Akshay Sonawane (apsonawane)]

Please also register and implement HyperConnectionLowRankMix in the Python symbolic-shape inference tool, alongside the existing hyper-connection handlers. This C++ schema callback does not provide inference support to that tool.

I reproduced a failure using the PR-head implementation with a valid node: streams=["batch",16], num_branches=2, norm_scale=[16], down_weight=[16,8], up_weight=[8,16], and inject_weight=[16,2]. When the output shapes are unspecified, SymbolicShapeInference.infer_shapes(model, auto_merge=True) raises Incomplete symbolic shape inference. An existing HyperConnectionPreMix control successfully infers ["batch",8].

Please add the dispatcher entry and handler to infer pre_output=[...,H], post_gate=[...,C], and optional streams_out=[...,C*H], preserving the input element type. Add regression coverage with symbolic leading dimensions and both two-output and three-output forms. Otherwise valid graphs containing this new operator fail the repository’s symbolic shape-inference preprocessing.

[Copilot]

🔵 Needs a closer look

BF16 test hardware guards need correction, and asynchronous CUDA execution and precision guarantees warrant GPU-maintainer review.

1 open issue

• Guard BF16 CUDA tests by device capability, not provider availability · New

2 resolved since last review

• Prevent reduced-precision cuBLAS reductions on FP16 general path
• Avoid unused scratch allocation and writes on general norm path

_{Review effort: Balanced}