From b3bcf853afdfd977673d57f8f3e471272c899b26 Mon Sep 17 00:00:00 2001
From: Umar Arshad <umar.arshad@intel.com>
Date: Tue, 17 Dec 2024 17:38:32 -0800
Subject: [PATCH] xe: sdpa: Fix KQ gemm alignment for the K tensor

---
 src/gpu/intel/ocl/micro_sdpa.cpp | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/src/gpu/intel/ocl/micro_sdpa.cpp b/src/gpu/intel/ocl/micro_sdpa.cpp
index 0555ee2d1cc..e1076623805 100644
--- a/src/gpu/intel/ocl/micro_sdpa.cpp
+++ b/src/gpu/intel/ocl/micro_sdpa.cpp
@@ -264,7 +264,10 @@ 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 = static_cast<int>(
+            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());