From 36a0e2ffc0b352b774f8d71c43668b75b3061ef5 Mon Sep 17 00:00:00 2001 From: Michael Taylor <162068037+mctylr-gh@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:37:28 -0300 Subject: [PATCH] sycl: FLASH_ATTN_EXT failures (#30761) --- ggml/src/ggml-sycl/fattn-mkl.cpp | 39 +++++++++++------------------ ggml/src/ggml-sycl/fattn-onednn.cpp | 32 ++--------------------- 2 files changed, 17 insertions(+), 54 deletions(-) diff --git a/ggml/src/ggml-sycl/fattn-mkl.cpp b/ggml/src/ggml-sycl/fattn-mkl.cpp index 497815b05..e3b3553de 100644 --- a/ggml/src/ggml-sycl/fattn-mkl.cpp +++ b/ggml/src/ggml-sycl/fattn-mkl.cpp @@ -269,19 +269,21 @@ struct mkl_fa_kv_desc { int64_t D = 0; // ne[0] int64_t nb1 = 0; // byte stride, seq dim int64_t nb2 = 0; // byte stride, head dim + int64_t nb3 = 0; // byte stride, batch 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) { +static mkl_fa_kv_desc mkl_fa_make_desc(const ggml_tensor * T, bool interleaved) { 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.nb3 = (int64_t)T->nb[3]; d.ts = (int64_t)ggml_type_size(T->type); if (T->type == GGML_TYPE_F16) { @@ -293,20 +295,8 @@ static mkl_fa_kv_desc mkl_fa_make_desc(const ggml_tensor * T, bool 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; - } + d.s01 = d.nb1 / d.ts; + d.s02 = d.nb2 / d.ts; } return d; } @@ -314,19 +304,20 @@ static mkl_fa_kv_desc mkl_fa_make_desc(const ggml_tensor * T, bool interleaved, // 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) { + sycl::half * out, int ikvh, int ib, int chunk_start, int this_chunk) { const int64_t D = d.D; switch (d.mode) { case MKL_FA_KV_MODE_F16_DENSE: { + // ib * nb3 selects this batch's KV (nb[3] != ne[1]*nb[1] for batched KV views) const char * base = d.data + (int64_t)ikvh * d.nb2 - + (int64_t)chunk_start * d.nb1; + + (int64_t)chunk_start * d.nb1 + (int64_t)ib * d.nb3; 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; + + (int64_t)chunk_start * d.nb1 + (int64_t)ib * d.nb3; const int64_t row_halfs = d.nb1 / (int64_t)sizeof(sycl::half); const sycl::half * src = (const sycl::half *)base; stream->parallel_for( @@ -340,7 +331,7 @@ static void mkl_fa_dequant_chunk( } case MKL_FA_KV_MODE_QUANT_CONTIG: { const char * base = d.data + (int64_t)ikvh * d.nb2 - + (int64_t)chunk_start * d.nb1; + + (int64_t)chunk_start * d.nb1 + (int64_t)ib * d.nb3; 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; @@ -348,7 +339,7 @@ static void mkl_fa_dequant_chunk( 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; + + (int64_t)chunk_start * d.s01 + (int64_t)ib * d.nb3 / d.ts; 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); @@ -477,9 +468,9 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * 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 K_desc = mkl_fa_make_desc(K, k_interleaved); const mkl_fa_kv_desc V_desc = V_is_K_view - ? K_desc : mkl_fa_make_desc(V, v_interleaved, n_kv_heads); + ? K_desc : mkl_fa_make_desc(V, v_interleaved); MKL_ACCUM(dequant_time_us, t_deq); @@ -574,10 +565,10 @@ void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * { MKL_TAKE_TIME(t0); mkl_fa_dequant_chunk(stream, K_desc, KQV, - K_chunk_f16_ptr, ikvh, chunk_start, this_chunk); + K_chunk_f16_ptr, ikvh, ib, 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); + V_chunk_f16_ptr, ikvh, ib, chunk_start, this_chunk); } MKL_ACCUM(dequant_time_us, t0); } diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index 6e030940a..4691c8600 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -306,29 +306,15 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso { const char * K_data = (const char *)K->data; const bool k_non_dense = ((int64_t)K->ne[1] * K->nb[1] != K->nb[2]) && K->ne[2] > 1; - const bool k_gemma = k_non_dense && - ((int64_t)K->nb[2] < (int64_t)K->ne[1] * (int64_t)K->nb[1]); if (ggml_is_contiguously_allocated(K) && !k_non_dense) { to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst); to_fp16(K_data, K_ptr, ggml_nelements(K), stream); } else { - const size_t bs = ggml_blck_size(K->type); const size_t ts = ggml_type_size(K->type); to_fp16_nc_sycl_t to_fp16 = ggml_get_to_fp16_nc_sycl(K->type); - int64_t s01, s02, s03; - if (k_gemma) { - const int64_t blk_per_row = (int64_t)K->ne[0] / bs; - s01 = (int64_t)Hkv * blk_per_row; - s02 = blk_per_row; - s03 = (int64_t)K->ne[1] * s01; - } else { - s01 = (int64_t)K->nb[1] / ts; - s02 = (int64_t)K->nb[2] / ts; - s03 = (int64_t)K->nb[3] / ts; - } to_fp16(K_data, K_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], - s01, s02, s03, stream); + (int64_t)K->nb[1] / ts, (int64_t)K->nb[2] / ts, (int64_t)K->nb[3] / ts, stream); } } // Quantized V: always dequant separately. Even when K and V share @@ -340,29 +326,15 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso { const char * V_data = (const char *)V->data; const bool v_non_dense = ((int64_t)V->ne[1] * V->nb[1] != V->nb[2]) && V->ne[2] > 1; - const bool v_gemma = v_non_dense && - ((int64_t)V->nb[2] < (int64_t)V->ne[1] * (int64_t)V->nb[1]); if (ggml_is_contiguously_allocated(V) && !v_non_dense) { to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst); to_fp16(V_data, V_ptr, ggml_nelements(V), stream); } else { - const size_t bs = ggml_blck_size(V->type); const size_t ts = ggml_type_size(V->type); to_fp16_nc_sycl_t to_fp16 = ggml_get_to_fp16_nc_sycl(V->type); - int64_t s01, s02, s03; - if (v_gemma) { - const int64_t blk_per_row = (int64_t)V->ne[0] / bs; - s01 = (int64_t)V->ne[2] * blk_per_row; - s02 = blk_per_row; - s03 = (int64_t)V->ne[1] * s01; - } else { - s01 = (int64_t)V->nb[1] / ts; - s02 = (int64_t)V->nb[2] / ts; - s03 = (int64_t)V->nb[3] / ts; - } to_fp16(V_data, V_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], - s01, s02, s03, stream); + (int64_t)V->nb[1] / ts, (int64_t)V->nb[2] / ts, (int64_t)V->nb[3] / ts, stream); } } } else {