This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 17be2d71ad [Unity] Support causal mask for `R.nn.attention` (#14907)
17be2d71ad is described below
commit 17be2d71ad6052f085e22739235cd60c109337b3
Author: Yaxing Cai <[email protected]>
AuthorDate: Tue May 23 18:07:46 2023 -0700
[Unity] Support causal mask for `R.nn.attention` (#14907)
* [Unity]Support causal mask for `R.nn.attention`
This PR introduces the causal mask for `R.nn.attention` and its cutlass
dispatch and legalization. The `causal_mask` argument accepts 2 types of causal
masks - "TopLeft" and "BottomRight".
For example, with `seq_len = 4`, `seq_len_kv = 2`,
mask for "TopLeft":
[[1, 0],
[1, 1],
[1, 1],
[1, 1]]
mask for "BottomRight":
[[1, 1],
[1, 1],
[1, 1],
[1, 1]]
with `seq_len = 2`, `seq_len_kv = 4`,
mask for "TopLeft":
[[1, 0, 0, 0],
[1, 1, 0, 0]]
mask for "BottomRight":
[[1, 1, 1, 0],
[1, 1, 1, 1]]
* update fx_translator and legalize_ops unit test
* apply code review suggestions
---
include/tvm/relax/attrs/nn.h | 3 +
python/tvm/contrib/cutlass/attention_operation.py | 1 +
python/tvm/contrib/cutlass/build.py | 10 ++
python/tvm/contrib/cutlass/gen_tensor_op.py | 2 +
python/tvm/relax/frontend/torch/fx_translator.py | 9 +-
python/tvm/relax/op/nn/nn.py | 28 +++++-
python/tvm/relax/transform/legalize_ops/nn.py | 26 ++++-
src/relax/op/nn/attention.cc | 7 +-
src/relax/op/nn/attention.h | 3 +-
tests/python/relax/test_codegen_cutlass.py | 59 ++++++++++--
tests/python/relax/test_frontend_from_fx.py | 36 +++++++
.../python/relax/test_transform_legalize_ops_nn.py | 105 ++++++++++++---------
12 files changed, 229 insertions(+), 60 deletions(-)
diff --git a/include/tvm/relax/attrs/nn.h b/include/tvm/relax/attrs/nn.h
index c1ca468fc9..b759ce6c26 100644
--- a/include/tvm/relax/attrs/nn.h
+++ b/include/tvm/relax/attrs/nn.h
@@ -313,10 +313,13 @@ struct DropoutAttrs : public tvm::AttrsNode<DropoutAttrs>
{
/*! \brief Attributes used in dropout operator */
struct AttentionAttrs : public tvm::AttrsNode<AttentionAttrs> {
Optional<FloatImm> scale;
+ Optional<String> causal_mask;
TVM_DECLARE_ATTRS(AttentionAttrs, "relax.attrs.AttentionAttrs") {
TVM_ATTR_FIELD(scale).describe(
"The custom scale applied before the softmax. The default value is 1 /
sqrt(head_dim).");
+ TVM_ATTR_FIELD(causal_mask)
+ .describe("The type of the causal mask, i.e. 'TopLeft' and
'BottomRight'.");
}
}; // struct AttentionAttrs
diff --git a/python/tvm/contrib/cutlass/attention_operation.py
b/python/tvm/contrib/cutlass/attention_operation.py
index c728f7fe4b..8a96e70fe4 100644
--- a/python/tvm/contrib/cutlass/attention_operation.py
+++ b/python/tvm/contrib/cutlass/attention_operation.py
@@ -108,6 +108,7 @@ def instantiate_attention_template(attrs):
p.num_queries = ${num_queries}; // S
p.num_keys = ${num_keys}; // S'
p.scale = ${scale};
+ p.custom_mask_type = ${custom_mask_type};
p.o_strideM = p.head_dim_value * p.num_heads; // H' * N
diff --git a/python/tvm/contrib/cutlass/build.py
b/python/tvm/contrib/cutlass/build.py
index f263c9c22f..b8b65d27bd 100644
--- a/python/tvm/contrib/cutlass/build.py
+++ b/python/tvm/contrib/cutlass/build.py
@@ -831,6 +831,15 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator):
_, _, _, head_dim_value = v_shape
scale = op_attrs.scale
+ if op_attrs.causal_mask is None:
+ custom_mask_type = 0
+ if op_attrs.causal_mask == "TopLeft":
+ custom_mask_type = 1
+ elif op_attrs.causal_mask == "BottomRight":
+ custom_mask_type = 2
+ else:
+ raise NotImplementedError()
+
return f.with_attrs(
{
"op_type": op_type,
@@ -845,6 +854,7 @@ class CutlassRelaxFunctionAnnotator(relax.PyExprMutator):
"scale": scale,
"arch": self.options["sm"],
"qkv_layout": qkv_layout,
+ "custom_mask_type": custom_mask_type,
**arg,
}
)
diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py
b/python/tvm/contrib/cutlass/gen_tensor_op.py
index f94d7ef467..abc13d3570 100644
--- a/python/tvm/contrib/cutlass/gen_tensor_op.py
+++ b/python/tvm/contrib/cutlass/gen_tensor_op.py
@@ -721,6 +721,8 @@ def instantiate_template(func_name, annotations, func_args):
attrs["scale"] = (
float(1 / math.sqrt(h.value)) if annotations["scale"] is None else
annotations["scale"]
)
+ attrs["custom_mask_type"] = annotations["custom_mask_type"]
+
assert (
attrs["scale"] > 0 or attrs["scale"] < 0
), "Cutlass may generate nan occasionally when scale == 0.0"
diff --git a/python/tvm/relax/frontend/torch/fx_translator.py
b/python/tvm/relax/frontend/torch/fx_translator.py
index c5d65e2f0d..ab6a707cb0 100644
--- a/python/tvm/relax/frontend/torch/fx_translator.py
+++ b/python/tvm/relax/frontend/torch/fx_translator.py
@@ -1015,19 +1015,22 @@ class TorchFXImporter:
)
def _scaled_dot_product_attention(self, node: fx.node.Node) -> relax.Var:
- assert len(node.args) <= 4, "Dropout, and causal masking are not
supported."
+ assert (
+ len(node.args) <= 4
+ ), "Dropout is not supported, and is_causal should be called by
kwargs."
transpose_S_H = lambda tensor: relax.op.permute_dims(tensor, [0, 2, 1,
3])
query = transpose_S_H(self.env[node.args[0]])
key = transpose_S_H(self.env[node.args[1]])
value = transpose_S_H(self.env[node.args[2]])
+ causal_mask = "TopLeft" if node.kwargs.get("is_causal", False) else
None
if len(node.args) == 4:
mask = self.env[node.args[3]]
msg = "Only a float mask is supported for the attn_mask input."
assert "float" in mask.struct_info.dtype, msg
- attn = relax.op.nn.attention(query, key, value, bias=mask)
+ attn = relax.op.nn.attention(query, key, value, bias=mask,
causal_mask=causal_mask)
else:
- attn = relax.op.nn.attention(query, key, value)
+ attn = relax.op.nn.attention(query, key, value,
causal_mask=causal_mask)
return self.block_builder.emit(attn)
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
index 601a0c8439..92bb7c2042 100644
--- a/python/tvm/relax/op/nn/nn.py
+++ b/python/tvm/relax/op/nn/nn.py
@@ -1005,6 +1005,7 @@ def attention(
value: Expr,
bias: Optional[Expr] = None,
scale: Optional[FloatImm] = None,
+ causal_mask: Optional[str] = None,
) -> Expr:
r"""Computes fused multi head attention.
@@ -1035,8 +1036,29 @@ def attention(
a 4-D tensor ending with seq_len_kv, and broadcastable to
(batch_size, num_head, seq_len, seq_len_kv).
- scale: Optional[FloatImm]
- The custom scale applied before the softmax. The default value is 1 /
sqrt(head_dim).
+ causal_mask: Optional[str]
+ The optional causal mask, i.e. 'TopLeft' and 'BottomRight'.
+ For 'TopLeft', the mask matrix is as `np.tril(*, k=0)`,
+ while for 'BottomRight', the mask matrix is as `np.tril(*,
k=abs(seq_len - seq_len_kv))`
+ For example, with seq_len = 4, seq_len_kv = 2,
+ mask for 'TopLeft':
+ [[1, 0],
+ [1, 1],
+ [1, 1],
+ [1, 1]]
+ mask for 'BottomRight':
+ [[1, 1],
+ [1, 1],
+ [1, 1],
+ [1, 1]]
+ with seq_len = 2, seq_len_kv = 4,
+ mask for 'TopLeft':
+ [[1, 0, 0, 0],
+ [1, 1, 0, 0]]
+ mask for 'BottomRight':
+ [[1, 1, 1, 0],
+ [1, 1, 1, 1]]
+
Returns
-------
@@ -1044,4 +1066,4 @@ def attention(
The computed result. The layout of the output should be
(batch_size, seq_len, num_head, head_dim_v).
"""
- return _ffi_api.attention(query, key, value, bias, scale) # type: ignore
+ return _ffi_api.attention(query, key, value, bias, scale, causal_mask) #
type: ignore
diff --git a/python/tvm/relax/transform/legalize_ops/nn.py
b/python/tvm/relax/transform/legalize_ops/nn.py
index 9eea40b3fb..514c3d0782 100644
--- a/python/tvm/relax/transform/legalize_ops/nn.py
+++ b/python/tvm/relax/transform/legalize_ops/nn.py
@@ -18,6 +18,7 @@
"""Default legalization function for neural network operators."""
import logging
import math
+from typing import Optional
from tvm import topi, tir, te
from ...block_builder import BlockBuilder
@@ -335,7 +336,12 @@ def _nn_dropout(bb: BlockBuilder, call: Call) -> Expr:
def _te_attention(
- q: te.Tensor, k: te.Tensor, v: te.Tensor, bias: te.Tensor, scale:
tir.FloatImm
+ q: te.Tensor,
+ k: te.Tensor,
+ v: te.Tensor,
+ bias: te.Tensor,
+ scale: tir.FloatImm,
+ causal_mask: Optional[str],
) -> te.Tensor:
batch_size, seq_len, num_head, head_dim = q.shape
_, seq_len_kv, _, head_dim_v = v.shape
@@ -354,7 +360,21 @@ def _te_attention(
p = topi.reshape(p, [batch_size, num_head, seq_len, seq_len_kv])
p = topi.add(p, bias)
p = topi.reshape(p, [batch_size * num_head, seq_len, seq_len_kv])
- s = topi.nn.softmax(p)
+ if causal_mask is None:
+ s = topi.nn.softmax(p)
+ else:
+ if causal_mask == "TopLeft":
+ offset = tir.IntImm("int32", 0)
+ elif causal_mask == "BottomRight":
+ offset = tir.IntImm("int32", abs(seq_len - seq_len_kv))
+ else:
+ raise NotImplementedError()
+ p_masked = topi.trilu(p, k=offset, upper=False)
+ p_masked_exp = topi.trilu(
+ topi.exp(p_masked - topi.max(p_masked, axis=-1, keepdims=True)),
k=offset, upper=False
+ )
+ p_masked_sum = topi.sum(p_masked_exp, axis=-1, keepdims=True)
+ s = topi.divide(p_masked_exp, p_masked_sum)
o = topi.nn.batch_matmul(s, v, transpose_b=False)
o = topi.reshape(o, [batch_size, num_head, seq_len, head_dim_v])
return topi.transpose(o, [0, 2, 1, 3])
@@ -369,6 +389,7 @@ def _nn_attention(bb: BlockBuilder, call: Call) -> Expr:
call.args[2],
None,
call.attrs.scale,
+ call.attrs.causal_mask,
primfunc_name_hint="attention",
)
@@ -382,6 +403,7 @@ def _nn_attention_bias(bb: BlockBuilder, call: Call) ->
Expr:
call.args[2],
call.args[3],
call.attrs.scale,
+ call.attrs.causal_mask,
primfunc_name_hint="attention_bias",
)
diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc
index c83e49c70c..55757552db 100644
--- a/src/relax/op/nn/attention.cc
+++ b/src/relax/op/nn/attention.cc
@@ -28,9 +28,11 @@ namespace relax {
/* relax.nn.attention */
TVM_REGISTER_NODE_TYPE(AttentionAttrs);
-Expr attention(Expr query, Expr key, Expr value, Optional<Expr> bias,
Optional<FloatImm> scale) {
+Expr attention(Expr query, Expr key, Expr value, Optional<Expr> bias,
Optional<FloatImm> scale,
+ Optional<String> causal_mask) {
ObjectPtr<AttentionAttrs> attrs = make_object<AttentionAttrs>();
attrs->scale = scale;
+ attrs->causal_mask = causal_mask;
if (bias.defined()) {
return Call(Op::Get("relax.nn.attention_bias"),
{std::move(query), std::move(key), std::move(value),
std::move(bias.value())},
@@ -110,7 +112,8 @@ StructInfo InferStructInfoAttention(const Call& call, const
BlockBuilder& ctx) {
}
Call InferMixedPrecisionAttention(const Call& call, const DataType& out_dtype)
{
- return Downcast<Call>(attention(call->args[0], call->args[1], call->args[2],
NullOpt, NullOpt));
+ return Downcast<Call>(
+ attention(call->args[0], call->args[1], call->args[2], NullOpt, NullOpt,
NullOpt));
}
TVM_REGISTER_OP("relax.nn.attention")
diff --git a/src/relax/op/nn/attention.h b/src/relax/op/nn/attention.h
index 7eda30b408..8bbf2596ce 100644
--- a/src/relax/op/nn/attention.h
+++ b/src/relax/op/nn/attention.h
@@ -33,7 +33,8 @@ namespace tvm {
namespace relax {
/*! \brief fused multi head attention */
-Expr attention(Expr query, Expr key, Expr value, Optional<Expr> bias,
Optional<FloatImm> scale);
+Expr attention(Expr query, Expr key, Expr value, Optional<Expr> bias,
Optional<FloatImm> scale,
+ Optional<String> causal_mask);
} // namespace relax
} // namespace tvm
diff --git a/tests/python/relax/test_codegen_cutlass.py
b/tests/python/relax/test_codegen_cutlass.py
index 000bea7a0c..c0a72fffb9 100644
--- a/tests/python/relax/test_codegen_cutlass.py
+++ b/tests/python/relax/test_codegen_cutlass.py
@@ -552,7 +552,7 @@ def attention_size(request):
return request.param
-def get_relax_attention_module(q, k, v, bias=None, qk_scale=None):
+def get_relax_attention_module(q, k, v, bias=None, qk_scale=None, causal=None):
dtype = str(q.dtype)
from tvm.script.ir_builder import IRBuilder
@@ -571,7 +571,7 @@ def get_relax_attention_module(q, k, v, bias=None,
qk_scale=None):
if bias is not None:
bias = R.arg("bias", R.Tensor(bias.shape, dtype))
with R.dataflow() as frame:
- result = R.emit(R.nn.attention(q, k, v, bias, qk_scale))
+ result = R.emit(R.nn.attention(q, k, v, bias, qk_scale,
causal))
R.output(result)
R.func_ret_value(frame.output_vars[0])
@@ -581,7 +581,7 @@ def get_relax_attention_module(q, k, v, bias=None,
qk_scale=None):
@memoize("topi.tests.test_codegen_cutlass.test_attention_offload")
-def get_numpy_attention_ref(b, s, s_kv, n, h, h_v, bias_shape, qk_scale,
dtype):
+def get_numpy_attention_ref(b, s, s_kv, n, h, h_v, bias_shape, qk_scale,
causal, dtype):
q = np.random.randn(b, s, n, h).astype(dtype)
k = np.random.randn(b, s_kv, n, h).astype(dtype)
v = np.random.randn(b, s_kv, n, h_v).astype(dtype)
@@ -596,7 +596,21 @@ def get_numpy_attention_ref(b, s, s_kv, n, h, h_v,
bias_shape, qk_scale, dtype):
score = score + bias # b, n, s, s_kv
else:
bias = None
- attn = tvm.topi.testing.softmax_python(score, -1)
+ if causal == "none":
+ attn = tvm.topi.testing.softmax_python(score, -1)
+ else:
+ if causal == "TopLeft":
+ offset = 0
+ elif causal == "BottomRight":
+ offset = abs(s - s_kv)
+ else:
+ raise NotImplementedError()
+ score_masked = np.tril(score, k=offset)
+ score_masked_exp = np.tril(
+ np.exp(score_masked - np.max(score_masked, axis=-1,
keepdims=True)), k=offset
+ )
+ score_masked_sum = np.sum(score_masked_exp, axis=-1, keepdims=True)
+ attn = np.divide(score_masked_exp, score_masked_sum)
vt = v.transpose(0, 2, 1, 3) # b, n, s_kv, h_v
ref = attn @ vt # b, n, s, h_v
return q, k, v, bias, ref.transpose(0, 2, 1, 3) # b, s, n, h_v
@@ -605,7 +619,7 @@ def get_numpy_attention_ref(b, s, s_kv, n, h, h_v,
bias_shape, qk_scale, dtype):
def test_attention_offload(attention_size, attention_dtype):
b, (s, s_kv), n, (h, h_v) = attention_size
q, k, v, _, ref = get_numpy_attention_ref(
- b, s, s_kv, n, h, h_v, "none", "none", attention_dtype
+ b, s, s_kv, n, h, h_v, "none", "none", "none", attention_dtype
)
mod = get_relax_attention_module(q, k, v)
@@ -634,7 +648,7 @@ def attention_bias_size(request):
def test_attention_bias_offload(attention_bias_size):
b, (s, s_kv), n, (h, h_v), bias_shape = attention_bias_size
q, k, v, bias, ref = get_numpy_attention_ref(
- b, s, s_kv, n, h, h_v, bias_shape, "none", "float32"
+ b, s, s_kv, n, h, h_v, bias_shape, "none", "none", "float32"
)
mod = get_relax_attention_module(q, k, v, bias)
@@ -662,7 +676,7 @@ def attention_scale(request):
def test_attention_scale_offload(attention_scale_size, attention_scale):
b, (s, s_kv), n, (h, h_v), bias_shape = attention_scale_size
q, k, v, bias, ref = get_numpy_attention_ref(
- b, s, s_kv, n, h, h_v, bias_shape, attention_scale, "float32"
+ b, s, s_kv, n, h, h_v, bias_shape, attention_scale, "none", "float32"
)
mod = get_relax_attention_module(q, k, v, bias, attention_scale)
@@ -673,6 +687,37 @@ def test_attention_scale_offload(attention_scale_size,
attention_scale):
tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)
[email protected](
+ params=[
+ # B, S, N, H, bias_shape
+ (2, (16, 8), 4, (8, 16), "none"),
+ (2, (8, 16), 4, (8, 16), "none"),
+ (2, (16, 8), 4, (8, 16), (2, 4, 16, 8)),
+ ]
+)
+def attention_causal_size(request):
+ return request.param
+
+
[email protected](params=["TopLeft", "BottomRight"])
+def attention_causal(request):
+ return request.param
+
+
+def test_attention_causal_offload(attention_causal_size, attention_causal):
+ b, (s, s_kv), n, (h, h_v), bias_shape = attention_causal_size
+ q, k, v, bias, ref = get_numpy_attention_ref(
+ b, s, s_kv, n, h, h_v, bias_shape, "none", attention_causal, "float32"
+ )
+
+ mod = get_relax_attention_module(q, k, v, bias, None, attention_causal)
+ if bias is None:
+ out = get_result_with_relax_cutlass_offload(mod, q, k, v,
num_final_bindings=3)
+ else:
+ out = get_result_with_relax_cutlass_offload(mod, q, k, v, bias,
num_final_bindings=3)
+ tvm.testing.assert_allclose(out, ref, rtol=1e-2, atol=1e-2)
+
+
@memoize("topi.tests.test_codegen_cutlass.test_stacked_attention_offload")
def get_numpy_stacked_attention_ref(b, s, n, h, h_v, bias_shape, qk_scale,
dtype):
qkv = np.random.randn(b, s, n * h + n * h + n * h_v).astype(dtype)
diff --git a/tests/python/relax/test_frontend_from_fx.py
b/tests/python/relax/test_frontend_from_fx.py
index 40b9519386..ba8f776aa6 100644
--- a/tests/python/relax/test_frontend_from_fx.py
+++ b/tests/python/relax/test_frontend_from_fx.py
@@ -2644,6 +2644,31 @@ def test_attention():
R.output(gv)
return gv
+ @I.ir_module
+ class Expected3:
+ @R.function
+ def main(
+ inp_0: R.Tensor((32, 8, 128, 64), dtype="float32"),
+ inp_1: R.Tensor((32, 8, 128, 64), dtype="float32"),
+ inp_2: R.Tensor((32, 8, 128, 64), dtype="float32"),
+ ) -> R.Tensor((32, 128, 8, 64), dtype="float32"):
+ with R.dataflow():
+ lv: R.Tensor((32, 128, 8, 64), dtype="float32") =
R.permute_dims(
+ inp_0, axes=[0, 2, 1, 3]
+ )
+ lv1: R.Tensor((32, 128, 8, 64), dtype="float32") =
R.permute_dims(
+ inp_1, axes=[0, 2, 1, 3]
+ )
+ lv2: R.Tensor((32, 128, 8, 64), dtype="float32") =
R.permute_dims(
+ inp_2, axes=[0, 2, 1, 3]
+ )
+ lv3: R.Tensor((32, 128, 8, 64), dtype="float32") =
R.nn.attention(
+ lv, lv1, lv2, scale=None, causal_mask="TopLeft"
+ )
+ gv: R.Tensor((32, 128, 8, 64), dtype="float32") = lv3
+ R.output(gv)
+ return gv
+
verify_model(
lambda q, k, v: F.scaled_dot_product_attention(q, k, v),
[
@@ -2667,6 +2692,17 @@ def test_attention():
Expected2,
)
+ verify_model(
+ lambda q, k, v: F.scaled_dot_product_attention(q, k, v,
is_causal=True),
+ [
+ ([32, 8, 128, 64], "float32"),
+ ([32, 8, 128, 64], "float32"),
+ ([32, 8, 128, 64], "float32"),
+ ],
+ {},
+ Expected3,
+ )
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_transform_legalize_ops_nn.py
b/tests/python/relax/test_transform_legalize_ops_nn.py
index 1ff0569629..cbbacbabda 100644
--- a/tests/python/relax/test_transform_legalize_ops_nn.py
+++ b/tests/python/relax/test_transform_legalize_ops_nn.py
@@ -2676,14 +2676,14 @@ def test_attention():
@R.function
def main(q: R.Tensor((4, 16, 32, 8), "float32"), k: R.Tensor((4, 8,
32, 8), "float32"), v: R.Tensor((4, 8, 32, 16), "float32"), bias: R.Tensor((4,
32, 16, 8), "float32")):
scale = T.FloatImm("float32", 0.1)
- gv: R.Tensor((4, 16, 32, 16), "float32") = R.nn.attention(q, k, v,
bias, scale)
+ gv: R.Tensor((4, 16, 32, 16), "float32") = R.nn.attention(q, k, v,
bias, scale=scale, causal_mask="TopLeft")
return gv
@tvm.script.ir_module
class Expected:
@T.prim_func
- def attention_bias(rxplaceholder: T.Buffer((T.int64(4), T.int64(16),
T.int64(32), T.int64(8)), "float32"), rxplaceholder_1: T.Buffer((T.int64(4),
T.int64(8), T.int64(32), T.int64(8)), "float32"), rxplaceholder_2:
T.Buffer((T.int64(4), T.int64(8), T.int64(32), T.int64(16)), "float32"),
rxplaceholder_3: T.Buffer((T.int64(4), T.int64(32), T.int64(16), T.int64(8)),
"float32"), T_transpose: T.Buffer((T.int64(4), T.int64(16), T.int64(32),
T.int64(16)), "float32")):
- T.func_attr({"tir.noalias": True})
+ def attention_bias(A: T.Buffer((T.int64(4), T.int64(16), T.int64(32),
T.int64(8)), "float32"), B: T.Buffer((T.int64(4), T.int64(8), T.int64(32),
T.int64(8)), "float32"), C: T.Buffer((T.int64(4), T.int64(8), T.int64(32),
T.int64(16)), "float32"), D: T.Buffer((T.int64(4), T.int64(32), T.int64(16),
T.int64(8)), "float32"), T_transpose: T.Buffer((T.int64(4), T.int64(16),
T.int64(32), T.int64(16)), "float32")):
+ T.func_attr({"tir.noalias": T.bool(True)})
# with T.block("root"):
T_transpose_1 = T.alloc_buffer((T.int64(4), T.int64(32),
T.int64(16), T.int64(8)))
T_reshape = T.alloc_buffer((T.int64(128), T.int64(16), T.int64(8)))
@@ -2694,10 +2694,13 @@ def test_attention():
T_reshape_2 = T.alloc_buffer((T.int64(4), T.int64(32),
T.int64(16), T.int64(8)))
T_add = T.alloc_buffer((T.int64(4), T.int64(32), T.int64(16),
T.int64(8)))
T_reshape_3 = T.alloc_buffer((T.int64(128), T.int64(16),
T.int64(8)))
- T_softmax_maxelem = T.alloc_buffer((T.int64(128), T.int64(16)))
- T_softmax_exp = T.alloc_buffer((T.int64(128), T.int64(16),
T.int64(8)))
- T_softmax_expsum = T.alloc_buffer((T.int64(128), T.int64(16)))
- T_softmax_norm = T.alloc_buffer((T.int64(128), T.int64(16),
T.int64(8)))
+ trilu = T.alloc_buffer((T.int64(128), T.int64(16), T.int64(8)))
+ trilu_red = T.alloc_buffer((T.int64(128), T.int64(16), T.int64(1)))
+ T_subtract = T.alloc_buffer((T.int64(128), T.int64(16),
T.int64(8)))
+ compute = T.alloc_buffer((T.int64(128), T.int64(16), T.int64(8)))
+ trilu_1 = T.alloc_buffer((T.int64(128), T.int64(16), T.int64(8)))
+ trilu_red_1 = T.alloc_buffer((T.int64(128), T.int64(16),
T.int64(1)))
+ T_divide = T.alloc_buffer((T.int64(128), T.int64(16), T.int64(8)))
T_transpose_3 = T.alloc_buffer((T.int64(4), T.int64(32),
T.int64(8), T.int64(16)))
T_reshape_4 = T.alloc_buffer((T.int64(128), T.int64(8),
T.int64(16)))
T_batch_matmul_NN = T.alloc_buffer((T.int64(128), T.int64(16),
T.int64(16)))
@@ -2705,9 +2708,9 @@ def test_attention():
for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(32),
T.int64(16), T.int64(8)):
with T.block("T_transpose"):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0,
ax1, ax2, ax3])
- T.reads(rxplaceholder[v_ax0, v_ax2, v_ax1, v_ax3])
+ T.reads(A[v_ax0, v_ax2, v_ax1, v_ax3])
T.writes(T_transpose_1[v_ax0, v_ax1, v_ax2, v_ax3])
- T_transpose_1[v_ax0, v_ax1, v_ax2, v_ax3] =
rxplaceholder[v_ax0, v_ax2, v_ax1, v_ax3]
+ T_transpose_1[v_ax0, v_ax1, v_ax2, v_ax3] = A[v_ax0,
v_ax2, v_ax1, v_ax3]
for ax0, ax1, ax2 in T.grid(T.int64(128), T.int64(16), T.int64(8)):
with T.block("T_reshape"):
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
@@ -2717,9 +2720,9 @@ def test_attention():
for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(32),
T.int64(8), T.int64(8)):
with T.block("T_transpose_1"):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0,
ax1, ax2, ax3])
- T.reads(rxplaceholder_1[v_ax0, v_ax2, v_ax1, v_ax3])
+ T.reads(B[v_ax0, v_ax2, v_ax1, v_ax3])
T.writes(T_transpose_2[v_ax0, v_ax1, v_ax2, v_ax3])
- T_transpose_2[v_ax0, v_ax1, v_ax2, v_ax3] =
rxplaceholder_1[v_ax0, v_ax2, v_ax1, v_ax3]
+ T_transpose_2[v_ax0, v_ax1, v_ax2, v_ax3] = B[v_ax0,
v_ax2, v_ax1, v_ax3]
for ax0, ax1, ax2 in T.grid(T.int64(128), T.int64(8), T.int64(8)):
with T.block("T_reshape_1"):
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
@@ -2750,50 +2753,67 @@ def test_attention():
for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(32),
T.int64(16), T.int64(8)):
with T.block("T_add"):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0,
ax1, ax2, ax3])
- T.reads(T_reshape_2[v_ax0, v_ax1, v_ax2, v_ax3],
rxplaceholder_3[v_ax0, v_ax1, v_ax2, v_ax3])
+ T.reads(T_reshape_2[v_ax0, v_ax1, v_ax2, v_ax3], D[v_ax0,
v_ax1, v_ax2, v_ax3])
T.writes(T_add[v_ax0, v_ax1, v_ax2, v_ax3])
- T_add[v_ax0, v_ax1, v_ax2, v_ax3] = T_reshape_2[v_ax0,
v_ax1, v_ax2, v_ax3] + rxplaceholder_3[v_ax0, v_ax1, v_ax2, v_ax3]
+ T_add[v_ax0, v_ax1, v_ax2, v_ax3] = T_reshape_2[v_ax0,
v_ax1, v_ax2, v_ax3] + D[v_ax0, v_ax1, v_ax2, v_ax3]
for ax0, ax1, ax2 in T.grid(T.int64(128), T.int64(16), T.int64(8)):
with T.block("T_reshape_3"):
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
T.reads(T_add[((v_ax2 // T.int64(8) + v_ax1) //
T.int64(16) + v_ax0) % T.int64(128) // T.int64(32), ((v_ax2 // T.int64(8) +
v_ax1) // T.int64(16) + v_ax0) % T.int64(32), (v_ax2 // T.int64(8) + v_ax1) %
T.int64(16), v_ax2 % T.int64(8)])
T.writes(T_reshape_3[v_ax0, v_ax1, v_ax2])
T_reshape_3[v_ax0, v_ax1, v_ax2] = T_add[((v_ax2 //
T.int64(8) + v_ax1) // T.int64(16) + v_ax0) % T.int64(128) // T.int64(32),
((v_ax2 // T.int64(8) + v_ax1) // T.int64(16) + v_ax0) % T.int64(32), (v_ax2 //
T.int64(8) + v_ax1) % T.int64(16), v_ax2 % T.int64(8)]
- for i0, i1, k in T.grid(T.int64(128), T.int64(16), T.int64(8)):
- with T.block("T_softmax_maxelem"):
- v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
- T.reads(T_reshape_3[v_i0, v_i1, v_k])
- T.writes(T_softmax_maxelem[v_i0, v_i1])
- with T.init():
- T_softmax_maxelem[v_i0, v_i1] =
T.float32(-3.4028234663852886e+38)
- T_softmax_maxelem[v_i0, v_i1] =
T.max(T_softmax_maxelem[v_i0, v_i1], T_reshape_3[v_i0, v_i1, v_k])
for i0, i1, i2 in T.grid(T.int64(128), T.int64(16), T.int64(8)):
- with T.block("T_softmax_exp"):
+ with T.block("trilu"):
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
- T.reads(T_reshape_3[v_i0, v_i1, v_i2],
T_softmax_maxelem[v_i0, v_i1])
- T.writes(T_softmax_exp[v_i0, v_i1, v_i2])
- T_softmax_exp[v_i0, v_i1, v_i2] = T.exp(T_reshape_3[v_i0,
v_i1, v_i2] - T_softmax_maxelem[v_i0, v_i1])
- for i0, i1, k in T.grid(T.int64(128), T.int64(16), T.int64(8)):
- with T.block("T_softmax_expsum"):
- v_i0, v_i1, v_k = T.axis.remap("SSR", [i0, i1, k])
- T.reads(T_softmax_exp[v_i0, v_i1, v_k])
- T.writes(T_softmax_expsum[v_i0, v_i1])
+ T.reads(T_reshape_3[v_i0, v_i1, v_i2])
+ T.writes(trilu[v_i0, v_i1, v_i2])
+ trilu[v_i0, v_i1, v_i2] = T.Select(v_i2 <= v_i1,
T_reshape_3[v_i0, v_i1, v_i2], T.float32(0))
+ for ax0, ax1, ax2, k2 in T.grid(T.int64(128), T.int64(16),
T.int64(1), T.int64(8)):
+ with T.block("trilu_red"):
+ v_ax0, v_ax1, v_ax2, v_k2 = T.axis.remap("SSSR", [ax0,
ax1, ax2, k2])
+ T.reads(trilu[v_ax0, v_ax1, v_k2])
+ T.writes(trilu_red[v_ax0, v_ax1, v_ax2])
with T.init():
- T_softmax_expsum[v_i0, v_i1] = T.float32(0)
- T_softmax_expsum[v_i0, v_i1] = T_softmax_expsum[v_i0,
v_i1] + T_softmax_exp[v_i0, v_i1, v_k]
+ trilu_red[v_ax0, v_ax1, v_ax2] =
T.float32(-3.4028234663852886e+38)
+ trilu_red[v_ax0, v_ax1, v_ax2] = T.max(trilu_red[v_ax0,
v_ax1, v_ax2], trilu[v_ax0, v_ax1, v_k2])
+ for ax0, ax1, ax2 in T.grid(T.int64(128), T.int64(16), T.int64(8)):
+ with T.block("T_subtract"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(trilu[v_ax0, v_ax1, v_ax2], trilu_red[v_ax0,
v_ax1, T.int64(0)])
+ T.writes(T_subtract[v_ax0, v_ax1, v_ax2])
+ T_subtract[v_ax0, v_ax1, v_ax2] = trilu[v_ax0, v_ax1,
v_ax2] - trilu_red[v_ax0, v_ax1, T.int64(0)]
for i0, i1, i2 in T.grid(T.int64(128), T.int64(16), T.int64(8)):
- with T.block("T_softmax_norm"):
+ with T.block("compute"):
v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
- T.reads(T_softmax_exp[v_i0, v_i1, v_i2],
T_softmax_expsum[v_i0, v_i1])
- T.writes(T_softmax_norm[v_i0, v_i1, v_i2])
- T.block_attr({"axis": 2})
- T_softmax_norm[v_i0, v_i1, v_i2] = T_softmax_exp[v_i0,
v_i1, v_i2] / T_softmax_expsum[v_i0, v_i1]
+ T.reads(T_subtract[v_i0, v_i1, v_i2])
+ T.writes(compute[v_i0, v_i1, v_i2])
+ compute[v_i0, v_i1, v_i2] = T.exp(T_subtract[v_i0, v_i1,
v_i2])
+ for i0, i1, i2 in T.grid(T.int64(128), T.int64(16), T.int64(8)):
+ with T.block("trilu_1"):
+ v_i0, v_i1, v_i2 = T.axis.remap("SSS", [i0, i1, i2])
+ T.reads(compute[v_i0, v_i1, v_i2])
+ T.writes(trilu_1[v_i0, v_i1, v_i2])
+ trilu_1[v_i0, v_i1, v_i2] = T.Select(v_i2 <= v_i1,
compute[v_i0, v_i1, v_i2], T.float32(0))
+ for ax0, ax1, ax2, k2 in T.grid(T.int64(128), T.int64(16),
T.int64(1), T.int64(8)):
+ with T.block("trilu_red_1"):
+ v_ax0, v_ax1, v_ax2, v_k2 = T.axis.remap("SSSR", [ax0,
ax1, ax2, k2])
+ T.reads(trilu_1[v_ax0, v_ax1, v_k2])
+ T.writes(trilu_red_1[v_ax0, v_ax1, v_ax2])
+ with T.init():
+ trilu_red_1[v_ax0, v_ax1, v_ax2] = T.float32(0)
+ trilu_red_1[v_ax0, v_ax1, v_ax2] = trilu_red_1[v_ax0,
v_ax1, v_ax2] + trilu_1[v_ax0, v_ax1, v_k2]
+ for ax0, ax1, ax2 in T.grid(T.int64(128), T.int64(16), T.int64(8)):
+ with T.block("T_divide"):
+ v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
+ T.reads(trilu_1[v_ax0, v_ax1, v_ax2], trilu_red_1[v_ax0,
v_ax1, T.int64(0)])
+ T.writes(T_divide[v_ax0, v_ax1, v_ax2])
+ T_divide[v_ax0, v_ax1, v_ax2] = trilu_1[v_ax0, v_ax1,
v_ax2] / trilu_red_1[v_ax0, v_ax1, T.int64(0)]
for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(32),
T.int64(8), T.int64(16)):
with T.block("T_transpose_2"):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0,
ax1, ax2, ax3])
- T.reads(rxplaceholder_2[v_ax0, v_ax2, v_ax1, v_ax3])
+ T.reads(C[v_ax0, v_ax2, v_ax1, v_ax3])
T.writes(T_transpose_3[v_ax0, v_ax1, v_ax2, v_ax3])
- T_transpose_3[v_ax0, v_ax1, v_ax2, v_ax3] =
rxplaceholder_2[v_ax0, v_ax2, v_ax1, v_ax3]
+ T_transpose_3[v_ax0, v_ax1, v_ax2, v_ax3] = C[v_ax0,
v_ax2, v_ax1, v_ax3]
for ax0, ax1, ax2 in T.grid(T.int64(128), T.int64(8), T.int64(16)):
with T.block("T_reshape_4"):
v_ax0, v_ax1, v_ax2 = T.axis.remap("SSS", [ax0, ax1, ax2])
@@ -2803,12 +2823,12 @@ def test_attention():
for b, i, j, k in T.grid(T.int64(128), T.int64(16), T.int64(16),
T.int64(8)):
with T.block("T_batch_matmul_NN"):
v_b, v_i, v_j, v_k = T.axis.remap("SSSR", [b, i, j, k])
- T.reads(T_softmax_norm[v_b, v_i, v_k], T_reshape_4[v_b,
v_k, v_j])
+ T.reads(T_divide[v_b, v_i, v_k], T_reshape_4[v_b, v_k,
v_j])
T.writes(T_batch_matmul_NN[v_b, v_i, v_j])
T.block_attr({"layout_free_placeholders": [T_reshape_4]})
with T.init():
T_batch_matmul_NN[v_b, v_i, v_j] = T.float32(0)
- T_batch_matmul_NN[v_b, v_i, v_j] = T_batch_matmul_NN[v_b,
v_i, v_j] + T_softmax_norm[v_b, v_i, v_k] * T_reshape_4[v_b, v_k, v_j]
+ T_batch_matmul_NN[v_b, v_i, v_j] = T_batch_matmul_NN[v_b,
v_i, v_j] + T_divide[v_b, v_i, v_k] * T_reshape_4[v_b, v_k, v_j]
for ax0, ax1, ax2, ax3 in T.grid(T.int64(4), T.int64(32),
T.int64(16), T.int64(16)):
with T.block("T_reshape_5"):
v_ax0, v_ax1, v_ax2, v_ax3 = T.axis.remap("SSSS", [ax0,
ax1, ax2, ax3])
@@ -2824,7 +2844,8 @@ def test_attention():
@R.function
def main(q: R.Tensor((4, 16, 32, 8), dtype="float32"), k: R.Tensor((4,
8, 32, 8), dtype="float32"), v: R.Tensor((4, 8, 32, 16), dtype="float32"),
bias: R.Tensor((4, 32, 16, 8), dtype="float32")) -> R.Tensor((4, 16, 32, 16),
dtype="float32"):
- gv = R.call_tir(Expected.attention_bias, (q, k, v, bias),
out_sinfo=R.Tensor((4, 16, 32, 16), dtype="float32"))
+ cls = Expected
+ gv = R.call_tir(cls.attention_bias, (q, k, v, bias),
out_sinfo=R.Tensor((4, 16, 32, 16), dtype="float32"))
return gv
# fmt: on
