From 748b84bfa91454caa5793ecb05f10518b1151ef9 Mon Sep 17 00:00:00 2001 From: Ozymandias_EBON <112784549+johnkarlhill@users.noreply.github.com> Date: Fri, 31 Jul 2026 02:43:16 -0500 Subject: [PATCH] =?UTF-8?q?SYCL:=20add=20oneMKL=20GEMM=20flash=20attention?= =?UTF-8?q?=20for=20XMX-accelerated=20prompt=20proc=E2=80=A6=20(#25025)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * SYCL: add oneMKL GEMM flash attention for XMX-accelerated prompt processing * fattn-mkl: fix interleaved dst layout in normalize kernel - Fix mkl_fa_normalize_head: use interleaved dst layout ((query * n_q_heads + head) * DV) matching TILE's flash_attn_combine_results. Previously used dense head-major layout which wrote head outputs to wrong addresses, corrupting attention for all models except Qwen3.6-27B (where GQA=6 heads were sparse enough to avoid visible overlap). - Remove 7 redundant stream->wait() calls — SYCL in-order queue already serializes pure SYCL kernel dependencies. Retain only the 4 MKL GEMM ↔ SYCL handshake barriers (oneMKL GEMM uses its own internal queue that does not respect SYCL in-order). - Remove unused dst_row_stride, diagnostic clutter, and dead K/V hex dump (fa_diag block in fattn-mkl.cpp). - Add MKL_FA_DISABLE=1 env var for A/B testing. - Add FA-DISP watchdog (MKL_FA_DEBUG=1) and FA-DIAG output fingerprint (MKL_FA_DIAG=1) in fattn.cpp. Tested: Gemma-4-26B, Gemma-4-31B, Qwen3.6-27B, Qwen3.6-35B-A3B Perf (B70/Battlemage, 32K, q8_0 KV): Gemma-4-26B: 1473 t/s MKL vs 746 TILE (1.97x) Qwen3.6-27B: 609 t/s MKL vs 330 TILE (1.85x) Co-Authored-By: Claude Code on DeepSeek-v4-Pro * Thank you for the review feedback: rename env vars, use GGML_LOG_INFO, document in SYCL.md Completed the following: - Rename MKL_FA_DISABLE → GGML_SYCL_ENABLE_MKL_FA (inverted: 0 to disable) - Rename MKL_FA_DEBUG → GGML_SYCL_MKL_FA_DEBUG - Rename MKL_FA_DIAG → GGML_SYCL_MKL_FA_DIAG - Replace fprintf(stderr, ...) / fflush(stderr) with GGML_LOG_INFO() macro - Document all three env vars in docs/backend/SYCL.md under Runtime - Add comment explaining MKL FA activation trigger (flash-attn + quantized KV cache + batch-size >= 1024 + n_kv >= 1024) Resolves review feedback from arthw. Again, thank you!!! Co-Authored-By: Claude Code on DeepSeek-v4-Pro * Thank you for the review feedback round 2: use ggml_sycl_get_env, remove dup waits, gate perf macros - Replace raw getenv() with ggml_sycl_get_env() in all 4 env-var checks (fattn.cpp: GGML_SYCL_ENABLE_MKL_FA, GGML_SYCL_MKL_FA_DEBUG, GGML_SYCL_MKL_FA_DIAG; fattn-mkl.cpp: GGML_SYCL_MKL_FA_DEBUG) - Remove duplicated stream->wait() before ev.wait_and_throw() in GEMM KQ and GEMM VKQ — ev.wait_and_throw() already waits for completion - Gate MKL_ACCUM macro behind do_print so timing accumulators are no-ops in normal operation - Remove redundant MIT/Intel copyright header from fattn-mkl.cpp - Remove unused #include - Expand SYCL.md MKL FA docs with step-by-step activation trigger and example llama-cli command Again, thank you!!! Co-Authored-By: Claude Code on DeepSeek-v4-Pro * fattn-mkl: enable MKL FA for all KV cache types Remove the quantized-only restriction on MKL activation — the MKL kernel converts any non-F16 K/V to F16 via to_fp16_sycl before GEMM, so F16 (default), BF16, and F32 caches all benefit from XMX hardware acceleration. The type restriction was an unnecessary gate. Before (F16/BF16 default cache + FA on at 32K prefill): ~356 t/s (TILE path) After: ~670 t/s (MKL path, matching quantized-cache baseline) Minimal change: two conditions removed, one comment updated in fattn.cpp. No kernel or conversion code changes — the dequant pipeline already covers all types. * fattn-mkl: rename mkl_disable -> mkl_enable for clarity * fattn-mkl: refine MKL FA dispatch gates Three changes: 1. Remove quantized-only restriction - MKL FA activates for all KV cache types (F16 default, BF16, F32, quantized). The MKL kernel converts non-F16 K/V via to_fp16_sycl before GEMM. 2. Rename mkl_disable -> mkl_enable to match env var (GGML_SYCL_ENABLE_MKL_FA). 3. Replace batch-size threshold with Q->ne[1] >= 32 gate. Keeps TG (Q=1) and MTP drafts (Q=3-8) on VEC path where fused kernel beats MKL launch overhead. Routes all multi-token prefill through XMX-accelerated GEMM. Production data confirms Q patterns: 1-8 TG, 32-127 cache reuse, 128+ full reprocess. At 32K F16/BF16 FA-on: 356 -> 670 t/s. * ggml-sycl: fix F16 cache + MKL FA multi-turn corruption; add gate guards Two changes: 1. Always copy F16 K/V to dense row-major buffers before MKL GEMM. Previously F16 was read in-place with raw tensor strides. During multi-turn conversations, the accumulated KV cache had different stride properties than a fresh prefill, producing corrupted outputs. Now dense F16 gets a fast memcpy; interleaved (Gemma) gets a strided copy kernel. This matches what the quantized paths already did through to_fp16_sycl. 2. Gate MKL FA on unsupported op params (max_bias, logit_softcap, batch dim mismatch) and pathological F16 strides (nb[1] not a multiple of ne[0]*2). These conditions would previously crash inside the MKL kernel. Pathological strides (test-only) and ALiBi/softcap fall through to TILE/VEC which handle them correctly. The stride check uses modulo rather than equality, so both dense (nb1 == ne0*2) and interleaved (nb1 == H * ne0*2) pass — all real models use these layouts. Only test cases with overlapping rows (nb1=32 or nb1=75 for ne0=40) are blocked. Thanks to hmscider for the oneDNN FA PR (#25222) which surfaced the same insight: always normalize inputs to contiguous F16 before GEMM. Co-Authored-By: Claude Code using DeepSeek-V4-Pro * fattn-mkl: fix quant+GQA KV strides, tighten MKL gate, add K>=1024 tests Adding K>=1024 flash-attn test cases surfaced several MKL bugs: - Quant K/V with a padded seq-view (real KV cache) used the wrong strides in the dequant path... only the true Gemma interleave layout should reconstruct strides. nb[2] vs ne[1]*nb[1] - Gate was firing on shapes the kernel doesn't handle: head_dim < 64 or not a multiple of 64, MHA, attention sinks, and bf16 decode... fell through to vec which no bf16 case. Gate MKL to the validated envelope: gqa>=2, head_dim 64 through 512 (has to be a multiple of 64) with matching K/V head size, mask, no sinks/alibi/softcap... everything else falls back to tile. Covers Qwen Dense/MoE and Gemma4 Dense/MoE Ran test-backend-ops -o FLASH_ATTN_EXT: 3641/3641 pass. Perplexity unchanged... 6.7267 MKL vs 6.7290 stock using Qwen 27b q5_k_xl * Update ggml/src/ggml-sycl/fattn.cpp Co-authored-by: Neo Zhang * Update ggml/src/ggml-sycl/fattn.cpp Co-authored-by: Neo Zhang * Update ggml/src/ggml-sycl/fattn.cpp Co-authored-by: Neo Zhang * fattn-mkl: bound attention scratch so it doesn't grow with batch or context... also dropped the bf16 comment in fattn.cpp per arthw review. * Update ggml/src/ggml-sycl/fattn-mkl.cpp Co-authored-by: Neo Zhang * Update ggml/src/ggml-sycl/fattn-mkl.cpp Co-authored-by: Neo Zhang * apply arthw suggestions: enum for dequant modes, macro for wg_size, env-var one-liners --------- Co-authored-by: Claude Code using DeepSeek-V4-Pro Co-authored-by: Neo Zhang --- ggml/src/ggml-sycl/fattn-mkl.cpp | 690 +++++++++++++++++++++++++++++++ ggml/src/ggml-sycl/fattn.cpp | 121 ++++++ ggml/src/ggml-sycl/fattn.hpp | 2 + 3 files changed, 813 insertions(+) create mode 100644 ggml/src/ggml-sycl/fattn-mkl.cpp diff --git a/ggml/src/ggml-sycl/fattn-mkl.cpp b/ggml/src/ggml-sycl/fattn-mkl.cpp new file mode 100644 index 000000000..fc22b7bdb --- /dev/null +++ b/ggml/src/ggml-sycl/fattn-mkl.cpp @@ -0,0 +1,690 @@ +// 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 +#include +#include + +#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::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 KQ_f32(pool); // [q_tile_rows x chunk] + ggml_sycl_pool_alloc S_f16(pool); // [q_tile_rows x chunk] + ggml_sycl_pool_alloc VKQ_chunk(pool); // [q_tile_rows x DV] + ggml_sycl_pool_alloc VKQ_accum(pool); // [n_query_rows x DV] (full) + ggml_sycl_pool_alloc KQ_max(pool); // [n_query_rows] (full) + ggml_sycl_pool_alloc KQ_sum(pool); // [n_query_rows] (full) + ggml_sycl_pool_alloc Q_head_f16(pool); // [n_query_rows x DKQ] (full) + ggml_sycl_pool_alloc K_chunk_f16(pool); // [chunk x DKQ] (per-chunk dequant) + ggml_sycl_pool_alloc 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); + } +} diff --git a/ggml/src/ggml-sycl/fattn.cpp b/ggml/src/ggml-sycl/fattn.cpp index 1772b9c85..f8efa53ac 100644 --- a/ggml/src/ggml-sycl/fattn.cpp +++ b/ggml/src/ggml-sycl/fattn.cpp @@ -99,8 +99,10 @@ enum best_fattn_kernel { 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 @@ -115,6 +117,7 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const 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); @@ -122,7 +125,49 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const 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; @@ -216,6 +261,37 @@ static best_fattn_kernel ggml_sycl_get_best_fattn_kernel(const int device, const 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"); @@ -232,6 +308,51 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst 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]); + } + } } } diff --git a/ggml/src/ggml-sycl/fattn.hpp b/ggml/src/ggml-sycl/fattn.hpp index f2a8ffc97..c093970a3 100644 --- a/ggml/src/ggml-sycl/fattn.hpp +++ b/ggml/src/ggml-sycl/fattn.hpp @@ -19,4 +19,6 @@ void ggml_sycl_flash_attn_ext(ggml_backend_sycl_context & ctx, ggml_tensor * dst bool ggml_sycl_flash_attn_ext_supported(int device, const ggml_tensor * dst); +void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst); + #endif // GGML_SYCL_FATTN_HPP