[Native] Fuse broadcast RoPE frequencies and lower tanh in ggml

backends/native/passes/fuse_rope.py

@@ -333,8 +333,31 @@ def _trig_source(
     return phase, scale
 
 
+def _broadcast_frequency_product(product: Node) -> "tuple[Node, Node] | None":
+    """Match ``position_ids[..., None].float() * inv_freq``, the outer product
+    some HF models (e.g. Muse Glimmer) build without a matmul."""
+    t = _fake_tensor(product)
+    if t is None or t.dtype != torch.float32 or t.dim() != 3:
+        return None
+    a, b = product.args[:2]
+    for pos, inv in ((a, b), (b, a)):
+        if not isinstance(pos, Node) or not isinstance(inv, Node):
+            continue
+        pos = _float_source(pos)
+        position_ids = _unsqueeze_source(pos, 2) if pos is not None else None
+        inv_freq = _float_source(inv)
+        if (
+            position_ids is not None
+            and _ndim(position_ids) == 2
+            and inv_freq is not None
+            and _ndim(inv_freq) == 1
+        ):
+            return position_ids, inv_freq
+    return None
+
+
 def _frequency_product(phase: Node) -> "tuple[Node, Node] | None":  # noqa: C901
-    """Recover positions/inverse frequencies from HF's duplicated batched outer product."""
+    """Recover positions/inverse frequencies from HF's duplicated outer product."""
     if not _is_op(phase, torch.ops.aten.cat.default):
         return None
     tensors = phase.args[0]
@@ -346,6 +369,8 @@ def _frequency_product(phase: Node) -> "tuple[Node, Node] | None":  # noqa: C901
         return None
     if (phase.args[1] if len(phase.args) > 1 else 0) not in (-1, 2):
         return None
+    if _is_op(tensors[0], torch.ops.aten.mul.Tensor):
+        return _broadcast_frequency_product(tensors[0])
     transposed = tensors[0]
     if _is_op(
         transposed, torch.ops.aten.permute_copy.default, torch.ops.aten.permute.default

backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp

@@ -688,6 +688,8 @@ void GraphBuilder::lower(const Node& node) {
     lower_unary(node, ggml_log);
   } else if (key == "aten.neg.default") {
     lower_unary(node, ggml_neg);
+  } else if (key == "aten.tanh.default") {
+    lower_unary(node, ggml_tanh);
   } else if (key == "aten.relu.default" || key == "aten.relu_.default") {
     lower_unary(node, ggml_relu);
   } else if (

backends/native/test/runtime/ggml/test_ggml_ops.py

@@ -620,6 +620,9 @@ def test_log(self):
     def test_neg(self):
         self._check(Lambda(torch.neg), (torch.randn(2, 3, 4),))
 
+    def test_tanh(self):
+        self._check(Lambda(torch.tanh), (torch.randn(2, 3, 4) * 3,))
+
     def test_relu(self):
         self._check(Lambda(F.relu), (torch.randn(2, 3, 4),))
 

backends/native/test/test_passes.py

@@ -302,6 +302,30 @@ def test_hf_rope_fusion_preserves_numerics(self):
                     actual[..., rotary:], x[..., rotary:], rtol=0, atol=0
                 )
 
+    def test_broadcast_frequency_product_fuses(self):
+        # Muse Glimmer builds the angles as position_ids[..., None] * inv_freq.
+        class BroadcastRoPE(nn.Module):
+            def forward(self, x, position_ids, inv_freq):
+                freqs = position_ids[..., None].float() * inv_freq.float()
+                emb = torch.cat((freqs, freqs), dim=-1)
+                cos = emb.cos().to(x.dtype).unsqueeze(1)
+                sin = emb.sin().to(x.dtype).unsqueeze(1)
+                return _apply_rope_tables(x, cos, sin)
+
+        torch.manual_seed(0)
+        for dtype, rotary in itertools.product((torch.float32, torch.float16), (4, 8)):
+            with self.subTest(dtype=dtype, rotary=rotary):
+                inputs = (
+                    torch.randn(1, 3, 4, 8, dtype=dtype),
+                    torch.arange(4).reshape(1, 4) + 2049,
+                    torch.linspace(0.01, 0.9, rotary // 2),
+                )
+                model = BroadcastRoPE().eval()
+                ep = _transformed(model, inputs, [FuseRoPEPass()])
+                targets = [n.target for n in ep.graph.nodes if n.op == "call_function"]
+                self.assertEqual(targets, [torch.ops.native.rope.default])
+                torch.testing.assert_close(ep.module()(*inputs), model(*inputs))
+
     def test_hf_rope_serializes_five_arguments(self):
         for rotary, scale in ((8, 1.0), (4, 1.25)):
             with self.subTest(rotary=rotary, scale=scale):