Yes, do_rotary=0 could work for 3D mRoPE at build time — the AttentionToGQA rewrite rule already does exactly this today. Making it work directly in the build path is feasible but requires a few pieces that are currently missing:
-
Different payload in context: GQAContext currently carries cos_cache/sin_cache (full tables [max_seq, rotary_dim]) specifically for do_rotary=1 — GQA indexes them internally. For do_rotary=0, we need the already-gathered, per-token (cos, sin) tensors (batch, seq, rotary_dim) produced by _MRopeBase.forward(), not raw tables. These are a completely different shape/purpose and would need to be plumbed through (either by extending GQAContext or creating a separate context type).
-
Apply RoPE inside _forward_gqa(): When do_rotary=0, apply_rotary_pos_emb() must be called on Q/K before emitting the GroupQueryAttention node (no cos/sin inputs to GQA in this mode). This branch doesn't exist in _forward_gqa() yet.
-
Remove the _MRopeBase guard in TextModel.forward(), or replace it with a condition that chooses do_rotary=0 for _MRopeBase vs do_rotary=1 for standard BaseRope.
Worth noting: the supports_fused_rope guard would also need revisiting — do_rotary=0 doesn't need fused RoPE inside the GQA kernel, so it could be applied even on EPs where supports_fused_rope=False (as long as dtype in caps.gqa_dtypes).
Happy to add this if you'd like to extend the PR. It would fully eliminate AttentionToGQA as a fallback for mRoPE models on GQA-capable EPs.
Originally posted by @Copilot in #134 (comment)
Yes,
do_rotary=0could work for 3D mRoPE at build time — theAttentionToGQArewrite rule already does exactly this today. Making it work directly in the build path is feasible but requires a few pieces that are currently missing:Different payload in context:
GQAContextcurrently carriescos_cache/sin_cache(full tables[max_seq, rotary_dim]) specifically fordo_rotary=1— GQA indexes them internally. Fordo_rotary=0, we need the already-gathered, per-token(cos, sin)tensors(batch, seq, rotary_dim)produced by_MRopeBase.forward(), not raw tables. These are a completely different shape/purpose and would need to be plumbed through (either by extendingGQAContextor creating a separate context type).Apply RoPE inside
_forward_gqa(): Whendo_rotary=0,apply_rotary_pos_emb()must be called on Q/K before emitting theGroupQueryAttentionnode (no cos/sin inputs to GQA in this mode). This branch doesn't exist in_forward_gqa()yet.Remove the
_MRopeBaseguard inTextModel.forward(), or replace it with a condition that choosesdo_rotary=0for_MRopeBasevsdo_rotary=1for standardBaseRope.Worth noting: the
supports_fused_ropeguard would also need revisiting —do_rotary=0doesn't need fused RoPE inside the GQA kernel, so it could be applied even on EPs wheresupports_fused_rope=False(as long asdtype in caps.gqa_dtypes).Happy to add this if you'd like to extend the PR. It would fully eliminate
AttentionToGQAas a fallback for mRoPE models on GQA-capable EPs.Originally posted by @Copilot in #134 (comment)