onnxscript 0.1.0__py3-none-any.whl

onnxscript/rewriter/ort_fusions/cos_sin_cache.py

@@ -0,0 +1,154 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import numpy as np
+
+import onnxscript.ir as ir
+from onnxscript.optimizer import remove_unused_nodes
+from onnxscript.rewriter import _ir_utils, pattern
+
+# Rewrite the computation of cos/sin cache into the form expected by ORT's custom ops.
+
+# We match against the following code pattern:
+# Original code (from transformers) for computing cos/sin cache for RoPE:
+# https://github.com/huggingface/transformers/blob/0ade1caa356dce6b70ef8293addeb0898f177206/src/transformers/models/llama/modeling_llama.py#L135
+# position_ids_expanded = position_ids[:, None, :].float()
+# freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
+# emb = torch.cat((freqs, freqs), dim=-1)
+# cos = emb.cos()
+# sin = emb.sin()
+#
+# We rewrite this pattern into the following form:
+# inv_freq_values = inv_freq_expanded.reshape(1, -1)
+# pos_id_range = np.arange(max_pos_id, dtype=np.float32).reshape(-1, 1)
+# angles = np.matmul(pos_id_range, inv_freq_values)
+# cos_value = np.cos(angles)
+# sin_value = np.sin(angles)
+# cos_2d = op.Constant(value=ir.tensor(cos_value))
+# sin_2d = op.Constant(value=ir.tensor(sin_value))
+#
+# This produces cos/sin values in a form that can be used by ORT's custom ops.
+
+# TODO: To apply the pattern-rewrite, we need to know the maximum position id.
+# Need to find a way to get this information from the model or its config.
+
+
+class CosSinCacheFusion(pattern.RewriteRuleClassBase):
+    def __init__(
+        self,
+        name: str,
+        max_pos_id: int,
+        *,
+        cast: bool = False,
+        reshape: bool = False,
+        const_freqs: bool = False,
+    ):
+        # This pattern makes use of shared Cos/Sin values. So, we can't remove the
+        # matched nodes as part of the rewrite-step. We apply a separate final
+        # pass to remove unused nodes.
+        super().__init__(name, remove_nodes=False)
+        self._max_pos_id = max_pos_id
+        # map from inv_freq to (cos, sin) values for transformed graph
+        self._inv_freq_cos_sin_cache: dict[ir.Value, tuple[ir.Value, ir.Value]] = {}
+        self._reshape = reshape
+        self._cast = cast
+        self._const_freqs = const_freqs
+
+    def cleanup(self):
+        self._inv_freq_cos_sin_cache.clear()
+
+    def pattern(self, op, x, inv_freq, position_ids, interleaved, num_heads, freqs, dtype):
+        if not self._const_freqs:
+            # Compute freqs from inv_freq and position_ids. In the _const_freqs case,
+            # this computation has been constant-folded away and freqs is a constant.
+            # B: batch size, S: sequence length, E: embedding dimension
+            # position_ids: [B, S]
+            # inv_freq: [1, E, 1]
+            position_ids_expanded = op.Unsqueeze(position_ids, 1)  # [B, S] => [B, 1, S]
+            position_ids_expanded = op.Cast(position_ids_expanded, to=ir.DataType.FLOAT)
+            # if self._reshape:
+            #     position_ids_expanded = op.Expand(position_ids_expanded, _allow_other_inputs=True)
+            #     position_ids_expanded = op.Reshape(position_ids_expanded, _allow_other_inputs=True)
+            freqs = op.MatMul(inv_freq, position_ids_expanded)  # [B, E, S]
+        # if self._reshape:
+        #     freqs = op.Reshape(freqs, freqs_3d_shape)  # redundant reshape
+        freqs = op.Transpose(freqs, perm=[0, 2, 1])  # [B, S, E]
+        emb = op.Concat(freqs, freqs, axis=-1)
+        cos = op.Cos(emb)
+        if self._cast:
+            cos = op.Cast(cos, to=dtype)
+        sin = op.Sin(emb)
+        if self._cast:
+            sin = op.Cast(sin, to=dtype)
+        cos_4d = op.Unsqueeze(cos, 1)  # convert
+        sin_4d = op.Unsqueeze(sin, 1)
+        return op.RotaryEmbedding(
+            x,
+            cos_4d,
+            sin_4d,
+            interleaved=interleaved,
+            num_heads=num_heads,
+            _domain="ai.onnxruntime.fusion",
+        )
+
+    def check(self, context, inv_freq, position_ids, freqs, **_):
+        # TODO(rama): handle redundant reshape/expand
+        if self._const_freqs:
+            return (freqs.const_value is not None) and _ir_utils.has_rank(freqs, 3)
+        if not _ir_utils.has_rank(position_ids, 2):
+            return False
+        if not _ir_utils.has_rank(inv_freq, 3):
+            return False
+        inv_freq_shape = inv_freq.shape
+        if inv_freq.const_value is None:  # TODO: should this be inv_freq_shape?
+            return False
+        return inv_freq_shape[0] == 1 and inv_freq_shape[2] == 1
+
+    def rewrite(
+        self, op, x, inv_freq, position_ids, interleaved, num_heads, freqs, dtype, **_
+    ):
+        if inv_freq in self._inv_freq_cos_sin_cache:
+            cos_2d, sin_2d = self._inv_freq_cos_sin_cache[inv_freq]
+        else:
+            if self._const_freqs:
+                angles = freqs.const_value.numpy()
+            else:
+                inv_freq_values = inv_freq.const_value.numpy().reshape(1, -1)
+                pos_id_range = np.arange(self._max_pos_id, dtype=np.float32).reshape(-1, 1)
+                angles = np.matmul(pos_id_range, inv_freq_values)
+            cos_value = np.cos(angles)
+            sin_value = np.sin(angles)
+            cos_2d = op.Constant(value=ir.tensor(cos_value))
+            sin_2d = op.Constant(value=ir.tensor(sin_value))
+            if self._cast:
+                cos_2d = op.Cast(cos_2d, to=dtype)
+                sin_2d = op.Cast(sin_2d, to=dtype)
+            self._inv_freq_cos_sin_cache[inv_freq] = (cos_2d, sin_2d)
+        return op.RotaryEmbedding(
+            x,
+            position_ids,
+            cos_2d,
+            sin_2d,
+            interleaved=interleaved,
+            num_heads=num_heads,
+            _domain="com.microsoft",
+        )
+
+
+_cast = CosSinCacheFusion.rule("CosSinCache", 2048, cast=True, const_freqs=True)
+_no_cast = CosSinCacheFusion.rule("CosSinCache", 2048, cast=False)
+
+cos_sin_cache_rules = pattern.RewriteRuleSet([_cast, _no_cast])
+
+debug: bool = True
+
+
+def fuse_cos_sin_cache(model: ir.Model) -> int:
+    count = cos_sin_cache_rules.apply_to_model(model)
+    if count == 0 and debug:
+        cos_sin_cache_rules.apply_to_model(model, debug=True)
+    else:
+        print(f"CosSinCache count: {count}")
+        remove_unused_nodes(model)
+    return count

