fix msc testcase

Author: mshr-h

python/tvm/contrib/msc/core/transform/pattern.py

@@ -330,15 +330,17 @@ def make_relax_attention_pattern() -> (
     q_trans = relax_pattern.is_op("relax.permute_dims")(weight_q)
     k_trans = relax_pattern.is_op("relax.permute_dims")(weight_k)
     v_trans = relax_pattern.is_op("relax.permute_dims")(weight_v)
-    out = relax_pattern.is_op("relax.nn.attention")(q_trans, k_trans, v_trans)
+    attention = relax_pattern.is_op("relax.nn.attention")(q_trans, k_trans, v_trans)
+    out = relax_pattern.is_op("relax.permute_dims")(attention)
     annotations = {
         "weight_q": weight_q,
         "weight_k": weight_k,
         "weight_v": weight_v,
         "q_trans": q_trans,
         "k_trans": k_trans,
         "v_trans": v_trans,
-        "attention": out,
+        "attention": attention,
+        "out": out,
     }
     return out, annotations
 
@@ -378,7 +380,8 @@ def make_relax_mask_attention_pattern() -> (
     q_trans = relax_pattern.is_op("relax.permute_dims")(weight_q)
     k_trans = relax_pattern.is_op("relax.permute_dims")(weight_k)
     v_trans = relax_pattern.is_op("relax.permute_dims")(weight_v)
-    out = relax_pattern.is_op("relax.nn.attention_bias")(q_trans, k_trans, v_trans, mask)
+    attention = relax_pattern.is_op("relax.nn.attention_bias")(q_trans, k_trans, v_trans, mask)
+    out = relax_pattern.is_op("relax.permute_dims")(attention)
     annotations = {
         "weight_q": weight_q,
         "weight_k": weight_k,
@@ -387,7 +390,8 @@ def make_relax_mask_attention_pattern() -> (
         "q_trans": q_trans,
         "k_trans": k_trans,
         "v_trans": v_trans,
-        "attention": out,
+        "attention": attention,
+        "out": out,
     }
     return out, annotations
 

src/contrib/msc/framework/tvm/relax_opcode.cc

@@ -107,6 +107,7 @@ class RelaxAttentionCodeGen : public RelaxOpCode {
           .op_list_arg<int>(axes_key, "axes");
     }
     stack_.op_call().op_inputs_arg(false).op_arg<float>("scale").op_str_arg("causal_mask");
+    stack_.op_call("relax.op.permute_dims").op_output_arg().op_list_arg<int>("axes_3", "axes");
   }
 };
 

tests/python/contrib/test_msc/test_graph_build.py

@@ -2362,12 +2362,7 @@ def forward(self, q_data, k_data, v_data):
             {"name": "inp_2", "shape": [1, 8, seq, 64], "dtype": "float32", "layout": "ACBD"},
         ],
         "outputs": [
-            {
-                "name": "attention",
-                "shape": [1, seq, 8, 64],
-                "dtype": "float32",
-                "layout": "ABCD",
-            }
+            {"name": "attention", "shape": [1, 8, seq, 64], "dtype": "float32", "layout": "ABCD"}
         ],
         "nodes": {"total": 4, "input": 3, "msc.attention": 1},
     }
@@ -2396,7 +2391,7 @@ def forward(self, q_data, k_data, v_data, mask):
         "outputs": [
             {
                 "name": "attention_bias",
-                "shape": [1, seq, 8, 64],
+                "shape": [1, 8, seq, 64],
                 "dtype": "float32",
                 "layout": "ABCD",
             }