diff --git a/backends/native/passes/fuse_rope.py b/backends/native/passes/fuse_rope.py index 8e7c670d25d..1215486ed15 100644 --- a/backends/native/passes/fuse_rope.py +++ b/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 diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp index c5522173893..040cbfb0c22 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp +++ b/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 ( diff --git a/backends/native/test/runtime/ggml/test_ggml_ops.py b/backends/native/test/runtime/ggml/test_ggml_ops.py index 2a7169ffecd..dd22a735767 100644 --- a/backends/native/test/runtime/ggml/test_ggml_ops.py +++ b/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),)) diff --git a/backends/native/test/test_passes.py b/backends/native/test/test_passes.py index 4c5b81b3d90..6806f6b8ee3 100644 --- a/backends/native/test/test_passes.py +++ b/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):