onnxscript/rewriter/ort_fusions/gqa.py

@@ -0,0 +1,156 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import onnxscript.ir as ir
+from onnxscript.optimizer import remove_unused_nodes
+from onnxscript.rewriter import pattern
+
+
+class GroupQueryAttention(pattern.RewriteRuleClassBase):
+    def __init__(self, name: str, *, use_2d_matmul: bool):
+        super().__init__(name, remove_nodes=False)
+        self._use_2d_matmul = use_2d_matmul
+
+    def _compute_packed_QKV(self, op, input, weight):
+        if self._use_2d_matmul:
+            # Convert batched input of shape (B, S, D) to 2D input (B*S, D)
+            input = op.Reshape(input, _allow_other_inputs=True)
+        projected = op.MatMul(input, weight)
+        if self._use_2d_matmul:
+            # Convert 2D output back to batched output of shape (B, S, D)
+            projected = op.Reshape(projected, _allow_other_inputs=True)
+        # Split combined QKV into Q, K, and V
+        query_3d = op.Slice(projected, _allow_other_inputs=True)
+        key_3d = op.Slice(projected, _allow_other_inputs=True)
+        value_3d = op.Slice(projected, _allow_other_inputs=True)
+        # Reshape from (B, S, D) to (B, S, H, D/H)
+        query_4d = op.Reshape(
+            query_3d,
+            _allow_other_inputs=True,
+            _allow_other_attributes=True,
+            _outputs=["query_mm_reshaped"],
+        )
+        # Transpose from (B, S, H, D/H) to (B, H, S, D/H)
+        query = op.Transpose(query_4d, perm=[0, 2, 1, 3])
+        key_4d = op.Reshape(
+            key_3d,
+            _allow_other_inputs=True,
+            _allow_other_attributes=True,
+            _outputs=["key_mm_reshaped"],
+        )
+        key = op.Transpose(key_4d, perm=[0, 2, 1, 3])
+        value_4d = op.Reshape(
+            value_3d,
+            _allow_other_inputs=True,
+            _allow_other_attributes=True,
+            _outputs=["value_mm_reshaped"],
+        )
+        value = op.Transpose(value_4d, perm=[0, 2, 1, 3])
+
+        return query, key, value
+
+    def pattern(
+        self,
+        op,
+        input,
+        qkv_weight,
+        mask,
+        cos,
+        sin,
+        past_key,
+        past_value,
+        position_ids,
+    ):
+        query, key, value = self._compute_packed_QKV(op, input, qkv_weight)
+
+        query_rope = op.RotaryEmbedding(query, position_ids, cos, sin, _domain="com.microsoft")
+
+        key_rope = op.RotaryEmbedding(key, position_ids, cos, sin, _domain="com.microsoft")
+        present_key = op.Concat(past_key, key_rope, axis=-2)
+        # Transpose last two axes of present_key to compute dot-product via matmul.
+        present_key = op.Transpose(present_key, perm=[0, 1, 3, 2])
+
+        present_value = op.Concat(past_value, value, axis=-2)
+
+        attention = op.SDPA(
+            query_rope, present_key, present_value, mask, _domain="ai.onnxruntime.fusion"
+        )
+        # Transpose back to (B, S, H, D/H)
+        attention_transposed = op.Transpose(attention, perm=[0, 2, 1, 3])
+        # Reshape back to (B, S, D)
+        attention_reshaped = op.Reshape(
+            attention_transposed, _allow_other_inputs=True, _outputs=["attention_reshaped"]
+        )
+        return attention_reshaped, present_key, present_value
+
+    def check(
+        self,
+        op,
+        # query_mm_reshaped,
+        # key_mm_reshaped,
+        # value_mm_reshaped,
+        # key_reshaped,
+        # key_transposed,
+        # attention_reshaped,
+        **_,
+    ):
+        # bindings: dict[str, int] = {}
+        # status = (
+        #     _check_shape(bindings, query_mm_reshaped, ["B", "S", "H", "d_h"])
+        #     and _check_shape(bindings, key_mm_reshaped, ["B", "S", "H", "d_h"])
+        #     and _check_shape(bindings, value_mm_reshaped, ["B", "S", "H", "d_h"])
+        #     and _check_shape(bindings, key_reshaped, ["B*H", "KVS", "d_h"])
+        #     and _check_shape(bindings, key_transposed, ["B", "H", "d_h", "KVS"])
+        #     and _check_shape(bindings, attention_reshaped, ["B", "S", "H*d_h"])
+        # )
+        # if not status:
+        #     return False
+        # if bindings["B"] * bindings["H"] != bindings["B*H"]:
+        #     return False
+        # if bindings["H"] * bindings["d_h"] != bindings["H*d_h"]:
+        #     return False
+        return True
+
+    def rewrite(
+        self,
+        op,
+        input,
+        qkv_weight,
+        mask,
+        cos,
+        sin,
+        past_key,
+        past_value,
+        position_ids,
+        query_mm_reshaped,
+        **_,
+    ):
+        num_heads = query_mm_reshaped.shape[2]
+        qkv = op.MatMul(input, qkv_weight)
+        return op.GroupQueryAttention(
+            qkv,
+            None,  # key
+            None,  # value
+            past_key,
+            past_value,
+            # seqlens_k,
+            # total_sequence_length,
+            cos,
+            sin,
+            num_heads=num_heads,
+            _domain="com.microsoft",
+            _outputs=3,
+        )
+
+
+_rule1 = GroupQueryAttention.rule("MHA_2dmm", use_2d_matmul=False)
+
+gqa_rules = pattern.RewriteRuleSet([_rule1])
+
+
+def fuse_gqa(model: ir.Model) -> int:
+    count = gqa_rules.apply_to_model(model)
+    print(f"GQA count: {count}")
+    remove_unused_nodes(model)
+    return count

