onnx2fx 0.0.0__py3-none-any.whl

onnx2fx/ops/attention.py

@@ -0,0 +1,1055 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Attention and Transformer related operators (standard ONNX domain).
+
+Microsoft domain operators are in attention_msft.py.
+"""
+
+from typing import TYPE_CHECKING
+
+import onnx
+import torch
+
+from ..op_registry import register
+from ..utils.attributes import get_attribute
+from ..utils.op_helpers import get_optional_input
+
+if TYPE_CHECKING:
+    from ..graph_builder import GraphBuilder
+
+
+# =============================================================================
+# Embedding and LayerNorm variants
+# =============================================================================
+
+
+@register("SkipLayerNormalization")
+def skip_layer_normalization(
+    builder: "GraphBuilder", node: onnx.NodeProto
+) -> torch.fx.Node:
+    """Skip connection + LayerNorm (common in transformers)."""
+    x = builder.get_value(node.input[0])
+    skip = builder.get_value(node.input[1])
+    gamma = builder.get_value(node.input[2])
+    beta = get_optional_input(builder, node, 3)
+    bias = get_optional_input(builder, node, 4)
+
+    epsilon = get_attribute(node, "epsilon", 1e-5)
+
+    def _skip_layer_norm(
+        inp: torch.Tensor,
+        sk: torch.Tensor,
+        g: torch.Tensor,
+        b: torch.Tensor | None,
+        bi: torch.Tensor | None,
+        eps: float,
+    ) -> torch.Tensor:
+        hidden = inp + sk
+        if bi is not None:
+            hidden = hidden + bi
+        return torch.nn.functional.layer_norm(
+            hidden, hidden.shape[-1:], weight=g, bias=b, eps=eps
+        )
+
+    return builder.call_function(
+        _skip_layer_norm, args=(x, skip, gamma, beta, bias, epsilon)
+    )
+
+
+@register("EmbedLayerNormalization")
+def embed_layer_normalization(
+    builder: "GraphBuilder", node: onnx.NodeProto
+) -> torch.fx.Node:
+    """Embedding + LayerNorm (common in BERT-like models)."""
+    input_ids = builder.get_value(node.input[0])
+    segment_ids = get_optional_input(builder, node, 1)
+    word_embedding = builder.get_value(node.input[2])
+    position_embedding = builder.get_value(node.input[3])
+    segment_embedding = get_optional_input(builder, node, 4)
+    gamma = get_optional_input(builder, node, 5)
+    beta = get_optional_input(builder, node, 6)
+
+    epsilon = get_attribute(node, "epsilon", 1e-5)
+
+    def _embed_layer_norm(
+        ids: torch.Tensor,
+        seg_ids: torch.Tensor | None,
+        word_emb: torch.Tensor,
+        pos_emb: torch.Tensor,
+        seg_emb: torch.Tensor | None,
+        g: torch.Tensor | None,
+        b: torch.Tensor | None,
+        eps: float,
+    ) -> torch.Tensor:
+        # Word embedding lookup
+        word_embed = torch.nn.functional.embedding(ids, word_emb)
+
+        # Position embedding (assume sequential positions)
+        seq_len = ids.shape[1]
+        pos_embed = pos_emb[:seq_len].unsqueeze(0).expand(ids.shape[0], -1, -1)
+
+        hidden = word_embed + pos_embed
+
+        # Add segment embedding if present
+        if seg_emb is not None and seg_ids is not None:
+            seg_embed = torch.nn.functional.embedding(seg_ids, seg_emb)
+            hidden = hidden + seg_embed
+
+        # Layer normalization
+        if g is not None:
+            hidden = torch.nn.functional.layer_norm(
+                hidden, hidden.shape[-1:], weight=g, bias=b, eps=eps
+            )
+
+        return hidden
+
+    return builder.call_function(
+        _embed_layer_norm,
+        args=(
+            input_ids,
+            segment_ids,
+            word_embedding,
+            position_embedding,
+            segment_embedding,
+            gamma,
+            beta,
+            epsilon,
+        ),
+    )
+
+
+# =============================================================================
+# Attention operator (ONNX standard domain, since opset 24)
+# =============================================================================
+
+
+@register("Attention")
+def attention(builder: "GraphBuilder", node: onnx.NodeProto) -> torch.fx.Node:
+    """ONNX Attention operator (standard domain, since opset 24).
+
+    Inputs:
+        Q: Query tensor
+           - 4D: (batch_size, q_num_heads, q_sequence_length, head_size)
+           - 3D: (batch_size, q_sequence_length, q_num_heads * head_size)
+        K: Key tensor
+           - 4D: (batch_size, kv_num_heads, kv_sequence_length, head_size)
+           - 3D: (batch_size, kv_sequence_length, kv_num_heads * head_size)
+        V: Value tensor (same format as K)
+        attn_mask (optional): Attention mask, broadcastable to
+            (batch_size, q_num_heads, q_sequence_length, total_sequence_length)
+        past_key (optional): Past key cache
+        past_value (optional): Past value cache
+        nonpad_kv_seqlen (optional): Non-padding KV sequence lengths
+
+    Attributes:
+        is_causal: If 1, use causal (lower triangular) mask
+        scale: Scaling factor for Q*K^T (default: 1/sqrt(head_size))
+        softcap: Softcap value for attention weights
+        q_num_heads: Number of query heads (required for 3D inputs)
+        kv_num_heads: Number of key/value heads (required for 3D inputs)
+        qk_matmul_output_mode: Output mode for QK matmul (0, 1, or 2)
+
+    Outputs:
+        Y: Output tensor (same format as Q)
+        present_key (optional): Updated key cache
+        present_value (optional): Updated value cache
+        qk_matmul_output (optional): QK matmul output
+    """
+    # Get inputs
+    query = builder.get_value(node.input[0])
+    key = builder.get_value(node.input[1])
+    value = builder.get_value(node.input[2])
+
+    attn_mask = get_optional_input(builder, node, 3)
+    past_key = get_optional_input(builder, node, 4)
+    past_value = get_optional_input(builder, node, 5)
+    nonpad_kv_seqlen = get_optional_input(builder, node, 6)
+
+    # Get attributes
+    is_causal = get_attribute(node, "is_causal", 0)
+    scale = get_attribute(node, "scale", None)
+    softcap = get_attribute(node, "softcap", 0.0)
+    q_num_heads = get_attribute(node, "q_num_heads", None)
+    kv_num_heads = get_attribute(node, "kv_num_heads", None)
+    qk_matmul_output_mode = get_attribute(node, "qk_matmul_output_mode", 0)
+
+    # Determine which outputs are needed
+    # Output positions: 0=Y, 1=present_key, 2=present_value, 3=qk_matmul_output
+    num_outputs = len(node.output)
+    has_present_key = num_outputs > 1 and node.output[1]
+    has_present_value = num_outputs > 2 and node.output[2]
+    has_qk_matmul_output = num_outputs > 3 and node.output[3]
+
+    # Check if we need multiple outputs (even if some are empty)
+    _needs_multiple_outputs = num_outputs > 1
+
+    # Use manual attention computation when:
+    # 1. We need qk_matmul_output
+    # 2. We have softcap
+    # 3. We have both is_causal and attn_mask
+    needs_manual_attention = (
+        has_qk_matmul_output or softcap != 0.0 or (is_causal and attn_mask is not None)
+    )
+
+    if needs_manual_attention:
+        # Manual attention implementation for advanced features
+
+        def _attention_manual(
+            q: torch.Tensor,
+            k: torch.Tensor,
+            v: torch.Tensor,
+            mask: torch.Tensor | None,
+            past_k: torch.Tensor | None,
+            past_v: torch.Tensor | None,
+            is_causal: int,
+            scale: float | None,
+            softcap: float,
+            q_num_heads: int | None,
+            kv_num_heads: int | None,
+            num_outputs: int,
+            qk_matmul_output_mode: int,
+        ) -> (
+            torch.Tensor
+            | tuple[torch.Tensor, torch.Tensor]
+            | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+            | tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+        ):
+            is_3d = q.dim() == 3
+
+            if is_3d:
+                batch_size = q.shape[0]
+                q_seq_len = q.shape[1]
+                kv_seq_len = k.shape[1]
+                n_q_heads = q_num_heads if q_num_heads is not None else 1
+                n_kv_heads = kv_num_heads if kv_num_heads is not None else n_q_heads
+                q_head_size = q.shape[2] // n_q_heads
+                kv_head_size = k.shape[2] // n_kv_heads
+                v_head_size = v.shape[2] // n_kv_heads
+                q_4d = q.view(batch_size, q_seq_len, n_q_heads, q_head_size).transpose(
+                    1, 2
+                )
+                k_4d = k.view(
+                    batch_size, kv_seq_len, n_kv_heads, kv_head_size
+                ).transpose(1, 2)
+                v_4d = v.view(
+                    batch_size, kv_seq_len, n_kv_heads, v_head_size
+                ).transpose(1, 2)
+            else:
+                q_4d = q
+                k_4d = k
+                v_4d = v
+                batch_size = q.shape[0]
+                n_q_heads = q.shape[1]
+                n_kv_heads = k.shape[1]
+                q_seq_len = q.shape[2]
+                q_head_size = q.shape[3]
+
+            # Handle past key/value (KV cache)
+            past_seq_len = 0
+            if past_k is not None and past_v is not None:
+                past_seq_len = past_k.shape[2]
+                k_4d = torch.cat([past_k, k_4d], dim=2)
+                v_4d = torch.cat([past_v, v_4d], dim=2)
+
+            # Save present_key and present_value before GQA expansion (if needed for output)
+            present_k = k_4d if num_outputs > 1 else None
+            present_v = v_4d if num_outputs > 2 else None
+
+            # Handle GQA: expand KV heads to match Q heads
+            if n_kv_heads != n_q_heads:
+                n_rep = n_q_heads // n_kv_heads
+                k_4d = k_4d.repeat_interleave(n_rep, dim=1)
+                v_4d = v_4d.repeat_interleave(n_rep, dim=1)
+
+            # Compute attention scale
+            if scale is None:
+                scale_val = 1.0 / (q_head_size**0.5)
+            else:
+                scale_val = scale
+
+            # Q @ K^T with scaling
+            # (batch, heads, q_seq, head) @ (batch, heads, head, kv_seq)
+            # -> (batch, heads, q_seq, kv_seq)
+            qk = torch.matmul(q_4d, k_4d.transpose(-2, -1)) * scale_val
+
+            # Save QK matmul output before applying mask/softmax (if needed for output)
+            # Mode 0: raw QK matmul output
+            qk_output = None
+            if num_outputs > 3 and qk_matmul_output_mode == 0:
+                qk_output = qk.clone()
+
+            # Build combined attention mask (applied BEFORE softcap per ONNX spec)
+            # The ONNX approach: create causal mask first, add attn_mask to it,
+            # then add combined mask to QK scores
+            kv_seq = k_4d.shape[2]
+            combined_mask = None
+
+            # Create causal mask if is_causal=1
+            # ONNX uses: Less(q_pos + past_len, k_pos) to determine masked positions
+            # This creates a strictly lower triangular mask where q_pos + past_len >= k_pos is allowed
+            if is_causal:
+                # Create causal mask: (q_pos + past_seq_len) < k_pos means masked
+                row = (
+                    torch.arange(q_seq_len, device=q.device).view(-1, 1) + past_seq_len
+                )
+                col = torch.arange(kv_seq, device=q.device).view(1, -1)
+                causal_bool = row < col  # True where masked
+                causal_mask = (
+                    torch.where(causal_bool, float("-inf"), 0.0)
+                    .to(q.dtype)
+                    .unsqueeze(0)
+                    .unsqueeze(0)
+                )
+                combined_mask = causal_mask
+
+            # Add attention mask to causal mask (or use just attn_mask if no causal)
+            if mask is not None:
+                if combined_mask is not None:
+                    combined_mask = mask + combined_mask  # attn_mask + causal_mask
+                else:
+                    combined_mask = mask
+
+            # Add combined mask to QK scores
+            if combined_mask is not None:
+                qk = qk + combined_mask
+
+            # Mode 1: after attention mask addition (including causal)
+            if num_outputs > 3 and qk_matmul_output_mode == 1:
+                qk_output = qk.clone()
+
+            # Apply softcap if specified (after mask, before softmax per ONNX spec)
+            if softcap != 0.0:
+                qk = softcap * torch.tanh(qk / softcap)
+
+            # Mode 2: after softcap
+            if num_outputs > 3 and qk_matmul_output_mode == 2:
+                qk_output = qk.clone()
+
+            # Softmax
+            attn_weights = torch.nn.functional.softmax(qk, dim=-1)
+
+            # Mode 3: after softmax
+            if num_outputs > 3 and qk_matmul_output_mode == 3:
+                qk_output = attn_weights.clone()
+
+            # Attention @ V
+            output = torch.matmul(attn_weights, v_4d)
+
+            if is_3d:
+                output = (
+                    output.transpose(1, 2).contiguous().view(batch_size, q_seq_len, -1)
+                )
+
+            # Return based on num_outputs (must match exactly)
+            # Output positions: 0=Y, 1=present_key, 2=present_value, 3=qk_matmul_output
+            if num_outputs == 1:
+                return output
+            elif num_outputs == 2:
+                return (output, present_k)
+            elif num_outputs == 3:
+                return (output, present_k, present_v)
+            else:  # num_outputs == 4
+                return (output, present_k, present_v, qk_output)
+
+        return builder.call_function(
+            _attention_manual,
+            args=(
+                query,
+                key,
+                value,
+                attn_mask,
+                past_key,
+                past_value,
+                is_causal,
+                scale,
+                softcap,
+                q_num_heads,
+                kv_num_heads,
+                num_outputs,
+                qk_matmul_output_mode,
+            ),
+        )
+    elif has_present_key or has_present_value:
+        # Use SDPA but also return present_key/present_value
+
+        def _attention_with_cache(
+            q: torch.Tensor,
+            k: torch.Tensor,
+            v: torch.Tensor,
+            mask: torch.Tensor | None,
+            past_k: torch.Tensor | None,
+            past_v: torch.Tensor | None,
+            is_causal: int,
+            scale: float | None,
+            q_num_heads: int | None,
+            kv_num_heads: int | None,
+            has_present_key: bool,
+            has_present_value: bool,
+        ) -> (
+            torch.Tensor
+            | tuple[torch.Tensor, torch.Tensor]
+            | tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+        ):
+            is_3d = q.dim() == 3
+
+            if is_3d:
+                batch_size = q.shape[0]
+                q_seq_len = q.shape[1]
+                kv_seq_len = k.shape[1]
+                n_q_heads = q_num_heads if q_num_heads is not None else 1
+                n_kv_heads = kv_num_heads if kv_num_heads is not None else n_q_heads
+                q_head_size = q.shape[2] // n_q_heads
+                kv_head_size = k.shape[2] // n_kv_heads
+                v_head_size = v.shape[2] // n_kv_heads
+                q_4d = q.view(batch_size, q_seq_len, n_q_heads, q_head_size).transpose(
+                    1, 2
+                )
+                k_4d = k.view(
+                    batch_size, kv_seq_len, n_kv_heads, kv_head_size
+                ).transpose(1, 2)
+                v_4d = v.view(
+                    batch_size, kv_seq_len, n_kv_heads, v_head_size
+                ).transpose(1, 2)
+            else:
+                q_4d = q
+                k_4d = k
+                v_4d = v
+                batch_size = q.shape[0]
+                n_q_heads = q.shape[1]
+                n_kv_heads = k.shape[1]
+                q_seq_len = q.shape[2]
+
+            # Handle past key/value (KV cache)
+            if past_k is not None and past_v is not None:
+                k_4d = torch.cat([past_k, k_4d], dim=2)
+                v_4d = torch.cat([past_v, v_4d], dim=2)
+
+            # Save present_key and present_value before GQA expansion
+            present_k = k_4d if has_present_key else None
+            present_v = v_4d if has_present_value else None
+
+            # Handle GQA: expand KV heads to match Q heads
+            if n_kv_heads != n_q_heads:
+                n_rep = n_q_heads // n_kv_heads
+                k_4d = k_4d.repeat_interleave(n_rep, dim=1)
+                v_4d = v_4d.repeat_interleave(n_rep, dim=1)
+
+            # Call SDPA
+            if mask is not None:
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    q_4d, k_4d, v_4d, attn_mask=mask, is_causal=False, scale=scale
+                )
+            elif is_causal:
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    q_4d, k_4d, v_4d, is_causal=True, scale=scale
+                )
+            else:
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    q_4d, k_4d, v_4d, scale=scale
+                )
+
+            if is_3d:
+                output = (
+                    output.transpose(1, 2).contiguous().view(batch_size, q_seq_len, -1)
+                )
+
+            # Build return tuple
+            results = [output]
+            if has_present_key:
+                results.append(present_k)
+            if has_present_value:
+                results.append(present_v)
+
+            if len(results) == 1:
+                return results[0]
+            return tuple(results)
+
+        return builder.call_function(
+            _attention_with_cache,
+            args=(
+                query,
+                key,
+                value,
+                attn_mask,
+                past_key,
+                past_value,
+                is_causal,
+                scale,
+                q_num_heads,
+                kv_num_heads,
+                has_present_key,
+                has_present_value,
+            ),
+        )
+    else:
+        # Simple case: just use SDPA
+
+        def _attention_standard(
+            q: torch.Tensor,
+            k: torch.Tensor,
+            v: torch.Tensor,
+            mask: torch.Tensor | None,
+            past_k: torch.Tensor | None,
+            past_v: torch.Tensor | None,
+            is_causal: int,
+            scale: float | None,
+            q_num_heads: int | None,
+            kv_num_heads: int | None,
+            nonpad_kv_seqlen: torch.Tensor | None,
+        ) -> torch.Tensor:
+            is_3d = q.dim() == 3
+
+            if is_3d:
+                batch_size = q.shape[0]
+                q_seq_len = q.shape[1]
+                kv_seq_len = k.shape[1]
+                n_q_heads = q_num_heads if q_num_heads is not None else 1
+                n_kv_heads = kv_num_heads if kv_num_heads is not None else n_q_heads
+                q_head_size = q.shape[2] // n_q_heads
+                kv_head_size = k.shape[2] // n_kv_heads
+                v_head_size = v.shape[2] // n_kv_heads
+                q_4d = q.view(batch_size, q_seq_len, n_q_heads, q_head_size).transpose(
+                    1, 2
+                )
+                k_4d = k.view(
+                    batch_size, kv_seq_len, n_kv_heads, kv_head_size
+                ).transpose(1, 2)
+                v_4d = v.view(
+                    batch_size, kv_seq_len, n_kv_heads, v_head_size
+                ).transpose(1, 2)
+            else:
+                q_4d = q
+                k_4d = k
+                v_4d = v
+                batch_size = q.shape[0]
+                n_q_heads = q.shape[1]
+                n_kv_heads = k.shape[1]
+                q_seq_len = q.shape[2]
+
+            # Handle past key/value (KV cache)
+            if past_k is not None and past_v is not None:
+                k_4d = torch.cat([past_k, k_4d], dim=2)
+                v_4d = torch.cat([past_v, v_4d], dim=2)
+
+            # Handle GQA: expand KV heads to match Q heads
+            if n_kv_heads != n_q_heads:
+                n_rep = n_q_heads // n_kv_heads
+                k_4d = k_4d.repeat_interleave(n_rep, dim=1)
+                v_4d = v_4d.repeat_interleave(n_rep, dim=1)
+
+            # Handle mask padding if mask is shorter than KV sequence
+            # Per ONNX spec: "The last dimension can also be shorter than
+            # total_sequence_length and will be padded with negative infinity"
+            kv_seq_len_actual = k_4d.shape[2]
+
+            if mask is not None:
+                # Pad mask if shorter than KV sequence (BEFORE adding nonpad mask)
+                if mask.shape[-1] < kv_seq_len_actual:
+                    pad_size = kv_seq_len_actual - mask.shape[-1]
+                    # For bool masks, pad with False; for float, pad with -inf
+                    if mask.dtype == torch.bool:
+                        mask = torch.nn.functional.pad(mask, (0, pad_size), value=False)
+                    else:
+                        mask = torch.nn.functional.pad(
+                            mask, (0, pad_size), value=float("-inf")
+                        )
+
+                # Expand mask for GQA if needed (mask might have n_kv_heads dimension)
+                if (
+                    mask.dim() == 4
+                    and mask.shape[1] == n_kv_heads
+                    and n_kv_heads != n_q_heads
+                ):
+                    n_rep = n_q_heads // n_kv_heads
+                    mask = mask.repeat_interleave(n_rep, dim=1)
+
+            # Handle nonpad_kv_seqlen: create a padding mask for each batch
+            # that masks out positions >= nonpad_kv_seqlen[batch]
+            if nonpad_kv_seqlen is not None:
+                # Create a position index tensor: (1, 1, 1, kv_seq_len)
+                positions = torch.arange(kv_seq_len_actual, device=q.device).view(
+                    1, 1, 1, -1
+                )
+                # Create mask: True where position < nonpad_kv_seqlen[batch]
+                # nonpad_kv_seqlen: (batch_size,) -> (batch_size, 1, 1, 1)
+                valid_mask = positions < nonpad_kv_seqlen.view(-1, 1, 1, 1)
+                # Convert to additive mask: 0 for valid, -inf for padding
+                pad_mask = torch.where(valid_mask, 0.0, float("-inf")).to(q.dtype)
+                # Combine with existing mask
+                if mask is not None:
+                    mask = mask + pad_mask
+                else:
+                    mask = pad_mask
+
+            # Call SDPA
+            if mask is not None:
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    q_4d, k_4d, v_4d, attn_mask=mask, is_causal=False, scale=scale
+                )
+            elif is_causal:
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    q_4d, k_4d, v_4d, is_causal=True, scale=scale
+                )
+            else:
+                output = torch.nn.functional.scaled_dot_product_attention(
+                    q_4d, k_4d, v_4d, scale=scale
+                )
+
+            if is_3d:
+                output = (
+                    output.transpose(1, 2).contiguous().view(batch_size, q_seq_len, -1)
+                )
+
+            return output
+
+        return builder.call_function(
+            _attention_standard,
+            args=(
+                query,
+                key,
+                value,
+                attn_mask,
+                past_key,
+                past_value,
+                is_causal,
+                scale,
+                q_num_heads,
+                kv_num_heads,
+                nonpad_kv_seqlen,
+            ),
+        )
+
+
+# =============================================================================
+# Simplified LayerNormalization variants (ONNX Runtime contrib ops)
+# =============================================================================
+
+
+@register("SimplifiedLayerNormalization")
+def simplified_layer_normalization(
+    builder: "GraphBuilder", node: onnx.NodeProto
+) -> torch.fx.Node:
+    """Simplified Layer Normalization (RMSNorm).
+
+    This is LayerNormalization without bias and mean subtraction.
+    Formula: output = x / sqrt(mean(x^2) + epsilon) * scale
+    """
+    x = builder.get_value(node.input[0])
+    scale = builder.get_value(node.input[1])
+
+    axis = get_attribute(node, "axis", -1)
+    epsilon = get_attribute(node, "epsilon", 1e-5)
+
+    def _simplified_layer_norm(x, scale, axis, epsilon):
+        # Simplified LayerNorm (RMSNorm)
+        # output = x * rsqrt(mean(x^2) + epsilon) * scale
+        if axis < 0:
+            axis_pos = x.dim() + axis
+        else:
+            axis_pos = axis
+
+        # Keep dims for broadcasting
+        dims = list(range(axis_pos, x.dim()))
+
+        # Compute RMS: sqrt(mean(x^2))
+        variance = x.pow(2).mean(dim=dims, keepdim=True)
+        x_normalized = x * torch.rsqrt(variance + epsilon)
+
+        return x_normalized * scale
+
+    return builder.call_function(_simplified_layer_norm, args=(x, scale, axis, epsilon))
+
+
+@register("SkipSimplifiedLayerNormalization")
+def skip_simplified_layer_normalization(
+    builder: "GraphBuilder", node: onnx.NodeProto
+) -> torch.fx.Node:
+    """Skip connection + Simplified Layer Normalization (RMSNorm)."""
+    x = builder.get_value(node.input[0])
+    skip = builder.get_value(node.input[1])
+    scale = builder.get_value(node.input[2])
+    bias = get_optional_input(builder, node, 3)
+
+    epsilon = get_attribute(node, "epsilon", 1e-5)
+
+    def _skip_simplified_layer_norm(x, skip, scale, bias, epsilon):
+        # Add skip connection
+        hidden = x + skip
+        if bias is not None:
+            hidden = hidden + bias
+
+        # Simplified LayerNorm (RMSNorm)
+        variance = hidden.pow(2).mean(dim=-1, keepdim=True)
+        hidden_normalized = hidden * torch.rsqrt(variance + epsilon)
+
+        return hidden_normalized * scale
+
+    return builder.call_function(
+        _skip_simplified_layer_norm, args=(x, skip, scale, bias, epsilon)
+    )
+
+
+@register("GroupQueryAttention")
+def group_query_attention(
+    builder: "GraphBuilder", node: onnx.NodeProto
+) -> torch.fx.Node:
+    """Group Query Attention (GQA) - used in LLaMA, Mistral, etc.
+
+    Inputs:
+        - query: [batch, seq_len, num_heads * head_size]
+        - key: [batch, kv_seq_len, num_kv_heads * head_size]
+        - value: [batch, kv_seq_len, num_kv_heads * head_size]
+        - past_key (optional): [batch, num_kv_heads, past_seq_len, head_size]
+        - past_value (optional): [batch, num_kv_heads, past_seq_len, head_size]
+        - seqlens_k (optional): cumulative sequence lengths for keys
+        - total_sequence_length (optional): total sequence length
+        - cos_cache (optional): [max_seq_len, head_size / 2] or [max_seq_len, head_size]
+        - sin_cache (optional): [max_seq_len, head_size / 2] or [max_seq_len, head_size]
+
+    Attributes:
+        - num_heads: number of attention heads
+        - kv_num_heads: number of key-value heads (for GQA)
+        - scale: scaling factor (default: 1/sqrt(head_size))
+        - local_window_size: for sliding window attention
+        - do_rotary: whether to apply rotary position embeddings
+        - rotary_interleaved: whether rotary is interleaved (GPT-NeoX style vs LLaMA)
+
+    Outputs:
+        - output: [batch, seq_len, num_heads * head_size]
+        - present_key: [batch, num_kv_heads, total_seq_len, head_size]
+        - present_value: [batch, num_kv_heads, total_seq_len, head_size]
+    """
+    # Get required inputs
+    query = builder.get_value(node.input[0])
+    key = builder.get_value(node.input[1])
+    value = builder.get_value(node.input[2])
+
+    # Get optional inputs
+    past_key = get_optional_input(builder, node, 3)
+    past_value = get_optional_input(builder, node, 4)
+    seqlens_k = get_optional_input(builder, node, 5)
+    total_seq_len = get_optional_input(builder, node, 6)
+    cos_cache = get_optional_input(builder, node, 7)
+    sin_cache = get_optional_input(builder, node, 8)
+
+    # Get attributes
+    num_heads = get_attribute(node, "num_heads", 1)
+    kv_num_heads = get_attribute(node, "kv_num_heads", num_heads)
+    scale = get_attribute(node, "scale", None)
+    local_window_size = get_attribute(node, "local_window_size", -1)
+    do_rotary = get_attribute(node, "do_rotary", 0)
+    rotary_interleaved = get_attribute(node, "rotary_interleaved", 0)
+
+    def _group_query_attention(
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        past_k: torch.Tensor | None,
+        past_v: torch.Tensor | None,
+        seqlens_k: torch.Tensor | None,
+        total_seq_len: torch.Tensor | None,
+        cos_cache: torch.Tensor | None,
+        sin_cache: torch.Tensor | None,
+        n_heads: int,
+        n_kv_heads: int,
+        attn_scale: float | None,
+        window_size: int,
+        do_rotary: int,
+        rotary_interleaved: int,
+    ):
+        batch_size, seq_len, _ = q.shape
+        head_size = q.shape[-1] // n_heads
+        kv_head_size = k.shape[-1] // n_kv_heads
+
+        # Reshape Q, K, V to [batch, num_heads, seq_len, head_size]
+        q = q.view(batch_size, seq_len, n_heads, head_size).transpose(1, 2)
+        k = k.view(batch_size, -1, n_kv_heads, kv_head_size).transpose(1, 2)
+        v = v.view(batch_size, -1, n_kv_heads, kv_head_size).transpose(1, 2)
+
+        # Calculate position offset from past cache
+        past_seq_len = 0
+        if past_k is not None and past_k.numel() > 0:
+            past_seq_len = past_k.shape[2]
+
+        # Apply rotary position embeddings if enabled
+        if do_rotary and cos_cache is not None and sin_cache is not None:
+            # Get the position indices
+            positions = torch.arange(
+                past_seq_len, past_seq_len + seq_len, device=q.device
+            )
+
+            # Get cos/sin values for current positions
+            cos = cos_cache[positions]  # [seq_len, rotary_dim]
+            sin = sin_cache[positions]  # [seq_len, rotary_dim]
+
+            # Expand for batch and heads: [1, 1, seq_len, rotary_dim]
+            cos = cos.unsqueeze(0).unsqueeze(0)
+            sin = sin.unsqueeze(0).unsqueeze(0)
+
+            rotary_dim = cos.shape[-1]
+
+            if rotary_interleaved:
+                # GPT-NeoX style: [x0, x1, x2, x3, ...] -> rotate pairs
+                q_rot = q[..., :rotary_dim]
+                q_pass = q[..., rotary_dim:]
+                k_rot = k[..., :rotary_dim]
+                k_pass = k[..., rotary_dim:]
+
+                # Apply rotation
+                q1, q2 = q_rot[..., ::2], q_rot[..., 1::2]
+                k1, k2 = k_rot[..., ::2], k_rot[..., 1::2]
+
+                cos_half = cos[..., ::2]
+                sin_half = sin[..., ::2]
+
+                q_rot_new = torch.stack(
+                    [q1 * cos_half - q2 * sin_half, q1 * sin_half + q2 * cos_half],
+                    dim=-1,
+                ).flatten(-2)
+                k_rot_new = torch.stack(
+                    [k1 * cos_half - k2 * sin_half, k1 * sin_half + k2 * cos_half],
+                    dim=-1,
+                ).flatten(-2)
+
+                q = torch.cat([q_rot_new, q_pass], dim=-1)
+                k = torch.cat([k_rot_new, k_pass], dim=-1)
+            else:
+                # LLaMA style: cos/sin are [seq, rotary_dim]
+                # rotary_dim is half the head_size in this format
+                # q/k first rotary_dim*2 elements are rotated:
+                # q1 = q[..., :rotary_dim], q2 = q[..., rotary_dim:rotary_dim*2]
+                # result = (q1*cos - q2*sin, q1*sin + q2*cos)
+
+                rotary_full = rotary_dim * 2  # total dims that get rotated
+                q_rot = q[..., :rotary_full]
+                q_pass = q[..., rotary_full:]
+                k_rot = k[..., :rotary_full]
+                k_pass = k[..., rotary_full:]
+
+                # Split into first half and second half
+                q1, q2 = q_rot[..., :rotary_dim], q_rot[..., rotary_dim:rotary_full]
+                k1, k2 = k_rot[..., :rotary_dim], k_rot[..., rotary_dim:rotary_full]
+
+                # cos/sin are already in the right shape [1, 1, seq_len, rotary_dim]
+                q_rot_new = torch.cat(
+                    [q1 * cos - q2 * sin, q1 * sin + q2 * cos], dim=-1
+                )
+                k_rot_new = torch.cat(
+                    [k1 * cos - k2 * sin, k1 * sin + k2 * cos], dim=-1
+                )
+
+                q = torch.cat([q_rot_new, q_pass], dim=-1)
+                k = torch.cat([k_rot_new, k_pass], dim=-1)
+
+        # Handle past key-value cache
+        if past_k is not None and past_k.numel() > 0:
+            k = torch.cat([past_k, k], dim=2)
+            v = torch.cat([past_v, v], dim=2)
+
+        # Present key-value for caching
+        present_k = k
+        present_v = v
+
+        # Expand K, V for GQA (repeat for each head group)
+        if n_kv_heads < n_heads:
+            n_rep = n_heads // n_kv_heads
+            k = k.repeat_interleave(n_rep, dim=1)
+            v = v.repeat_interleave(n_rep, dim=1)
+
+        # Compute attention scale
+        if attn_scale is None:
+            attn_scale = 1.0 / (head_size**0.5)
+
+        # Use scaled_dot_product_attention
+        # For autoregressive with past cache, don't use causal mask for new tokens
+        # since past_k/v already handled the causality
+        is_causal = seq_len > 1 and past_seq_len == 0
+        output = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, scale=attn_scale, is_causal=is_causal
+        )
+
+        # Reshape output: [batch, num_heads, seq_len, head_size] -> [batch, seq_len, hidden]
+        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, -1)
+
+        return output, present_k, present_v
+
+    # Build the call
+    result = builder.call_function(
+        _group_query_attention,
+        args=(
+            query,
+            key,
+            value,
+            past_key,
+            past_value,
+            seqlens_k,
+            total_seq_len,
+            cos_cache,
+            sin_cache,
+            num_heads,
+            kv_num_heads,
+            scale,
+            local_window_size,
+            do_rotary,
+            rotary_interleaved,
+        ),
+    )
+
+    # Return tuple output
+    return result
+
+
+# =============================================================================
+# Rotary Embedding (ONNX standard domain, since opset 23)
+# =============================================================================
+
+
+@register("RotaryEmbedding", since_version=23)
+def rotary_embedding_onnx(
+    builder: "GraphBuilder", node: onnx.NodeProto
+) -> torch.fx.Node:
+    """ONNX RotaryEmbedding operator (standard domain, since opset 23).
+
+    Applies rotary position embeddings (RoPE) to the input tensor based on
+    https://arxiv.org/pdf/2104.09864
+
+    Inputs:
+        - X: 4D tensor (batch_size, num_heads, sequence_length, head_size) or
+             3D tensor (batch_size, sequence_length, hidden_size)
+        - cos_cache: 2D tensor (max_position_id+1, head_size/2) when position_ids provided,
+                     or 3D tensor (batch_size, sequence_length, head_size/2) otherwise
+        - sin_cache: Same shape as cos_cache
+        - position_ids (optional): 2D tensor (batch_size, sequence_length)
+
+    Attributes:
+        - interleaved: Whether to use interleaved pattern. Default is 0 (False).
+        - num_heads: Number of attention heads (required for 3D input).
+        - rotary_embedding_dim: Partial rotary dimension. Default is 0 (full rotation).
+
+    Outputs:
+        - Y: Tensor with same shape as input.
+    """
+    # Get inputs
+    input_tensor = builder.get_value(node.input[0])
+    cos_cache = builder.get_value(node.input[1])
+    sin_cache = builder.get_value(node.input[2])
+    position_ids = get_optional_input(builder, node, 3)
+
+    # Get attributes
+    interleaved = get_attribute(node, "interleaved", 0)
+    num_heads = get_attribute(node, "num_heads", 0)
+    rotary_embedding_dim = get_attribute(node, "rotary_embedding_dim", 0)
+
+    def _rotary_embedding_onnx(
+        x: torch.Tensor,
+        cos_cache: torch.Tensor,
+        sin_cache: torch.Tensor,
+        position_ids: torch.Tensor | None,
+        interleaved: int,
+        num_heads: int,
+        rotary_dim: int,
+    ) -> torch.Tensor:
+        """Apply ONNX-standard rotary position embeddings."""
+        original_shape = x.shape
+        is_3d = x.dim() == 3
+
+        # First ensure input has shape [batch_size, seq_len, num_heads, head_size]
+        if x.dim() == 4:
+            # Input is (batch_size, num_heads, seq_len, head_size)
+            # Transpose to (batch_size, seq_len, num_heads, head_size)
+            x = x.transpose(1, 2)
+            batch_size, seq_len, n_heads, head_size = x.shape
+        else:
+            # Input is (batch_size, seq_len, hidden_size)
+            batch_size, seq_len, hidden_size = x.shape
+            assert num_heads != 0, "num_heads must be provided for 3D input"
+            head_size = hidden_size // num_heads
+            x = x.view(batch_size, seq_len, num_heads, head_size)
+            _n_heads = num_heads
+
+        # Determine rotary_embedding_dim
+        if rotary_dim == 0:
+            rot_dim = head_size
+        else:
+            rot_dim = rotary_dim
+
+        rotary_dim_half = rot_dim // 2
+
+        # Split into rotary and pass-through parts
+        x_rotate = x[..., :rot_dim]
+        x_not_rotate = x[..., rot_dim:] if rot_dim < head_size else None
+
+        # Retrieve sin and cos caches using position ids
+        if position_ids is not None:
+            # cos_cache/sin_cache shape: (max_pos+1, rotary_dim/2)
+            # position_ids shape: (batch_size, seq_len)
+            # Result shape: (batch_size, seq_len, rotary_dim/2)
+            cos = cos_cache[position_ids]
+            sin = sin_cache[position_ids]
+        else:
+            # cos_cache/sin_cache already have shape (batch_size, seq_len, rotary_dim/2)
+            cos = cos_cache
+            sin = sin_cache
+
+        # Validate cache dimensions
+        if cos.shape[-1] != rotary_dim_half:
+            raise ValueError(
+                f"Last dimension of cos cache ({cos.shape[-1]}) does not match "
+                f"rotary_embedding_dim/2 ({rotary_dim_half})."
+            )
+
+        # Add head dimension: (batch_size, seq_len, 1, rotary_dim/2)
+        cos = cos.unsqueeze(2)
+        sin = sin.unsqueeze(2)
+
+        # Apply rotation based on interleaved pattern
+        if interleaved:
+            # Interleaved: x_rotate[..., 0::2] and x_rotate[..., 1::2]
+            x1 = x_rotate[..., 0::2]
+            x2 = x_rotate[..., 1::2]
+
+            # Calculate real and imaginary values
+            real = (cos * x1) - (sin * x2)
+            imag = (sin * x1) + (cos * x2)
+
+            # Interleave back
+            real = real.unsqueeze(-1)
+            imag = imag.unsqueeze(-1)
+            x_rotate_result = torch.cat((real, imag), dim=-1).flatten(-2)
+        else:
+            # Non-interleaved: split in halves
+            x1 = x_rotate[..., :rotary_dim_half]
+            x2 = x_rotate[..., rotary_dim_half:rot_dim]
+
+            # Calculate real and imaginary values
+            real = (cos * x1) - (sin * x2)
+            imag = (sin * x1) + (cos * x2)
+
+            x_rotate_result = torch.cat((real, imag), dim=-1)
+
+        # Concatenate with non-rotated part
+        if x_not_rotate is not None:
+            output = torch.cat((x_rotate_result, x_not_rotate), dim=-1)
+        else:
+            output = x_rotate_result
+
+        # Reshape back to original shape
+        if is_3d:
+            output = output.view(original_shape)
+        else:
+            # Transpose back to (batch_size, num_heads, seq_len, head_size)
+            output = output.transpose(1, 2)
+
+        return output
+
+    return builder.call_function(
+        _rotary_embedding_onnx,
+        args=(
+            input_tensor,
+            cos_cache,
+            sin_cache,
+            position_ids,
+            interleaved,
+            num_heads,
+            rotary_embedding_dim,
+        ),
+    )