Resolving inconsistency between attention/attention_bias

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)

python/tvm/relax/transform/legalize_ops/nn.py

@@ -656,30 +656,13 @@ def _nn_attention(bb: BlockBuilder, call: Call) -> Expr:
         call.args[0],
         call.args[1],
         call.args[2],
-        None,
+        call.attrs.bias,
         call.attrs.scale,
         call.attrs.causal_mask,
         primfunc_name_hint="attention",
     )
 
 
-@register_legalize("relax.nn.attention_bias")
-def _nn_attention_bias(bb: BlockBuilder, call: Call) -> Expr:
-    assert (
-        call.attrs.window_size is None
-    ), "Legalization for sliding-window attention is not supported yet."
-    return bb.call_te(
-        _te_attention,
-        call.args[0],
-        call.args[1],
-        call.args[2],
-        call.args[3],
-        call.attrs.scale,
-        call.attrs.causal_mask,
-        primfunc_name_hint="attention_bias",
-    )
-
-
 @register_legalize("relax.nn.attention_var_len")
 def _nn_attention_var_len(bb: BlockBuilder, call: Call) -> Expr:
     raise RuntimeError("Legalization of attention_var_len op is not supported yet.")

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)