onnxscript/rewriter/ort_fusions/mha.py

@@ -0,0 +1,198 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+from typing import Sequence
+
+import onnxscript.ir as ir
+from onnxscript.rewriter import pattern
+
+"""
+The MultiHeadAttention pattern:
+
+B: Batch size
+S: Sequence length
+D: input embedding dimension
+H: number of heads
+d_h: head size (usually, D = H * d_h)
+
+thus, weights are usually of shape (D, D) and (D, D) and (D, D)
+
+for each of Q, K, and V, we have the following pattern:
+   MatMul (Input, W), producing output of shape (B, S, D)
+   Reshape to produce a matrix of shape (B, S, H, d_h)
+   Transpose middle two axes to produce a matrix of shape (B, H, S, d_h)
+
+This is followed by a RotaryEmbedding pattern for Q and K
+
+The last two axes of the key-embedding are then swapped (using a Reshape/Transpose/Reshape sequence)
+
+The dot-product attention is then computed using SDPA.
+Finally, the output is transposed and reshaped back to (B, S, D) shape
+"""
+
+
+def _check_shape(bindings: dict[str, int], val: ir.Value, shape: Sequence[str]) -> bool:
+    if val.shape is None:
+        return False
+    if val.shape.rank() != len(shape):
+        return False
+    for actual, expected in zip(val.shape, shape):
+        if expected not in bindings:
+            bindings[expected] = actual  # type: ignore[assignment]
+        elif actual != bindings[expected]:
+            return False
+    return True
+
+
+class MultiHeadAttention(pattern.RewriteRuleClassBase):
+    def __init__(self, name: str, *, use_2d_matmul: bool):
+        super().__init__(name)
+        self._use_2d_matmul = use_2d_matmul
+
+    def _compute_QKV(self, op, input, weight, reshape_var: str):
+        """Applied to generate each of Q, K, and V from input."""
+        if self._use_2d_matmul:
+            # Convert batched input of shape (B, S, D) to 2D input (B*S, D)
+            input = op.Reshape(input, _allow_other_inputs=True)
+        projected = op.MatMul(input, weight)
+        if self._use_2d_matmul:
+            # Convert 2D output back to batched output of shape (B, S, D)
+            projected = op.Reshape(projected, _allow_other_inputs=True)
+        # Reshape from (B, S, D) to (B, S, H, D/H)
+        reshaped = op.Reshape(
+            projected,
+            _allow_other_inputs=True,
+            _allow_other_attributes=True,
+            _outputs=[reshape_var],
+        )
+        # Transpose from (B, S, H, D/H) to (B, H, S, D/H)
+        transposed = op.Transpose(reshaped, perm=[0, 2, 1, 3])
+        return transposed
+
+    def pattern(
+        self,
+        op,
+        input,
+        query_weight,
+        key_weight,
+        value_weight,
+        qkv_weight,
+        mask,
+        cos,
+        sin,
+        past_key,
+        past_value,
+        position_ids,
+    ):
+        query = self._compute_QKV(op, input, query_weight, "query_mm_reshaped")
+        key = self._compute_QKV(op, input, key_weight, "key_mm_reshaped")
+        value = self._compute_QKV(op, input, value_weight, "value_mm_reshaped")
+
+        query_rope = op.RotaryEmbedding(query, position_ids, cos, sin, _domain="com.microsoft")
+
+        key_rope = op.RotaryEmbedding(key, position_ids, cos, sin, _domain="com.microsoft")
+        key_rope = op.Concat(past_key, key_rope, axis=-2)
+        # Transpose last two axes of key_rope to compute dot-product via matmul.
+        key_reshaped = op.Reshape(
+            key_rope, _allow_other_inputs=True, _outputs=["key_reshaped"]
+        )
+        key_reshaped_transposed = op.Transpose(key_reshaped, perm=[0, 2, 1])
+        key_transposed = op.Reshape(
+            key_reshaped_transposed, _allow_other_inputs=True, _outputs=["key_transposed"]
+        )
+
+        value = op.Concat(past_value, value, axis=-2)
+
+        attention = op.SDPA(
+            query_rope, key_transposed, value, mask, _domain="ai.onnxruntime.fusion"
+        )
+        # Transpose back to (B, S, H, D/H)
+        attention_transposed = op.Transpose(attention, perm=[0, 2, 1, 3])
+        # Reshape back to (B, S, D)
+        attention_reshaped = op.Reshape(
+            attention_transposed, _allow_other_inputs=True, _outputs=["attention_reshaped"]
+        )
+        return attention_reshaped, key_rope, value
+
+    def check(
+        self,
+        op,
+        query_mm_reshaped,
+        key_mm_reshaped,
+        value_mm_reshaped,
+        key_reshaped,
+        key_transposed,
+        attention_reshaped,
+        **_,
+    ):
+        bindings: dict[str, int] = {}
+        status = (
+            _check_shape(bindings, query_mm_reshaped, ["B", "S", "H", "d_h"])
+            and _check_shape(bindings, key_mm_reshaped, ["B", "S", "H", "d_h"])
+            and _check_shape(bindings, value_mm_reshaped, ["B", "S", "H", "d_h"])
+            and _check_shape(bindings, key_reshaped, ["B*H", "KVS", "d_h"])
+            and _check_shape(bindings, key_transposed, ["B", "H", "d_h", "KVS"])
+            and _check_shape(bindings, attention_reshaped, ["B", "S", "H*d_h"])
+        )
+        if not status:
+            return False
+        # if bindings["B"] * bindings["H"] != bindings["B*H"]:
+        #     return False
+        # if bindings["H"] * bindings["d_h"] != bindings["H*d_h"]:
+        #     return False
+        return True
+
+    def rewrite(
+        self,
+        op,
+        input,
+        query_weight,
+        key_weight,
+        value_weight,
+        mask,
+        cos,
+        sin,
+        past_key,
+        past_value,
+        position_ids,
+        query_mm_reshaped,
+        **_,
+    ):
+        num_heads = query_mm_reshaped.shape[2]
+        query = op.MatMul(input, query_weight)
+        key = op.MatMul(input, key_weight)
+        value = op.MatMul(input, value_weight)
+
+        query_rope = op.RotaryEmbedding(query, position_ids, cos, sin, _domain="com.microsoft")
+        key_rope = op.RotaryEmbedding(key, position_ids, cos, sin, _domain="com.microsoft")
+
+        return op.MultiHeadAttention(
+            query_rope,
+            key_rope,
+            value,
+            None,  # bias
+            None,  # key padding mask
+            mask,  # attention mask/bias
+            past_key,
+            past_value,
+            num_heads=num_heads,
+            _domain="com.microsoft",
+            _outputs=3,
+        )
+
+
+_rule1 = MultiHeadAttention.rule("MHA_2dmm", use_2d_matmul=False)
+
+mha_rules = pattern.RewriteRuleSet([_rule1])
+
+debug: bool = True
+
+
+def fuse_mha(model: ir.Model) -> int:
+    count = mha_rules.apply_to_model(model)
+    if count == 0 and debug:
+        mha_rules.apply_to_model(model, debug=True)
+    else:
+        print(f"MHA count: {count}")
+    return count

