graph: dnnl: refine check and layout propagation

Author: xiang1guo

src/graph/backend/dnnl/kernels/sdp_primitive_config.cpp

@@ -286,10 +286,15 @@ status_t sdp_primitive_config_t::initial_check(
 
     VCHECK_SDP_PRIMITIVE(q_id != -1 && k_id != -1 && v_id != -1,
             status::unimplemented, "Q, K, V are not found");
-    VCHECK_SDP_PRIMITIVE(ltw(inputs[q_id]).vdims().size() == 4
-                    && ltw(inputs[k_id]).vdims().size() == 4
-                    && ltw(inputs[v_id]).vdims().size() == 4,
-            status::unimplemented, "Q, K, V should be 4-dims");
+
+    // Note: sdpa_primitive_v1 kenrel accept 5D GQA pattern, and will reshape to
+    // 4D in later compilation pass.
+    if (!v1_kenrel) {
+        VCHECK_SDP_PRIMITIVE(ltw(inputs[q_id]).vdims().size() == 4
+                        && ltw(inputs[k_id]).vdims().size() == 4
+                        && ltw(inputs[v_id]).vdims().size() == 4,
+                status::unimplemented, "Q, K, V should be 4-dims");
+    }
 
     // sdp_primitive only supports single scale value.
     if (scale) {

src/graph/backend/dnnl/layout_propagator.cpp

@@ -1578,13 +1578,38 @@ status_t layout_propagator_for_sdpa(std::shared_ptr<op_t> &op,
 
     value_ptr dst_val = op->get_output_value(0);
     const logical_tensor_t &out_lt = dst_val->get_logical_tensor();
-
     dnnl::memory::desc expected_md;
-    // Set default output layout format for sdpa as acbd
+
     if (ltw(out_lt).is_any()) {
-        expected_md = {ltw(out_lt).vdims(),
-                static_cast<dnnl::memory::data_type>(ltw(out_lt).data_type()),
-                dnnl::memory::format_tag::acbd};
+        // For GQA, we need to check the layout of the dnnl_reshape output
+        // following dnnl_sdpa, which is given by the user.
+        if (!dst_val->get_consumers().empty()) {
+            const auto &consumer_op = dst_val->get_consumers()[0].get_op();
+            const logical_tensor_t &consumer_out
+                    = consumer_op.get_output_value(0)->get_logical_tensor();
+            if (consumer_op.get_kind() == op_kind::dnnl_reshape
+                    && ltw(consumer_out).ndims() == 5
+                    && ltw(consumer_out).is_strided()) {
+                const auto &ori_strides = ltw(consumer_out).vstrides();
+                std::vector<dim_t> strides = {ori_strides[0], ori_strides[2],
+                        ori_strides[3], ori_strides[4]};
+                dnnl::memory::desc tmp_md {ltw(out_lt).vdims(),
+                        static_cast<dnnl::memory::data_type>(
+                                ltw(out_lt).data_type()),
+                        strides};
+                expected_md = tmp_md;
+            } else {
+                expected_md = {ltw(out_lt).vdims(),
+                        static_cast<dnnl::memory::data_type>(
+                                ltw(out_lt).data_type()),
+                        dnnl::memory::format_tag::acbd};
+            }
+        } else {
+            expected_md = {ltw(out_lt).vdims(),
+                    static_cast<dnnl::memory::data_type>(
+                            ltw(out_lt).data_type()),
+                    dnnl::memory::format_tag::acbd};
+        }
     } else {
         expected_md = make_dnnl_memory_desc(out_lt);
     }