cyx-6 commented on code in PR #14150:
URL: https://github.com/apache/tvm/pull/14150#discussion_r1124896803
##########
python/tvm/relax/op/nn/nn.py:
##########
@@ -577,3 +577,41 @@ def cross_entropy_with_logits(predictions: Expr, labels:
Expr) -> Expr:
The computed result.
"""
return _ffi_api.cross_entropy_with_logits(predictions, labels) # type:
ignore
+
+
+def attention(query: Expr, key: Expr, value: Expr, bias: Optional[Expr] =
None) -> Expr:
+ r"""Computes fused multi head attention.
+
+ All input tensors are of 4-D tensors with BSNH layout.
+
+ .. math::
+ FMA(Q, K, V) = \text{Softmax}(Q @ K^T) @ V
+
+ .. note::
+ The input tensor is required to have float16 dtype
+
+ Parameters
+ ----------
+ query: relax.Expr
+ The input query to the operator. The layout of the input query should
be
+ (batch_size, seq_len, num_head, head_dim).
+
+ key: relax.Expr
+ The input key to the operator. The layout of the input key should be
+ (batch_size, seq_len_kv, num_head, head_dim).
+
+ value: relax.Expr
+ The input value to the operator. The layout of the input value should
be
+ (batch_size, seq_len_kv, num_head, head_dim_v).
+
+ bias: Optional[Expr]
+ The optional attention bias to the operator. The layout of the
attention bias should be
Review Comment:
Sure! I have updated the layout of the bias to `BNSS'`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