onnxscript/rewriter/ort_fusions/rms_normalization.py

@@ -0,0 +1,95 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import onnxscript.ir as ir
+from onnxscript.rewriter import _ir_utils, pattern
+
+"""
+RMS Normalization: This is referred to as SimplifiedLayerNormalization in the ORT codebase.
+See https://github.com/microsoft/onnxruntime/blob/6d9636f07cccdb6e4ac453087ad54c3bc9854d50/onnxruntime/core/graph/contrib_ops/contrib_defs.cc#L2981
+
+Key points for the fusion optimization:
+* Input and scale are allowed to be of different types.
+* The normalization of the input can be done in a different precision than the input type,
+which is also the precision of reciprocal_rms returned by operation.
+* Input (x) must be: float or double or float16 or bfloat16
+* Scale must be: float or double or float16 or bfloat16
+* Normalization precision must be float or double
+"""
+
+float_types = [
+    ir.DataType.FLOAT,
+    ir.DataType.FLOAT16,
+    ir.DataType.BFLOAT16,
+    ir.DataType.DOUBLE,
+]
+fp_float_types = [ir.DataType.FLOAT, ir.DataType.DOUBLE]
+
+
+class RmsNormFusion(pattern.RewriteRuleClassBase):
+    def __init__(self, name: str, *, cast_input: bool, cast_normalized: bool):
+        """
+        Args:
+            name: Name of the rule.
+            cast_input: Whether to cast input to do the normalization in a different precision.
+            cast_normalized: Whether to cast the normalized output to the target dtype (same as scale).
+        """
+        super().__init__(name=name)
+        self._cast_input = cast_input
+        self._cast_normalized = cast_normalized
+
+    def pattern(self, op, x, scale, epsilon, compute_dtype, target_dtype):
+        if self._cast_input:
+            x = op.Cast(x, to=compute_dtype)
+        x_square = op.Pow(x, 2.0)
+        mean_square = op.ReduceMean(x_square, [-1], keepdims=1, noop_with_empty_axes=0)
+        mean_square_plus_epsilon = op.Add(mean_square, epsilon)
+        rms = op.Sqrt(mean_square_plus_epsilon)
+        reciprocal_rms = op.Reciprocal(rms)
+        normalized = op.Mul(x, reciprocal_rms)
+        if self._cast_normalized:
+            normalized = op.Cast(normalized, to=target_dtype)
+        return op.Mul(scale, normalized)
+
+    def check(self, op, x, scale, epsilon, compute_dtype, target_dtype):
+        """Check if the pattern matches conditions for use of SimplifiedLayerNormalization op."""
+        # epsilon must be a scalar
+        epsilon_value = _ir_utils.get_singleton_value(epsilon)
+        if not isinstance(epsilon_value, float):  # TODO: support other types
+            return False
+        # input and output must be same dtype
+        if x.dtype not in float_types:
+            return False
+        if scale.dtype not in float_types:
+            return False
+        stash_dtype = compute_dtype.value if self._cast_input else x.dtype
+        if stash_dtype not in fp_float_types:
+            return False
+        return True
+
+    def rewrite(self, op, x, scale, epsilon, compute_dtype, target_dtype):
+        stash_dtype = compute_dtype.value if self._cast_input else x.dtype
+        # Note: ORT's SimplifiedLayerNormalization was placed in onnx domain by mistake.
+        # No need to use com.microsoft domain here.
+        return op.SimplifiedLayerNormalization(
+            x,
+            scale,
+            axis=-1,
+            epsilon=_ir_utils.get_singleton_value(epsilon),
+            stash_type=stash_dtype,
+        )
+
+
+_rule_0 = RmsNormFusion.rule("RmsNorm-0", cast_input=True, cast_normalized=True)
+_rule_1 = RmsNormFusion.rule("RmsNorm-1", cast_input=False, cast_normalized=True)
+_rule_2 = RmsNormFusion.rule("RmsNorm-2", cast_input=True, cast_normalized=False)
+_rule_3 = RmsNormFusion.rule("RmsNorm-3", cast_input=False, cast_normalized=False)
+
+rms_normalization_rules = [_rule_0, _rule_1, _rule_2, _rule_3]
+rms_normalization_ruleset = pattern.RewriteRuleSet(rms_normalization_rules)
+
+
+def fuse_rms_normalization(model: ir.Model) -> None:
+    count = rms_normalization_ruleset.apply_to_model(model)
+    print(f"RMS Normalization count: {count}")

