SYCL: add oneMKL GEMM flash attention for XMX-accelerated prompt proc… (#25025)
* 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 <cfloat> - 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 <noreply@anthropic.com> * 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 <zhang.jianyu@outlook.com> * Update ggml/src/ggml-sycl/fattn.cpp Co-authored-by: Neo Zhang <zhang.jianyu@outlook.com> * Update ggml/src/ggml-sycl/fattn.cpp Co-authored-by: Neo Zhang <zhang.jianyu@outlook.com> * 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 <zhang.jianyu@outlook.com> * Update ggml/src/ggml-sycl/fattn-mkl.cpp Co-authored-by: Neo Zhang <zhang.jianyu@outlook.com> * apply arthw suggestions: enum for dequant modes, macro for wg_size, env-var one-liners --------- Co-authored-by: Claude Code using DeepSeek-V4-Pro <noreply@anthropic.com> Co-authored-by: Neo Zhang <zhang.jianyu@outlook.com>
This commit is contained in:
Ozymandias_EBON
co-authored by
Neo Zhang
Claude Code using DeepSeek-V4-Pro
Neo Zhang
Claude Code using DeepSeek-V4-Pro
parent
d5d3e05bf8
commit
9d9a6d29f6
@@ -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]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user