Use more threads

Author: sshlyapn

src/plugins/intel_gpu/src/kernel_selector/cl_kernels/pa_sdpa_opt.cl

@@ -7,7 +7,7 @@
 #include "include/batch_headers/sub_group_block_write.cl"
 #include "include/batch_headers/sub_group_shuffle.cl"
 
-#define SUBGROUPS_PER_WG (HEAD_SIZE / SUBGROUP_SIZE)
+#define SUBGROUPS_PER_WG ((HEAD_SIZE / SUBGROUP_SIZE) * SG_SCALE_FACTOR)
 #define PAGED_ATTENTION_BLOCKS_PER_PARTITION (SEQ_LEN_PARTITION_SIZE / PAGED_ATTENTION_BLOCK_SIZE)
 
 #if HEAD_SIZE > 128
@@ -75,7 +75,7 @@ KERNEL(pa_sdpa_opt)(
 
     const uint seq_idx = get_global_id(0);
     const uint head_num_idx = get_global_id(1);
-    const uint head_size_idx = get_global_id(2);
+    const uint head_size_idx = get_local_id(2);
     const uint sglid = get_sub_group_local_id();
     const uint sgid = get_sub_group_id();
     const uint total_partitions_num = get_num_groups(2);
@@ -93,7 +93,6 @@ KERNEL(pa_sdpa_opt)(
 #endif
 
     const uint partition_idx = get_group_id(2);
-    const uint block_start_idx = partition_idx * SEQ_LEN_PARTITION_SIZE / PAGED_ATTENTION_BLOCK_SIZE;
 
     if (partition_idx * SEQ_LEN_PARTITION_SIZE >= seq_len) {
         return;
@@ -336,6 +335,15 @@ KERNEL(pa_sdpa_opt)(
         OUTPUT_TYPE acc = OUTPUT_VAL_ZERO;
 
         const uint partition_seq_len = min(seq_len - partition_idx * SEQ_LEN_PARTITION_SIZE, (uint)SEQ_LEN_PARTITION_SIZE);
+
+#if SG_SCALE_FACTOR > 1
+        const uint block_start_idx = (sgid / (SUBGROUPS_PER_WG / SG_SCALE_FACTOR)) * (SEQ_LEN_PARTITION_SIZE / SG_SCALE_FACTOR / SUBGROUP_SIZE);
+        const uint block_end_idx = min(block_start_idx + (SEQ_LEN_PARTITION_SIZE / SG_SCALE_FACTOR / SUBGROUP_SIZE), partition_seq_len / SUBGROUP_SIZE);
+#else
+        const uint block_start_idx = 0;
+        const uint block_end_idx = partition_seq_len / SUBGROUP_SIZE;
+#endif
+
         uint blocks_num_per_partition = min(total_blocks_num - partition_idx * PAGED_ATTENTION_BLOCKS_PER_PARTITION, (uint)PAGED_ATTENTION_BLOCKS_PER_PARTITION);
 
         uint leftovers = blocks_num_per_partition * PAGED_ATTENTION_BLOCK_SIZE - partition_seq_len;
@@ -346,7 +354,7 @@ KERNEL(pa_sdpa_opt)(
 
         const uint start_block_idx = block_indices_begins[subsequence_idx] + partition_idx * PAGED_ATTENTION_BLOCKS_PER_PARTITION;
 
-        for (uint block_num = 0; block_num < blocks_num_per_partition; block_num++) {
+        for (uint block_num = block_start_idx; block_num < block_end_idx; block_num++) {
 #ifdef BROADCAST_GROUP_SIZE
             const uint head_idx = head_num_idx / BROADCAST_GROUP_SIZE;
 #else
@@ -389,6 +397,10 @@ KERNEL(pa_sdpa_opt)(
             }
         }
 
+
+#if SG_SCALE_FACTOR > 1
+        if (sgid >= SUBGROUPS_PER_WG / SG_SCALE_FACTOR) {
+#endif
         if (leftovers != 0) {
 #ifdef BROADCAST_GROUP_SIZE
             const uint head_idx = head_num_idx / BROADCAST_GROUP_SIZE;
@@ -429,6 +441,32 @@ KERNEL(pa_sdpa_opt)(
             }
         }
 
+
+#if SG_SCALE_FACTOR > 1
+        }
+#endif
+
+#if SG_SCALE_FACTOR > 1
+        if ((partition_seq_len > (SEQ_LEN_PARTITION_SIZE / SG_SCALE_FACTOR)) || (leftovers != 0)) {
+            barrier(CLK_LOCAL_MEM_FENCE);
+
+            if (sgid >= SUBGROUPS_PER_WG / SG_SCALE_FACTOR) {
+                // Reuse slm_qk_vals SLM to sum-up results between two groups of subgroups
+                slm_qk_vals[head_size_idx] = acc;
+            }
+
+            barrier(CLK_LOCAL_MEM_FENCE);
+            
+            if (sgid < SUBGROUPS_PER_WG / SG_SCALE_FACTOR) {
+                acc += slm_qk_vals[head_size_idx];
+            }
+        }
+#endif
+
+#if SG_SCALE_FACTOR > 1
+        if (sgid < SUBGROUPS_PER_WG / SG_SCALE_FACTOR) {
+#endif
+
         if (seq_len > SEQ_LEN_PARTITION_SIZE) {
             const uint tmp_out_offset = seq_idx * (HEADS_NUM * HEAD_SIZE * total_partitions_num) +
                                         head_num_idx * (HEAD_SIZE * total_partitions_num) +
@@ -446,6 +484,11 @@ KERNEL(pa_sdpa_opt)(
             output[output_offset] = acc;
         }
 
+#if SG_SCALE_FACTOR > 1
+        }
+#endif
+
+
     }
 }
 

src/plugins/intel_gpu/src/kernel_selector/kernels/sdpa/pa_sdpa_kernel_opt.cpp

@@ -24,6 +24,19 @@ constexpr size_t subgroup_size = 16;
 constexpr size_t seq_len_partition_size = 256;
 constexpr size_t paged_attention_block_size = 16;
 constexpr Datatype softmax_acc_dt = Datatype::F32;
+
+size_t get_sg_number_scale_factor(const Params& params, size_t head_size, size_t kernel_type) {
+    return 1;
+    const size_t optimal_scale_factor = 2;
+    if (kernel_type == KernelsTypes::SINGLE_TOKEN ||
+        kernel_type == KernelsTypes::MULTI_TOKENS) {
+        if (head_size * optimal_scale_factor <= params.engineInfo.maxWorkGroupSize) {
+            return optimal_scale_factor;
+        }
+    }
+
+    return 1;
+}
 }  // namespace
 
 static std::string GetKernelName(std::string base_name, KernelsTypes type) {
@@ -211,6 +224,7 @@ JitConstants PagedAttentionSDPAKernelOpt::GetJitConstants(const pa_sdpa_params&
     jit.AddConstant(MakeJitConstant("PAGED_ATTENTION_BLOCK_SIZE", paged_attention_block_size));
     jit.AddConstant(MakeJitConstant("SUBGROUP_SIZE", subgroup_size));
     jit.AddConstant(MakeJitConstant("IS_KV_COMPRESSED", params.conf.is_kv_compressed));
+    jit.AddConstant(MakeJitConstant("SG_SCALE_FACTOR", get_sg_number_scale_factor(params, config.head_size, kernel_idx)));
 
     if (params.conf.is_kv_compressed) {
         auto scales_zp_size = 2 * 2; // FP16 * (scale + zp)
@@ -272,10 +286,11 @@ CommonDispatchData PagedAttentionSDPAKernelOpt::SetDefault(const pa_sdpa_params&
         const size_t head_size = static_cast<size_t>(params.conf.head_size);
 
         if (kernel_idx == KernelsTypes::SINGLE_TOKEN || kernel_idx == KernelsTypes::MULTI_TOKENS) {
+            auto sg_scale = get_sg_number_scale_factor(params, head_size, kernel_idx);
             dispatch_data.gws = { total_tokens,
                                   heads_num,
-                                  head_size * num_of_partitions };
-            dispatch_data.lws = { 1, 1, head_size };
+                                  head_size * num_of_partitions * sg_scale };
+            dispatch_data.lws = { 1, 1, head_size * sg_scale };
         } else if (kernel_idx == KernelsTypes::SCORES_CALCULATION) {
             const auto& past_lens = params.inputs[3];
             const auto subsequences_number = past_lens.Batch().v;