xe: sdpa: Fix KQ gemm alignment for the K tensor

oneapi-src · Dec 18, 2024 · 02fd3ad · 02fd3ad
1 parent 7e450f8
commit 02fd3ad
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/gpu/intel/ocl/micro_sdpa.cpp b/src/gpu/intel/ocl/micro_sdpa.cpp
@@ -264,7 +264,9 @@ status_t micro_sdpa_t::pd_t::init_microkernels(impl::engine_t *engine) {
 
     problem_kq.B.layout = MatrixLayout::Pr;
     problem_kq.C.layout = MatrixLayout::T;
-    problem_kq.A.setAlignment(alignmentForLD(d->head_size() * problem.Ta));
+    const memory_desc_wrapper key_mdw(key_md());
+    auto ldk = gemm_desc_t::get_ld(*key_md()) * key_mdw.data_type_size();
+    problem_kq.A.setAlignment(alignmentForLD(ldk));
     problem_kq.B.setAlignment(64); // Q is packed in VNNI format in SLM
     problem_kq.B.crosspack = 2;
     problem_kq.B.tileR = into<uint16_t>(d_max());