Scaled dot product attention (#20492)

* Added experimental ScaledDotProductAttention operation in opset12. Supported in PT FE for aten::scaled_dot_product_attention translation. Decomposed in the common optimizations as functional reference.

* Better ScaledDotProductAttention

- Moved decomposition to the decomposing transformation
- Implemented more ctors for the op
- Renamed is_causal to causal
- Shape/type inference native code instead of using decomposition
- Moved the op from opset12 to opset13
- Added Python wrapper for ScaledDotProductAttention

* Fix test that counts ops in the opsets

* Update src/core/src/op/scaled_dot_product_attention.cpp

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>

* Update src/core/src/op/scaled_dot_product_attention.cpp

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>

* Move ScaledDotProductAttentionDecomposition from fusions to decompositions.

* Remove not used legacy shape inference in ScaledDotProductAttention

* Better namespace usage

* Register all nodes in ScaledDotProductDecomposition for correct tracking of nodes and running next mather passes on all new nodes.

* Don't use register_new_node_

* ScaledDotProductAttention specification (with an extra scale argument)

* Code style fix

* Scale input implementation for ScaledDotProductAttention

* Handle attention_mask=0 case in the op spec

* Better description of scale input

* N->M in scale description

* Code style fix, remove debug print.

* Apply suggestions from code review

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>
Co-authored-by: Mateusz Mikolajczyk <mateusz.mikolajczyk@intel.com>

* Fix for case when is_causal is not passed

* Extended description of ScaledDotProduct op

* Better description in py op wrapper

* Basic shape propagation tests for ScaledDotProductAttention

* Added ScaledDotProductAttention to toc.

* Add op impl check

---------

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>
Co-authored-by: Mateusz Mikolajczyk <mateusz.mikolajczyk@intel.com>

src/bindings/python/src/openvino/runtime/opset13/__init__.py

@@ -151,6 +151,7 @@ from openvino.runtime.opset9.ops import roi_align
 from openvino.runtime.opset2.ops import roi_pooling
 from openvino.runtime.opset7.ops import roll
 from openvino.runtime.opset5.ops import round
+from openvino.runtime.opset13.ops import scaled_dot_product_attention
 from openvino.runtime.opset12.ops import scatter_elements_update
 from openvino.runtime.opset3.ops import scatter_update
 from openvino.runtime.opset1.ops import select

src/bindings/python/src/openvino/runtime/opset13/ops.py

@@ -136,7 +136,8 @@ def multinomial(
     inputs = as_nodes(probs, num_samples)
 
     if global_seed < 0:
-        raise RuntimeError(f"global_seed should be positive or 0. Got: {global_seed}")
+        raise RuntimeError(
+            f"global_seed should be positive or 0. Got: {global_seed}")
 
     if op_seed < 0:
         raise RuntimeError(f"op_seed should be positive or 0. Got: {op_seed}")
@@ -178,7 +179,8 @@ def nms_rotated(
     :param clockwise: Flag that specifies direction of the box rotation.
     :return: The new node which performs NMSRotated
     """
-    inputs = as_nodes(boxes, scores, max_output_boxes_per_class, iou_threshold, score_threshold)
+    inputs = as_nodes(boxes, scores, max_output_boxes_per_class,
+                      iou_threshold, score_threshold)
 
     attributes = {
         "sort_result_descending": sort_result_descending,
@@ -187,3 +189,36 @@ def nms_rotated(
     }
 
     return _get_node_factory_opset13().create("NMSRotated", inputs, attributes)
+
+
+@nameable_op
+def scaled_dot_product_attention(
+    query: NodeInput,
+    key: NodeInput,
+    value: NodeInput,
+    attention_mask: Optional[NodeInput] = None,
+    scale: Optional[NodeInput] = None,
+    causal: bool = False,
+    name: Optional[str] = None,
+) -> Node:
+    """Return a node which implements Scaled Dot Product Attention.
+
+    :param query: Query tensor of shape [N, ..., L, E] and floating-point datatype.
+    :param key: Key tensor of shape [N, ..., S, E] and floating-point datatype.
+    :param value: Value tensor of shape [N, ..., S, Ev] and floating-point datatype.
+    :param attention_mask: Optional attention mask tensor of shape [N, ..., L, S] or scalar float type zero value.
+                           Refer to the operation specification for a complete description.
+    :param scale: Optional alternative scale, a floating-point type scalar.
+    :param causal: If true, then autogenerates causal attention mask instead of using attention_mask input.
+                   In this case attention_mask input is ignored.
+    :param name: The optional new name for output node.
+
+    :return: The new node performing Scaled Dot Product Attention operation.
+    """
+    inputs = as_nodes(query, key, value, attention_mask) if attention_mask is not None else as_nodes(
+        query, key, value, scale)
+
+    attributes = {
+        "causal": causal,
+    }
+    return _get_node_factory_opset13().create("ScaledDotProductAttention", inputs, attributes)