Skip to content

Add fused MRotaryEmbedding contrib op for Qwen mRoPE variants - #31728

Draft
kunal-vaishnavi with Copilot wants to merge 3 commits into
mainfrom
copilot/mrope-variants-support
Draft

Add fused MRotaryEmbedding contrib op for Qwen mRoPE variants#31728
kunal-vaishnavi with Copilot wants to merge 3 commits into
mainfrom
copilot/mrope-variants-support

Conversation

Copilot AI commented Aug 8, 2026

Copy link
Copy Markdown
Contributor

Description

Introduces a new com.microsoft contrib op, MRotaryEmbedding, that fuses the multimodal RoPE (mRoPE) subgraphs used by the Qwen vision-language model family into a single op, replacing the multi-node Shape/Gather/Unsqueeze/Expand/MatMul/Concat/Cos/Sin/Split/Gather/Slice/Neg/Concat chain currently built per Q/K tensor by the onnxruntime-genai model builder.

Schema (bert_defs.cc, ms_opset.h)

  • Generalizes RotaryEmbedding: position_ids becomes a required 3D (3, batch_size, sequence_length) tensor stacking Temporal/Height/Width position streams.
  • New required attr mrope_section (3 ints, sum == rotary_embedding_dim/2).
  • New optional attr mrope_layout (0 = Sectioned/Chunked — Qwen2-VL/2.5-VL, 1 = Interleaved — Qwen3-VL/3.5 family).
  • Retains scale, interleaved, rotary_embedding_dim, num_heads, is_packed_batching from RotaryEmbedding so a single full-width section reduces to standard RoPE.

CPU kernel (contrib_ops/cpu/bert/mrotary_embedding*)

  • ComputeDimAssignments precomputes, per cos/sin column, which T/H/W stream owns it (based on mrope_section/mrope_layout), independent of runtime data.
  • Per-token: gathers position id for all 3 streams, builds combined cos/sin buffers by column selection, applies scale, then rotates via MlasRotaryEmbedOneRow.

CUDA kernel (contrib_ops/cuda/bert/mrotary_embedding*)

  • Per-thread stream selection (contiguous-range for Sectioned, modulo-3 for Interleaved) avoids needing a lookup table on GPU.
  • Mirrors the existing RotaryEmbeddingBSNH launcher (grid/block/shared-memory-for-in-place strategy).

Both kernels are registered for float/MLFloat16 (CPU) and float/MLFloat16/BFloat16 (CUDA).

Motivation and Context

Qwen2-VL, Qwen2.5-VL, Qwen3-VL(-MoE), and Qwen3.5(-MoE) each require a distinct, hand-built mRoPE subgraph (~15-20 ops per Q/K rotary application) in the ONNX graph produced by the model builder. Fusing this into one op removes that per-layer overhead, reduces graph size, and gives a single, versioned contract for all current and future mRoPE layout variants.

…le attribute

Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

Comment thread onnxruntime/contrib_ops/cpu/bert/mrotary_embedding_helper.h Fixed

@github-actions github-actions Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You can commit the suggested changes from lintrunner.

if (position_ids_dims.size() != 3 || position_ids_dims[0] != 3) {
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
"MRotaryEmbedding: 'position_ids' is expected to have shape "
"(3, batch_size, sequence_length), got ", position_ids_dims.size(), " dimensions",

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
"(3, batch_size, sequence_length), got ", position_ids_dims.size(), " dimensions",
"(3, batch_size, sequence_length), got ",
position_ids_dims.size(), " dimensions",

Comment on lines +42 to +43
int mrope_section[3]; // section sizes for T, H, W (sum == rotary_embedding_dim / 2)
MRopeLayout mrope_layout; // how the 3 sections are combined

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
int mrope_section[3]; // section sizes for T, H, W (sum == rotary_embedding_dim / 2)
MRopeLayout mrope_layout; // how the 3 sections are combined
int mrope_section[3]; // section sizes for T, H, W (sum == rotary_embedding_dim / 2)
MRopeLayout mrope_layout; // how the 3 sections are combined

Copilot AI and others added 2 commits August 8, 2026 00:31
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Co-authored-by: kunal-vaishnavi <115581922+kunal-vaishnavi@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants