xe: sdpa: (tmp) verify graph api changes w/o innaccurate softmax alg

Author: pv-pterab-s

src/common/sdpa_utils.hpp

@@ -166,7 +166,7 @@ static inline status_t create_sdpa_pd(
         const memory_desc_t *q_md, const memory_desc_t *k_md,
         const memory_desc_t *v_md, const memory_desc_t *dst_md,
         const memory_desc_t *attn_mask_md, const memory_desc_t *scale_md,
-        data_type_t scale_dt, bool invert_scale, dim_t kv_head_number,
+        bool invert_scale, dim_t kv_head_number,
         attn_mask_type_t attn_mask_type, alg_kind_t softmax_alg,
         const primitive_attr_t *attr, const primitive_attr_t *kq_attr = nullptr,
         const primitive_attr_t *vs_attr = nullptr) {

src/gpu/intel/ocl/reusable_softmax.hpp

@@ -205,18 +205,18 @@ struct reusable_softmax_fwd_t : public gpu_primitive_t {
 
             conf.algorithm_number = [&]() { // -> int
                 if (arch != arch_t::xe_hpg) {
-                    if (rt_conf.softmax_axis_stride == 1
-                            && rt_conf.softmax_axis_size >= 128
-                            && nelems > (1 << 17)
-                            && dnnl::impl::utils::div_up(
-                                       rt_conf.softmax_axis_size,
-                                       conf.subgroup_size)
-                                    <= 1024)
-                        return vectorized;
-                    if (rt_conf.softmax_axis_stride == 1
-                            && rt_conf.softmax_axis_size <= conf.subgroup_size
-                            && nelems < (1 << 15))
-                        return small;
+                    // if (rt_conf.softmax_axis_stride == 1
+                    //         && rt_conf.softmax_axis_size >= 128
+                    //         && nelems > (1 << 17)
+                    //         && dnnl::impl::utils::div_up(
+                    //                    rt_conf.softmax_axis_size,
+                    //                    conf.subgroup_size)
+                    //                 <= 1024)
+                    //     return vectorized;
+                    // if (rt_conf.softmax_axis_stride == 1
+                    //         && rt_conf.softmax_axis_size <= conf.subgroup_size
+                    //         && nelems < (1 << 15))
+                    //     return small;
                 }
                 if (rt_conf.softmax_axis_size < 6 && nelems > 64000)
                     return many_reductions_per_workgroup;

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

@@ -351,8 +351,11 @@ status_t sdp_primitive_config_t::init(std::shared_ptr<subgraph_t> &sg,
     if (attn_mask_)
         md_mask = make_dnnl_memory_desc(attn_mask_->get_logical_tensor());
 
-    auto scale_dt = impl::data_type::undef;
-    if (scale_) scale_dt = scale_->get_logical_tensor().data_type;
+    dnnl::memory::desc scale_md;
+    if (scale_)
+        scale_md = {dims {1},
+                static_cast<data_type>(scale_->get_logical_tensor().data_type),
+                dnnl::memory::format_tag::a};
 
     dnnl::primitive_attr attr, qk_attr, vs_attr;
 
@@ -376,9 +379,9 @@ status_t sdp_primitive_config_t::init(std::shared_ptr<subgraph_t> &sg,
             ? alg_kind::softmax_accurate_inf_as_zero
             : alg_kind::softmax_accurate;
     CHECK(create_sdpa_pd(sdpa_pd_, p_engine.get(), md_q.get(), md_k.get(),
-            md_v.get(), md_dst.get(), md_mask.get(), dnnl::memory::desc().get(),
-            scale_dt, invert_scale_, kv_head_number_, mask_type_, softmax_alg,
-            attr.get(), qk_attr.get(), vs_attr.get()));
+            md_v.get(), md_dst.get(), md_mask.get(), scale_md.get(),
+            invert_scale_, kv_head_number_, mask_type_, softmax_alg, attr.get(),
+            qk_attr.get(), vs_attr.get()));
 
     auto status = sdpa_pd_->create_primitive(sdpa_prim_, p_engine.get());
 

src/graph/backend/dnnl/op_executable.hpp

@@ -2796,12 +2796,14 @@ struct sdpa_executable_t : public op_executable_t {
         auto md_dst = make_dnnl_memory_desc(
                 op->get_output_value(0)->get_logical_tensor());
 
-        auto scale_dt = impl::data_type::undef;
         size_t idx = 3;
+        dnnl::memory::desc scale_md;
         if (with_scale_)
-            scale_dt = op->get_input_value(idx++)
-                               ->get_logical_tensor()
-                               .data_type;
+            scale_md = {dims {1},
+                    static_cast<data_type>(op->get_input_value(idx++)
+                                    ->get_logical_tensor()
+                                    .data_type),
+                    dnnl::memory::format_tag::a};
 
         dnnl::memory::desc md_mask;
         with_explicit_mask_ = mask_type_ == attn_mask_type::buffer;
@@ -2826,8 +2828,8 @@ struct sdpa_executable_t : public op_executable_t {
                 : alg_kind::softmax_accurate;
         status_t s = create_sdpa_pd(sdpa_pd_, p_engine.get(), md_q.get(),
                 md_k.get(), md_v.get(), md_dst.get(), md_mask.get(),
-                dnnl::memory::desc().get(), scale_dt, is_invert_scale_,
-                kv_head_number, mask_type_, softmax_alg, attr.get());
+                scale_md.get(), is_invert_scale_, kv_head_number, mask_type_,
+                softmax_alg, attr.get());
         if (s != dnnl::impl::status::success) {
             is_initialized_ = false;
         } else {