metal : support non-contiguous KV
ggml-ci
This commit is contained in:
Georgi Gerganov
parent
a2e6b9dee1
commit
f8c96dfd97
2 changed files with 37 additions and 53 deletions
55
ggml-metal.m
55
ggml-metal.m
|
|
@ -2518,7 +2518,6 @@ static enum ggml_status ggml_metal_graph_compute(
|
||||||
GGML_ASSERT(src0->type == GGML_TYPE_F32);
|
GGML_ASSERT(src0->type == GGML_TYPE_F32);
|
||||||
|
|
||||||
GGML_ASSERT(ggml_are_same_shape (src1, src2));
|
GGML_ASSERT(ggml_are_same_shape (src1, src2));
|
||||||
GGML_ASSERT(ggml_are_same_stride(src1, src2));
|
|
||||||
|
|
||||||
struct ggml_tensor * src3 = gf->nodes[i]->src[3];
|
struct ggml_tensor * src3 = gf->nodes[i]->src[3];
|
||||||
|
|
||||||
|
|
@ -2530,6 +2529,11 @@ static enum ggml_status ggml_metal_graph_compute(
|
||||||
GGML_ASSERT(!src3 || src3->ne[1] >= GGML_PAD(src0->ne[1], 8) &&
|
GGML_ASSERT(!src3 || src3->ne[1] >= GGML_PAD(src0->ne[1], 8) &&
|
||||||
"the Flash-Attention Metal kernel requires the mask to be padded to 8 and at least n_queries big");
|
"the Flash-Attention Metal kernel requires the mask to be padded to 8 and at least n_queries big");
|
||||||
|
|
||||||
|
const uint64_t nb20 = src2 ? src2->nb[0] : 0; GGML_UNUSED(nb20);
|
||||||
|
const uint64_t nb21 = src2 ? src2->nb[1] : 0;
|
||||||
|
const uint64_t nb22 = src2 ? src2->nb[2] : 0;
|
||||||
|
const uint64_t nb23 = src2 ? src2->nb[3] : 0;
|
||||||
|
|
||||||
const int64_t ne30 = src3 ? src3->ne[0] : 0; GGML_UNUSED(ne30);
|
const int64_t ne30 = src3 ? src3->ne[0] : 0; GGML_UNUSED(ne30);
|
||||||
//const int64_t ne31 = src3 ? src3->ne[1] : 0;
|
//const int64_t ne31 = src3 ? src3->ne[1] : 0;
|
||||||
const int64_t ne32 = src3 ? src3->ne[2] : 0; GGML_UNUSED(ne32);
|
const int64_t ne32 = src3 ? src3->ne[2] : 0; GGML_UNUSED(ne32);
|
||||||
|
|
@ -2598,32 +2602,29 @@ static enum ggml_status ggml_metal_graph_compute(
|
||||||
[encoder setBuffer:id_src0 offset:offs_src0 atIndex:3];
|
[encoder setBuffer:id_src0 offset:offs_src0 atIndex:3];
|
||||||
}
|
}
|
||||||
[encoder setBuffer:id_dst offset:offs_dst atIndex:4];
|
[encoder setBuffer:id_dst offset:offs_dst atIndex:4];
|
||||||
[encoder setBytes:&ne00 length:sizeof( int64_t) atIndex:5];
|
[encoder setBytes:&ne01 length:sizeof( int64_t) atIndex:5];
|
||||||
[encoder setBytes:&ne01 length:sizeof( int64_t) atIndex:6];
|
[encoder setBytes:&ne02 length:sizeof( int64_t) atIndex:6];
|
||||||
[encoder setBytes:&ne02 length:sizeof( int64_t) atIndex:7];
|
[encoder setBytes:&ne03 length:sizeof( int64_t) atIndex:7];
|
||||||
[encoder setBytes:&ne03 length:sizeof( int64_t) atIndex:8];
|
[encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:8];
|
||||||
[encoder setBytes:&nb00 length:sizeof(uint64_t) atIndex:9];
|
[encoder setBytes:&nb02 length:sizeof(uint64_t) atIndex:9];
|
||||||
[encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:10];
|
[encoder setBytes:&nb03 length:sizeof(uint64_t) atIndex:10];
|
||||||
[encoder setBytes:&nb02 length:sizeof(uint64_t) atIndex:11];
|
[encoder setBytes:&ne11 length:sizeof( int64_t) atIndex:11];
|
||||||
[encoder setBytes:&nb03 length:sizeof(uint64_t) atIndex:12];
|
[encoder setBytes:&ne12 length:sizeof( int64_t) atIndex:12];
|
||||||
[encoder setBytes:&ne10 length:sizeof( int64_t) atIndex:13];
|
[encoder setBytes:&ne13 length:sizeof( int64_t) atIndex:13];
|
||||||
[encoder setBytes:&ne11 length:sizeof( int64_t) atIndex:14];
|
[encoder setBytes:&nb11 length:sizeof(uint64_t) atIndex:14];
|
||||||
[encoder setBytes:&ne12 length:sizeof( int64_t) atIndex:15];
|
[encoder setBytes:&nb12 length:sizeof(uint64_t) atIndex:15];
|
||||||
[encoder setBytes:&ne13 length:sizeof( int64_t) atIndex:16];
|
[encoder setBytes:&nb13 length:sizeof(uint64_t) atIndex:16];
|
||||||
[encoder setBytes:&nb10 length:sizeof(uint64_t) atIndex:17];
|
[encoder setBytes:&nb21 length:sizeof(uint64_t) atIndex:17];
|
||||||
[encoder setBytes:&nb11 length:sizeof(uint64_t) atIndex:18];
|
[encoder setBytes:&nb22 length:sizeof(uint64_t) atIndex:18];
|
||||||
[encoder setBytes:&nb12 length:sizeof(uint64_t) atIndex:19];
|
[encoder setBytes:&nb23 length:sizeof(uint64_t) atIndex:19];
|
||||||
[encoder setBytes:&nb13 length:sizeof(uint64_t) atIndex:20];
|
[encoder setBytes:&nb31 length:sizeof(uint64_t) atIndex:20];
|
||||||
[encoder setBytes:&nb31 length:sizeof(uint64_t) atIndex:21];
|
[encoder setBytes:&ne1 length:sizeof( int64_t) atIndex:21];
|
||||||
[encoder setBytes:&ne0 length:sizeof( int64_t) atIndex:22];
|
[encoder setBytes:&ne2 length:sizeof( int64_t) atIndex:22];
|
||||||
[encoder setBytes:&ne1 length:sizeof( int64_t) atIndex:23];
|
[encoder setBytes:&scale length:sizeof( float) atIndex:23];
|
||||||
[encoder setBytes:&ne2 length:sizeof( int64_t) atIndex:24];
|
[encoder setBytes:&max_bias length:sizeof( float) atIndex:24];
|
||||||
[encoder setBytes:&ne3 length:sizeof( int64_t) atIndex:25];
|
[encoder setBytes:&m0 length:sizeof(m0) atIndex:25];
|
||||||
[encoder setBytes:&scale length:sizeof( float) atIndex:26];
|
[encoder setBytes:&m1 length:sizeof(m1) atIndex:26];
|
||||||
[encoder setBytes:&max_bias length:sizeof( float) atIndex:27];
|
[encoder setBytes:&n_head_log2 length:sizeof(n_head_log2) atIndex:27];
|
||||||
[encoder setBytes:&m0 length:sizeof(m0) atIndex:28];
|
|
||||||
[encoder setBytes:&m1 length:sizeof(m1) atIndex:29];
|
|
||||||
[encoder setBytes:&n_head_log2 length:sizeof(n_head_log2) atIndex:30];
|
|
||||||
|
|
||||||
if (!use_vec_kernel) {
|
if (!use_vec_kernel) {
|
||||||
// half8x8 kernel
|
// half8x8 kernel
|
||||||
|
|
|
||||||
|
|
@ -2049,27 +2049,24 @@ typedef void (flash_attn_ext_f16_t)(
|
||||||
device const char * v,
|
device const char * v,
|
||||||
device const char * mask,
|
device const char * mask,
|
||||||
device float * dst,
|
device float * dst,
|
||||||
constant int64_t & ne00,
|
|
||||||
constant int64_t & ne01,
|
constant int64_t & ne01,
|
||||||
constant int64_t & ne02,
|
constant int64_t & ne02,
|
||||||
constant int64_t & ne03,
|
constant int64_t & ne03,
|
||||||
constant uint64_t & nb00,
|
|
||||||
constant uint64_t & nb01,
|
constant uint64_t & nb01,
|
||||||
constant uint64_t & nb02,
|
constant uint64_t & nb02,
|
||||||
constant uint64_t & nb03,
|
constant uint64_t & nb03,
|
||||||
constant int64_t & ne10,
|
|
||||||
constant int64_t & ne11,
|
constant int64_t & ne11,
|
||||||
constant int64_t & ne12,
|
constant int64_t & ne12,
|
||||||
constant int64_t & ne13,
|
constant int64_t & ne13,
|
||||||
constant uint64_t & nb10,
|
|
||||||
constant uint64_t & nb11,
|
constant uint64_t & nb11,
|
||||||
constant uint64_t & nb12,
|
constant uint64_t & nb12,
|
||||||
constant uint64_t & nb13,
|
constant uint64_t & nb13,
|
||||||
|
constant uint64_t & nb21,
|
||||||
|
constant uint64_t & nb22,
|
||||||
|
constant uint64_t & nb23,
|
||||||
constant uint64_t & nb31,
|
constant uint64_t & nb31,
|
||||||
constant int64_t & ne0,
|
|
||||||
constant int64_t & ne1,
|
constant int64_t & ne1,
|
||||||
constant int64_t & ne2,
|
constant int64_t & ne2,
|
||||||
constant int64_t & ne3,
|
|
||||||
constant float & scale,
|
constant float & scale,
|
||||||
constant float & max_bias,
|
constant float & max_bias,
|
||||||
constant float & m0,
|
constant float & m0,
|
||||||
|
|
@ -2090,27 +2087,24 @@ kernel void kernel_flash_attn_ext_f16(
|
||||||
device const char * v,
|
device const char * v,
|
||||||
device const char * mask,
|
device const char * mask,
|
||||||
device float * dst,
|
device float * dst,
|
||||||
constant int64_t & ne00,
|
|
||||||
constant int64_t & ne01,
|
constant int64_t & ne01,
|
||||||
constant int64_t & ne02,
|
constant int64_t & ne02,
|
||||||
constant int64_t & ne03,
|
constant int64_t & ne03,
|
||||||
constant uint64_t & nb00,
|
|
||||||
constant uint64_t & nb01,
|
constant uint64_t & nb01,
|
||||||
constant uint64_t & nb02,
|
constant uint64_t & nb02,
|
||||||
constant uint64_t & nb03,
|
constant uint64_t & nb03,
|
||||||
constant int64_t & ne10,
|
|
||||||
constant int64_t & ne11,
|
constant int64_t & ne11,
|
||||||
constant int64_t & ne12,
|
constant int64_t & ne12,
|
||||||
constant int64_t & ne13,
|
constant int64_t & ne13,
|
||||||
constant uint64_t & nb10,
|
|
||||||
constant uint64_t & nb11,
|
constant uint64_t & nb11,
|
||||||
constant uint64_t & nb12,
|
constant uint64_t & nb12,
|
||||||
constant uint64_t & nb13,
|
constant uint64_t & nb13,
|
||||||
|
constant uint64_t & nb21,
|
||||||
|
constant uint64_t & nb22,
|
||||||
|
constant uint64_t & nb23,
|
||||||
constant uint64_t & nb31,
|
constant uint64_t & nb31,
|
||||||
constant int64_t & ne0,
|
|
||||||
constant int64_t & ne1,
|
constant int64_t & ne1,
|
||||||
constant int64_t & ne2,
|
constant int64_t & ne2,
|
||||||
constant int64_t & ne3,
|
|
||||||
constant float & scale,
|
constant float & scale,
|
||||||
constant float & max_bias,
|
constant float & max_bias,
|
||||||
constant float & m0,
|
constant float & m0,
|
||||||
|
|
@ -2180,10 +2174,6 @@ kernel void kernel_flash_attn_ext_f16(
|
||||||
const short ne22 = ne12;
|
const short ne22 = ne12;
|
||||||
const short ne23 = ne13;
|
const short ne23 = ne13;
|
||||||
|
|
||||||
const uint nb21 = nb11;
|
|
||||||
const uint nb22 = nb12;
|
|
||||||
const uint nb23 = nb13;
|
|
||||||
|
|
||||||
// broadcast
|
// broadcast
|
||||||
const short rk2 = ne02/ne12;
|
const short rk2 = ne02/ne12;
|
||||||
const short rk3 = ne03/ne13;
|
const short rk3 = ne03/ne13;
|
||||||
|
|
@ -2430,27 +2420,24 @@ kernel void kernel_flash_attn_ext_vec_f16(
|
||||||
device const char * v,
|
device const char * v,
|
||||||
device const char * mask,
|
device const char * mask,
|
||||||
device float * dst,
|
device float * dst,
|
||||||
constant int64_t & ne00,
|
|
||||||
constant int64_t & ne01,
|
constant int64_t & ne01,
|
||||||
constant int64_t & ne02,
|
constant int64_t & ne02,
|
||||||
constant int64_t & ne03,
|
constant int64_t & ne03,
|
||||||
constant uint64_t & nb00,
|
|
||||||
constant uint64_t & nb01,
|
constant uint64_t & nb01,
|
||||||
constant uint64_t & nb02,
|
constant uint64_t & nb02,
|
||||||
constant uint64_t & nb03,
|
constant uint64_t & nb03,
|
||||||
constant int64_t & ne10,
|
|
||||||
constant int64_t & ne11,
|
constant int64_t & ne11,
|
||||||
constant int64_t & ne12,
|
constant int64_t & ne12,
|
||||||
constant int64_t & ne13,
|
constant int64_t & ne13,
|
||||||
constant uint64_t & nb10,
|
|
||||||
constant uint64_t & nb11,
|
constant uint64_t & nb11,
|
||||||
constant uint64_t & nb12,
|
constant uint64_t & nb12,
|
||||||
constant uint64_t & nb13,
|
constant uint64_t & nb13,
|
||||||
|
constant uint64_t & nb21,
|
||||||
|
constant uint64_t & nb22,
|
||||||
|
constant uint64_t & nb23,
|
||||||
constant uint64_t & nb31,
|
constant uint64_t & nb31,
|
||||||
constant int64_t & ne0,
|
|
||||||
constant int64_t & ne1,
|
constant int64_t & ne1,
|
||||||
constant int64_t & ne2,
|
constant int64_t & ne2,
|
||||||
constant int64_t & ne3,
|
|
||||||
constant float & scale,
|
constant float & scale,
|
||||||
constant float & max_bias,
|
constant float & max_bias,
|
||||||
constant float & m0,
|
constant float & m0,
|
||||||
|
|
@ -2526,10 +2513,6 @@ kernel void kernel_flash_attn_ext_vec_f16(
|
||||||
const short ne22 = ne12;
|
const short ne22 = ne12;
|
||||||
const short ne23 = ne13;
|
const short ne23 = ne13;
|
||||||
|
|
||||||
const uint nb21 = nb11;
|
|
||||||
const uint nb22 = nb12;
|
|
||||||
const uint nb23 = nb13;
|
|
||||||
|
|
||||||
// broadcast
|
// broadcast
|
||||||
const short rk2 = ne02/ne12;
|
const short rk2 = ne02/ne12;
|
||||||
const short rk3 = ne03/ne13;
|
const short rk3 = ne03/ne13;
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue