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 dd02791e5a [Uniy][Op] Expand support of attention bias layout (#14737)
dd02791e5a is described below
commit dd02791e5a18715997ef68ff1d755770b6bee63f
Author: Yaxing Cai <[email protected]>
AuthorDate: Fri Apr 28 05:11:37 2023 -0700
[Uniy][Op] Expand support of attention bias layout (#14737)
This PR expands the support of attention bias layout. In the past, we only
support 3 types of layout, with `BNSS'`, `BSS'`, `BS'`, e.g. `BNSS'` for
(batch_size, num_head, seq_len, seq_len_kv). However, cutlass accepts 8 types
of layout, as long as the shape ends with `S` and is broadcastable to `BNSS'`.
So in this PR, we fixed the ndim of `bias` to 4 and accepts all 8 types of
layout as cutlass.
---
python/tvm/contrib/cutlass/attention_operation.py | 30 +++---------
python/tvm/contrib/cutlass/gen_tensor_op.py | 20 ++++++--
python/tvm/relax/op/nn/nn.py | 4 +-
python/tvm/relax/transform/legalize_ops/nn.py | 4 --
src/relax/op/nn/attention.cc | 29 ++++++------
tests/python/relax/test_codegen_cutlass.py | 57 ++++++++++++-----------
6 files changed, 67 insertions(+), 77 deletions(-)
diff --git a/python/tvm/contrib/cutlass/attention_operation.py
b/python/tvm/contrib/cutlass/attention_operation.py
index 1921edb397..0e507d3be1 100644
--- a/python/tvm/contrib/cutlass/attention_operation.py
+++ b/python/tvm/contrib/cutlass/attention_operation.py
@@ -23,32 +23,14 @@ def instantiate_attention_template(attrs):
"""Return CUTLASS host code for fused multi head attention
based on a template and the provided attribute map."""
- bias_template = {
- "B11S'": """
- CHECK(${bias}->ndim == 2); // B, 1, 1, S'
-
- p.attn_bias_ptr = reinterpret_cast<T *>(${bias}->data);
- p.bias_strideM = 0; // 0
- p.bias_strideH = 0; // 0
- p.bias_strideB = p.num_keys; // S'
-""",
- "B1SS'": """
- CHECK(${bias}->ndim == 3); // B, 1, S, S'
-
- p.attn_bias_ptr = reinterpret_cast<T *>(${bias}->data);
- p.bias_strideM = p.num_keys; // S'
- p.bias_strideH = 0; // 0
- p.bias_strideB = p.bias_strideM * p.num_queries; // S' * S
-""",
- "BNSS'": """
+ bias_template = """
CHECK(${bias}->ndim == 4); // B, N, S, S'
p.attn_bias_ptr = reinterpret_cast<T *>(${bias}->data);
- p.bias_strideM = p.num_keys; // S'
- p.bias_strideH = p.bias_strideM * p.num_queries; // S' * S
- p.bias_strideB = p.bias_strideH * p.num_heads; // S' * S * N
-""",
- }
+ p.bias_strideM = ${bias_strideM};
+ p.bias_strideH = ${bias_strideH};
+ p.bias_strideB = ${bias_strideB};
+"""
qkv_template = {
"default": """
@@ -159,7 +141,7 @@ def instantiate_attention_template(attrs):
template,
{
"qkv_template": qkv_template[attrs["qkv_layout"]],
- "bias_template": bias_template[attrs["bias_layout"]] if
"bias_layout" in attrs else "",
+ "bias_template": bias_template if "bias" in attrs else "",
},
)
diff --git a/python/tvm/contrib/cutlass/gen_tensor_op.py
b/python/tvm/contrib/cutlass/gen_tensor_op.py
index bb4d224329..392a7bf7e2 100644
--- a/python/tvm/contrib/cutlass/gen_tensor_op.py
+++ b/python/tvm/contrib/cutlass/gen_tensor_op.py
@@ -741,11 +741,21 @@ def instantiate_template(func_name, annotations,
func_args):
if "bias" in attrs:
attrs["kSupportsBias"] = True
if len(annotations["bias_shape"]) == 4:
- attrs["bias_layout"] = "BNSS'"
- elif len(annotations["bias_shape"]) == 3:
- attrs["bias_layout"] = "B1SS'"
- elif len(annotations["bias_shape"]) == 2:
- attrs["bias_layout"] = "B11S'"
+ strides = "p.num_keys"
+ if annotations["bias_shape"][2] == 1:
+ attrs["bias_strideM"] = 0
+ else:
+ attrs["bias_strideM"] = strides
+ strides = f"p.num_queries * {strides}"
+ if annotations["bias_shape"][1] == 1:
+ attrs["bias_strideH"] = 0
+ else:
+ attrs["bias_strideH"] = strides
+ strides = f"p.num_heads * {strides}"
+ if annotations["bias_shape"][0] == 1:
+ attrs["bias_strideB"] = 0
+ else:
+ attrs["bias_strideB"] = strides
else:
raise NotImplementedError()
else:
diff --git a/python/tvm/relax/op/nn/nn.py b/python/tvm/relax/op/nn/nn.py
index 5483e7c5ee..fb5e0736ff 100644
--- a/python/tvm/relax/op/nn/nn.py
+++ b/python/tvm/relax/op/nn/nn.py
@@ -1009,8 +1009,8 @@ def attention(
bias: Optional[Expr]
The optional attention bias to the operator. The layout of the
attention bias should be
- (batch_size, num_head, seq_len, seq_len_kv),
- (batch_size, seq_len, seq_len_kv) or (batch_size, seq_len_kv).
+ 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).
diff --git a/python/tvm/relax/transform/legalize_ops/nn.py
b/python/tvm/relax/transform/legalize_ops/nn.py
index efed7b586b..9c98682e32 100644
--- a/python/tvm/relax/transform/legalize_ops/nn.py
+++ b/python/tvm/relax/transform/legalize_ops/nn.py
@@ -336,10 +336,6 @@ def _te_attention(
p = topi.divide(p, tir.sqrt(tir.Cast(p.dtype, head_dim)))
if bias is not None:
p = topi.reshape(p, [batch_size, num_head, seq_len, seq_len_kv])
- if len(bias.shape) == 2:
- bias = topi.reshape(bias, [batch_size, 1, 1, seq_len_kv])
- elif len(bias.shape) == 3:
- bias = topi.reshape(bias, [batch_size, 1, 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)
diff --git a/src/relax/op/nn/attention.cc b/src/relax/op/nn/attention.cc
index 7f0a1d3aea..56e5a04e12 100644
--- a/src/relax/op/nn/attention.cc
+++ b/src/relax/op/nn/attention.cc
@@ -85,23 +85,24 @@ StructInfo InferStructInfoAttention(const Call& call, const
BlockBuilder& ctx) {
if (input_sinfo.size() == 4) {
TensorStructInfo bias_sinfo = input_sinfo[3];
const ShapeExprNode* bias_shape = bias_sinfo->shape.as<ShapeExprNode>();
- if (bias_sinfo->ndim == 4) {
- diag_equal(num_batches, bias_shape->values[0], "query", "bias", "batch
size");
- diag_equal(num_heads, bias_shape->values[1], "query", "bias", "number of
heads");
- diag_equal(num_queries, bias_shape->values[2], "query", "bias",
"sequence length");
- diag_equal(num_keys, bias_shape->values[3], "key", "bias", "sequence
length");
- } else if (bias_sinfo->ndim == 3) {
- diag_equal(num_batches, bias_shape->values[0], "query", "bias", "batch
size");
- diag_equal(num_queries, bias_shape->values[1], "query", "bias",
"sequence length");
- diag_equal(num_keys, bias_shape->values[2], "key", "bias", "sequence
length");
- } else if (bias_sinfo->ndim == 2) {
- diag_equal(num_batches, bias_shape->values[0], "query", "bias", "batch
size");
- diag_equal(num_keys, bias_shape->values[1], "key", "bias", "sequence
length");
- } else {
+ if (bias_sinfo->ndim != 4) {
ctx->ReportFatal(Diagnostic::Error(call)
- << "The bias should have 2, 3 or 4 dimensions."
+ << "The bias should have 4 dimensions."
<< "However, the bias input has " << bias_sinfo->ndim
<< " dimensions.");
}
+ auto diag_equal_or_broadcast = [&](PrimExpr v1, PrimExpr v2, String m1,
String m2, String dim) {
+ if (analyzer->CanProve(v1 != v2) && !tir::is_one(v2)) {
+ ctx->ReportFatal(Diagnostic::Error(call)
+ << "The " << m1 << " " << dim << " and the " << m2 <<
" " << dim
+ << " should be the same or broadcastable. However,
the " << dim << " of "
+ << m1 << " is " << v1 << " while the " << dim << " of
" << m2 << " is "
+ << v2);
+ }
+ };
+ diag_equal_or_broadcast(num_batches, bias_shape->values[0], "query",
"bias", "batch size");
+ diag_equal_or_broadcast(num_heads, bias_shape->values[1], "query", "bias",
"number of heads");
+ diag_equal_or_broadcast(num_queries, bias_shape->values[2], "query",
"bias", "sequence length");
+ diag_equal(num_keys, bias_shape->values[3], "key", "bias", "sequence
length");
}
Array<PrimExpr> output_shape = {num_batches, num_queries, num_heads,
head_dim_value};
diff --git a/tests/python/relax/test_codegen_cutlass.py
b/tests/python/relax/test_codegen_cutlass.py
index 3d0cc3a54c..d5d3142cab 100644
--- a/tests/python/relax/test_codegen_cutlass.py
+++ b/tests/python/relax/test_codegen_cutlass.py
@@ -549,7 +549,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, bias_reshape,
qk_scale, dtype):
+def get_numpy_attention_ref(b, s, s_kv, n, h, h_v, bias_shape, qk_scale,
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)
@@ -561,7 +561,7 @@ def get_numpy_attention_ref(b, s, s_kv, n, h, h_v,
bias_shape, bias_reshape, qk_
score = qt @ kt / np.sqrt(q.shape[-1]) # b, n, s, s_kv
if not bias_shape == "none":
bias = np.random.randn(*bias_shape).astype(dtype)
- score = score + bias.reshape(*bias_reshape) # b, n, s, s_kv
+ score = score + bias # b, n, s, s_kv
else:
bias = None
attn = tvm.topi.testing.softmax_python(score, -1)
@@ -573,7 +573,7 @@ def get_numpy_attention_ref(b, s, s_kv, n, h, h_v,
bias_shape, bias_reshape, qk_
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", "none", attention_dtype
+ b, s, s_kv, n, h, h_v, "none", "none", attention_dtype
)
mod = get_relax_attention_module(q, k, v)
@@ -584,10 +584,15 @@ def test_attention_offload(attention_size,
attention_dtype):
@pytest.fixture(
params=[
- # B, S, N, H, bias_shape, bias_reshape
- (4, (16, 8), 32, (8, 16), (4, 32, 16, 8), (4, 32, 16, 8)),
- (4, (16, 8), 32, (8, 16), (4, 16, 8), (4, 1, 16, 8)),
- (4, (16, 8), 32, (8, 16), (4, 8), (4, 1, 1, 8)),
+ # B, S, N, H, bias_shape
+ (4, (16, 8), 32, (8, 16), (4, 32, 16, 8)),
+ (4, (16, 8), 32, (8, 16), (4, 1, 16, 8)),
+ (4, (16, 8), 32, (8, 16), (4, 32, 1, 8)),
+ (4, (16, 8), 32, (8, 16), (4, 1, 1, 8)),
+ (4, (16, 8), 32, (8, 16), (1, 32, 16, 8)),
+ (4, (16, 8), 32, (8, 16), (1, 1, 16, 8)),
+ (4, (16, 8), 32, (8, 16), (1, 32, 1, 8)),
+ (4, (16, 8), 32, (8, 16), (1, 1, 1, 8)),
]
)
def attention_bias_size(request):
@@ -595,9 +600,9 @@ def attention_bias_size(request):
def test_attention_bias_offload(attention_bias_size):
- b, (s, s_kv), n, (h, h_v), bias_shape, bias_reshape = 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, bias_reshape, "none", "float32"
+ b, s, s_kv, n, h, h_v, bias_shape, "none", "float32"
)
mod = get_relax_attention_module(q, k, v, bias)
@@ -608,9 +613,9 @@ def test_attention_bias_offload(attention_bias_size):
@pytest.fixture(
params=[
- # B, S, N, H, bias_shape, bias_reshape
- (4, (16, 8), 32, (8, 16), (4, 32, 16, 8), (4, 32, 16, 8)),
- (4, (16, 8), 32, (8, 16), "none", "none"),
+ # B, S, N, H, bias_shape
+ (4, (16, 8), 32, (8, 16), (4, 32, 16, 8)),
+ (4, (16, 8), 32, (8, 16), "none"),
]
)
def attention_scale_size(request):
@@ -623,9 +628,9 @@ def attention_scale(request):
def test_attention_scale_offload(attention_scale_size, attention_scale):
- b, (s, s_kv), n, (h, h_v), bias_shape, bias_reshape = attention_scale_size
+ 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, bias_reshape, attention_scale,
"float32"
+ b, s, s_kv, n, h, h_v, bias_shape, attention_scale, "float32"
)
mod = get_relax_attention_module(q, k, v, bias, attention_scale)
@@ -637,7 +642,7 @@ def test_attention_scale_offload(attention_scale_size,
attention_scale):
@memoize("topi.tests.test_codegen_cutlass.test_stacked_attention_offload")
-def get_numpy_stacked_attention_ref(b, s, n, h, h_v, bias_shape, bias_reshape,
qk_scale, dtype):
+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)
split_qkv = np.split(qkv, [n * h, n * h * 2], axis=2)
q = np.reshape(split_qkv[0], (b, s, n, h))
@@ -651,7 +656,7 @@ def get_numpy_stacked_attention_ref(b, s, n, h, h_v,
bias_shape, bias_reshape, q
score = qt @ kt / np.sqrt(q.shape[-1]) # b, n, s, s
if not bias_shape == "none":
bias = np.random.randn(*bias_shape).astype(dtype)
- score = score + bias.reshape(*bias_reshape) # b, n, s, s
+ score = score + bias # b, n, s, s
else:
bias = None
attn = tvm.topi.testing.softmax_python(score, -1)
@@ -709,10 +714,10 @@ def get_relax_stacked_attention_module(
@pytest.fixture(
params=[
- # B, S, N, H, bias_shape, bias_reshape, scale, single_shape
- (4, 8, 32, (64, 32), "none", "none", "none", False),
- (4, 8, 32, (64, 32), (4, 32, 8, 8), (4, 32, 8, 8), 0.5, False),
- (4, 8, 32, (64, 64), "none", "none", "none", True),
+ # B, S, N, H, bias_shape, scale, single_shape
+ (4, 8, 32, (64, 32), "none", "none", False),
+ (4, 8, 32, (64, 32), (4, 32, 8, 8), 0.5, False),
+ (4, 8, 32, (64, 64), "none", "none", True),
]
)
def stacked_attention_size(request):
@@ -720,10 +725,8 @@ def stacked_attention_size(request):
def test_stacked_attention_split_offload(stacked_attention_size):
- b, s, n, (h, h_v), bias_shape, bias_reshape, scale, single_shape =
stacked_attention_size
- qkv, bias, ref = get_numpy_stacked_attention_ref(
- b, s, n, h, h_v, bias_shape, bias_reshape, scale, "float32"
- )
+ b, s, n, (h, h_v), bias_shape, scale, single_shape = stacked_attention_size
+ qkv, bias, ref = get_numpy_stacked_attention_ref(b, s, n, h, h_v,
bias_shape, scale, "float32")
if scale == "none":
mod = get_relax_stacked_attention_module(
qkv, b, s, n, h, h_v, "split", bias, single_shape=single_shape
@@ -741,10 +744,8 @@ def
test_stacked_attention_split_offload(stacked_attention_size):
def test_stacked_attention_strided_slice_offload(stacked_attention_size):
- b, s, n, (h, h_v), bias_shape, bias_reshape, scale, single_shape =
stacked_attention_size
- qkv, bias, ref = get_numpy_stacked_attention_ref(
- b, s, n, h, h_v, bias_shape, bias_reshape, scale, "float32"
- )
+ b, s, n, (h, h_v), bias_shape, scale, single_shape = stacked_attention_size
+ qkv, bias, ref = get_numpy_stacked_attention_ref(b, s, n, h, h_v,
bias_shape, scale, "float32")
if scale == "none":
mod = get_relax_stacked_attention_module(
qkv, b, s, n, h, h_v, "strided_slice", bias,
single_shape=single_shape
