From ddb4ff29d5a24760d8e8b202a2ee23a940efc2eb Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 11 Nov 2026 13:47:38 -0800 Subject: [PATCH] Fuse broadcast RoPE frequencies and lower tanh in the native ggml engine Muse Glimmer's Hugging Face rotary embedding builds its angles as position_ids[..., None] * inv_freq rather than HF's batched matmul, so FuseRoPEPass left the cos/sin tables in the graph, which the ggml engine cannot run. The pass now recognizes that outer product too. The ggml engine also lowers aten.tanh, which Muse Glimmer's final logit soft cap uses. --- backends/native/passes/fuse_rope.py | 27 ++++++++++++++++++- .../runtime/engine/ggml/GgmlGraphBuilder.cpp | 2 ++ .../native/test/runtime/ggml/test_ggml_ops.py | 3 +++ backends/native/test/test_passes.py | 24 +++++++++++++++++ 4 files changed, 55 insertions(+), 1 deletion(-) 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):