sycl: FLASH_ATTN_EXT failures (#30761)
This commit is contained in:
1 parent
418f905a4e
commit
36a0e2ffc0
2 files changed
+17
-54
No files matched your search
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
Reference in new issue
Block a user