Resolving inconsistency between attention/attention_bias

Author: parsifal-47

include/tvm/relax/attrs/nn.h

@@ -546,11 +546,13 @@ struct DropoutAttrs : public tvm::AttrsNode<DropoutAttrs> {
 
 /*! \brief Attributes used in Attention operator */
 struct AttentionAttrs : public tvm::AttrsNode<AttentionAttrs> {
+  Optional<Expr> bias;
   Optional<FloatImm> scale;
   Optional<String> causal_mask;
   Optional<IntImm> window_size;
 
   TVM_DECLARE_ATTRS(AttentionAttrs, "relax.attrs.AttentionAttrs") {
+    TVM_ATTR_FIELD(bias).describe("The input bias tensor.");
     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)

src/relax/op/nn/attention.cc

@@ -34,12 +34,8 @@ Expr attention(Expr query, Expr key, Expr value, Optional<Expr> bias, Optional<F
   attrs->scale = scale;
   attrs->causal_mask = causal_mask;
   attrs->window_size = window_size;
+  attrs->bias = bias;
 
-  if (bias) {
-    return Call(Op::Get("relax.nn.attention_bias"),
-                {std::move(query), std::move(key), std::move(value), std::move(bias.value())},
-                Attrs(attrs), {});
-  }
   return Call(Op::Get("relax.nn.attention"), {std::move(query), std::move(key), std::move(value)},
               Attrs(attrs), {});
 }
@@ -152,18 +148,6 @@ TVM_REGISTER_OP("relax.nn.attention")
     .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoAttention)
     .set_attr<Bool>("FPurity", Bool(true));
 
-TVM_REGISTER_OP("relax.nn.attention_bias")
-    .set_attrs_type<AttentionAttrs>()
-    .set_num_inputs(4)
-    .add_argument("query", "Tensor", "The input queries tensor.")
-    .add_argument("key", "Tensor", "The input keys tensor.")
-    .add_argument("value", "Tensor", "The input values tensor.")
-    .add_argument("bias", "Tensor", "The input bias tensor.")
-    .set_attr<TMixedPrecisionPolicy>("TMixedPrecisionPolicy", MixedPrecisionPolicyKind::kAlways)
-    .set_attr<FInferMixedPrecision>("FInferMixedPrecision", InferMixedPrecisionAttention)
-    .set_attr<FInferStructInfo>("FInferStructInfo", InferStructInfoAttention)
-    .set_attr<Bool>("FPurity", Bool(true));
-
 TVM_REGISTER_OP("relax.nn.attention_var_len")
     .set_attrs_type<AttentionAttrs>()
     .set_num_inputs(7)