[GPU] Fix accuracy issue in PagedAttention kernel for large prompts (…

…4K/8K tokens) by correcting index calculation in sub_group_broadcast function to ensure accurate data broadcasting within the subgroup
sshlyapn · Nov 6, 2024 · f471153 · f471153
1 parent 5833781
commit f471153
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/pa_sdpa_opt.cl b/src/plugins/intel_gpu/src/kernel_selector/cl_kernels/pa_sdpa_opt.cl
@@ -436,7 +436,7 @@ KERNEL(pa_sdpa_finalization_stage)(
                                         partition_num * HEAD_SIZE +
                                         head_size_idx;
             OUTPUT_TYPE out_val = tmp_out[tmp_out_offset];
-            acc += TO_SOFTMAX_ACCUMULATOR_TYPE(out_val) * TO_SOFTMAX_ACCUMULATOR_TYPE(sub_group_broadcast(exp_sum[partition_num / SUBGROUP_SIZE], partition_num)) / TO_SOFTMAX_ACCUMULATOR_TYPE(global_sum);
+            acc += TO_SOFTMAX_ACCUMULATOR_TYPE(out_val) * TO_SOFTMAX_ACCUMULATOR_TYPE(sub_group_broadcast(exp_sum[partition_num / SUBGROUP_SIZE], partition_num % SUBGROUP_SIZE)) / TO_SOFTMAX_ACCUMULATOR_TYPE(global_sum);
         }
         const uint out_offset = seq_idx * (HEADS_NUM * HEAD_SIZE) +
                                 head_num_idx * HEAD_SIZE +