Add rotary embedding onnx domain support

[Xiaoyu Z (xiaoyu-work)]

Description

Mobius exports standard ONNX rotary embedding op. Adding support for this.

Motivation and Context

[Copilot]

Pull request overview

This PR updates the GroupQueryAttention fusion optimizer pass to recognize standard (ONNX-domain) RotaryEmbedding nodes when extracting the cos_cache/sin_cache inputs needed to fuse rotary embedding into the com.microsoft.GroupQueryAttention node.

Changes:

• Adds a helper to retrieve cos_cache/sin_cache NodeArgs for both com.microsoft.RotaryEmbedding and ONNX-domain RotaryEmbedding (different input ordering).
• Updates the fusion pattern-matching logic to use that helper and to require that rotary cache inputs were successfully identified before fusing.

[Ti-Tai Wang (titaiwangms)]

cc Tianlei Wu (@tianleiwu)

[Tianlei Wu (tianleiwu)]

Summary

The change extends GroupQueryAttentionFusion to also match the standard ONNX-domain RotaryEmbedding (X, cos_cache, sin_cache, position_ids) in addition to com.microsoft.RotaryEmbedding, plumbing position_ids through to GQA input 9 and setting the input-arg-count. The approach is sound:

• Requiring position_ids to be present for the ONNX-domain path (and bailing otherwise) correctly avoids the 3D per-batch cos/sin cache form that GQA's 2D rotary cache validation cannot consume.
• The position_ids_arg_mismatch guard and the added cos_cache_arg == nullptr || sin_cache_arg == nullptr checks make the fusion safely skip ambiguous/mixed cases.
• MutableInputArgsCount()[9] is in-bounds because the GQA schema declares formal inputs up to index 11, so UpdateInputArgCount() sizes the vector accordingly.
• Both the fused (with position_ids) and non-fused (omitted position_ids) cases are covered by new tests.

Main concern

Rotary interleaved / rotary_embedding_dim attributes are not validated or propagated. GQA's do_rotary path runs non-interleaved, full-width RoPE (rotary_interleaved defaults to 0 and the fusion never sets it). A standard ONNX RotaryEmbedding with interleaved=1, or a partial-rotary node with rotary_embedding_dim > 0 (and a correspondingly narrower cos/sin cache), is silently fused into a GQA that applies a different rotation — producing incorrect results with no error. Since this PR specifically targets a standard-ONNX export path where interleaved RoPE is common, the fusion should either verify interleaved == 0 and rotary_embedding_dim == 0 (full rotary) before matching, or propagate interleaved to GQA's rotary_interleaved. Inline comment below. (Note: the pre-existing com.microsoft.RotaryEmbedding path has the same latent gap.)

Minor

• Only position_ids is checked for consistency between the two rotary nodes; cos_cache/sin_cache are taken from whichever rotary node is visited first without verifying the second uses the same caches. In practice Q/K share caches, but an explicit equality guard (mirroring the position_ids_arg_mismatch check) would make the fusion robust to malformed graphs.

[Tianlei Wu (tianleiwu)]

Thanks for updating the ONNX-domain RotaryEmbedding handling and adding the targeted fusion tests. I found one remaining correctness issue around explicit 2D position_ids: the fused GQA path can use different prompt-time position semantics than ONNX RotaryEmbedding, so I am requesting changes for that case.

[Tianlei Wu (tianleiwu)]

Summary

The ONNX-domain RotaryEmbedding handling is now correctly disabled, and propagating rotary_interleaved plus the new Q/K mismatch guard address the earlier interleaved/attribute concerns — thanks for those. Two smaller follow-ups below, plus a note left on the existing position_ids thread.

• Positive-path test coverage. All five new tests exercise only the not-fused branches (ONNX-domain skip and interleaved/position-id mismatch). The substantive new MS-domain behavior this PR adds — forwarding position_ids as GQA input 9 and setting the rotary_interleaved attribute on the fused node — has no test. The existing TestGQAFusion model tests only assert op counts, so they won't catch a missing/incorrect position_ids input or a wrong rotary_interleaved value. Consider a positive test that fuses a com.microsoft.RotaryEmbedding graph and asserts the fused GQA has 10 inputs (position_ids at index 9) and the expected rotary_interleaved attribute (including the interleaved=1 case).
• MS-domain position_ids forwarding — see the reply on the existing thread; the full 2D position_ids is only partially honored on the first prompt.
• Minor cleanup noted inline in TryGetRotaryEmbeddingArgs.