onnxruntime

245f59b7 - [CUDA] Add HyperConnectionLowRankMix fused low-rank hyper-connection op (#33384)

[CUDA] Add HyperConnectionLowRankMix fused low-rank hyper-connection op (#33384) ## 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` dc89460336): * `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 #33312. * No graph fusion transformer yet (follow-up). ## Checklist - [x] Tests added/updated - [x] No breaking changes (new op, no existing behavior touched) - [x] Documentation updated (schema doc, `docs/ContribOperators.md`, `docs/OperatorKernels.md`) --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>