--- /dev/null
+// Flash attention via oneMKL GEMM (XMX-accelerated).
+// Uses column_major::gemm for Q*K^T and S*V matmuls
+// with an online softmax SYCL kernel.
+//
+// All GQA query heads sharing a KV head are batched into single
+// GEMM calls, amortizing MKL launch overhead across K and V reuse.
+//
+
+#include "common.hpp"
+#include "fattn-common.hpp"
+#include "fattn-buffers.hpp"
+#include "convert.hpp"
+#include "fattn.hpp"
+
+#include <oneapi/mkl.hpp>
+#include <cstdio>
+#include <chrono>
+
+#define MKL_FA_CHUNK_SIZE_KV 8192
+
+// Number of query rows processed per tile. The score buffers (KQ_f32, S_f16)
+// are sized q_tile_rows * chunk_size, so this bounds their footprint
+// regardless of batch size (n_query_rows = n_queries * gqa_ratio). A typical
+// single-ubatch prefill (e.g. ubatch 1024 * gqa 8 = 8192 rows) is exactly one
+// tile, so it runs with no extra iterations. Larger batches tile and stay
+// bounded. Override with GGML_SYCL_MKL_FA_Q_TILE.
+#define MKL_FA_Q_TILE 8192
+
+#define MKL_FA_WG_SIZE 256
+
+using oneapi::mkl::transpose;
+using oneapi::mkl::blas::column_major::gemm;
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+// Pack all GQA Q heads for one KV head into fp16, applying q_scale.
+// Launches one kernel per GQA group — each kernel copies exactly
+// n_queries * DKQ elements using the per-group dst offset and
+// per-head source stride.
+static void mkl_fa_pack_q_fp16(
+ dpct::queue_ptr stream,
+ sycl::half * __restrict dst,
+ const float * __restrict q_src,
+ int n_queries, int n_query_rows, int DKQ,
+ int gqa_ratio, int kvh_base_head,
+ float q_scale, int64_t q_row_stride, int64_t q_head_stride,
+ int64_t wg_size) {
+
+ for (int iqg = 0; iqg < gqa_ratio; iqg++) {
+ int iqh = kvh_base_head + iqg;
+ sycl::half * dst_g = dst + (int64_t)iqg * n_queries * DKQ;
+
+ const int64_t n_elem = (int64_t)n_queries * DKQ;
+ const int64_t wg = ((n_elem + wg_size - 1) / wg_size) * wg_size;
+
+ stream->submit([&](sycl::handler & cgh) {
+ cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+ [=](sycl::nd_item<1> item) {
+ int64_t e = item.get_global_id(0);
+ if (e >= n_elem) return;
+
+ int64_t q = e / DKQ;
+ int64_t d = e - q * DKQ;
+
+ // Stride-aware source offset: handles permuted,
+ // sliced, or contiguous Q tensor layouts.
+ int64_t src_off = d
+ + q * q_row_stride
+ + (int64_t)iqh * q_head_stride;
+
+ dst_g[e] = sycl::half(
+ q_src[src_off] * q_scale);
+ });
+ });
+ }
+}
+
+// Zero-initialize the online softmax state arrays.
+// KQ_max → -inf, KQ_sum → 0, VKQ_accum → 0.
+// Merged into one kernel to avoid per-array launch overhead.
+static void mkl_fa_init_softmax_state(
+ dpct::queue_ptr stream,
+ float * kmax, float * ksum, float * vacc,
+ int n_query_rows, int DV, int64_t wg_size) {
+
+ const float neg_inf = -1e30f;
+ const int64_t n_maxsum = n_query_rows;
+ const int64_t n_vacc = (int64_t)n_query_rows * DV;
+ const int64_t total = (n_vacc > n_maxsum) ? n_vacc : n_maxsum;
+ const int64_t wg = ((total + wg_size - 1) / wg_size) * wg_size;
+
+ stream->submit([&](sycl::handler & cgh) {
+ cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+ [=](sycl::nd_item<1> item) {
+ int64_t i = item.get_global_id(0);
+ if (i < n_maxsum) {
+ kmax[i] = neg_inf;
+ ksum[i] = 0.0f;
+ }
+ if (i < n_vacc) {
+ vacc[i] = 0.0f;
+ }
+ });
+ });
+}
+
+// Online softmax over one KV chunk for a tile of GQA query rows.
+// The tile spans absolute rows [q0, q0 + q_rows). Score buffers
+// (KQ_f32/S_f16) are indexed RELATIVE to the tile; the persistent state
+// (VKQ_accum/KQ_max/KQ_sum) and mask are indexed by ABSOLUTE row.
+// For each row: find local max → rescale previous VKQ_accum →
+// compute exp(s - max) → write S_f16 → update running max/sum.
+static void mkl_fa_online_softmax_chunk(
+ dpct::queue_ptr stream,
+ float * __restrict KQ_f32,
+ sycl::half * __restrict S_f16,
+ float * __restrict KQ_max,
+ float * __restrict KQ_sum,
+ float * __restrict VKQ_accum,
+ int q0, int q_rows, int n_queries, int DV,
+ int chunk_size, int chunk_start,
+ int kvh_head, int gqa_ratio,
+ const sycl::half * mask_data, int64_t mask_head_stride,
+ int64_t mask_row_stride, int mask_n_heads,
+ float logit_softcap, int64_t wg_size) {
+
+ const int64_t wg = ((q_rows + wg_size - 1) / wg_size) * wg_size;
+
+ stream->submit([&](sycl::handler & cgh) {
+ cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+ [=](sycl::nd_item<1> item) {
+ int jc_rel = item.get_global_id(0);
+ if (jc_rel >= q_rows) return;
+ int jc_abs = q0 + jc_rel;
+
+ const int gqa_group = jc_abs / n_queries;
+ const int q_row = jc_abs % n_queries;
+
+ // Score buffers are tile-local (relative index).
+ const float * __restrict KQ_row = KQ_f32
+ + jc_rel * (int64_t)chunk_size;
+ // Persistent accumulator is full-sized (absolute index).
+ float * __restrict vkq = VKQ_accum
+ + jc_abs * (int64_t)DV;
+
+ const sycl::half * mask_h = nullptr;
+ int64_t m_stride = 0;
+ if (mask_data) {
+ int m_head = (mask_n_heads > 1)
+ ? (kvh_head + gqa_group) : 0;
+ mask_h = mask_data + (int64_t)m_head * mask_head_stride;
+ m_stride = mask_row_stride;
+ }
+
+ // Row-wise local maximum (softcap before mask)
+ float local_max = -1e30f;
+ for (int i = 0; i < chunk_size; i++) {
+ float s = KQ_row[i];
+ if (logit_softcap != 0.0f) {
+ s = logit_softcap * sycl::tanh(s);
+ }
+ if (mask_h) {
+ s += (float)mask_h[q_row * m_stride
+ + (chunk_start + i)];
+ }
+ if (s > local_max) local_max = s;
+ }
+
+ // Rescale previous accumulator by exp(old_max - new_max)
+ float old_max = KQ_max[jc_abs];
+ float new_max = (old_max > local_max) ? old_max : local_max;
+ float rescale = (old_max < -1e29f) ? 1.0f
+ : sycl::native::exp(old_max - new_max);
+
+ for (int v = 0; v < DV; v++) {
+ vkq[v] *= rescale;
+ }
+
+ // Softmax and write S_f16 (tile-local index)
+ float local_sum = 0.0f;
+ sycl::half * __restrict S_row = S_f16
+ + jc_rel * (int64_t)chunk_size;
+
+ for (int i = 0; i < chunk_size; i++) {
+ float s = KQ_row[i];
+ if (logit_softcap != 0.0f) {
+ s = logit_softcap * sycl::tanh(s);
+ }
+ if (mask_h) {
+ s += (float)mask_h[q_row * m_stride
+ + (chunk_start + i)];
+ }
+ float val = sycl::native::exp(s - new_max);
+ S_row[i] = sycl::half(val);
+ local_sum += val;
+ }
+
+ KQ_sum[jc_abs] = KQ_sum[jc_abs] * rescale + local_sum;
+ KQ_max[jc_abs] = new_max;
+ });
+ });
+}
+
+// Write one GQA group's normalized output to its destination head.
+static void mkl_fa_normalize_head(
+ dpct::queue_ptr stream,
+ float * __restrict dst_batch,
+ const float * __restrict VKQ_accum,
+ const float * __restrict KQ_sum,
+ int iqh, int n_queries, int DV, int n_q_heads,
+ int64_t src_offset, int64_t wg_size) {
+
+ const int64_t wg = ((n_queries + wg_size - 1) / wg_size) * wg_size;
+
+ stream->submit([&](sycl::handler & cgh) {
+ cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+ [=](sycl::nd_item<1> item) {
+ int jc = item.get_global_id(0);
+ if (jc >= n_queries) return;
+
+ int ksum_idx = (int)(src_offset / DV) + jc;
+ float inv_sum = 1.0f / KQ_sum[ksum_idx];
+ const float * __restrict src = VKQ_accum
+ + src_offset + jc * (int64_t)DV;
+ // Interleaved dst layout (matching TILE):
+ // rows alternate between heads, then increment query.
+ // offset = (query * n_q_heads + head) * DV
+ float * __restrict dst_row = dst_batch
+ + ((int64_t)jc * n_q_heads + iqh) * (int64_t)DV;
+
+ for (int v = 0; v < DV; v++) {
+ dst_row[v] = src[v] * inv_sum;
+ }
+ });
+ });
+}
+
+// ---------------------------------------------------------------------------
+// Per-chunk dequant
+//
+// Rather than dequantizing all of K/V up front (footprint scales with
+// context), we dequant one KV-head chunk at a time into a dense
+// [this_chunk x D] fp16 buffer (row-major, lda = D). The source address of
+// element (head=ikvh, row=chunk_start+r, col=c) decomposes into independent
+// linear terms head_off(ikvh) + row_off(chunk_start) + (r,c), so slicing a
+// chunk is a clean pointer offset in every layout case. The true-Gemma-
+// interleave vs padded-seq-view distinction is resolved once when the
+// descriptor is built; slicing does not reintroduce it.
+// ---------------------------------------------------------------------------
+enum mkl_fa_kv_desc_mode {
+ MKL_FA_KV_MODE_F16_DENSE = 0,
+ MKL_FA_KV_MODE_F16_INTERLEAVED = 1,
+ MKL_FA_KV_MODE_QUANT_CONTIG = 2,
+ MKL_FA_KV_MODE_QUANT_NC = 3,
+};
+
+struct mkl_fa_kv_desc {
+ const char * data = nullptr;
+ ggml_type type = GGML_TYPE_F16;
+ int64_t D = 0; // ne[0]
+ int64_t nb1 = 0; // byte stride, seq dim
+ int64_t nb2 = 0; // byte stride, head dim
+ mkl_fa_kv_desc_mode mode = MKL_FA_KV_MODE_F16_DENSE;
+ int64_t ts = 0; // type size (mode 3 base offset)
+ int64_t s01 = 0; // nc row stride in blocks (mode 3)
+ int64_t s02 = 0; // nc head stride in blocks (mode 3)
+};
+
+static mkl_fa_kv_desc mkl_fa_make_desc(const ggml_tensor * T, bool interleaved, int n_kv_heads) {
+ mkl_fa_kv_desc d;
+ d.data = (const char *)T->data;
+ d.type = T->type;
+ d.D = T->ne[0];
+ d.nb1 = (int64_t)T->nb[1];
+ d.nb2 = (int64_t)T->nb[2];
+ d.ts = (int64_t)ggml_type_size(T->type);
+
+ if (T->type == GGML_TYPE_F16) {
+ d.mode = interleaved ? MKL_FA_KV_MODE_F16_INTERLEAVED
+ : MKL_FA_KV_MODE_F16_DENSE;
+ } else if (ggml_is_contiguously_allocated(T) && !interleaved) {
+ d.mode = MKL_FA_KV_MODE_QUANT_CONTIG;
+ } else {
+ d.mode = MKL_FA_KV_MODE_QUANT_NC;
+ const int64_t bs = (int64_t)ggml_blck_size(T->type);
+ const int64_t blk_per_row = T->ne[0] / bs;
+ // True Gemma interleave packs heads within a row (nb[2] < ne[1]*nb[1])
+ // → reconstruct physical strides. Padded seq-views (nb[2] > ne[1]*nb[1])
+ // already have correct physical strides.
+ const bool gemma = interleaved &&
+ ((int64_t)T->nb[2] < (int64_t)T->ne[1] * (int64_t)T->nb[1]);
+ if (gemma) {
+ d.s01 = (int64_t)n_kv_heads * blk_per_row;
+ d.s02 = blk_per_row;
+ } else {
+ d.s01 = d.nb1 / d.ts;
+ d.s02 = d.nb2 / d.ts;
+ }
+ }
+ return d;
+}
+
+// Dequant one KV-head chunk into a dense [this_chunk x D] fp16 buffer.
+static void mkl_fa_dequant_chunk(
+ dpct::queue_ptr stream, const mkl_fa_kv_desc & d, ggml_tensor * dst_ctx,
+ sycl::half * out, int ikvh, int chunk_start, int this_chunk) {
+
+ const int64_t D = d.D;
+ switch (d.mode) {
+ case MKL_FA_KV_MODE_F16_DENSE: {
+ const char * base = d.data + (int64_t)ikvh * d.nb2
+ + (int64_t)chunk_start * d.nb1;
+ stream->memcpy(out, base, (size_t)this_chunk * D * sizeof(sycl::half));
+ break;
+ }
+ case MKL_FA_KV_MODE_F16_INTERLEAVED: {
+ const char * base = d.data + (int64_t)ikvh * d.nb2
+ + (int64_t)chunk_start * d.nb1;
+ const int64_t row_halfs = d.nb1 / (int64_t)sizeof(sycl::half);
+ const sycl::half * src = (const sycl::half *)base;
+ stream->parallel_for(
+ sycl::range<2>((size_t)this_chunk, (size_t)D),
+ [=](sycl::item<2> it) {
+ int64_t r = it.get_id(0);
+ int64_t c = it.get_id(1);
+ out[r * D + c] = src[r * row_halfs + c];
+ });
+ break;
+ }
+ case MKL_FA_KV_MODE_QUANT_CONTIG: {
+ const char * base = d.data + (int64_t)ikvh * d.nb2
+ + (int64_t)chunk_start * d.nb1;
+ to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(d.type, dst_ctx);
+ to_fp16(base, out, (int64_t)this_chunk * D, stream);
+ break;
+ }
+ default: { // MKL_FA_KV_MODE_QUANT_NC
+ to_fp16_nc_sycl_t to_fp16 = ggml_get_to_fp16_nc_sycl(d.type);
+ const int64_t base_blocks = (int64_t)ikvh * d.s02
+ + (int64_t)chunk_start * d.s01;
+ const char * base = d.data + base_blocks * d.ts;
+ // ne02 = ne03 = 1 → s02/s03 inert; head+chunk offset carried by base.
+ to_fp16(base, out, D, this_chunk, 1, 1, d.s01, d.s02, d.s02, stream);
+ break;
+ }
+ }
+}
+
+// ---------------------------------------------------------------------------
+// MKL Flash Attention orchestrator
+//
+// Pipeline: dequantize K/V → for each KV head:
+// pack GQA Q heads → MKL GEMM KQ → online softmax →
+// MKL GEMM VKQ → accumulate → normalize → scatter to dst
+// ---------------------------------------------------------------------------
+void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
+
+ const ggml_tensor * Q = dst->src[0];
+ const ggml_tensor * K = dst->src[1];
+ const ggml_tensor * V = dst->src[2];
+ const ggml_tensor * mask = dst->src[3];
+ ggml_tensor * KQV = dst;
+
+ GGML_ASSERT(Q->type == GGML_TYPE_F32);
+ GGML_ASSERT(KQV->type == GGML_TYPE_F32);
+
+ // --- Op params ---
+ float scale = 1.0f, max_bias = 0.0f, logit_softcap = 0.0f;
+ memcpy(&scale, (const float *)KQV->op_params + 0, sizeof(float));
+ memcpy(&max_bias, (const float *)KQV->op_params + 1, sizeof(float));
+ memcpy(&logit_softcap, (const float *)KQV->op_params + 2, sizeof(float));
+
+ const float q_scale = scale;
+
+ // --- Dimensions ---
+ const int DKQ = (int)K->ne[0];
+ const int DV = (int)V->ne[0];
+ const int n_queries = (int)Q->ne[1];
+ const int n_q_heads = (int)Q->ne[2];
+ const int n_kv_heads = (int)K->ne[2];
+ const int n_batch = (int)Q->ne[3];
+ const int n_kv = (int)K->ne[1];
+ const int gqa_ratio = n_q_heads / n_kv_heads;
+ const int n_query_rows = n_queries * gqa_ratio;
+
+ GGML_ASSERT(n_q_heads % n_kv_heads == 0);
+ GGML_ASSERT(max_bias == 0.0f); // ALiBi not supported
+ GGML_ASSERT(Q->ne[3] == K->ne[3] || K->ne[3] == 1);
+
+ const int chunk_size = std::min(MKL_FA_CHUNK_SIZE_KV, n_kv);
+
+ // Query rows are processed in tiles of q_tile_rows so the score buffers
+ // (KQ_f32/S_f16 = q_tile_rows * chunk_size) stay bounded regardless of
+ // batch size. n_query_rows <= Q_TILE is a single tile (no extra work).
+ static int q_tile_env = ggml_sycl_get_env("GGML_SYCL_MKL_FA_Q_TILE", MKL_FA_Q_TILE);
+ const int q_tile_rows = std::max(1, std::min(q_tile_env, n_query_rows));
+
+ const int64_t wg_size = MKL_FA_WG_SIZE;
+
+ // --- Debug output (gated by GGML_SYCL_MKL_FA_DEBUG=1) ---
+ static int mkl_call_count = 0;
+ mkl_call_count++;
+ static int mkl_debug = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DEBUG", 0);
+ const bool do_print = (mkl_debug == 1);
+
+ const int64_t q_row_stride = Q->nb[1] / sizeof(float);
+ const int64_t q_head_stride = Q->nb[2] / sizeof(float);
+
+ const bool V_is_K_view = V->view_src
+ && (V->view_src == K || (V->view_src == K->view_src
+ && V->view_offs == K->view_offs));
+
+ // Early interleaved detection for debug output.
+ // True interleaved detection happens after dequant (nb12_fp16 == nb11_fp16),
+ // but we can pre-detect on the original tensor strides.
+ const bool k_early_interleaved =
+ ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]);
+ const bool v_early_interleaved =
+ !V_is_K_view && ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]);
+
+ if (do_print) {
+ GGML_LOG_INFO("[MKL-FA] #%d D=%d DV=%d n_q=%d n_kv=%d "
+ "n_qh=%d n_kvh=%d gqa=%d batch=%d K=%s V=%s "
+ "chunk=%d buf=%.1fMB%s%s\n",
+ mkl_call_count, DKQ, DV, n_queries, n_kv,
+ n_q_heads, n_kv_heads, gqa_ratio, n_batch,
+ ggml_type_name(K->type), ggml_type_name(V->type),
+ chunk_size,
+ (double)((int64_t)n_query_rows * chunk_size * sizeof(float))
+ / (1024.0 * 1024.0),
+ k_early_interleaved ? " K_ILV" : "",
+ v_early_interleaved ? " V_ILV" : "");
+ GGML_LOG_INFO("[MKL-FA] #%d Q-nb1=%lld Q-nb2=%lld "
+ "q_rs=%lld q_hs=%lld dst_rs=%lld dst_hs=%lld\n",
+ mkl_call_count,
+ (long long)Q->nb[1], (long long)Q->nb[2],
+ (long long)q_row_stride, (long long)q_head_stride,
+ (long long)(KQV->nb[1] / sizeof(float)),
+ (long long)(KQV->nb[2] / sizeof(float)));
+ }
+
+ // --- Stream and allocators ---
+ dpct::queue_ptr stream = ctx.stream();
+
+#define MKL_TAKE_TIME(t0) auto t0 = std::chrono::steady_clock::now()
+#define MKL_ACCUM(acc, t0) do { if (do_print) { \
+ acc += (int64_t)std::chrono::duration_cast \
+ <std::chrono::microseconds>(std::chrono::steady_clock::now() - (t0)).count(); \
+} } while(0)
+
+ int64_t gemm_kq_time_us = 0;
+ int64_t gemm_vkq_time_us = 0;
+ int64_t softmax_time_us = 0;
+ int64_t dequant_time_us = 0;
+
+ MKL_TAKE_TIME(t_deq);
+
+ // --- K/V dequant descriptors ---
+ // Dequant is done per-chunk inside the KV loop (footprint independent of
+ // context). Output is always dense row-major fp16 [this_chunk x D], lda=D.
+ // Interleaved detection: ne[1]*nb[1] != nb[2] means heads are interleaved.
+ const bool k_interleaved =
+ ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]) && K->ne[2] > 1;
+ const bool v_interleaved =
+ ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]) && V->ne[2] > 1;
+
+ const mkl_fa_kv_desc K_desc = mkl_fa_make_desc(K, k_interleaved, n_kv_heads);
+ const mkl_fa_kv_desc V_desc = V_is_K_view
+ ? K_desc : mkl_fa_make_desc(V, v_interleaved, n_kv_heads);
+
+ MKL_ACCUM(dequant_time_us, t_deq);
+
+ // --- Resolve mask pointers ---
+ const sycl::half * mask_data = nullptr;
+ int64_t mask_head_stride = 0;
+ int64_t mask_row_stride = 0;
+ int mask_n_heads = 0;
+
+ if (mask) {
+ // Use actual fp16 device size (2 bytes), NOT sizeof(sycl::half)
+ // which may be 4 on the host in oneAPI.
+ mask_head_stride = mask->nb[2] / 2;
+ mask_row_stride = mask->nb[1] / 2;
+ mask_n_heads = (int)mask->ne[2];
+ }
+
+ // --- Allocate intermediates from pool ---
+ ggml_sycl_pool & pool = ctx.pool();
+
+ ggml_sycl_pool_alloc<float> KQ_f32(pool); // [q_tile_rows x chunk]
+ ggml_sycl_pool_alloc<sycl::half> S_f16(pool); // [q_tile_rows x chunk]
+ ggml_sycl_pool_alloc<float> VKQ_chunk(pool); // [q_tile_rows x DV]
+ ggml_sycl_pool_alloc<float> VKQ_accum(pool); // [n_query_rows x DV] (full)
+ ggml_sycl_pool_alloc<float> KQ_max(pool); // [n_query_rows] (full)
+ ggml_sycl_pool_alloc<float> KQ_sum(pool); // [n_query_rows] (full)
+ ggml_sycl_pool_alloc<sycl::half> Q_head_f16(pool); // [n_query_rows x DKQ] (full)
+ ggml_sycl_pool_alloc<sycl::half> K_chunk_f16(pool); // [chunk x DKQ] (per-chunk dequant)
+ ggml_sycl_pool_alloc<sycl::half> V_chunk_f16(pool); // [chunk x DV] (per-chunk dequant)
+
+ KQ_f32.alloc((size_t)q_tile_rows * chunk_size);
+ S_f16.alloc((size_t)q_tile_rows * chunk_size);
+ VKQ_chunk.alloc((size_t)q_tile_rows * DV);
+ VKQ_accum.alloc((size_t)n_query_rows * DV);
+ KQ_max.alloc(n_query_rows);
+ KQ_sum.alloc(n_query_rows);
+ Q_head_f16.alloc((size_t)n_query_rows * DKQ);
+ K_chunk_f16.alloc((size_t)chunk_size * DKQ);
+
+ sycl::half * V_chunk_f16_ptr;
+ if (V_is_K_view) {
+ V_chunk_f16_ptr = K_chunk_f16.ptr; // V aliases K (DV == DKQ)
+ } else {
+ V_chunk_f16.alloc((size_t)chunk_size * DV);
+ V_chunk_f16_ptr = V_chunk_f16.ptr;
+ }
+
+ sycl::half * Q_head_f16_ptr = Q_head_f16.ptr;
+ float * KQ_f32_ptr = KQ_f32.ptr;
+ sycl::half * S_f16_ptr = S_f16.ptr;
+ float * VKQ_chunk_ptr = VKQ_chunk.ptr;
+ float * VKQ_accum_ptr = VKQ_accum.ptr;
+ float * KQ_max_ptr = KQ_max.ptr;
+ float * KQ_sum_ptr = KQ_sum.ptr;
+ sycl::half * K_chunk_f16_ptr = K_chunk_f16.ptr;
+
+ const float alpha = 1.0f;
+ const float beta = 0.0f;
+
+ for (int ib = 0; ib < n_batch; ib++) {
+ const float * Q_batch = (const float *)Q->data
+ + ib * (Q->nb[3] / sizeof(float));
+ float * dst_batch = (float *)KQV->data
+ + ib * (KQV->nb[3] / sizeof(float));
+
+ const sycl::half * mask_batch = nullptr;
+ if (mask) {
+ int m_batch = (mask->ne[3] > 1) ? ib : 0;
+ mask_batch = (const sycl::half *)mask->data
+ + m_batch * (mask->nb[3] / 2); // 2 = actual fp16 device size
+ }
+
+ for (int ikvh = 0; ikvh < n_kv_heads; ikvh++) {
+ int kvh_base_head = ikvh * gqa_ratio;
+
+ // 1. Pack all GQA Q heads into fp16 (full n_query_rows)
+ mkl_fa_pack_q_fp16(stream,
+ Q_head_f16_ptr, Q_batch,
+ n_queries, n_query_rows, DKQ,
+ gqa_ratio, kvh_base_head,
+ q_scale, q_row_stride, q_head_stride, wg_size);
+
+ // 2. Initialize softmax state (full n_query_rows)
+ mkl_fa_init_softmax_state(stream,
+ KQ_max_ptr, KQ_sum_ptr, VKQ_accum_ptr,
+ n_query_rows, DV, wg_size);
+
+ // Sync before MKL GEMM (MKL may use an internal queue)
+ stream->wait();
+
+ // 3. KV chunk loop (OUTER): dequant each chunk once, then tile queries.
+ for (int chunk_start = 0; chunk_start < n_kv; chunk_start += chunk_size) {
+ int this_chunk = std::min(chunk_size, n_kv - chunk_start);
+
+ // 3a. Dequant this KV chunk to dense fp16 (once per chunk)
+ {
+ MKL_TAKE_TIME(t0);
+ mkl_fa_dequant_chunk(stream, K_desc, KQV,
+ K_chunk_f16_ptr, ikvh, chunk_start, this_chunk);
+ if (!V_is_K_view) {
+ mkl_fa_dequant_chunk(stream, V_desc, KQV,
+ V_chunk_f16_ptr, ikvh, chunk_start, this_chunk);
+ }
+ stream->wait(); // dequant must be ready before MKL GEMM
+ MKL_ACCUM(dequant_time_us, t0);
+ }
+
+ // 3b. Query tile loop (INNER) — bounds KQ_f32/S_f16 footprint.
+ for (int q0 = 0; q0 < n_query_rows; q0 += q_tile_rows) {
+ int q_rows = std::min(q_tile_rows, n_query_rows - q0);
+
+ // GEMM: KQ = Q_tile × K_chunk^T
+ {
+ MKL_TAKE_TIME(t0);
+ sycl::event ev = gemm(*stream,
+ transpose::trans, transpose::nontrans,
+ this_chunk, q_rows, DKQ,
+ alpha,
+ K_chunk_f16_ptr, DKQ,
+ Q_head_f16_ptr + (int64_t)q0 * DKQ, DKQ,
+ beta,
+ KQ_f32_ptr, this_chunk);
+ try { ev.wait_and_throw(); } catch (sycl::exception & e) {
+ GGML_LOG_INFO("[MKL-FA] GEMM KQ: %s\n", e.what());
+ GGML_ABORT("MKL GEMM KQ failed");
+ }
+ MKL_ACCUM(gemm_kq_time_us, t0);
+ }
+ // Online softmax over this chunk for this query tile
+ {
+ MKL_TAKE_TIME(t0);
+ mkl_fa_online_softmax_chunk(stream,
+ KQ_f32_ptr, S_f16_ptr,
+ KQ_max_ptr, KQ_sum_ptr, VKQ_accum_ptr,
+ q0, q_rows, n_queries, DV,
+ this_chunk, chunk_start,
+ kvh_base_head, gqa_ratio,
+ mask_batch, mask_head_stride,
+ mask_row_stride, mask_n_heads,
+ logit_softcap, wg_size);
+ stream->wait(); // S_f16 must be ready for GEMM
+ MKL_ACCUM(softmax_time_us, t0);
+ }
+
+ // GEMM: VKQ_chunk = S × V_chunk
+ {
+ MKL_TAKE_TIME(t0);
+ sycl::event ev = gemm(*stream,
+ transpose::nontrans, transpose::nontrans,
+ DV, q_rows, this_chunk,
+ alpha,
+ V_chunk_f16_ptr, DV,
+ S_f16_ptr, this_chunk,
+ beta,
+ VKQ_chunk_ptr, DV);
+ try { ev.wait_and_throw(); } catch (sycl::exception & e) {
+ GGML_LOG_INFO("[MKL-FA] GEMM VKQ: %s\n", e.what());
+ GGML_ABORT("MKL GEMM VKQ failed");
+ }
+ MKL_ACCUM(gemm_vkq_time_us, t0);
+ }
+ // VKQ_accum[q0..] += VKQ_chunk
+ {
+ const int64_t n_total = (int64_t)q_rows * DV;
+ const int64_t wg = ((n_total + wg_size - 1) / wg_size)
+ * wg_size;
+ float * accum = VKQ_accum_ptr + (int64_t)q0 * DV;
+ stream->submit([&](sycl::handler & cgh) {
+ cgh.parallel_for(sycl::nd_range<1>(wg, wg_size),
+ [=](sycl::nd_item<1> item) {
+ int64_t i = item.get_global_id(0);
+ if (i < n_total) {
+ accum[i] += VKQ_chunk_ptr[i];
+ }
+ });
+ });
+ }
+ }
+ }
+
+ // 4. Normalize and scatter each GQA head to dst
+ for (int iqg = 0; iqg < gqa_ratio; iqg++) {
+ int iqh = kvh_base_head + iqg;
+ int64_t src_offset = (int64_t)iqg * n_queries * DV;
+ mkl_fa_normalize_head(stream,
+ dst_batch, VKQ_accum_ptr, KQ_sum_ptr,
+ iqh, n_queries, DV, n_q_heads,
+ src_offset, wg_size);
+ }
+ }
+ }
+
+#undef MKL_TAKE_TIME
+#undef MKL_ACCUM
+
+ if (do_print) {
+ const int64_t v_chunk_elems = V_is_K_view ? 0 : (int64_t)chunk_size * DV;
+ double total_mb = (double)(
+ (int64_t)q_tile_rows * chunk_size * sizeof(float) // KQ_f32
+ + (int64_t)q_tile_rows * chunk_size * sizeof(sycl::half) // S_f16
+ + (int64_t)q_tile_rows * DV * sizeof(float) // VKQ_chunk
+ + (int64_t)n_query_rows * DV * sizeof(float) // VKQ_accum
+ + (int64_t)n_query_rows * sizeof(float) // KQ_max
+ + (int64_t)n_query_rows * sizeof(float) // KQ_sum
+ + (int64_t)n_query_rows * DKQ * sizeof(sycl::half) // Q_head_f16
+ + (int64_t)chunk_size * DKQ * sizeof(sycl::half) // K_chunk_f16
+ + v_chunk_elems * (int64_t)sizeof(sycl::half) // V_chunk_f16
+ ) / (1024.0 * 1024.0);
+ GGML_LOG_INFO("[MKL-FA] #%d n_kv=%d n_q=%d q_tile=%d time_us: "
+ "dequant=%lld GEMM_KQ=%lld softmax=%lld GEMM_VKQ=%lld "
+ "buf_mb=%.1f\n",
+ mkl_call_count, n_kv, n_queries, q_tile_rows,
+ (long long)dequant_time_us,
+ (long long)gemm_kq_time_us,
+ (long long)softmax_time_us,
+ (long long)gemm_vkq_time_us,
+ total_mb);
+ }
+}
BEST_FATTN_KERNEL_VEC = 100,
BEST_FATTN_KERNEL_ONEDNN = 150, // added enum for onednn==150
BEST_FATTN_KERNEL_TILE = 200,
+ BEST_FATTN_KERNEL_MKL = 300,
};
+
static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const ggml_tensor * dst) {
GGML_UNUSED(device);
#ifndef SYCL_FLASH_ATTN
const ggml_tensor * K = dst->src[1];
const ggml_tensor * V = dst->src[2];
const ggml_tensor * mask = dst->src[3];
+ const ggml_tensor * sinks = dst->src[4];
const int gqa_ratio = Q->ne[2] / K->ne[2];
GGML_ASSERT(Q->ne[2] % K->ne[2] == 0);
float max_bias = 0.0f;
memcpy(&max_bias, (const float *) KQV->op_params + 1, sizeof(float));
+ float logit_softcap = 0.0f;
+ memcpy(&logit_softcap, (const float *) KQV->op_params + 2, sizeof(float));
+
bool gqa_opt_applies = gqa_ratio >= 2 && mask && max_bias == 0.0f && K->ne[1] % FATTN_KQ_STRIDE == 0;
+
+ // MKL path: XMX-accelerated GEMM for prompt processing (all KV cache types).
+ // The MKL kernel converts non-F16 K/V to F16 via to_fp16_sycl before GEMM,
+ // so quantized, F16, BF16, and F32 caches all benefit from XMX acceleration.
+ // Activates automatically when flash-attn is enabled (--flash-attn on or -fa)
+ // and n_kv >= 1024. Falls through to TILE/VEC for ALiBi, logit softcap,
+ // and mismatched batch dimensions (unsupported by the MKL kernel).
+ // Set GGML_SYCL_ENABLE_MKL_FA=0 to force TILE/VEC path for A/B testing.
+ // Example: GGML_SYCL_ENABLE_MKL_FA=0 llama-cli -m model.gguf -fa -ngl 99 ...
+ // Note: MKL GEMM calls are incompatible with SYCL graph capture replay.
+ static int mkl_enable = ggml_sycl_get_env("GGML_SYCL_ENABLE_MKL_FA", 1);
+ // MKL is validated for the mainstream GQA envelope: grouped-query
+ // (gqa_ratio >= 2), head_dim a multiple of 64 in [64,512] with matching
+ // K/V head size, mask, no sinks/ALiBi/softcap. Gemma's global layers use
+ // head_dim 512, so the cap must include it. Head sizes not a multiple of
+ // 64 (72/80/96), MHA (gqa_ratio == 1), and MLA (DKQ != DV, e.g. 576/512)
+ // fall through to TILE/VEC; see follow-up work.
+ if (mkl_enable == 1 && mask && !sinks && gqa_ratio >= 2 &&
+ Q->ne[0] >= 64 && Q->ne[0] <= 512 && Q->ne[0] % 64 == 0 &&
+ Q->ne[0] == V->ne[0] &&
+ Q->ne[1] >= 32 && K->ne[1] >= 1024 &&
+ max_bias == 0.0f && logit_softcap == 0.0f &&
+ (Q->ne[3] == K->ne[3] || K->ne[3] == 1)) {
+ // F16 K/V strides must be a multiple of ne[0]*2 (the natural row size
+ // in bytes). This passes both dense (nb1 == ne0*2) and interleaved
+ // (nb1 == H * ne0*2). Only pathological test strides like nb1=32 or
+ // nb1=75 for ne0=40 fall through to TILE.
+ bool kv_strides_ok = true;
+ for (const ggml_tensor * t : {K, V}) {
+ if (t->type == GGML_TYPE_F16 && t->nb[1] % (t->ne[0] * 2) != 0) {
+ kv_strides_ok = false;
+ break;
+ }
+ }
+ if (kv_strides_ok) {
+ return BEST_FATTN_KERNEL_MKL;
+ }
+ }
+
for (const ggml_tensor * t : {Q, K, V, mask}) {
if (t == nullptr || ggml_is_quantized(t->type)) {
continue;
void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst) {
ggml_sycl_set_device(ctx.device);
+
+ // n_kv watchdog: log when n_kv differs from the last FA call with
+ // the same D — helps detect cache-truncation issues.
+ static int nkv_debug = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DEBUG", 0);
+ if (nkv_debug == 1) {
+ const ggml_tensor * K_dbg = dst->src[1];
+ const ggml_tensor * V_dbg = dst->src[2];
+ static int64_t last_nkv_d256 = 0, last_nkv_d512 = 0;
+ static int fa_call_seq = 0;
+ fa_call_seq++;
+ int64_t cur_nkv = K_dbg->ne[1];
+ int Dk = (int)K_dbg->ne[0];
+ const char * kname = "TILE";
+ best_fattn_kernel k = ggml_sycl_get_best_fattn_kernel(ctx.device, dst);
+ if (k == BEST_FATTN_KERNEL_MKL) kname = "MKL";
+ if (k == BEST_FATTN_KERNEL_VEC) kname = "VEC";
+ int64_t delta = 0;
+ if (Dk == 256) {
+ delta = cur_nkv - last_nkv_d256;
+ last_nkv_d256 = cur_nkv;
+ } else if (Dk == 512) {
+ delta = cur_nkv - last_nkv_d512;
+ last_nkv_d512 = cur_nkv;
+ }
+ GGML_LOG_INFO("[FA-DISP] #%d %s D=%d n_kv=%lld delta=%lld "
+ "V_ne1=%lld\n",
+ fa_call_seq, kname, Dk,
+ (long long)cur_nkv, (long long)delta,
+ (long long)V_dbg->ne[1]);
+ }
+
switch (ggml_sycl_get_best_fattn_kernel(ggml_sycl_get_device(), dst)) {
case BEST_FATTN_KERNEL_NONE:
GGML_ABORT("Not support Flash-Attention");
case BEST_FATTN_KERNEL_VEC:
ggml_sycl_flash_attn_ext_vec(ctx, dst);
break;
+ case BEST_FATTN_KERNEL_MKL:
+ ggml_sycl_flash_attn_ext_mkl(ctx, dst);
+ break;
+ }
+
+ // --- Output fingerprint (GGML_SYCL_MKL_FA_DIAG=1) ---
+ // Copy first 64 float output values to host for fingerprinting.
+ // Compare MKL vs TILE (GGML_SYCL_ENABLE_MKL_FA=0) to detect divergence.
+ // Only fingerprints the first 6 FA calls with n_kv >= 1024.
+ static int fa_diag = ggml_sycl_get_env("GGML_SYCL_MKL_FA_DIAG", 0);
+ static int fa_diag_count = 0;
+ if (fa_diag == 1 && fa_diag_count < 6) {
+ const ggml_tensor * K_diag = dst->src[1];
+ const ggml_tensor * V_diag = dst->src[2];
+ const ggml_tensor * Q_diag = dst->src[0];
+ if (K_diag->ne[1] >= 1024) {
+ fa_diag_count++;
+ float diag_buf[64];
+ dpct::queue_ptr q = ctx.stream();
+ q->memcpy(diag_buf, dst->data, 64 * sizeof(float));
+ q->wait();
+ const char * kname = "???";
+ best_fattn_kernel kb = ggml_sycl_get_best_fattn_kernel(ctx.device, dst);
+ if (kb == BEST_FATTN_KERNEL_MKL) kname = "MKL";
+ if (kb == BEST_FATTN_KERNEL_TILE) kname = "TILE";
+ if (kb == BEST_FATTN_KERNEL_VEC) kname = "VEC";
+ GGML_LOG_INFO("[FA-DIAG] #%d %s D=%d n_kv=%lld n_q=%lld "
+ "n_qh=%lld n_kvh=%lld K=%s V=%s "
+ "nb1=%zu nb2=%zu first 64 floats:\n",
+ fa_diag_count, kname,
+ (int)K_diag->ne[0], (long long)K_diag->ne[1],
+ (long long)Q_diag->ne[1],
+ (long long)Q_diag->ne[2], (long long)K_diag->ne[2],
+ ggml_type_name(K_diag->type),
+ ggml_type_name(V_diag->type),
+ K_diag->nb[1], K_diag->nb[2]);
+ for (int i = 0; i < 64; i += 8) {
+ GGML_LOG_INFO(" [%2d] %08x %08x %08x %08x %08x %08x %08x %08x\n",
+ i,
+ *(unsigned *)&diag_buf[i+0], *(unsigned *)&diag_buf[i+1],
+ *(unsigned *)&diag_buf[i+2], *(unsigned *)&diag_buf[i+3],
+ *(unsigned *)&diag_buf[i+4], *(unsigned *)&diag_buf[i+5],
+ *(unsigned *)&diag_buf[i+6], *(unsigned *)&diag_buf[i+7]);
+ }
+ }
}
}
overview / pkg / ggml / sources / llama.cpp / commitdiff