[CUDA] Add QKNormRotaryEmbedding and InverseRotaryRegroup contrib operators

[Tianlei Wu (tianleiwu)]

Stacked PR. Based on tlwu/20260817/sparse_kv_select (#32124) because it shares quant_sim_common.cuh. Please review that PR first; only the last commit here is new. I will retarget this to main once the base merges.

Description

Adds two CUDA contrib operators for the query/KV preparation around MLA-style attention. QKNormRotaryEmbedding takes the raw query and latent KV projections up to the attention kernel — head split, RMS norm, rotary embedding and optional FP8 cache simulation — in one launch. InverseRotaryRegroup undoes that rotation on the attention output and regroups the heads for a grouped output projection.

Both are ~10-node primitive subgraphs whose per-node cost is dominated by launch overhead at decode, and the norms have precision contracts (which intermediate is rounded to T, and where) that are easy to get subtly wrong when the subgraph is rebuilt by hand. The schemas state those contracts explicitly.

Summary of Changes

QKNormRotaryEmbedding

File	Change
onnxruntime/contrib_ops/cuda/math/qk_norm_rotary.{h,cc}	OpKernel — attributes, validation, output shapes
onnxruntime/contrib_ops/cuda/math/qk_norm_rotary_impl.{h,cu}	Fused norm + rotary + FP8-simulation kernels

• Query: weightless RMS norm per head. The reciprocal is rounded to T before the multiply, and the multiply is done in T, matching the unfused subgraph.
• Latent KV row (shared by every head): weighted RMS norm, computed in float then rounded to T.
• cos/sin are stored pre-interleaved, so the rotation is the signed swap of each adjacent pair:

  x[nope + t] = x[nope + t] * cos[t] + (t odd ? x[nope + t - 1] : -x[nope + t + 1]) * sin[t]
  

• With simulate_fp8, the leading head_dim - rope_head_dim channels of the KV row take a simulated FP8-E4M3 round trip in blocks of 64 with a power-of-two scale — what a quantized attention cache will store.

InverseRotaryRegroup

File	Change
onnxruntime/contrib_ops/cuda/math/inverse_rotary_regroup.{h,cc}	OpKernel
onnxruntime/contrib_ops/cuda/math/inverse_rotary_regroup_impl.{h,cu}	Inverse rotation + regroup kernel

• The inverse rotation is the forward one with the signed swap flipped.
• The regroup is the reshape/transpose/reshape trio turning (tokens, num_heads * head_dim) into (num_groups, tokens, group_dim). Both views index the same flat channel, so this is only a change of addressing and is folded into the same kernel.

Shared

File	Change
onnxruntime/core/graph/contrib_ops/contrib_defs.cc	Schemas, shape inference, formulas
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
docs/ContribOperators.md, docs/OperatorKernels.md	Regenerated entries

Testing

• onnxruntime/test/python/transformers/test_qk_norm_rotary.py — PyTorch references for both operators over float32/float16/bfloat16 and several head/dim configurations, with and without simulate_fp8, plus a forward/inverse rotary round-trip test.

python -m pytest onnxruntime/test/python/transformers/test_qk_norm_rotary.py
# 7 passed, 30 subtests passed

Additive only — no existing operator, schema or kernel is modified.

Checklist

• Tests added
• No breaking changes
• Documentation updated

[kunal-vaishnavi]

Instead of naming it QKNormRotaryEmbedding, could we do something more generic like AttentionQK, QKProjection, or QKSubgraph? The QK subgraph keeps changing per model and this op name would only reflect certain fused ops.

[kunal-vaishnavi]

Can we not fuse this op inside the QK op?