Fix micro_sdpa kernel arguments

Author: sshlyapn

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

@@ -152,12 +152,12 @@ KERNEL(micro_sdpa)(OPTIONAL_SHAPE_INFO_ARG
 #if WITH_SCALE
         global SCALE_DATA_T *scale_ptr,
 #endif
-        int d,
 #if IS_PAGED_ATTENTION
         const __global int* blocked_indexes_start,
         const __global int* blocked_indexes_end,
         const __global int* gws_seq_indexes_correspondence
 #else
+        int d,
         int k,
         int q
 #endif
@@ -172,11 +172,18 @@ KERNEL(micro_sdpa)(OPTIONAL_SHAPE_INFO_ARG
     const uint q_tile_idx = get_group_id(0);
     const uint block_start_pos = blocked_indexes_start[q_tile_idx];
     const uint block_end_pos = blocked_indexes_end[q_tile_idx];
-    const uint subsequence_q_tile_idx = block_start_pos - subsequence_begins[gws_seq_indexes_correspondence[q_tile_idx]];
+    const uint gws_mapping = gws_seq_indexes_correspondence[q_tile_idx];
+    const uint subsequence_begin = subsequence_begins[gws_mapping];
+    const uint subsequence_begin_next = subsequence_begins[gws_mapping + 1];
+    const uint subsequence_q_tile_idx = block_start_pos - subsequence_begin;
     // const uint sequence_idx_end = block_end_pos - block_start_pos;
-    const uint subsequence_begin = subsequence_begins[gws_seq_indexes_correspondence[q_tile_idx]];
-    const int k = subsequence_begins[gws_seq_indexes_correspondence[q_tile_idx] + 1] - subsequence_begins[gws_seq_indexes_correspondence[q_tile_idx]];
+    const int k = subsequence_begins[gws_mapping + 1] - subsequence_begin;
     const int q = k;
+    const int d = HEAD_SIZE;
+    if (get_global_id(0) == 0 && get_global_id(1) == 0 && get_global_id(2) == 0) {
+        printf("q_tile_idx=%d block_start_pos=%d block_end_pos=%d gws_mapping=%d subsequence_begin=%d subsequence_begin_next=%d subsequence_q_tile_idx=%d k=%d d=%d\n",
+                q_tile_idx, block_start_pos, block_end_pos, gws_mapping, subsequence_begin, subsequence_begin_next, subsequence_q_tile_idx, k, d);
+    }
 #endif
     uint sg_ij = sub_group_broadcast(get_local_id(1), 0);
     uint b0 = get_group_id(1);

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

@@ -508,7 +508,9 @@ bool SDPAKernelMicro::Validate(const Params& p) const {
         print_arr(q_dims, q_dims.size(), "q_dims");
         print_arr(k_dims, k_dims.size(), "k_dims");
         print_arr(v_dims, v_dims.size(), "v_dims");
-        std::cout << Q_num_heads_dim.is_dynamic << " " << K_num_heads_dim.is_dynamic << " " << V_num_heads_dim.is_dynamic << " " << K_num_heads_dim.v << " " << V_num_heads_dim.v << "\n";
+        std::cout << Q_num_heads_dim.is_dynamic << " "
+                  << K_num_heads_dim.is_dynamic << " "
+                  << V_num_heads_dim.is_dynamic << " " << K_num_heads_dim.v << " " << V_num_heads_dim.v << "\n";
         return false;
     }
 
@@ -555,6 +557,7 @@ JitConstants SDPAKernelMicro::GetJitConstants(const sdpa_params& params, const m
     jit.AddConstant(MakeJitConstant("SUBGROUP_SIZE", subgroup_size(prim_params.engineInfo.arch)));
     jit.AddConstant(MakeJitConstant("INVERT_SCALE", false));
     jit.AddConstant(MakeJitConstant("SCALE_DATA_T", "half"));
+    jit.AddConstant(MakeJitConstant("HEAD_SIZE", head_size));
 
     jit.AddConstant(MakeJitConstant("WITH_ATTN_MASK", sdpa_inputs > 3));
     jit.AddConstant(MakeJitConstant("WITH_SCALE", sdpa_inputs > 4));
@@ -733,6 +736,8 @@ CommonDispatchData SDPAKernelMicro::SetDefault(const sdpa_params& params, const
     auto seq_length = get_seq_length(params.inputs[0], params.input0_order).v;
     if (params.conf.is_paged_attention) {
         seq_length = params.conf.paged_attention_aligned_seq_len;
+        GPU_DEBUG_TRACE_DETAIL << "seq_len=" << seq_length << "\n";
+        GPU_DEBUG_TRACE_DETAIL << "wg_tile_q=" << wg_tile_q << "\n";
     }
 
     dispatch_data.gws[0] *= CeilDiv(seq_length, wg_tile_q);
@@ -767,6 +772,14 @@ clKernelData SDPAKernelMicro::get_kernel_data(const sdpa_params& params, bool is
     kernel.params.arguments.push_back({ArgumentDescriptor::Types::INPUT, 2}); // V
     kernel.params.arguments.push_back({ArgumentDescriptor::Types::OUTPUT, 0}); // A
 
+
+    if (params.conf.is_paged_attention) {
+        kernel.params.arguments.push_back({ArgumentDescriptor::Types::INPUT, 3}); // subsequence_begins
+
+        kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 0});
+        kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 1});
+        kernel.params.arguments.push_back({ArgumentDescriptor::Types::INTERNAL_BUFFER, 2});
+    }
     // if (params.inputs.size() >= 4)
     //     kernel.params.arguments.push_back({ArgumentDescriptor::Types::INPUT, 3}); // mask
     // if (params.inputs.size() >= 5)