Add fused rope pass (pytorch#22918)

Summary:

FuseRoPEPass folds the HuggingFace rotate_half rotary-embedding apply into a
single native::rope op. This diff adds the op as well as the pass.

custom_ops.py registers `native::rope(Tensor input, Tensor cos, Tensor sin, bool
interleaved=False) -> Tensor`, with a fake (meta) impl used by AOT
export/serialization and a CompositeExplicitAutograd reference body used for
eager execution. Registration is an import side effect, so the module is imported
from the backend entry points to guarantee the op exists before lowering.

HF apply_rotary_pos_emb lowers each rotated tensor to
`x*cos + rotate_half(x)*sin`, where `rotate_half(x) = cat([-x2, x1], -1)` splits
the last dim in half. The pass matches that ~7-op chain (an add over two muls,
over cat / neg / slice_copy) and rewrites it to one
`native::rope(x, cos, sin, interleaved=False)`.

Partial rotary is folded too: when only the leading rotary_dim channels are
rotated and the apply result is concatenated with the pass-through
`x[..., rotary_dim:]`, the whole slice/apply/cat is replaced by a single
native::rope over the full tensor. The op rotates the first `cos.shape[-1]`
channels and passes the remainder through, so the rope input stays contiguous
with no leftover strided slice.

cos/sin are captured as the already-scaled tensors feeding the multiplies, so any
rope variant (default, YaRN, long-rope) that ends in this apply is handled
without the pass modelling frequency construction. Only the split-half layout is
matched; the interleaved (GPT-NeoX "traditional") layout is left decomposed,
though the op implements it for completeness. Every intermediate node must be
single-use so the decomposition DCEs away after the rewrite.

The pass is added to get_default_passes(), so it runs pre-partition on the
edge-dialect graph, and the partitioner claims native::rope via
_SUPPORTED_NON_CORE_OPS.

Reviewed By: digantdesai

Differential Revision: D111753295

Pull Request resolved: pytorch#22918

Author: metascroy

backends/native/BUCK

@@ -15,8 +15,11 @@ fbcode_target(
     name = "lib",
     srcs = [
         "__init__.py",
+        "custom_ops.py",
         "partitioner.py",
         "passes/__init__.py",
+        "passes/_utils.py",
+        "passes/fuse_rope.py",
         "passes/reinplace.py",
         "passes/replace_copy_with_alias.py",
         "preprocess.py",

backends/native/custom_ops.py

@@ -0,0 +1,107 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+"""
+Custom ops for the native backend.
+
+Registers ``native::rope``, the fused rotary position embedding op produced by
+``FuseRoPEPass``. Registration happens as an import side effect, so this module is
+imported from the backend's entry points (``passes`` / ``partitioner.py``) to
+guarantee the op exists before lowering. The reference ``CompositeExplicitAutograd``
+body is only used for eager execution; AOT export/serialization uses the fake
+(meta) impl.
+"""
+
+import torch
+from torch import Tensor
+
+lib: torch.library.Library = torch.library.Library("native", "DEF")
+
+
+lib.define(
+    "rope(Tensor input, Tensor position_ids, Tensor inv_freq, "
+    "bool interleaved=False, float attention_scale=1.0) -> Tensor"
+)
+
+
+def _check_rope_shapes(input: Tensor, position_ids: Tensor, inv_freq: Tensor) -> None:
+    torch._check(input.is_floating_point(), lambda: "rope input must be floating point")
+    torch._check(input.dim() == 4, lambda: "rope input must have shape [B, H, T, D]")
+    torch._check(
+        position_ids.dim() == 2, lambda: "rope position_ids must have shape [Bpos, T]"
+    )
+    torch._check(inv_freq.dim() == 1, lambda: "rope inv_freq must have shape [R / 2]")
+    torch._check(
+        (position_ids.shape[0] == 1) | (position_ids.shape[0] == input.shape[0]),
+        lambda: "rope position batch must be 1 or match the input batch",
+    )
+    torch._check(
+        position_ids.shape[1] == input.shape[2],
+        lambda: "rope position sequence length must match the input",
+    )
+    torch._check(
+        2 * inv_freq.shape[0] <= input.shape[3],
+        lambda: "rope rotary width must not exceed the input width",
+    )
+
+
+@torch.library.register_fake("native::rope", lib=lib)
+def _rope_fake(
+    input: Tensor,
+    position_ids: Tensor,
+    inv_freq: Tensor,
+    interleaved: bool = False,
+    attention_scale: float = 1.0,
+) -> Tensor:
+    _check_rope_shapes(input, position_ids, inv_freq)
+    return input.new_empty(input.shape, dtype=input.dtype)
+
+
+@torch.library.impl("native::rope", "CompositeExplicitAutograd", lib=lib)
+def _rope_impl(
+    input: Tensor,
+    position_ids: Tensor,
+    inv_freq: Tensor,
+    interleaved: bool = False,
+    attention_scale: float = 1.0,
+) -> Tensor:
+    """Rotate the leading R channels of [B, H, T, D], preserving the tail.
+
+    Positions [Bpos, T] broadcast over heads and, when Bpos is 1, batches.
+    inv_freq [R / 2] determines the rotary width. Phase, trig and attention
+    scaling use FP32; the tables are then cast to the input dtype.
+    """
+    _check_rope_shapes(input, position_ids, inv_freq)
+    # Elementwise multiplication keeps phase construction in FP32 under autocast,
+    # unlike the equivalent batched matmul used by HF's source pattern.
+    phase = position_ids.float().unsqueeze(-1) * inv_freq.float()
+    if interleaved:
+        phase = phase.repeat_interleave(2, dim=-1)
+    else:
+        phase = torch.cat((phase, phase), dim=-1)
+    cos = (phase.cos() * attention_scale).to(input.dtype).unsqueeze(1)
+    sin = (phase.sin() * attention_scale).to(input.dtype).unsqueeze(1)
+    rotary = 2 * inv_freq.shape[0]
+    x_rot = input[..., :rotary]
+    x_pass = input[..., rotary:]
+    if interleaved:
+        x1 = x_rot[..., 0::2]
+        x2 = x_rot[..., 1::2]
+        rotated = torch.stack((-x2, x1), dim=-1).flatten(-2)
+    else:
+        half = rotary // 2
+        x1 = x_rot[..., :half]
+        x2 = x_rot[..., half:]
+        rotated = torch.cat((-x2, x1), dim=-1)
+    out = x_rot * cos + rotated * sin
+    if x_pass.shape[-1] == 0:
+        return out
+    return torch.cat((out, x_pass), dim=-1)
+
+
+rope_op: torch._ops.OpOverload = torch.ops.native.rope.default

backends/native/partitioner.py

@@ -18,6 +18,7 @@
 # Registers torch.ops.torchao.dequantize_gguf, referenced in _SUPPORTED_NON_CORE_OPS.
 import executorch.extension.llm.export.gguf  # noqa: F401
 import torch
+from executorch.backends.native.custom_ops import rope_op
 from executorch.backends.native.passes import backend_inplace_aten_variants
 
 from executorch.exir.backend.compile_spec_schema import CompileSpec
@@ -43,6 +44,7 @@
     # GGUF weight dequantize stays in the delegate; the serializer folds it into a
     # PackedQuant weight on the consuming op.
     torch.ops.torchao.dequantize_gguf.default,
+    rope_op,
     torch.ops.aten.rms_norm.default,
 ]
 

backends/native/passes/__init__.py

@@ -29,6 +29,7 @@
 
 from typing import List, Union
 
+from executorch.backends.native.passes.fuse_rope import FuseRoPEPass
 from executorch.backends.native.passes.reinplace import (
     backend_inplace_aten_variants,
     BACKEND_INPLACE_OPS,
@@ -54,6 +55,7 @@
     "CollapseViewCopyPass",
     "FuseGQAWithSDPAPass",
     "FuseRMSNormPass",
+    "FuseRoPEPass",
     "get_default_passes",
     "NativeReinplacePass",
     "NormalizeSDPAInputRankPass",
@@ -69,6 +71,7 @@ def get_default_passes() -> List[Union[ExportPass, ExportedProgramPassBase]]:
     their in-place edge forms.
     """
     return [
+        FuseRoPEPass(),
         FuseGQAWithSDPAPass(),
         NormalizeSDPAInputRankPass(),
         FuseRMSNormPass(fold_dtype_casts=True, allow_lossy_weight_casts=True),

backends/native/passes/_utils.py

@@ -0,0 +1,34 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Shared helpers for native backend graph passes."""
+
+import torch
+from torch.fx import Node
+
+
+def _resolve_aten(target: object) -> "torch._ops.OpOverload | None":
+    """Return the underlying aten OpOverload for a node target, or None.
+
+    Handles both edge-dialect ops (EdgeOpOverload, whose ``_op`` is the aten op)
+    and plain aten OpOverloads (whose ``_op`` is the C++ builtin, so it must not
+    be unwrapped). See graph_serialize._resolve_op_overload for the same logic.
+    """
+    inner = getattr(target, "_op", None)
+    if isinstance(inner, torch._ops.OpOverload):
+        return inner
+    if isinstance(target, torch._ops.OpOverload):
+        return target
+    return None
+
+
+def _single_user(node: object) -> bool:
+    return isinstance(node, Node) and len(node.users) == 1
+
+
+def _fake_tensor(node: object) -> "torch.Tensor | None":
+    val = node.meta.get("val") if isinstance(node, Node) else None
+    return val if isinstance(val, torch.Tensor) else None