Repository: microsoft/onnxscript

Issue:
aten::scaled_dot_product_attention
https://github.com/pytorch/pytorch/issues/97262


Discussion:

Resolving patch:
diff --git a/onnxscript/function_libs/torch_aten/ops/core.py b/onnxscript/function_libs/torch_aten/ops/core.py
--- a/onnxscript/function_libs/torch_aten/ops/core.py
+++ b/onnxscript/function_libs/torch_aten/ops/core.py
@@ -5754,10 +5754,11 @@ def aten_triangular_solve(
     raise NotImplementedError()
 
 
-def aten_tril(self: TensorType, diagonal: int = 0) -> TensorType:
+@torch_op("aten::tril")
+def aten_tril(self: TTensor, diagonal: int = 0) -> TTensor:
     """tril(Tensor self, int diagonal=0) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.Trilu(self, diagonal, upper=0)
 
 
 def aten_tril_indices(row: int, col: int, offset: int = 0) -> TensorType:
@@ -5781,10 +5782,11 @@ def aten_triplet_margin_loss(
     raise NotImplementedError()
 
 
+@torch_op("aten::triu")
 def aten_triu(self: TensorType, diagonal: int = 0) -> TensorType:
     """triu(Tensor self, int diagonal=0) -> Tensor"""
 
-    raise NotImplementedError()
+    return op.Trilu(self, diagonal, upper=1)
 
 
 def aten_triu_indices(row: int, col: int, offset: int = 0) -> TensorType:
diff --git a/onnxscript/function_libs/torch_aten/ops/nn.py b/onnxscript/function_libs/torch_aten/ops/nn.py
--- a/onnxscript/function_libs/torch_aten/ops/nn.py
+++ b/onnxscript/function_libs/torch_aten/ops/nn.py
@@ -17,6 +17,7 @@
 import math
 from typing import Optional, Sequence
 
+import onnxscript
 from onnxscript import FLOAT, INT64
 from onnxscript.function_libs.torch_aten.registration import torch_op
 from onnxscript.function_libs.torch_aten.tensor_typing import (
@@ -25,7 +26,7 @@
     TReal,
 )
 from onnxscript.onnx_opset import opset18 as op
-from onnxscript.onnx_types import TensorType
+from onnxscript.onnx_types import BOOL, TensorType
 
 _MATH_PI = math.pi
 
@@ -1059,6 +1060,253 @@ def aten_rrelu_with_noise_backward(
     raise NotImplementedError()
 
 
+@onnxscript.script()
+def _causal_attention_mask(query: TFloat, key: TFloat) -> TFloat:
+    """Create a causal mask for the given query and key tensors.
+
+    Equivalent to::
+        mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
+        attn_mask = torch.zeros(L, S, dtype=torch.float)
+        attn_mask = attn_mask.masked_fill(not mask, -float('inf'))
+
+    Args:
+        query: Tensor of shape [..., L, E]
+        key: Tensor of shape [..., S, E]
+
+    Returns:
+        Tensor of shape [L, S]
+    """
+    target_length = op.Shape(query)[-2:-1]
+    source_length = op.Shape(key)[-2:-1]
+    # attn_mask = torch.ones(L, S) := {
+    size = op.Concat(target_length, source_length, axis=0)
+    attn_mask = op.Expand(1.0, size)
+    # }
+    attn_mask = op.Trilu(attn_mask, upper=0)
+    # The causal mask has 0s in the lower triangle and -inf in the upper triangle.
+    attn_mask = op.Where(op.Equal(attn_mask, 0.0), op.Constant(value_float=-float("inf")), 0.0)
+    return attn_mask
+
+
+@onnxscript.script()
+def _attention_scale(query: TFloat) -> TFloat:
+    """Calculate the scale factor for the attention result.
+
+    Args:
+        query: Tensor of shape [..., L, E]
+
+    Returns:
+        Scalar scale factor := 1 / math.sqrt(query.size(-1))
+    """
+    embedding_size = op.CastLike(op.Shape(query)[-1], query)
+    scale = op.Div(1.0, op.Sqrt(embedding_size))
+    return scale
+
+
+@torch_op("aten::scaled_dot_product_attention", trace_only=True)
+def aten_scaled_dot_product_attention(
+    query: TFloat,
+    key: TFloat,
+    value: TFloat,
+    attn_mask: Optional[TFloat] = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+):
+    """scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None) -> Tensor
+
+    Reference: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
+
+    Equivalent to the PyTorch code::
+        scale_factor = 1 / math.sqrt(Q.size(-1)) if scale is None else scale
+        attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) if is_causal else attn_mask
+        attn_mask = attn_mask.masked_fill(not attn_mask, -float('inf')) if attn_mask.dtype==torch.bool else attn_mask
+        attn_weight = torch.softmax((Q @ K.transpose(-2, -1) * scale_factor) + attn_mask, dim=-1)
+        attn_weight = torch.dropout(attn_weight, dropout_p)
+        return attn_weight @ V
+
+    where Q, K, V are the query, key, and value tensors, respectively.
+    L is the target sequence length, S is the source sequence length, and E is the embedding size.
+    """
+    # Use trace_only to handle optional inputs
+    assert (not is_causal) or (
+        is_causal and attn_mask is None
+    ), "is_causal and attn_mask cannot be set at the same time"
+
+    # Reference: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
+    if scale is None:
+        scale = _attention_scale(query)
+
+    if is_causal:
+        attn_mask = _causal_attention_mask(query, key)
+
+    if attn_mask is None:
+        return _aten_scaled_dot_product_attention_no_mask_onnx(
+            query, key, value, scale, dropout_p
+        )
+
+    return _aten_scaled_dot_product_attention_float_mask_onnx(
+        query, key, value, attn_mask, scale, dropout_p
+    )
+
+
+@torch_op("aten::scaled_dot_product_attention", trace_only=True, overload=True)
+def aten_scaled_dot_product_attention_bool_mask(
+    query: TFloat,
+    key: TFloat,
+    value: TFloat,
+    attn_mask: Optional[BOOL] = None,
+    dropout_p: float = 0.0,
+    is_causal: bool = False,
+    scale: Optional[float] = None,
+):
+    """scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None) -> Tensor
+
+    Reference: https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
+
+    Equivalent to the PyTorch code::
+        scale_factor = 1 / math.sqrt(Q.size(-1)) if scale is None else scale
+        attn_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) if is_causal else attn_mask
+        attn_mask = attn_mask.masked_fill(not attn_mask, -float('inf')) if attn_mask.dtype==torch.bool else attn_mask
+        attn_weight = torch.softmax((Q @ K.transpose(-2, -1) * scale_factor) + attn_mask, dim=-1)
+        attn_weight = torch.dropout(attn_weight, dropout_p)
+        return attn_weight @ V
+
+    where Q, K, V are the query, key, and value tensors, respectively.
+    L is the target sequence length, S is the source sequence length, and E is the embedding size.
+    """
+    # Use trace_only to handle optional inputs
+    assert (not is_causal) or (
+        is_causal and attn_mask is None
+    ), "is_causal and attn_mask cannot be set at the same time"
+
+    if scale is None:
+        scale = _attention_scale(query)
+
+    if is_causal:
+        attn_mask = _causal_attention_mask(query, key)
+        # The causal mask is always float
+        return _aten_scaled_dot_product_attention_float_mask_onnx(
+            query, key, value, attn_mask, scale, dropout_p
+        )
+
+    if attn_mask is None:
+        return _aten_scaled_dot_product_attention_no_mask_onnx(
+            query, key, value, scale, dropout_p
+        )
+
+    return _aten_scaled_dot_product_attention_bool_mask_onnx(
+        query, key, value, attn_mask, scale, dropout_p
+    )
+
+
+@torch_op("aten::scaled_dot_product_attention", private=True)
+def _aten_scaled_dot_product_attention_no_mask_onnx(
+    query: TFloat,
+    key: TFloat,
+    value: TFloat,
+    scale: FLOAT,
+    dropout_p: float,
+):
+    # Swap the last two axes of key
+    key_shape = op.Shape(key)
+    key_last_dim = key_shape[-1:]
+    key_second_last_dim = key_shape[-2:-1]
+    key_first_dims = key_shape[:-2]
+    # Contract the dimensions that are not the last two so we can transpose
+    # with a static permutation.
+    key_squeezed_shape = op.Concat(
+        op.Constant(value_ints=[-1]), key_second_last_dim, key_last_dim, axis=0
+    )
+    key_squeezed = op.Reshape(key, key_squeezed_shape)
+    key_squeezed_transposed = op.Transpose(key_squeezed, perm=[0, 2, 1])
+    key_transposed_shape = op.Concat(key_first_dims, key_last_dim, key_second_last_dim, axis=0)
+    key_transposed = op.Reshape(key_squeezed_transposed, key_transposed_shape)
+
+    # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653
+    # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math
+    query_scaled = op.Mul(query, op.Sqrt(scale))
+    key_transposed_scaled = op.Mul(key_transposed, op.Sqrt(scale))
+    attn_weight = op.Softmax(
+        op.MatMul(query_scaled, key_transposed_scaled),
+        axis=-1,
+    )
+    attn_weight, _ = op.Dropout(attn_weight, dropout_p)
+    return op.MatMul(attn_weight, value)
+
+
+@torch_op("aten::scaled_dot_product_attention", private=True)
+def _aten_scaled_dot_product_attention_bool_mask_onnx(
+    query: TFloat,
+    key: TFloat,
+    value: TFloat,
+    attn_mask: BOOL,
+    scale: FLOAT,
+    dropout_p: float,
+):
+    # Swap the last two axes of key
+    key_shape = op.Shape(key)
+    key_last_dim = key_shape[-1:]
+    key_second_last_dim = key_shape[-2:-1]
+    key_first_dims = key_shape[:-2]
+    # Contract the dimensions that are not the last two so we can transpose
+    # with a static permutation.
+    key_squeezed_shape = op.Concat(
+        op.Constant(value_ints=[-1]), key_second_last_dim, key_last_dim, axis=0
+    )
+    key_squeezed = op.Reshape(key, key_squeezed_shape)
+    key_squeezed_transposed = op.Transpose(key_squeezed, perm=[0, 2, 1])
+    key_transposed_shape = op.Concat(key_first_dims, key_last_dim, key_second_last_dim, axis=0)
+    key_transposed = op.Reshape(key_squeezed_transposed, key_transposed_shape)
+
+    # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653
+    # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math
+    query_scaled = op.Mul(query, op.Sqrt(scale))
+    key_transposed_scaled = op.Mul(key_transposed, op.Sqrt(scale))
+    attn_weight = op.Softmax(
+        op.Add(op.MatMul(query_scaled, key_transposed_scaled), attn_mask),
+        axis=-1,
+    )
+    attn_weight, _ = op.Dropout(attn_weight, dropout_p)
+    return op.MatMul(attn_weight, value)
+
+
+@torch_op("aten::scaled_dot_product_attention", private=True)
+def _aten_scaled_dot_product_attention_float_mask_onnx(
+    query: TFloat,
+    key: TFloat,
+    value: TFloat,
+    attn_mask: TFloat,
+    scale: TFloat,
+    dropout_p: float,
+):
+    # Swap the last two axes of key
+    key_shape = op.Shape(key)
+    key_last_dim = key_shape[-1:]
+    key_second_last_dim = key_shape[-2:-1]
+    key_first_dims = key_shape[:-2]
+    # Contract the dimensions that are not the last two so we can transpose
+    # with a static permutation.
+    key_squeezed_shape = op.Concat(
+        op.Constant(value_ints=[-1]), key_second_last_dim, key_last_dim, axis=0
+    )
+    key_squeezed = op.Reshape(key, key_squeezed_shape)
+    key_squeezed_transposed = op.Transpose(key_squeezed, perm=[0, 2, 1])
+    key_transposed_shape = op.Concat(key_first_dims, key_last_dim, key_second_last_dim, axis=0)
+    key_transposed = op.Reshape(key_squeezed_transposed, key_transposed_shape)
+
+    # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653
+    # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math
+    query_scaled = op.Mul(query, op.Sqrt(scale))
+    key_transposed_scaled = op.Mul(key_transposed, op.Sqrt(scale))
+    attn_weight = op.Softmax(
+        op.Add(op.MatMul(query_scaled, key_transposed_scaled), attn_mask),
+        axis=-1,
+    )
+    attn_weight, _ = op.Dropout(attn_weight, dropout_p)
+    return op.MatMul(attn_weight, value)
+
+
 def aten_sigmoid_backward(grad_output: TensorType, output: TensorType) -> TensorType:
     """sigmoid_backward(Tensor grad_output, Tensor output) -> Tensor"""
 


Test patch:
diff --git a/onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py b/onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py
--- a/onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py
+++ b/onnxscript/tests/function_libs/torch_aten/ops_correctness_test.py
@@ -26,6 +26,7 @@
 
 import copy
 import dataclasses
+import os
 import pprint
 import unittest
 import warnings
@@ -74,6 +75,7 @@
 )
 
 TEST_OPSET_VERSION = 18
+IS_WINDOWS = os.name == "nt"
 
 
 def dtypes_except(*dtypes: torch.dtype) -> Sequence[torch.dtype]:
@@ -538,6 +540,8 @@ def _where_input_wrangler(
     "tan": core_ops.aten_tan,
     "tanh": core_ops.aten_tanh,
     "topk": core_ops.aten_topk,
+    "tril": core_ops.aten_tril,
+    "triu": core_ops.aten_triu,
     "trunc": core_ops.aten_trunc,
     "unsqueeze": core_ops.aten_unsqueeze,
     "view": core_ops.aten_view,
@@ -573,6 +577,8 @@ def _where_input_wrangler(
     "nn.functional.conv3d": core_ops.aten_conv3d,
     "nn.functional.gelu": nn_ops.aten_gelu,
     "nn.functional.linear": nn_ops.aten_linear,
+    "nn.functional.scaled_dot_product_attention": nn_ops.aten_scaled_dot_product_attention,
+    "nn.functional.scaled_dot_product_attention_bool_mask": nn_ops.aten_scaled_dot_product_attention_bool_mask,
     "nn.functional.upsample_nearest2d": (
         nn_ops.aten_upsample_nearest2d,
         _upsample_input_wrangler,
@@ -681,6 +687,16 @@ def _where_input_wrangler(
         reason="fixme: Onnx [ShapeInferenceError] Inferred shape and existing shape differ in rank: (0) vs (1)",
         test_class_name="TestOutputConsistencyFullGraph",
     ),
+    skip(
+        "nn.functional.scaled_dot_product_attention",
+        reason="fixme: ORT crashes on Windows",
+        enabled_if=IS_WINDOWS,
+    ),
+    skip(
+        "nn.functional.scaled_dot_product_attention_bool_mask",
+        reason="fixme: ORT crashes on Windows",
+        enabled_if=IS_WINDOWS,
+    ),
     xfail(
         "nn.functional.upsample_nearest2d",
         reason="fixme: ORT fails with invalid model: 'INVALID_ARGUMENT : Failed to load model with error: vector::_M_range_check: __n (which is 1) >= this->size() (which is 1)'",
@@ -863,6 +879,28 @@ def _where_input_wrangler(
         matcher=lambda sample: "weight" not in sample.kwargs,
         reason="this Aten overload need weight as kwargs",
     ),
+    skip(
+        "nn.functional.scaled_dot_product_attention",
+        matcher=lambda sample: (attn_mask := sample.kwargs.get("attn_mask")) is not None
+        and attn_mask.dtype == torch.bool,
+        reason="this overload takes a non-boolean mask",
+    ),
+    skip(
+        "nn.functional.scaled_dot_product_attention",
+        matcher=lambda sample: sample.kwargs.get("dropout_p") != 0.0,
+        reason="dropout is random so the results do not match",
+    ),
+    skip(
+        "nn.functional.scaled_dot_product_attention_bool_mask",
+        matcher=lambda sample: (attn_mask := sample.kwargs.get("attn_mask")) is not None
+        and attn_mask.dtype != torch.bool,
+        reason="this overload takes a boolean mask",
+    ),
+    skip(
+        "nn.functional.scaled_dot_product_attention_bool_mask",
+        matcher=lambda sample: sample.kwargs.get("dropout_p") != 0.0,
+        reason="dropout is random so the results do not match",
+    ),
     skip(
         "nn.functional.upsample_nearest2d",
         # Shape should be [N, C, H, W]
@@ -920,6 +958,12 @@ def _where_input_wrangler(
 
 duplicate_opinfo(OPS_DB, "nn.functional.nll_loss", ("nn.functional.nll_loss_weight",))
 
+duplicate_opinfo(
+    OPS_DB,
+    "nn.functional.scaled_dot_product_attention",
+    ("nn.functional.scaled_dot_product_attention_bool_mask",),
+)
+
 duplicate_opinfo(
     OPS_DB,
     "min",
@@ -1143,8 +1187,14 @@ def _capture_graph_and_evaluate_torch_script_evaluator(function: Callable, args,
 
         onnx_model = onnxscript_graph.to_model_proto(TEST_OPSET_VERSION)
         # Make sure the model is valid
-        onnx.checker.check_model(onnx_model, full_check=True)
-
+        try:
+            onnx.checker.check_model(onnx_model, full_check=True)
+        except onnx.checker.ValidationError as e:
+            raise AssertionError(
+                f"ONNX model is invalid: {e}. "
+                f"Model:\n"
+                f"{onnxscript.proto2text(onnx_model)}"
+            ) from e
         # Disable all ORT optimizations
         session_options = onnxruntime.SessionOptions()
         session_options.graph_optimization_level = (
@@ -1158,6 +1208,7 @@ def _capture_graph_and_evaluate_torch_script_evaluator(function: Callable, args,
             onnxruntime.capi.onnxruntime_pybind11_state.RuntimeException,  # pylint: disable=c-extension-no-member
             onnxruntime.capi.onnxruntime_pybind11_state.InvalidArgument,  # pylint: disable=c-extension-no-member
             onnxruntime.capi.onnxruntime_pybind11_state.InvalidGraph,  # pylint: disable=c-extension-no-member
+            onnxruntime.capi.onnxruntime_pybind11_state.NotImplemented,  # pylint: disable=c-extension-no-member
         ) as e:
             raise AssertionError(
                 f"ONNX Runtime failed to evaluate:\n"
@@ -1223,7 +1274,14 @@ def run_test_output_match(
         # Provide the repr to subtest because tensors are not serializable in parallel test runs
         with test_suite.subTest(
             sample_num=i,
-            inputs=repr(inputs),
+            inputs=repr(
+                [
+                    f"Tensor<{inp.shape}, dtype={inp.dtype}>"
+                    if isinstance(inp, torch.Tensor)
+                    else inp
+                    for inp in inputs
+                ]
+            ),
             kwargs=repr(cpu_sample.kwargs),
         ):
             skip_reason = _should_skip_test_sample(op.name, cpu_sample)