onnxscript/rewriter/ort_fusions/rotary_embedding.py

@@ -0,0 +1,64 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+from __future__ import annotations
+
+import onnxscript.ir as ir
+from onnxscript.rewriter import _ir_utils, pattern
+
+# Add first version of the RotaryEmbeddingFusion rule. This considers only one simple pattern
+# for full rotation without interleaving.
+# TODO(rama): Add pattern variations to handle other cases (interleaved, as well as partial rotation).
+
+# Note: This targets the new op being proposed to ONNX. This version does not exist in ORT yet.
+# so it can't be tested by running against ORT. See cos_sin_cache.py for a transformation that
+# rewrites the pattern into one that can be run against ORT.
+
+
+def _rotate_half_pattern(op, x, start1, end1, start2, end2):
+    # Slice(input, starts, ends, axes, steps)
+    x1 = op.Slice(x, start1, end1, [3], [1])
+    x2 = op.Slice(x, start2, end2, [3], [1])
+    minus_x2 = op.Neg(x2)
+    rotated_x = op.Concat(minus_x2, x1, axis=-1)
+    return rotated_x
+
+
+class RotaryEmbeddingFusion(pattern.RewriteRuleClassBase):
+    def pattern(self, op, x, cos, sin, start1, end1, start2, end2):
+        return x * cos + _rotate_half_pattern(op, x, start1, end1, start2, end2) * sin
+
+    def check(self, op, x, start1, end1, start2, end2, **_):
+        # x needs to be a 4D tensor with known last dimension size (== head_size) and known second dimension (num_heads)
+        if x is None or x.shape is None or len(x.shape) != 4:
+            return False
+        if not isinstance(x.shape[1], int):
+            return False
+        head_size = x.shape[3]
+        if not isinstance(head_size, int):
+            return False
+        half_head_size = head_size // 2
+
+        # Check that x is being split into two equal halves of size half_head_size
+        return (
+            _ir_utils.is_singleton_value(start1, 0)
+            and _ir_utils.is_singleton_value(end1, half_head_size)
+            and _ir_utils.is_singleton_value(start2, half_head_size)
+            and _ir_utils.is_singleton_value(end2, lambda x: x >= head_size)
+        )
+
+    def rewrite(self, op, x, cos, sin, **_):
+        num_heads = x.shape[1]
+        return op.RotaryEmbedding(
+            x, cos, sin, interleaved=0, num_heads=num_heads, _domain="ai.onnxruntime.fusion"
+        )
+
+
+_rule = RotaryEmbeddingFusion.rule()
+
+rotary_embedding_rules = pattern.RewriteRuleSet([_rule])
+
+
+def fuse_rotary_embedding(model: ir.Model) -> int:
+    count = rotary_embedding_rules.apply_to_model(model)
+    print(f"Rotary Embedding count: {count}")
+    return count