mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-17 20:31:47 +02:00
metal : support n_kv_max sparse mask hint in flash attention vec kernel
- add kernel_flash_attn_ext_vec_idx: compacts finite mask entries into a per-row index list (Hillis-Steele scan, one threadgroup per row) - extend vec FA kernel with optional sparse index gathering (FC slot 5) - add host-side gate: sparse path when n_kv_max > 0, mask present, supported head sizes / KV types, n_kv_max <= 4096 - new buffer region extra_idx for the index list - pipeline getter extended with has_sparse param - add test cases: head sizes, quant types, nb>1, nr23 variants, sinks, ALiBi, softcap, permute, v_view_of_k, no-mask fallback Note: multi-row (nb*nr23[1] > 1) cases still failing - rid mapping in the store phase needs revisiting for the sparse path. Assisted-by: pi:llama.cpp/Qwen3.8-27B
This commit is contained in:
Georgi Gerganov
@@ -1577,6 +1577,26 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
|
||||
ggml_metal_library_t lib,
|
||||
const ggml_tensor * op) {
|
||||
assert(op->op == GGML_OP_FLASH_ATTN_EXT);
|
||||
assert(op->src[3]);
|
||||
|
||||
char name[256];
|
||||
|
||||
snprintf(name, 256, "kernel_flash_attn_ext_vec_idx");
|
||||
|
||||
ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
|
||||
if (!res.pipeline) {
|
||||
res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
|
||||
}
|
||||
|
||||
GGML_UNUSED(op);
|
||||
|
||||
return res;
|
||||
}
|
||||
|
||||
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec(
|
||||
ggml_metal_library_t lib,
|
||||
const ggml_tensor * op,
|
||||
@@ -1585,6 +1605,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
|
||||
bool has_bias,
|
||||
bool has_scap,
|
||||
bool has_kvpad,
|
||||
bool has_sparse,
|
||||
int32_t nqpsg,
|
||||
int32_t ne,
|
||||
int32_t nsg,
|
||||
@@ -1614,13 +1635,14 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
|
||||
dv,
|
||||
qne_suffix);
|
||||
|
||||
snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
|
||||
snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
|
||||
base,
|
||||
has_mask,
|
||||
has_sinks,
|
||||
has_bias,
|
||||
has_scap,
|
||||
has_kvpad,
|
||||
has_sparse,
|
||||
ns10,
|
||||
ns20,
|
||||
nsg, nwg);
|
||||
@@ -1633,7 +1655,8 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v
|
||||
ggml_metal_cv_set_bool(cv, has_sinks, FC_FLASH_ATTN_EXT_VEC + 1);
|
||||
ggml_metal_cv_set_bool(cv, has_bias, FC_FLASH_ATTN_EXT_VEC + 2);
|
||||
ggml_metal_cv_set_bool(cv, has_scap, FC_FLASH_ATTN_EXT_VEC + 3);
|
||||
ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT_VEC + 4);
|
||||
ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT_VEC + 4);
|
||||
ggml_metal_cv_set_bool(cv, has_sparse, FC_FLASH_ATTN_EXT_VEC + 5);
|
||||
|
||||
ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20);
|
||||
ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21);
|
||||
|
||||
@@ -201,6 +201,10 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
|
||||
int32_t ns10,
|
||||
int32_t ns20);
|
||||
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
|
||||
ggml_metal_library_t lib,
|
||||
const struct ggml_tensor * op);
|
||||
|
||||
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec(
|
||||
ggml_metal_library_t lib,
|
||||
const struct ggml_tensor * op,
|
||||
@@ -209,6 +213,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_att
|
||||
bool has_bias,
|
||||
bool has_scap,
|
||||
bool has_kvpad,
|
||||
bool has_sparse,
|
||||
int32_t nqpsg,
|
||||
int32_t ne,
|
||||
int32_t nsg,
|
||||
|
||||
@@ -107,6 +107,7 @@
|
||||
#define FC_SUM_ROWS 1400
|
||||
#define FC_UPSCALE 1500
|
||||
#define FC_GATED_DELTA_NET 1600
|
||||
#define FC_FLASH_ATTN_EXT_VEC_IDX 1700
|
||||
|
||||
// op-specific constants
|
||||
#define OP_FLASH_ATTN_EXT_NQPSG 8
|
||||
@@ -458,8 +459,21 @@ typedef struct {
|
||||
float m1;
|
||||
int32_t n_head_log2;
|
||||
float logit_softcap;
|
||||
int32_t n_kv_max_padded;
|
||||
} ggml_metal_kargs_flash_attn_ext_vec;
|
||||
|
||||
typedef struct {
|
||||
int32_t ne30;
|
||||
int32_t ne31;
|
||||
int32_t ne32;
|
||||
int32_t ne33;
|
||||
uint64_t nb31;
|
||||
uint64_t nb32;
|
||||
uint64_t nb33;
|
||||
int32_t n_kv_max;
|
||||
int32_t n_kv_max_padded;
|
||||
} ggml_metal_kargs_flash_attn_ext_vec_idx;
|
||||
|
||||
typedef struct {
|
||||
int32_t nrows;
|
||||
} ggml_metal_kargs_flash_attn_ext_vec_reduce;
|
||||
|
||||
@@ -2841,6 +2841,70 @@ static bool ggml_metal_op_flash_attn_ext_use_kv_f16(const ggml_tensor * op) {
|
||||
}
|
||||
}
|
||||
|
||||
// returns the n_kv_max hint if the sparse path is available for this op, or 0 otherwise
|
||||
// the mask (src[3]) remains the single source of truth: finite entries are the valid KV positions,
|
||||
// n_kv_max is only an upper bound on their number per mask row, used to size the index lists
|
||||
static int ggml_metal_op_flash_attn_ext_n_kv_max_sparse(const ggml_tensor * op) {
|
||||
assert(op->op == GGML_OP_FLASH_ATTN_EXT);
|
||||
|
||||
int32_t n_kv_max = 0;
|
||||
memcpy(&n_kv_max, ((const int32_t *) op->op_params) + 4, sizeof(n_kv_max));
|
||||
|
||||
if (n_kv_max <= 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// the sparse indices are gathered from the mask
|
||||
if (!op->src[3]) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// the sparse path is implemented for the vec kernels only
|
||||
if (!ggml_metal_op_flash_attn_ext_use_vec(op)) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// bound the size of the index lists
|
||||
if (n_kv_max > 4096) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
// vec kernel instantiations exist for these (type, dk, dv) combinations only
|
||||
const int64_t dk = op->src[1]->ne[0];
|
||||
const int64_t dv = op->src[2]->ne[0];
|
||||
|
||||
const bool dk_dv_ok = (dk == 32 && dv == 32) ||
|
||||
(dk == 64 && dv == 64) ||
|
||||
(dk == 96 && dv == 96) ||
|
||||
(dk == 128 && dv == 128) ||
|
||||
(dk == 192 && dv == 128) ||
|
||||
(dk == 192 && dv == 192) ||
|
||||
(dk == 256 && dv == 256) ||
|
||||
(dk == 320 && dv == 256) ||
|
||||
(dk == 512 && dv == 512) ||
|
||||
(dk == 576 && dv == 512);
|
||||
|
||||
if (!dk_dv_ok) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
switch (op->src[1]->type) {
|
||||
case GGML_TYPE_F16:
|
||||
case GGML_TYPE_BF16:
|
||||
case GGML_TYPE_F32:
|
||||
case GGML_TYPE_Q4_0:
|
||||
case GGML_TYPE_Q4_1:
|
||||
case GGML_TYPE_Q5_0:
|
||||
case GGML_TYPE_Q5_1:
|
||||
case GGML_TYPE_Q8_0:
|
||||
break;
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
|
||||
return n_kv_max;
|
||||
}
|
||||
|
||||
// in some models (e.g. MLA-based), V is a view of K (the first ne20 elements of each K row);
|
||||
// the dequantized V is then a view of the dequantized K and does not need its own dequant or scratch
|
||||
// - ref: https://github.com/ggml-org/llama.cpp/pull/13435
|
||||
@@ -3011,6 +3075,24 @@ size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const ggml_tensor * op) {
|
||||
return k_size + v_size;
|
||||
}
|
||||
|
||||
// size of the sparse index lists: one list of KV indices per mask row,
|
||||
// padded with -1 up to a multiple of OP_FLASH_ATTN_EXT_VEC_NCPSG
|
||||
size_t ggml_metal_op_flash_attn_ext_extra_idx(const ggml_tensor * op) {
|
||||
assert(op->op == GGML_OP_FLASH_ATTN_EXT);
|
||||
|
||||
GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne);
|
||||
|
||||
const int n_kv_max = ggml_metal_op_flash_attn_ext_n_kv_max_sparse(op);
|
||||
|
||||
if (n_kv_max <= 0) {
|
||||
return 0;
|
||||
}
|
||||
|
||||
const int n_kv_max_padded = GGML_PAD(n_kv_max, OP_FLASH_ATTN_EXT_VEC_NCPSG);
|
||||
|
||||
return GGML_PAD(sizeof(int32_t)*(size_t) n_kv_max_padded*ne31*ne32*ne33, 16);
|
||||
}
|
||||
|
||||
int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_tensor * op = ctx->node(idx);
|
||||
|
||||
@@ -3088,7 +3170,16 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_metal_buffer_id bid_kv_f16 = bid_tmp;
|
||||
bid_kv_f16.offs += ggml_metal_op_flash_attn_ext_extra_tmp(op);
|
||||
|
||||
const bool use_kv_f16 = ggml_metal_op_flash_attn_ext_use_kv_f16(op);
|
||||
// sparse path: gather the finite mask entries into index lists and run the vec kernels over them
|
||||
const int n_kv_max_sparse = ggml_metal_op_flash_attn_ext_n_kv_max_sparse(op);
|
||||
const bool use_sparse = n_kv_max_sparse > 0;
|
||||
const int n_kv_max_padded = use_sparse ? GGML_PAD(n_kv_max_sparse, OP_FLASH_ATTN_EXT_VEC_NCPSG) : 0;
|
||||
|
||||
// the vec kernels dequantize the KV inline; no need for the F16 dequant pass in the sparse path
|
||||
const bool use_kv_f16 = !use_sparse && ggml_metal_op_flash_attn_ext_use_kv_f16(op);
|
||||
|
||||
ggml_metal_buffer_id bid_idx = bid_kv_f16;
|
||||
bid_idx.offs += ggml_metal_op_flash_attn_ext_extra_kv_f16(op);
|
||||
|
||||
ggml_metal_buffer_id bid_k = bid_src1;
|
||||
ggml_metal_buffer_id bid_v = bid_src2;
|
||||
@@ -3362,12 +3453,15 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
|
||||
#undef FATTN_SMEM
|
||||
} else {
|
||||
// half4x4 kernel
|
||||
auto cfg = ggml_metal_tuning::fa_vec_pick(
|
||||
props_dev->device_id,
|
||||
props_dev->gpu_family,
|
||||
(int) op->src[1]->type,
|
||||
(int) ne00, (int) ne20, // dk, dv (ne00 == dk for FA)
|
||||
ne11, ne01);
|
||||
// sparse: the index lists are per query row, so a threadgroup can share KV with Q == 1 only
|
||||
auto cfg = use_sparse
|
||||
? ggml_metal_tuning::fa_vec_baseline_cfg((int) ne00, (int) ne20)
|
||||
: ggml_metal_tuning::fa_vec_pick(
|
||||
props_dev->device_id,
|
||||
props_dev->gpu_family,
|
||||
(int) op->src[1]->type,
|
||||
(int) ne00, (int) ne20, // dk, dv (ne00 == dk for FA)
|
||||
ne11, ne01);
|
||||
int nqptg = cfg.Q; // queries per threadgroup
|
||||
const int ncpsg = OP_FLASH_ATTN_EXT_VEC_NCPSG; // cache values per simdgroup !! sync with kernel template arguments !!
|
||||
const int nhptg = 1; // heads per threadgroup
|
||||
@@ -3378,7 +3472,37 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
|
||||
|
||||
bool need_sync = false;
|
||||
|
||||
const bool has_kvpad = ne11 % ncpsg != 0;
|
||||
const bool has_kvpad = !use_sparse && ne11 % ncpsg != 0;
|
||||
|
||||
if (use_sparse) {
|
||||
assert(ggml_metal_op_flash_attn_ext_extra_idx(op) != 0);
|
||||
|
||||
ggml_metal_kargs_flash_attn_ext_vec_idx args0 = {
|
||||
/*.ne30 =*/ ne30,
|
||||
/*.ne31 =*/ ne31,
|
||||
/*.ne32 =*/ ne32,
|
||||
/*.ne33 =*/ ne33,
|
||||
/*.nb31 =*/ nb31,
|
||||
/*.nb32 =*/ nb32,
|
||||
/*.nb33 =*/ nb33,
|
||||
/*.n_kv_max =*/ n_kv_max_sparse,
|
||||
/*.n_kv_max_padded =*/ n_kv_max_padded,
|
||||
};
|
||||
|
||||
auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(lib, op);
|
||||
|
||||
ggml_metal_encoder_set_pipeline(enc, pipeline0);
|
||||
ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_src3, 1);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_idx, 2);
|
||||
|
||||
int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline0), 256);
|
||||
nth = std::max(32, (nth/32)*32);
|
||||
|
||||
ggml_metal_encoder_dispatch_threadgroups(enc, ne31, ne32, ne33, nth, 1, 1);
|
||||
|
||||
need_sync = true;
|
||||
}
|
||||
|
||||
if (has_kvpad) {
|
||||
assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0);
|
||||
@@ -3444,6 +3568,11 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
|
||||
// however, this does not lead to significant improvement, so disabled
|
||||
nwg = 1;
|
||||
nsg = 4;
|
||||
} else if (use_sparse) {
|
||||
// the reduce kernel sums over one workgroup per lane (32), so nwg must be 32;
|
||||
// workgroups beyond the number of chunks emit empty partials which the reduce ignores
|
||||
nsg = 1;
|
||||
nwg = 32;
|
||||
} else {
|
||||
nwg = 32;
|
||||
nsg = 1;
|
||||
@@ -3468,7 +3597,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
|
||||
/*.nb01 =*/ nb01,
|
||||
/*.nb02 =*/ nb02,
|
||||
/*.nb03 =*/ nb03,
|
||||
/*.ne11 =*/ ne11,
|
||||
/*.ne11 =*/ use_sparse ? n_kv_max_padded : ne11,
|
||||
/*.ne_12_2 =*/ ne12,
|
||||
/*.ne_12_3 =*/ ne13,
|
||||
/*.ns10 =*/ ns10,
|
||||
@@ -3494,9 +3623,10 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
|
||||
/*.m1 =*/ m1,
|
||||
/*.n_head_log2 =*/ n_head_log2,
|
||||
/*.logit_softcap =*/ logit_softcap,
|
||||
/*.n_kv_max_padded =*/ n_kv_max_padded,
|
||||
};
|
||||
|
||||
auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
|
||||
auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, use_sparse, nqptg, cfg.NE, nsg, nwg, use_kv_f16, ns10, ns20);
|
||||
|
||||
GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline));
|
||||
|
||||
@@ -3507,6 +3637,7 @@ int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) {
|
||||
ggml_metal_encoder_set_buffer (enc, bid_v, 3);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_src3, 4);
|
||||
ggml_metal_encoder_set_buffer (enc, bid_src4, 5);
|
||||
ggml_metal_encoder_set_buffer (enc, use_sparse ? bid_idx : bid_src0, 8);
|
||||
|
||||
const size_t smem = FATTN_SMEM(nsg);
|
||||
|
||||
|
||||
@@ -43,6 +43,7 @@ size_t ggml_metal_op_flash_attn_ext_extra_pad(const struct ggml_tensor * op);
|
||||
size_t ggml_metal_op_flash_attn_ext_extra_blk(const struct ggml_tensor * op);
|
||||
size_t ggml_metal_op_flash_attn_ext_extra_tmp(const struct ggml_tensor * op);
|
||||
size_t ggml_metal_op_flash_attn_ext_extra_kv_f16(const struct ggml_tensor * op);
|
||||
size_t ggml_metal_op_flash_attn_ext_extra_idx(const struct ggml_tensor * op);
|
||||
|
||||
int ggml_metal_op_concat (ggml_metal_op_t ctx, int idx);
|
||||
int ggml_metal_op_repeat (ggml_metal_op_t ctx, int idx);
|
||||
|
||||
@@ -232,6 +232,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_
|
||||
res += ggml_metal_op_flash_attn_ext_extra_blk(tensor);
|
||||
res += ggml_metal_op_flash_attn_ext_extra_tmp(tensor);
|
||||
res += ggml_metal_op_flash_attn_ext_extra_kv_f16(tensor);
|
||||
res += ggml_metal_op_flash_attn_ext_extra_idx(tensor);
|
||||
} break;
|
||||
case GGML_OP_CUMSUM:
|
||||
case GGML_OP_ARGSORT:
|
||||
|
||||
@@ -1071,6 +1071,93 @@ constant int32_t FC_flash_attn_ext_vec_ns10 [[function_constant(FC_FLASH_ATTN_EX
|
||||
constant int32_t FC_flash_attn_ext_vec_ns20 [[function_constant(FC_FLASH_ATTN_EXT_VEC + 21)]];
|
||||
constant int32_t FC_flash_attn_ext_vec_nsg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 22)]];
|
||||
constant int32_t FC_flash_attn_ext_vec_nwg [[function_constant(FC_FLASH_ATTN_EXT_VEC + 23)]];
|
||||
constant bool FC_flash_attn_ext_vec_has_sparse [[function_constant(FC_FLASH_ATTN_EXT_VEC + 5)]];
|
||||
|
||||
// compress the finite entries of each KQ mask row into a list of KV indices (ascending order),
|
||||
// padded with -1 up to n_kv_max_padded (a multiple of OP_FLASH_ATTN_EXT_VEC_NCPSG)
|
||||
// one threadgroup per mask row; the mask remains the single source of truth for the values
|
||||
kernel void kernel_flash_attn_ext_vec_idx(
|
||||
constant ggml_metal_kargs_flash_attn_ext_vec_idx & args,
|
||||
device const half * mask,
|
||||
device int * idx,
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiitg[[thread_index_in_threadgroup]],
|
||||
ushort3 ntg[[threads_per_threadgroup]]) {
|
||||
constexpr short NW = N_SIMDWIDTH;
|
||||
|
||||
const int i1 = tgpig[0];
|
||||
const int i2 = tgpig[1];
|
||||
const int i3 = tgpig[2];
|
||||
|
||||
device const half * pm = (device const half *) (mask + i1*args.nb31 + i2*args.nb32 + i3*args.nb33);
|
||||
device int * pidx = idx + ((i3*args.ne32 + i2)*args.ne31 + i1)*args.n_kv_max_padded;
|
||||
|
||||
const int n = args.ne30;
|
||||
const int q = n/ntg.x;
|
||||
const int r = n%ntg.x;
|
||||
|
||||
// each thread handles a contiguous slice of the mask row
|
||||
const int r0 = q*tiitg + min((int) tiitg, r);
|
||||
const int r1 = r0 + q + (tiitg < r ? 1 : 0);
|
||||
|
||||
// count the finite entries in the slice
|
||||
int cnt = 0;
|
||||
for (int i = r0; i < r1; ++i) {
|
||||
cnt += isfinite((float) pm[i]) ? 1 : 0;
|
||||
}
|
||||
|
||||
const short sgitg = tiitg/NW;
|
||||
const short tiisg = tiitg%NW;
|
||||
|
||||
threadgroup int tcount[8];
|
||||
|
||||
// simd_sum is a collective: all lanes must evaluate it
|
||||
const int sg_sum = simd_sum(cnt);
|
||||
if (tiisg == 0) {
|
||||
tcount[sgitg] = sg_sum;
|
||||
}
|
||||
|
||||
threadgroup_barrier(mem_flags::mem_threadgroup);
|
||||
|
||||
int total = 0;
|
||||
for (short s = 0; s < ntg.x/NW; ++s) {
|
||||
total += tcount[s];
|
||||
}
|
||||
|
||||
// base offset of this thread's slice in the output list (exclusive scan within the simdgroup)
|
||||
int sg_base = 0;
|
||||
for (short s = 0; s < sgitg; ++s) {
|
||||
sg_base += tcount[s];
|
||||
}
|
||||
|
||||
// exclusive prefix scan of the per-thread counts within the simdgroup
|
||||
int incl = cnt;
|
||||
for (int d = 1; d < NW; d <<= 1) {
|
||||
const int v = simd_shuffle_up(incl, d);
|
||||
if (tiisg >= d) {
|
||||
incl += v;
|
||||
}
|
||||
}
|
||||
const int base = sg_base + (incl - cnt);
|
||||
|
||||
// write the finite positions in order; if the hint is violated, keep only the first n_kv_max entries
|
||||
int j = 0;
|
||||
for (int i = r0; i < r1; ++i) {
|
||||
if (base + j >= args.n_kv_max) {
|
||||
break;
|
||||
}
|
||||
if (isfinite((float) pm[i])) {
|
||||
pidx[base + j] = i;
|
||||
j++;
|
||||
}
|
||||
}
|
||||
|
||||
// pad the tail of the list with -1
|
||||
const int count = min(total, args.n_kv_max);
|
||||
for (int i = count + tiitg; i < args.n_kv_max_padded; i += ntg.x) {
|
||||
pidx[i] = -1;
|
||||
}
|
||||
}
|
||||
|
||||
template<
|
||||
typename q4_t, // query types in shared memory
|
||||
@@ -1091,6 +1178,7 @@ template<
|
||||
short NE = 4, // head elements per thread
|
||||
short Q = OP_FLASH_ATTN_EXT_VEC_NQPSG, // queries per threadgroup
|
||||
short C = OP_FLASH_ATTN_EXT_VEC_NCPSG> // cache items per threadgroup
|
||||
|
||||
kernel void kernel_flash_attn_ext_vec(
|
||||
constant ggml_metal_kargs_flash_attn_ext_vec & args,
|
||||
device const char * q,
|
||||
@@ -1100,6 +1188,7 @@ kernel void kernel_flash_attn_ext_vec(
|
||||
device const char * sinks,
|
||||
device const char * pad,
|
||||
device char * dst,
|
||||
device const char * idx,
|
||||
threadgroup half * shmem_f16 [[threadgroup(0)]],
|
||||
uint3 tgpig[[threadgroup_position_in_grid]],
|
||||
ushort tiisg[[thread_index_in_simdgroup]],
|
||||
@@ -1205,7 +1294,15 @@ kernel void kernel_flash_attn_ext_vec(
|
||||
const short ty = tiisg/NL;
|
||||
|
||||
// pointer to the mask
|
||||
device const half * pm_base = (device const half *) (mask + iq1*Q*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
|
||||
device const half * pm_base = (device const half *) (mask + (FC_flash_attn_ext_vec_has_sparse ? (iq1%args.ne31) : (iq1*Q))*args.nb31 + (iq2%args.ne32)*args.nb32 + (iq3%args.ne33)*args.nb33);
|
||||
|
||||
// sparse indices: the list of finite mask entries per query row
|
||||
// the sparse path requires Q == 1 (enforced by the host)
|
||||
device const int * pidx = nullptr;
|
||||
if (FC_flash_attn_ext_vec_has_sparse) {
|
||||
pidx = (device const int *) (idx + ((iq3%args.ne33)*args.ne32 + (iq2%args.ne32))*args.ne31*args.n_kv_max_padded
|
||||
+ (iq1%args.ne31)*args.n_kv_max_padded);
|
||||
}
|
||||
|
||||
float slope = 1.0f;
|
||||
|
||||
@@ -1265,11 +1362,22 @@ kernel void kernel_flash_attn_ext_vec(
|
||||
}
|
||||
|
||||
if (FC_flash_attn_ext_vec_has_mask) {
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
if ((iq1*Q + qq) < args.ne01) {
|
||||
sm[qq*C + tiisg] = pm[qq][ic + tiisg];
|
||||
} else {
|
||||
sm[qq*C + tiisg] = -MAXHALF;
|
||||
if (FC_flash_attn_ext_vec_has_sparse) {
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
const int i11 = pidx[ic + tiisg];
|
||||
if ((iq1*Q + qq) < args.ne01 && i11 >= 0) {
|
||||
sm[qq*C + tiisg] = pm[qq][i11];
|
||||
} else {
|
||||
sm[qq*C + tiisg] = -MAXHALF;
|
||||
}
|
||||
}
|
||||
} else {
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
if ((iq1*Q + qq) < args.ne01) {
|
||||
sm[qq*C + tiisg] = pm[qq][ic + tiisg];
|
||||
} else {
|
||||
sm[qq*C + tiisg] = -MAXHALF;
|
||||
}
|
||||
}
|
||||
}
|
||||
} else {
|
||||
@@ -1294,9 +1402,13 @@ kernel void kernel_flash_attn_ext_vec(
|
||||
|
||||
// Q*K^T
|
||||
{
|
||||
device const k4_t * pk4 = (device const k4_t *) (k + ic*args.nb11);
|
||||
device const k4_t * pk4 = nullptr;
|
||||
|
||||
pk4 += ty*NS10/4 + tx;
|
||||
if (!FC_flash_attn_ext_vec_has_sparse) {
|
||||
pk4 = (device const k4_t *) (k + ic*args.nb11);
|
||||
|
||||
pk4 += ty*NS10/4 + tx;
|
||||
}
|
||||
|
||||
qk_t mqk[Q][C/NE];
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
@@ -1307,7 +1419,35 @@ kernel void kernel_flash_attn_ext_vec(
|
||||
|
||||
// each simdgroup processes Q queries and NE (NW/NL) cache elements
|
||||
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
|
||||
if (is_same<kd4_t, k4_t>::value) {
|
||||
if (FC_flash_attn_ext_vec_has_sparse) {
|
||||
// the KV rows are gathered from the index list; -1 entries are padding
|
||||
const int i11 = pidx[ic + NE*cc + ty];
|
||||
if (i11 >= 0) {
|
||||
if (is_same<kd4_t, k4_t>::value) {
|
||||
device const k4_t * pk4s = (device const k4_t *) (k + i11*args.nb11) + tx;
|
||||
FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) {
|
||||
const k4_t k_elem = pk4s[ii*NL];
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
mqk[qq][cc] += dot((float4) k_elem, (float4) sq4[qq*PK4 + ii*NL + tx]);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
device const kd4_t * pk = (device const kd4_t *) (k + i11*args.nb11);
|
||||
|
||||
k4_t mk;
|
||||
|
||||
FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) {
|
||||
const short i = ii*NL + tx;
|
||||
|
||||
deq_k_t4(pk + i/nl_k, i%nl_k, mk);
|
||||
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
mqk[qq][cc] += dot((float4) mk, (float4) sq4[qq*PK4 + i]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (is_same<kd4_t, k4_t>::value) {
|
||||
FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) {
|
||||
const k4_t k_elem = pk4[cc*NE*NS10/4 + ii*NL];
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
@@ -1422,7 +1562,40 @@ kernel void kernel_flash_attn_ext_vec(
|
||||
}
|
||||
}
|
||||
|
||||
if (is_same<vd4_t, v4_t>::value) {
|
||||
if (FC_flash_attn_ext_vec_has_sparse) {
|
||||
FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) {
|
||||
// the KV rows are gathered from the index list; -1 entries are padding
|
||||
const int i11 = pidx[ic + NE*cc + ty];
|
||||
if (i11 >= 0) {
|
||||
if (is_same<vd4_t, v4_t>::value) {
|
||||
device const v4_t * pv4 = (device const v4_t *) (v + i11*args.nb21);
|
||||
|
||||
pv4 += tx;
|
||||
|
||||
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
||||
const v4_t v_elem = pv4[ii*NL];
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
lo[qq][ii] += o4_t(float4(v_elem)*float4(ss[qq*C + cc*NE + ty]));
|
||||
}
|
||||
}
|
||||
} else {
|
||||
device const vd4_t * pv4 = (device const vd4_t *) (v + i11*args.nb21);
|
||||
|
||||
FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) {
|
||||
const short i = ii*NL + tx;
|
||||
|
||||
v4_t mv;
|
||||
|
||||
deq_v_t4(pv4 + i/nl_v, i%nl_v, mv);
|
||||
|
||||
FOR_UNROLL (short qq = 0; qq < Q; ++qq) {
|
||||
lo[qq][ii] += o4_t(float4(mv)*float4(ss[qq*C + cc*NE + ty]));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (is_same<vd4_t, v4_t>::value) {
|
||||
device const v4_t * pv4 = (device const v4_t *) (v + ic*args.nb21);
|
||||
|
||||
pv4 += ty*NS20/4 + tx;
|
||||
|
||||
@@ -1503,11 +1503,6 @@ struct test_case {
|
||||
double err = ud->tc->err(f1.data(), f2.data(), f1.size());
|
||||
if (err > ud->tc->max_err(ud->backend1)) {
|
||||
printf("[%s] ERR = %.9f > %.9f ", ggml_op_desc(t1), err, ud->tc->max_err(ud->backend1));
|
||||
//for (int i = 0; i < (int) f1.size(); i++) {
|
||||
// printf("%5d %9.6f %9.6f, diff = %9.6f\n", i, f1[i], f2[i], f1[i] - f2[i]);
|
||||
//}
|
||||
//printf("\n");
|
||||
//exit(1);
|
||||
ud->ok = false;
|
||||
}
|
||||
return true;
|
||||
@@ -10192,8 +10187,55 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
|
||||
test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {16, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 768));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512, false));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 2304));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: more head sizes
|
||||
test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(256, 256, 2, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(192, 128, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext( 64, 64, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext( 96, 96, 4, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: quantized KV
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q4_0, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_1, GGML_TYPE_Q4_1, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q5_0, GGML_TYPE_Q5_0, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q5_1, GGML_TYPE_Q5_1, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_BF16, GGML_TYPE_BF16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: no mask (dense fallback, hint ignored)
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: with sinks
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, true, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: with ALiBi bias
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 1.0f, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: with logit softcap
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 50.0f, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: larger batch
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 8, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: V is a view of K
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true, 512));
|
||||
|
||||
// Sparse mask hint: permuted KV layout
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 2, 1, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: batch > 1
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 1}, 4096, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(128, 128, 4, { 8, 1}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// Sparse mask hint: nr23 != [8,1]
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 4, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, {16, 1}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 8, 2}, 4096, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 4, 2}, 4096, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
test_cases.emplace_back(new test_flash_attn_ext(512, 512, 1, { 2, 1}, 4096, 3, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, false, 512));
|
||||
|
||||
// more V-is-sub-view-of-K cases: other head shapes, and full views with equal head sizes
|
||||
test_cases.emplace_back(new test_flash_attn_ext(320, 256, 1, {32, 1}, 512, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16, {0, 1, 2, 3}, true, true));
|
||||
|
||||
Reference in New Issue
Block a user