From e43559f5a02889cda1d9c77287af68ef31b62542 Mon Sep 17 00:00:00 2001 From: meatposes Date: Tue, 28 Jul 2026 03:37:25 -0500 Subject: [PATCH] sycl: fix use-after-return of the SDPA scale in the oneDNN flash-attention path (llama/25880) * sycl: fix use-after-return of the SDPA scale in the oneDNN flash-attention path The scale was uploaded with an async memcpy sourced from a stack local. On the in-order queue that copy is ordered behind the K/V staging kernels; once n_kv is large enough (>= ~26k observed on Arc Pro B70) the staging outlives the host stack frame and the copy reads recycled memory, feeding the SDPA a garbage scale. Output then collapses to a single repeated token and the KV cache is poisoned for the rest of the session. Short contexts win the race by accident, and test-backend-ops caps FLASH_ATTN_EXT at kv=1024, which is why CI never caught it. The previous device_count > 1 wait_and_throw() gate (and reverting it, PR #25741) fixes the symptom only by keeping the frame alive across the copy at the cost of a host sync on every FA call. Fix: cache one device scalar per (device, value) -- the scale is constant per model -- and upload it synchronously once. The single-device fast path (no per-call host sync) is then safe: every device-side hazard already serializes on the in-order queue. The multi-GPU conservative wait is kept unchanged. Also: - GGML_SYCL_FA_ONEDNN_MAX_KV env (0 = unlimited): optional n_kv ceiling that routes very long sequences to the native FA kernel. - test-backend-ops: FLASH_ATTN_EXT F16 cases up to kv=65536 (Qwen3.6-27B geometry hsk=hsv=256 GQA 6, and hsk=128 GQA 4), closing the kv=1024 blind spot. Note the race itself needs a live multi-op pipeline to reproduce; single-op runs pass even on broken builds. Verified on Arc Pro B70 (bmg_g31), Qwen3.6-27B Q4_K, -c 131072: output byte-identical at temp 0 to the native FA path through 32k-deep prefill, with prefill depth-flat at 820-840 t/s (vs 340-350 native at 32k depth). Assisted-by: Claude Fable 5 * sycl: handle GGML_SYCL_FA_ONEDNN_MAX_KV like the other runtime env vars and document it Review feedback on #25880: - read the variable once at backend init into g_ggml_sycl_fa_onednn_max_kv via ggml_sycl_get_env, and print it in the startup env listing (-lv 4 shows it) - document GGML_SYCL_FA_ONEDNN and GGML_SYCL_FA_ONEDNN_MAX_KV in the SYCL.md runtime table Also trim the added FLASH_ATTN_EXT cases to kv={4096,16384}: the 32768/65536 shapes exceed the legacy NMSE threshold on both the oneDNN and native kernels (long-sequence fp16 accumulation drift, present before this PR) and would fail CI for an unrelated reason. Assisted-by: Claude Fable 5 * sycl: clarify GGML_SYCL_FA_ONEDNN_MAX_KV default is disabled Assisted-by: Claude Fable 5 * sycl: state default behavior of GGML_SYCL_FA_ONEDNN_MAX_KV explicitly Assisted-by: Claude Fable 5 * Update ggml/src/ggml-sycl/fattn-onednn.cpp Co-authored-by: Neo Zhang * sycl: write the SDPA scale from a kernel instead of caching it The per-(device, value) scale cache was a function-local static unordered_map with no synchronization, so concurrent backend instances could access and rehash it at the same time. Write the scalar with a single_task instead. The value is captured into the command, so no host memory has to outlive the call -- which is what the use-after-return fix needed in the first place. That removes the shared container, the leaked device allocation and the string key, and it also closes the remaining async-memcpy-from-a-stack-local on the first flash-attention call. Ordering does not rely on timing: the queue is created with sycl::property::queue::in_order and the dnnl stream wraps that same queue, so the write completes before the SDPA reads the scalar. The multi-GPU wait_and_throw() branch is unchanged. Also drop the include, which is unused. Assisted-by: Claude Opus 5 --------- Co-authored-by: Neo Zhang --- ggml/src/ggml-sycl/common.hpp | 1 + ggml/src/ggml-sycl/fattn-onednn.cpp | 30 ++++++++++++++++++++--------- ggml/src/ggml-sycl/ggml-sycl.cpp | 3 +++ 3 files changed, 25 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-sycl/common.hpp b/ggml/src/ggml-sycl/common.hpp index e5d9ee89d..f27ec5dd6 100644 --- a/ggml/src/ggml-sycl/common.hpp +++ b/ggml/src/ggml-sycl/common.hpp @@ -65,6 +65,7 @@ extern int g_ggml_sycl_prioritize_dmmv; extern int g_ggml_sycl_enable_flash_attention; extern int g_ggml_sycl_dev2dev_memcpy; extern int g_ggml_sycl_fa_onednn; +extern int g_ggml_sycl_fa_onednn_max_kv; #if defined(__clang__) && __has_builtin(__builtin_expect) diff --git a/ggml/src/ggml-sycl/fattn-onednn.cpp b/ggml/src/ggml-sycl/fattn-onednn.cpp index f2e12ef1a..8465e1224 100644 --- a/ggml/src/ggml-sycl/fattn-onednn.cpp +++ b/ggml/src/ggml-sycl/fattn-onednn.cpp @@ -38,6 +38,12 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) { if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) { return false; } + // Optional KV-length ceiling (GGML_SYCL_FA_ONEDNN_MAX_KV, 0 = unlimited). Escape hatch: + // very long sequences make the fused SDPA slow enough to risk the xe driver watchdog on + // some stacks; past the cap we fall back to the native FA kernel instead. + if (g_ggml_sycl_fa_onednn_max_kv > 0 && K->ne[1] > g_ggml_sycl_fa_onednn_max_kv) { + return false; + } // gate for the following cases // 1. if the oneDNN graph Add node has no input --> skip // 2. types other than f16 need different logical_tensor declaration @@ -208,9 +214,17 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso cont_to_f16_sycl((const char *) V->data, Vf.get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream); // divide-by-(1/scale) reproduces ggml's score *= kq_scale on the proven probe graph. + // + // The scale must not be uploaded with an async memcpy from a stack local: on the in-order + // queue that copy waits behind the K/V staging kernels, and once those take long enough + // (n_kv >= ~26k on B70) the host frame is recycled before the copy runs, feeding the SDPA a + // garbage scale (output collapses to a repeated token). Write the scalar from a kernel + // instead -- the value is captured into the command, so no host memory has to outlive the + // call, and the enqueue stays async. const sycl::half scale_h = (sycl::half) (1.0f / kq_scale); ggml_sycl_pool_alloc scbuf(ctx.pool(), 1); - stream->memcpy(scbuf.get(), &scale_h, sizeof(sycl::half)); + sycl::half * const scale_dev = scbuf.get(); + stream->single_task([=]() { *scale_dev = scale_h; }); ggml_sycl_pool_alloc outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d] @@ -232,7 +246,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso if (r == E.id_q) return Qf.get(); if (r == E.id_k) return Kf.get(); if (r == E.id_v) return Vf.get(); - if (r == E.id_scale) return scbuf.get(); + if (r == E.id_scale) return scale_dev; if (r == E.id_mask) return (void *) mask->data; return nullptr; }; @@ -245,14 +259,12 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso E.cp.execute(strm, ti, {to}); permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream); - // Single device: no sync is required, and actually PP perf is ~6% > wait_and_throw() (tested on llama-3.1-8b & qwen3.6-27b, both Q8_0, with Arc B70). - // Any future multi-GPU refactor MUST re-measure this single-device path and keep the best - // single-device PP speed. Otherwise (multiple devices/streams can race the reuse): + // Single device needs no sync: the dnnl stream wraps this same in-order queue, so the SDPA + // serializes with the staging kernels before it and the permute/pool reuse after it. The + // garbage output formerly blamed on the missing sync here was the scale use-after-return + // fixed above. Keep the conservative wait for multi-GPU, where other devices' streams can + // race the pool: if (ggml_sycl_info().device_count > 1) { - // cont_to_f16 -> oneDNN execute -> permute is async on this stream, but the - // pool_alloc*s above free their device buffers at host return. Without this wait the next - // scheduler op re-acquires those bytes while the GPU is still computing the SDPA, turning - // it into garbage and collapsing multi-turn output to a single repeated token ("GGGGG..."). stream->wait_and_throw(); } } diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp index cb8974eed..3b807c7cb 100644 --- a/ggml/src/ggml-sycl/ggml-sycl.cpp +++ b/ggml/src/ggml-sycl/ggml-sycl.cpp @@ -85,6 +85,7 @@ int g_ggml_sycl_enable_optimize = 1; int g_ggml_sycl_enable_graph = 0; int g_ggml_sycl_enable_dnn = 1; int g_ggml_sycl_fa_onednn = 1; +int g_ggml_sycl_fa_onednn_max_kv = 0; int g_ggml_sycl_enable_vmm = 1; int g_ggml_sycl_enable_fusion = 1; int g_ggml_sycl_prioritize_dmmv = 0; @@ -287,6 +288,7 @@ static void ggml_check_sycl() try { g_ggml_sycl_enable_graph = ggml_sycl_get_env("GGML_SYCL_ENABLE_GRAPH", 0); g_ggml_sycl_enable_dnn = ggml_sycl_get_env("GGML_SYCL_ENABLE_DNN", 1); g_ggml_sycl_fa_onednn = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN", 1); + g_ggml_sycl_fa_onednn_max_kv = ggml_sycl_get_env("GGML_SYCL_FA_ONEDNN_MAX_KV", 0); g_ggml_sycl_enable_vmm = ggml_sycl_get_env("GGML_SYCL_ENABLE_VMM", 1); g_ggml_sycl_enable_fusion = ggml_sycl_get_env("GGML_SYCL_ENABLE_FUSION", 1); g_ggml_sycl_prioritize_dmmv = ggml_sycl_get_env("GGML_SYCL_PRIORITIZE_DMMV", 0); @@ -359,6 +361,7 @@ static void ggml_check_sycl() try { GGML_LOG_INFO(" GGML_SYCL_ENABLE_DNN: DNN disabled by compile flag\n"); GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN: %d\n", g_ggml_sycl_fa_onednn); #endif + GGML_LOG_INFO(" GGML_SYCL_FA_ONEDNN_MAX_KV: %d\n", g_ggml_sycl_fa_onednn_max_kv); #ifdef SYCL_FLASH_ATTN GGML_LOG_INFO(" GGML_SYCL_ENABLE_FLASH_ATTN: %d\n", g_ggml_sycl_enable_flash_attention); #else