wip

Author: masahi

src/relax/op/nn/attention.cc

@@ -79,8 +79,8 @@ StructInfo InferStructInfoAttention(const Call& call, const BlockBuilder& ctx) {
   };
   diag_equal(num_batches, k_shape->values[0], "query", "key", "batch size");
   diag_equal(num_batches, v_shape->values[0], "query", "value", "batch size");
-  diag_equal(num_heads, k_shape->values[2], "query", "key", "number of heads");
-  diag_equal(num_heads, v_shape->values[2], "query", "value", "number of heads");
+  // diag_equal(num_heads, k_shape->values[2], "query", "key", "number of heads");
+  // diag_equal(num_heads, v_shape->values[2], "query", "value", "number of heads");
   diag_equal(num_keys, v_shape->values[1], "key", "value", "sequence length");
   diag_equal(head_dim, k_shape->values[3], "query", "key", "dimension of heads");
 

tests/python/relax/test_codegen_cutlass.py

@@ -2006,7 +2006,8 @@ def main(
             return lv6_1
 
     Module["main"] = rewrite_attention(Module["main"])
-    print(Module)
+    mod = partition_for_cutlass(Module)
+    print(mod)
 
 if __name__ == "__main__":
     # tvm.testing.main()