Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Titaniumtown merged 1 commit into from
Aug 29, 2026
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions ggml/src/ggml-sycl/fattn-common.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#include "convert.hpp"
#include "vecdotq.hpp"
#include "fattn-buffers.hpp"
#include "fattn.hpp"

#include "ggml.h"

Expand Down Expand Up @@ -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);
Expand All @@ -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;
Expand All @@ -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) {
Expand All @@ -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;
Expand All @@ -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;
}
}

Expand Down
84 changes: 54 additions & 30 deletions ggml/src/ggml-sycl/fattn-onednn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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}) {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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.
Expand All @@ -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]);
Expand All @@ -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;
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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);
}
Expand All @@ -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();
}
Comment thread
nicois marked this conversation as resolved.

// 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.
Expand All @@ -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;
Expand All @@ -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
Expand Down
6 changes: 5 additions & 1 deletion ggml/src/ggml-sycl/fattn-onednn.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
73 changes: 73 additions & 0 deletions ggml/src/ggml-sycl/fattn.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 changes: 18 additions & 0 deletions ggml/src/ggml-sycl/fattn.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Comment thread
nicois marked this conversation as resolved.
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
5 changes: 4 additions & 1 deletion ggml/src/ggml-sycl/ggml-sycl.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)) {
Expand Down
Loading