sycl: make --fit respect --fit-target better (#27629)

improve the --fit algorithm to take into account the actual peak
required VRAM for a given context size on a SYCL backend.

This includes both properly accounting for how much VRAM is required
when the allocated context is fully used (which makes the reported
context drop below what it did before, but stop it OOMing) as well
as preventing some overly-conservative calculations which meant too much
VRAM was being reserved.

Tested on a Arc b70 with unsloth's qwen3.8 (Q4_K_XL), able to get 262144 context,
fully usable, with q8_0 KV and MTP and 4k ubatch size using --fit-target 1
This commit is contained in:
Nick Farrell
2026-08-29 19:00:09 +10:00
committed by GitHub
parent c9ca51c1f6
commit cc83d7b482
6 changed files with 167 additions and 41 deletions
+13 -9
View File
@@ -6,6 +6,7 @@
#include "convert.hpp"
#include "vecdotq.hpp"
#include "fattn-buffers.hpp"
#include "fattn.hpp"
#include "ggml.h"
@@ -926,6 +927,7 @@ void launch_fattn(
ggml_sycl_fattn_alloc K_f16(fbuf.K);
ggml_sycl_fattn_alloc V_f16(fbuf.V);
const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst);
ggml_sycl_pool_alloc<int> KV_max(pool);
ggml_sycl_pool_alloc<float> dst_tmp(pool);
ggml_sycl_pool_alloc<sycl::float2> dst_tmp_meta(pool);
@@ -944,10 +946,11 @@ void launch_fattn(
const size_t bs = ggml_blck_size(K->type);
const size_t ts = ggml_type_size(K->type);
K_f16.alloc(ggml_nelements(K));
sycl::half * K_f16_ptr = extra.K_buffer_ptr ? (sycl::half *) extra.K_buffer_ptr
: K_f16.alloc(ggml_nelements(K));
if (ggml_is_contiguously_allocated(K)) {
to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(K->type, dst);
to_fp16(K_data, K_f16.ptr, ggml_nelements(K), main_stream);
to_fp16(K_data, K_f16_ptr, ggml_nelements(K), main_stream);
nb11 = nb11 * bs * sizeof(sycl::half) / ts;
nb12 = nb12 * bs * sizeof(sycl::half) / ts;
@@ -958,13 +961,13 @@ void launch_fattn(
const int64_t s01 = nb11 / ts;
const int64_t s02 = nb12 / ts;
const int64_t s03 = nb13 / ts;
to_fp16(K_data, K_f16.ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream);
to_fp16(K_data, K_f16_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3], s01, s02, s03, main_stream);
nb11 = K->ne[0] * sizeof(sycl::half);
nb12 = K->ne[1] * nb11;
nb13 = K->ne[2] * nb12;
}
K_data = (char *) K_f16.ptr;
K_data = (char *) K_f16_ptr;
}
if (need_f16_V && V->type != GGML_TYPE_F16) {
@@ -977,11 +980,12 @@ void launch_fattn(
const size_t bs = ggml_blck_size(V->type);
const size_t ts = ggml_type_size(V->type);
V_f16.alloc(ggml_nelements(V));
sycl::half * V_f16_ptr = extra.V_buffer_ptr ? (sycl::half *) extra.V_buffer_ptr
: V_f16.alloc(ggml_nelements(V));
if (ggml_is_contiguously_allocated(V)) {
to_fp16_sycl_t to_fp16 = ggml_get_to_fp16_sycl(V->type, dst);
to_fp16(V_data, V_f16.ptr, ggml_nelements(V), main_stream);
V_data = (char *) V_f16.ptr;
to_fp16(V_data, V_f16_ptr, ggml_nelements(V), main_stream);
V_data = (char *) V_f16_ptr;
nb21 = nb21 * bs * sizeof(sycl::half) / ts;
nb22 = nb22 * bs * sizeof(sycl::half) / ts;
@@ -992,13 +996,13 @@ void launch_fattn(
const int64_t s01 = nb21 / ts;
const int64_t s02 = nb22 / ts;
const int64_t s03 = nb23 / ts;
to_fp16(V_data, V_f16.ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream);
to_fp16(V_data, V_f16_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3], s01, s02, s03, main_stream);
nb21 = V->ne[0] * sizeof(sycl::half);
nb22 = V->ne[1] * nb21;
nb23 = V->ne[2] * nb22;
}
V_data = (char *) V_f16.ptr;
V_data = (char *) V_f16_ptr;
}
}
+54 -30
View File
@@ -14,9 +14,21 @@
// set minimum query length to treat as prefill (32)
#define GGML_SYCL_FA_ONEDNN_MIN_Q 32
bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) {
bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V) {
if (K->type != GGML_TYPE_F16 || V->type != GGML_TYPE_F16) {
return false;
}
auto bindable = [](const ggml_tensor * t) {
return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 &&
t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0;
};
return bindable(K) && bindable(V);
}
bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit) {
#if !GGML_SYCL_DNNL
GGML_UNUSED(dst);
GGML_UNUSED(use_shape_limit);
return false;
#else
if (!g_ggml_sycl_fa_onednn) {
@@ -44,7 +56,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) {
if (!k_ok || !v_ok) {
return false;
}
if (Q->ne[1] < 32 || K->ne[1] < 1024) {
if (use_shape_limit && (Q->ne[1] < 32 || K->ne[1] < 1024)) {
return false;
}
for (const ggml_tensor * t : {K, V}) {
@@ -94,7 +106,7 @@ bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst) {
return false;
}
// Prefill only.
if (Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) {
if (use_shape_limit && Q->ne[1] < GGML_SYCL_FA_ONEDNN_MIN_Q) {
return false;
}
return true;
@@ -240,9 +252,16 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
dnnl::engine eng = ctx.engine_dnnl(stream);
dnnl::stream strm = ctx.stream_dnnl(stream);
const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst);
// Q: always f32 -- copy to dense f16.
ggml_sycl_pool_alloc<sycl::half> Qf(ctx.pool(), (size_t) H * q * d);
cont_to_f16_sycl<float>((const char *) Q->data, Qf.get(), d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);
std::optional<ggml_sycl_pool_alloc<sycl::half>> Qf_pool;
sycl::half * Qf_ptr = (sycl::half *) extra.Q_buffer_ptr;
if (!Qf_ptr) {
Qf_pool.emplace(ctx.pool(), (size_t) H * q * d);
Qf_ptr = Qf_pool->get();
}
cont_to_f16_sycl<float>((const char *) Q->data, Qf_ptr, d, q, H, mb, Q->nb[1], Q->nb[2], Q->nb[3], stream);
// K/V: bind the f16 cache in place. llama.cpp permutes it to [token][head][dim], so its head
// plane is strided rather than dense, which is what an explicit stride vector expresses.
@@ -253,11 +272,12 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
std::array<int64_t, 5> v_str = k_str;
std::optional<ggml_sycl_pool_alloc<sycl::half>> Kf_pool;
std::optional<ggml_sycl_pool_alloc<sycl::half>> Vf_pool;
// Helper: hand out reserved space, or fall back to the pool.
auto stage_k = [&](size_t n) { if (extra.K_buffer_ptr) { return (sycl::half *) extra.K_buffer_ptr; }
Kf_pool.emplace(ctx.pool(), n); return Kf_pool->get(); };
auto stage_v = [&](size_t n) { if (extra.V_buffer_ptr) { return (sycl::half *) extra.V_buffer_ptr; }
Vf_pool.emplace(ctx.pool(), n); return Vf_pool->get(); };
auto bindable = [](const ggml_tensor * t) {
return t->nb[0] == sizeof(sycl::half) && t->nb[1] % sizeof(sycl::half) == 0 &&
t->nb[2] % sizeof(sycl::half) == 0 && t->nb[3] % sizeof(sycl::half) == 0;
};
auto elem_strides = [](const ggml_tensor * t) {
const int64_t s1 = (int64_t) (t->nb[1] / t->nb[0]);
const int64_t s2 = (int64_t) (t->nb[2] / t->nb[0]);
@@ -266,22 +286,19 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
return std::array<int64_t, 5>{ s3, s2, s2, s1, 1 };
};
if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && bindable(K) && bindable(V)) {
if (ggml_sycl_fattn_onednn_binds_kv(K, V)) {
K_ptr = (sycl::half *) K->data;
V_ptr = (sycl::half *) V->data;
k_str = elem_strides(K);
v_str = elem_strides(V);
} else if (K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16) {
Kf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
Vf_pool.emplace(ctx.pool(), (size_t) Hkv * seq * d);
cont_to_f16_sycl<sycl::half>((const char *) K->data, Kf_pool->get(), d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
cont_to_f16_sycl<sycl::half>((const char *) V->data, Vf_pool->get(), d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
K_ptr = Kf_pool->get();
V_ptr = Vf_pool->get();
K_ptr = stage_k((size_t) Hkv * seq * d);
V_ptr = stage_v((size_t) Hkv * seq * d);
cont_to_f16_sycl<sycl::half>((const char *) K->data, K_ptr, d, seq, Hkv, mb, K->nb[1], K->nb[2], K->nb[3], stream);
cont_to_f16_sycl<sycl::half>((const char *) V->data, V_ptr, d, seq, Hkv, mb, V->nb[1], V->nb[2], V->nb[3], stream);
} else if (ggml_is_quantized(K->type)) {
// Quantized K/V: dequant to dense F16 using pool, same lifetime as F16 path.
Kf_pool.emplace(ctx.pool(), ggml_nelements(K));
K_ptr = Kf_pool->get();
K_ptr = stage_k((size_t) ggml_nelements(K));
{
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;
@@ -315,8 +332,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
// data pointer), their logical values differ because the quantized
// elements at different positions/offsets represent different K/V
// data. Master's F16 path also never aliases K and V.
Vf_pool.emplace(ctx.pool(), ggml_nelements(V));
V_ptr = Vf_pool->get();
V_ptr = stage_v((size_t) ggml_nelements(V));
{
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;
@@ -347,12 +363,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
}
} else {
// F32: strided copy to dense F16 via cont_to_f16_sycl<float>.
Kf_pool.emplace(ctx.pool(), ggml_nelements(K));
K_ptr = Kf_pool->get();
K_ptr = stage_k((size_t) ggml_nelements(K));
cont_to_f16_sycl<float>((const char *) K->data, K_ptr, K->ne[0], K->ne[1], K->ne[2], K->ne[3],
K->nb[1], K->nb[2], K->nb[3], stream);
Vf_pool.emplace(ctx.pool(), ggml_nelements(V));
V_ptr = Vf_pool->get();
V_ptr = stage_v((size_t) ggml_nelements(V));
cont_to_f16_sycl<float>((const char *) V->data, V_ptr, V->ne[0], V->ne[1], V->ne[2], V->ne[3],
V->nb[1], V->nb[2], V->nb[3], stream);
}
@@ -366,11 +380,21 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
// 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<sycl::half> scbuf(ctx.pool(), 1);
sycl::half * const scale_dev = scbuf.get();
std::optional<ggml_sycl_pool_alloc<sycl::half>> scbuf;
sycl::half * scale_dev = (sycl::half *) extra.scale_buffer_ptr;
if (!scale_dev) {
scbuf.emplace(ctx.pool(), 1);
scale_dev = scbuf->get();
}
stream->single_task([=]() { *scale_dev = scale_h; });
ggml_sycl_pool_alloc<sycl::half> outf(ctx.pool(), (size_t) H * q * d); // f16 contiguous SDPA out [mb,H,q,d]
// f16 contiguous SDPA out [mb,H,q,d]
std::optional<ggml_sycl_pool_alloc<sycl::half>> outf_pool;
sycl::half * outf_ptr = (sycl::half *) extra.out_buffer_ptr;
if (!outf_ptr) {
outf_pool.emplace(ctx.pool(), (size_t) H * q * d);
outf_ptr = outf_pool->get();
}
// compile once per (device, shape, KV strides), reuse across layers/calls. Stride 2 always
// repeats stride 1 and stride 4 is always 1, so the key covers every entry that can differ.
@@ -392,7 +416,7 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
}
auto id2ptr = [&](size_t r) -> void * {
if (r == E.id_q) return Qf.get();
if (r == E.id_q) return Qf_ptr;
if (r == E.id_k) return K_ptr;
if (r == E.id_v) return V_ptr;
if (r == E.id_scale) return scale_dev;
@@ -404,10 +428,10 @@ void ggml_sycl_flash_attn_ext_onednn(ggml_backend_sycl_context & ctx, ggml_tenso
for (auto & lt : E.ins) {
ti.emplace_back(lt, eng, id2ptr(lt.get_id()));
}
tensor to(E.out, eng, outf.get());
tensor to(E.out, eng, outf_ptr);
E.cp.execute(strm, ti, {to});
permute_sdpa_out_sycl(outf.get(), (float *) dst->data, mb, H, q, d, stream);
permute_sdpa_out_sycl(outf_ptr, (float *) dst->data, mb, H, q, d, stream);
// 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
+5 -1
View File
@@ -5,7 +5,11 @@
// Static-only check: fused-XMX oneDNN Graph SDPA path==flash-attn op
// (f16 KV, no softcap/ALiBi, single stream, tuned head_dim, prefill-sized q.)
bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst);
bool ggml_sycl_flash_attn_ext_onednn_supported(const ggml_tensor * dst, bool use_shape_limit = true);
// True when the oneDNN path binds an F16 KV cache in place instead of staging a dense copy of
// it. Depends only on the types and strides of K and V, so the answer holds for every call.
bool ggml_sycl_fattn_onednn_binds_kv(const ggml_tensor * K, const ggml_tensor * V);
// Run flash attention through oneDNN's fused xmx SDPA
// execute the cached SDPA partition, write the f32 dst. Falls back to the TILE kernel on any failure.
+73
View File
@@ -378,3 +378,76 @@ 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) {
return ggml_sycl_get_best_fattn_kernel(device, dst) != BEST_FATTN_KERNEL_NONE;
}
static uintptr_t ggml_sycl_fattn_reserve_halves(ggml_sycl_fattn_extra & extra, size_t n_halves) {
if (n_halves == 0) {
return 0;
}
extra.end = GGML_PAD(extra.end, SYCL_BUFFER_ALIGNMENT);
const uintptr_t block = extra.end;
extra.end += n_halves * sizeof(sycl::half);
return block;
}
ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst) {
ggml_sycl_fattn_extra extra;
extra.end = (uintptr_t) dst->data + ggml_nbytes(dst);
if (dst->op != GGML_OP_FLASH_ATTN_EXT) {
return extra;
}
const ggml_tensor * Q = dst->src[0];
const ggml_tensor * K = dst->src[1];
const ggml_tensor * V = dst->src[2];
if (!Q || !K || !V) {
return extra;
}
const int64_t d = K->ne[0];
const int64_t H = Q->ne[2];
const int64_t q = Q->ne[1];
// calculate the worst-case memory consumption across all kernels
const bool onednn_supported = ggml_sycl_flash_attn_ext_onednn_supported(dst, /* use_shape_limit */ false);
const bool tile_needs_K = K->type != GGML_TYPE_F16;
const bool tile_needs_V = V->type != GGML_TYPE_F16;
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));
size_t need_K = 0, need_V = 0, need_Q = 0, need_out = 0, need_scale = 0;
if (onednn_supported) {
need_Q = (size_t) H * q * d;
need_out = (size_t) H * q * d;
need_scale = 1;
// an f16 cache is bound in place, so it needs no staging copy
if (!ggml_sycl_fattn_onednn_binds_kv(K, V)) {
need_K = (size_t) ggml_nelements(K);
need_V = (size_t) ggml_nelements(V);
}
}
if (tile_needs_K) {
need_K = std::max(need_K, (size_t) ggml_nelements(K));
}
if (tile_needs_V) {
need_V = std::max(need_V, (size_t) ggml_nelements(V));
}
extra.Q_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_Q);
extra.K_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_K);
extra.V_buffer_ptr = (V_is_K_view && !onednn_supported && need_V)
? extra.K_buffer_ptr
: ggml_sycl_fattn_reserve_halves(extra, need_V);
extra.scale_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_scale);
extra.out_buffer_ptr = ggml_sycl_fattn_reserve_halves(extra, need_out);
return extra;
}
size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst) {
const ggml_sycl_fattn_extra extra = ggml_sycl_fattn_get_extra(dst);
return (size_t) (extra.end - (uintptr_t) dst->data);
}
+18
View File
@@ -19,6 +19,24 @@ 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);
// Scratch that flash attention needs beyond the output tensor
struct ggml_sycl_fattn_extra {
uintptr_t K_buffer_ptr = 0; // F16 copy of the K cache
uintptr_t V_buffer_ptr = 0; // F16 copy of the V cache
uintptr_t Q_buffer_ptr = 0; // dense F16 copy of Q, oneDNN only
uintptr_t scale_buffer_ptr = 0; // the softmax scale as an F16 scalar, oneDNN only
uintptr_t out_buffer_ptr = 0; // F16 SDPA output before conversion to F32, oneDNN only
uintptr_t end = 0; // one past the last reserved byte; sizes the allocation
};
// ggml_sycl_fattn_get_extra() is the single source of truth for the layout: it both sizes
// the reservation and hands out the pointers, so the two cannot disagree.
// Each field is the address of one reserved block, or 0 if that block was not reserved,
// in which case the caller allocates from the scratch pool instead.
ggml_sycl_fattn_extra ggml_sycl_fattn_get_extra(const ggml_tensor * dst);
size_t ggml_sycl_flash_attn_ext_get_alloc_size(const ggml_tensor * dst);
void ggml_sycl_flash_attn_ext_mkl(ggml_backend_sycl_context & ctx, ggml_tensor * dst);
#endif // GGML_SYCL_FATTN_HPP
+4 -1
View File
@@ -955,7 +955,10 @@ static size_t ggml_backend_sycl_buffer_type_get_max_size(ggml_backend_buffer_typ
}
static size_t ggml_backend_sycl_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
size_t size = ggml_nbytes(tensor);
// Reserve the additional scratch so it's visible to the graph allocator
size_t size = tensor->op == GGML_OP_FLASH_ATTN_EXT
? ggml_sycl_flash_attn_ext_get_alloc_size(tensor)
: ggml_nbytes(tensor);
int64_t ne0 = tensor->ne[0];
if (ggml_is_quantized(tensor->type)) {