xe: sdpa: compute scale/iscale on host for host-side-scale

Author: pv-pterab-s

src/gpu/intel/ocl/micro_sdpa.cl

@@ -212,7 +212,7 @@ __attribute__((intel_reqd_sub_group_size(SUBGROUP_SIZE))) kernel void
 micro_sdpa(const global KEY_DATA_T *K, const global QRY_DATA_T *Q,
         const global VAL_DATA_T *V, global DST_DATA_T *A,
 #if HOST_SIDE_SCALE
-        const SCALE_DATA_T scale_value,
+        const float scale, const float iscale,
 #else
         const global SCALE_DATA_T *scale_ptr,
 #endif
@@ -339,29 +339,23 @@ micro_sdpa(const global KEY_DATA_T *K, const global QRY_DATA_T *Q,
 #endif
     }
 
+#if HOST_SIDE_SCALE == 0
     /* Load scale */
     float scale = 1.0f;
     float iscale = 1.0f;
     if (k0end > 0) {
 #if WITH_ATTN_SCALE
 #if INVERT_SCALE
-#if HOST_SIDE_SCALE
-        iscale = SCALES_TO_FLOAT(scale_value);
-#else
         iscale = SCALES_TO_FLOAT(*scale_ptr);
-#endif
         scale = native_recip(iscale);
-#else
-#if HOST_SIDE_SCALE
-        scale = SCALES_TO_FLOAT(scale_value);
 #else
         scale = SCALES_TO_FLOAT(*scale_ptr);
-#endif
         iscale = native_recip(scale);
 #endif
 #endif
         scale *= 1.442695f; // log2(e)
     }
+#endif
 
 #ifdef PREFETCH_K0
     if (k0end > 0) {

src/gpu/intel/ocl/micro_sdpa.cpp

@@ -564,15 +564,31 @@ status_t micro_sdpa_t::execute(const exec_ctx_t &ctx) const {
     arg_list.append(dst);
 
     if (pd()->scale_md()->format_kind == format_kind::host_side_scalar) {
-        void *handle = scale.data_handle();
-        switch (pd()->scale_md()->data_type) {
-            case data_type::f16: arg_list.append(*(float16_t *)handle); break;
-            case data_type::bf16:
-                arg_list.append(*(unsigned short *)handle);
-                break;
-            case data_type::f32: arg_list.append(*(float *)handle); break;
-            default: break;
+        const void *handle = scale.data_handle();
+
+        const float scale_value = [&]() {
+            switch (pd()->scale_md()->data_type) {
+                case data_type::f16: return float(*(float16_t *)handle);
+                case data_type::bf16: return float(*(bfloat16_t *)handle);
+                case data_type::f32: return *(float *)handle;
+                default: return 0.f;
+            }
+        }();
+
+        float scale = 1.f, iscale = 1.f;
+        if (pd()->with_attn_scale()) {
+            if (pd()->desc()->invert_scale) {
+                iscale = scale_value;
+                scale = 1.f / iscale;
+            } else {
+                scale = scale_value;
+                iscale = 1.f / scale;
+            }
         }
+        scale *= 1.442695f;
+
+        arg_list.append(scale);
+        arg_list.append(iscale);
     } else {
         arg_list.append(scale);
     }

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;