diff --git a/onnxoptimizer/pass_registry.h b/onnxoptimizer/pass_registry.h index 3fdfa4978..94c06f7c6 100644 --- a/onnxoptimizer/pass_registry.h +++ b/onnxoptimizer/pass_registry.h @@ -39,6 +39,7 @@ #include "onnxoptimizer/passes/eliminate_unused_initializer.h" #include "onnxoptimizer/passes/extract_constant_to_initializer.h" #include "onnxoptimizer/passes/fuse_add_bias_into_conv.h" +#include "onnxoptimizer/passes/fuse_attention.h" #include "onnxoptimizer/passes/fuse_bn_into_conv.h" #include "onnxoptimizer/passes/fuse_concat_into_reshape.h" #include "onnxoptimizer/passes/fuse_consecutive_concats.h" @@ -94,6 +95,7 @@ struct GlobalPassRegistry { registerPass(); registerPass(); registerPass(); + registerPass(); registerPass(); registerPass(); registerPass(); diff --git a/onnxoptimizer/passes/fuse_attention.h b/onnxoptimizer/passes/fuse_attention.h new file mode 100644 index 000000000..7a64a40ed --- /dev/null +++ b/onnxoptimizer/passes/fuse_attention.h @@ -0,0 +1,258 @@ +// Copyright (c) ONNX Project Contributors +// +// SPDX-License-Identifier: Apache-2.0 + +// ATTENTION: The code in this file is highly EXPERIMENTAL. +// Adventurous users should note that the APIs will probably change. + +#pragma once + +#include "onnx/defs/tensor_util.h" +#include "onnxoptimizer/pass.h" +#include "onnxoptimizer/passes/pass_util.h" + +namespace ONNX_NAMESPACE { +namespace optimization { + +// Fuse the sub-graph that `torch.onnx.export` (the TorchScript exporter) +// produces for `torch.nn.functional.scaled_dot_product_attention` into a +// single ONNX `Attention` operator (opset 23+). +// +// The exported pattern (working forward) looks like: +// +// key_t = Transpose(K, perm=[..., -1, -2]) // swap last two axes +// q_s = Mul(Q, scale_q) // scale_q == sqrt(scale) +// k_s = Mul(key_t, scale_k) // scale_k == sqrt(scale) +// qk = MatMul(q_s, k_s) +// biased = Add(qk, attn_mask) // optional (mask / causal) +// weights = Softmax(biased, axis=-1) +// out = MatMul(weights, V) +// +// `scale_q` and `scale_k` are either +// * two constant scalars (when an explicit `scale=` was passed), or +// * two `Sqrt` nodes fed from the default `1 / sqrt(head_size)` sub-graph +// that torch derives from `Shape(Q)` (when `scale` was left to default). +// +// Because `Attention` internally computes `softmax(scale * (Q @ K^T) + mask) @ +// V` with a default `scale` of `1 / sqrt(head_size)`, the whole sub-graph is +// equivalent to `Attention(Q, K, V, attn_mask)` where the `scale` attribute is +// set to `scale_q * scale_k` for the explicit-scale case and left unset (i.e. +// the default) for the derived-scale case. +struct FuseAttention final : public PredicateBasedPass { + explicit FuseAttention() + : PredicateBasedPass(PassType::Fuse, PassEfficiency::Complete, + PassOptimizationType::Compute) {} + + std::string getPassName() const override { + return "fuse_attention"; + } + + // The anchor is the final `MatMul(Softmax(...), V)`. + bool patternMatchPredicate(Node* node) override { + return CheckKind(node, kMatMul) && CheckKind(node->input(0), kSoftmax); + } + + // Returns true when `axis` refers to the last dimension of a `rank`-D tensor. + static bool IsLastAxis(int64_t axis, int64_t rank) { + if (rank <= 0) { + // Rank is unknown: only the canonical `-1` is accepted. + return axis == -1; + } + if (axis < 0) { + axis += rank; + } + return axis == rank - 1; + } + + // Checks that `perm` swaps the last two axes and leaves the rest untouched, + // e.g. [0, 1, 3, 2] for a rank-4 tensor. + static bool IsLastTwoAxesSwap(const std::vector& perm) { + const int64_t rank = static_cast(perm.size()); + if (rank < 2) { + return false; + } + for (int64_t i = 0; i < rank - 2; ++i) { + if (perm[i] != i) { + return false; + } + } + return perm[rank - 2] == rank - 1 && perm[rank - 1] == rank - 2; + } + + // Recognizes the default `1 / sqrt(head_size)` scale sub-graph that torch + // derives from the last dimension of `q`: + // Cast(Div(Constant(1), Sqrt(Cast(Slice(Shape(q), starts=[-1], ...))))) + // The leading/trailing `Cast`s are optional (they only appear for some + // dtypes). `q` must be the query fed into the fused attention. + static bool IsDefaultScaleSource(const Value* x, const Value* q) { + const Node* n = x->node(); + if (CheckKind(n, kCast)) { + n = n->input(0)->node(); + } + if (!CheckKind(n, kDiv)) { + return false; + } + float numerator; + if (!FetchSoleValueOfTensor(n->input(0), numerator) || numerator != 1.0f) { + return false; + } + const Node* sqrt = n->input(1)->node(); + if (!CheckKind(sqrt, kSqrt)) { + return false; + } + const Node* dim = sqrt->input(0)->node(); + if (CheckKind(dim, kCast)) { + dim = dim->input(0)->node(); + } + if (!CheckKind(dim, kSlice)) { + return false; + } + // The slice must select the last dimension of the shape. + int64_t start; + if (dim->inputs().size() < 2 || + !FetchSoleIntValueOfTensor(dim->input(1), start) || start != -1) { + return false; + } + const Node* shape = dim->input(0)->node(); + return CheckKind(shape, "Shape") && shape->input(0) == q; + } + + bool runTransform(Node* n, Graph& graph, + NodeDestroyType& destroy_current) override { + destroy_current = NodeDestroyType::DestroyZero; + + // `Attention` was introduced in opset 23; refuse to emit it otherwise. + if (getOpsetVersion(graph) < 23) { + return false; + } + + // out = MatMul(weights, V) + Node* softmax = n->input(0)->node(); + if (n->input(0)->uses().size() != 1) { + return false; + } + Value* value = n->input(1); + + // weights = Softmax(biased, axis=-1) + { + int64_t rank = + softmax->input(0)->has_sizes() + ? static_cast(softmax->input(0)->sizes().size()) + : -1; + int64_t axis = GetValueFromAttrWithDefault(softmax, kaxis, (int64_t)-1); + if (!IsLastAxis(axis, rank)) { + return false; + } + } + if (softmax->input(0)->uses().size() != 1) { + return false; + } + + // biased = Add(qk, attn_mask) (optional) + Node* score = softmax->input(0)->node(); + Value* attn_mask = nullptr; + Node* qk = nullptr; + if (CheckKind(score, kMatMul)) { + qk = score; + } else if (CheckKind(score, kAdd)) { + for (size_t i = 0; i < 2; ++i) { + if (CheckKind(score->input(i), kMatMul)) { + qk = score->input(i)->node(); + attn_mask = score->input(1 - i); + break; + } + } + if (qk == nullptr || score->output()->uses().size() != 1 || + qk->output()->uses().size() != 1) { + return false; + } + } else { + return false; + } + + // qk = MatMul(q_scaled, k_scaled) + Node* mul_q = qk->input(0)->node(); + Node* mul_k = qk->input(1)->node(); + if (!CheckKind(mul_q, kMul) || !CheckKind(mul_k, kMul) || + qk->input(0)->uses().size() != 1 || qk->input(1)->uses().size() != 1) { + return false; + } + + // k_scaled = Mul(Transpose(K), scale_k) + Node* transpose = nullptr; + Value* scale_k = nullptr; + for (size_t i = 0; i < 2; ++i) { + if (CheckKind(mul_k->input(i), kTranspose)) { + transpose = mul_k->input(i)->node(); + scale_k = mul_k->input(1 - i); + break; + } + } + if (transpose == nullptr || transpose->output()->uses().size() != 1) { + return false; + } + std::vector perm; + if (!GetValueFromAttr(transpose, kperm, perm) || !IsLastTwoAxesSwap(perm)) { + return false; + } + Value* key = transpose->input(0); + + // q_scaled = Mul(Q, scale_q): the scale operand is a constant scalar or a + // Sqrt node, the other operand is the query. + auto is_scale_operand = [](const Value* v) { + float tmp; + return CheckKind(v, kSqrt) || FetchSoleValueOfTensor(v, tmp); + }; + Value* query = nullptr; + Value* scale_q = nullptr; + if (is_scale_operand(mul_q->input(1))) { + query = mul_q->input(0); + scale_q = mul_q->input(1); + } else if (is_scale_operand(mul_q->input(0))) { + query = mul_q->input(1); + scale_q = mul_q->input(0); + } else { + return false; + } + + // Resolve the effective scale applied to Q @ K^T. + bool has_scale = false; + float scale = 0.0f; + float scale_q_val; + float scale_k_val; + if (FetchSoleValueOfTensor(scale_q, scale_q_val) && + FetchSoleValueOfTensor(scale_k, scale_k_val)) { + // Explicit scale: torch emits sqrt(scale) constants on both branches. + has_scale = true; + scale = scale_q_val * scale_k_val; + } else if (CheckKind(scale_q, kSqrt) && CheckKind(scale_k, kSqrt) && + scale_q->node()->input(0) == scale_k->node()->input(0) && + IsDefaultScaleSource(scale_q->node()->input(0), query)) { + // Default scale: 1 / sqrt(head_size), which is `Attention`'s default. + has_scale = false; + } else { + return false; + } + + Node* attention = graph.create("Attention"_sym, 1); + attention->addInput(query); + attention->addInput(key); + attention->addInput(value); + if (attn_mask != nullptr) { + attention->addInput(attn_mask); + } + if (has_scale) { + attention->f_(kscale, scale); + } + attention->insertBefore(n); + + if (!tryReplacingAllUsesWith(n->output(), attention->output())) { + return false; + } + destroy_current = NodeDestroyType::DestroyOne; + return true; + } +}; + +} // namespace optimization +} // namespace ONNX_NAMESPACE diff --git a/onnxoptimizer/test/optimizer_test.py b/onnxoptimizer/test/optimizer_test.py index 2ee7d3003..ba394c47e 100644 --- a/onnxoptimizer/test/optimizer_test.py +++ b/onnxoptimizer/test/optimizer_test.py @@ -18,6 +18,13 @@ except ImportError: has_tv = False +try: + import torch + + has_torch = True +except ImportError: + has_torch = False + import onnx import pytest from onnx import ( @@ -4938,6 +4945,222 @@ def test_fuse_qkv(self): # type: () -> None self._test_fuse_qkv_with_opset(opset_version) self._test_fuse_qkv_with_opset(LATEST_STABLE_OPSET_VERSION) + # ---- fuse_attention ------------------------------------------------- + # The graphs built here mirror what `torch.onnx.export` (the TorchScript + # exporter) produces for `scaled_dot_product_attention`: Q and the + # transposed K are each scaled by `sqrt(scale)`, multiplied together, + # optionally biased by an attention mask, passed through Softmax and + # finally multiplied by V. + ATTENTION_OPSET_VERSION = 23 + + def _make_sdpa_model(self, mask="none", scale=None, opset=None): + # type: (str, Optional[float], Optional[int]) -> onnx.ModelProto + if opset is None: + opset = self.ATTENTION_OPSET_VERSION + B, H, S, D = 2, 4, 8, 16 + INT64_MAX = 9223372036854775807 + + def vi(name, shape): + return helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + + nodes = [] + inits = [] + inputs = [vi("q", [B, H, S, D]), vi("k", [B, H, S, D]), vi("v", [B, H, S, D])] + + if scale is None: + # Default scale 1/sqrt(head_size), derived from Shape(q) exactly as + # torch emits it. + inits += [ + numpy_helper.from_array(np.array([-1], np.int64), "neg1"), + numpy_helper.from_array(np.array([INT64_MAX], np.int64), "imax"), + numpy_helper.from_array(np.array(1.0, np.float32), "one"), + ] + nodes += [ + helper.make_node("Shape", ["q"], ["shp"]), + helper.make_node("Slice", ["shp", "neg1", "imax"], ["lastdim"]), + helper.make_node("Cast", ["lastdim"], ["lastf"], to=TensorProto.FLOAT), + helper.make_node("Sqrt", ["lastf"], ["sqrt_hs"]), + helper.make_node("Div", ["one", "sqrt_hs"], ["scale_val"]), + helper.make_node("Cast", ["scale_val"], ["scale_c"], to=TensorProto.FLOAT), + helper.make_node("Sqrt", ["scale_c"], ["sqrt_q"]), + helper.make_node("Sqrt", ["scale_c"], ["sqrt_k"]), + ] + sq, sk = "sqrt_q", "sqrt_k" + else: + sf = float(np.sqrt(scale)) + inits += [ + numpy_helper.from_array(np.array(sf, np.float32), "sqf_q"), + numpy_helper.from_array(np.array(sf, np.float32), "sqf_k"), + ] + sq, sk = "sqf_q", "sqf_k" + + nodes += [ + helper.make_node("Transpose", ["k"], ["kt"], perm=[0, 1, 3, 2]), + helper.make_node("Mul", ["q", sq], ["mq"]), + helper.make_node("Mul", ["kt", sk], ["mk"]), + helper.make_node("MatMul", ["mq", "mk"], ["qk"]), + ] + + if mask == "none": + score = "qk" + elif mask == "additive": + inputs.append(vi("mask", [B, H, S, S])) + nodes.append(helper.make_node("Add", ["qk", "mask"], ["score"])) + score = "score" + elif mask == "causal": + causal = np.triu(np.full((S, S), -np.inf, np.float32), 1) + inits.append(numpy_helper.from_array(causal, "causal")) + nodes.append(helper.make_node("Add", ["qk", "causal"], ["score"])) + score = "score" + else: + raise ValueError(mask) + + nodes += [ + helper.make_node("Softmax", [score], ["weights"], axis=-1), + helper.make_node("MatMul", ["weights", "v"], ["out"]), + ] + graph = helper.make_graph( + nodes, "sdpa", inputs, [vi("out", [B, H, S, D])], inits + ) + return helper.make_model( + graph, + producer_name="onnx-test", + opset_imports=[helper.make_opsetid("", opset)], + ir_version=10, + ) + + def test_fuse_attention_default_scale(self): # type: () -> None + model = self._make_sdpa_model(mask="none", scale=None) + optimized_model = self._optimized( + model, ["fuse_attention", "eliminate_deadend", "eliminate_unused_initializer"] + ) + assert len(optimized_model.graph.node) == 1 + node = optimized_model.graph.node[0] + assert node.op_type == "Attention" + assert list(node.input) == ["q", "k", "v"] + # Default scale must be left implicit. + assert all(a.name != "scale" for a in node.attribute) + + def test_fuse_attention_explicit_scale(self): # type: () -> None + model = self._make_sdpa_model(mask="none", scale=0.125) + optimized_model = self._optimized( + model, ["fuse_attention", "eliminate_deadend", "eliminate_unused_initializer"] + ) + assert len(optimized_model.graph.node) == 1 + node = optimized_model.graph.node[0] + assert node.op_type == "Attention" + scale_attr = [a for a in node.attribute if a.name == "scale"] + assert len(scale_attr) == 1 + np.testing.assert_allclose(scale_attr[0].f, 0.125, rtol=1e-6) + + def test_fuse_attention_additive_mask(self): # type: () -> None + model = self._make_sdpa_model(mask="additive", scale=None) + optimized_model = self._optimized( + model, ["fuse_attention", "eliminate_deadend", "eliminate_unused_initializer"] + ) + assert len(optimized_model.graph.node) == 1 + node = optimized_model.graph.node[0] + assert node.op_type == "Attention" + assert list(node.input) == ["q", "k", "v", "mask"] + + def test_fuse_attention_causal_mask(self): # type: () -> None + model = self._make_sdpa_model(mask="causal", scale=None) + optimized_model = self._optimized( + model, ["fuse_attention", "eliminate_deadend", "eliminate_unused_initializer"] + ) + # The (constant) causal mask is folded into a single initializer that + # feeds the fused Attention node as attn_mask. + attention_nodes = [n for n in optimized_model.graph.node if n.op_type == "Attention"] + assert len(attention_nodes) == 1 + assert len(attention_nodes[0].input) == 4 + + def test_fuse_attention_opset_too_low_no_fuse(self): # type: () -> None + # Attention requires opset >= 23; the pass must be a no-op otherwise. + model = self._make_sdpa_model(mask="none", scale=0.125, opset=17) + optimized_model = self._optimized(model, ["fuse_attention"]) + assert all(n.op_type != "Attention" for n in optimized_model.graph.node) + + def test_fuse_attention_plain_softmax_no_fuse(self): # type: () -> None + # A bare Softmax -> MatMul (no scaled Q/K MatMul) must not be fused. + B, H, S, D = 2, 4, 8, 16 + + def vi(name, shape): + return helper.make_tensor_value_info(name, TensorProto.FLOAT, shape) + + graph = helper.make_graph( + [ + helper.make_node("Softmax", ["q"], ["s"], axis=-1), + helper.make_node("MatMul", ["s", "v"], ["out"]), + ], + "plain", + [vi("q", [B, H, S, D]), vi("v", [B, H, D, S])], + [vi("out", [B, H, S, S])], + ) + model = helper.make_model( + graph, + producer_name="onnx-test", + opset_imports=[helper.make_opsetid("", self.ATTENTION_OPSET_VERSION)], + ir_version=10, + ) + optimized_model = self._optimized(model, ["fuse_attention"]) + assert all(n.op_type != "Attention" for n in optimized_model.graph.node) + + @unittest.skipUnless(has_torch, "onnx test needs torch") + def test_fuse_attention_torch_sdpa_export(self): # type: () -> None + # Fuse graphs produced by the real torch.onnx exporter. + import torch.onnx._constants as _onnx_constants + from torch.nn.functional import scaled_dot_product_attention + + if _onnx_constants.ONNX_MAX_OPSET < self.ATTENTION_OPSET_VERSION: + self.skipTest("installed torch cannot export opset 23") + + B, H, S, D = 2, 4, 8, 16 + + class SDPA(torch.nn.Module): + def __init__(self, **kwargs): + super().__init__() + self._kwargs = kwargs + + def forward(self, *args): + if len(args) == 4: + return scaled_dot_product_attention( + args[0], args[1], args[2], attn_mask=args[3], **self._kwargs + ) + return scaled_dot_product_attention( + args[0], args[1], args[2], **self._kwargs + ) + + q = torch.randn(B, H, S, D) + k = torch.randn(B, H, S, D) + v = torch.randn(B, H, S, D) + mask = torch.randn(B, H, S, S) + + cases = [ + (SDPA(), (q, k, v), ["q", "k", "v"]), + (SDPA(scale=0.125), (q, k, v), ["q", "k", "v"]), + (SDPA(is_causal=True), (q, k, v), ["q", "k", "v"]), + (SDPA(), (q, k, v, mask), ["q", "k", "v", "mask"]), + ] + for module, args, names in cases: + buffer = io.BytesIO() + torch.onnx.export( + module, + args, + buffer, + opset_version=self.ATTENTION_OPSET_VERSION, + dynamo=False, + input_names=names, + ) + model = onnx.load_from_string(buffer.getvalue()) + optimized_model = self._optimized( + model, + ["fuse_attention", "eliminate_deadend", "eliminate_unused_initializer"], + ) + attention_nodes = [ + n for n in optimized_model.graph.node if n.op_type == "Attention" + ] + assert len(attention_nodes) == 1, [n.op_type for n in optimized_model.graph.node] + def test_fuse_consecutive_unsqueezes_opset13(self): # type: () -> None graph = parser.parse_graph(""" agraph (float[4, 64, 160, 160] X) => (float[1, 1, 1, 4, 64, 1, 160, 160, 1, 1, 1] Z)