return res;
}
+ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer(
+ ggml_metal_library_t lib,
+ const ggml_tensor * op) {
+ GGML_ASSERT(op->op == GGML_OP_LIGHTNING_INDEXER);
+
+ char name[256];
+
+ snprintf(name, 256, "kernel_lightning_indexer_%s", ggml_type_name(op->src[1]->type));
+
+ 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);
+ }
+
+ return res;
+}
+
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, ggml_op op) {
const char * name = nullptr;
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max (ggml_metal_library_t lib, const struct ggml_tensor * op);
+struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc (ggml_metal_library_t lib, enum ggml_op op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv (ggml_metal_library_t lib, const struct ggml_tensor * op);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched (ggml_metal_library_t lib, const struct ggml_tensor * op, int ssm_conv_bs);
#import "ggml-impl.h"
#import "ggml-backend-impl.h"
+#import "ggml-metal-impl.h"
#include <Foundation/Foundation.h>
return false;
}
return has_simdgroup_mm; // TODO: over-restricted for vec-kernels
+ case GGML_OP_LIGHTNING_INDEXER:
+ if (op->src[0]->ne[0] != OP_LIGHTNING_INDEXER_DK ||
+ op->src[0]->ne[1] != OP_LIGHTNING_INDEXER_NH) {
+ return false;
+ }
+ if (!has_simdgroup_mm ||
+ op->src[0]->type != GGML_TYPE_F32 ||
+ op->src[2]->type != GGML_TYPE_F32 ||
+ op->src[3]->type != GGML_TYPE_F16 ||
+ op->type != GGML_TYPE_F32 ||
+ !ggml_is_contiguous_rows(op->src[0]) ||
+ !ggml_is_contiguous_rows(op->src[1]) ||
+ !ggml_is_contiguous_rows(op->src[2]) ||
+ !ggml_is_contiguous_rows(op->src[3])) {
+ return false;
+ }
+ switch (op->src[1]->type) {
+ case GGML_TYPE_F32:
+ case GGML_TYPE_F16:
+ 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:
+ return true;
+ case GGML_TYPE_BF16:
+ return has_bfloat;
+ default:
+ return false;
+ }
case GGML_OP_DSV4_HC_COMB:
return has_simdgroup_reduction &&
op->src[0]->type == GGML_TYPE_F32 &&
#define OP_FLASH_ATTN_EXT_VEC_NQPSG 1
#define OP_FLASH_ATTN_EXT_VEC_NCPSG 32
+#define OP_LIGHTNING_INDEXER_DK 128
+#define OP_LIGHTNING_INDEXER_NH 64
+#define OP_LIGHTNING_INDEXER_NHPTG 8
+#define OP_LIGHTNING_INDEXER_NKPSG 8
+#define OP_LIGHTNING_INDEXER_NSG 8
+#define OP_LIGHTNING_INDEXER_NBPTG 8
+
#define OP_UNARY_NUM_SCALE 10
#define OP_UNARY_NUM_FILL 11
#define OP_UNARY_NUM_CLAMP 12
int64_t val;
} ggml_metal_kargs_memset;
+typedef struct {
+ int32_t n_kv;
+ int32_t n_batch;
+ int32_t mask_ne3;
+ uint64_t nb1;
+ uint64_t nb3;
+ uint64_t nbq1;
+ uint64_t nbq2;
+ uint64_t nbq3;
+ uint64_t nbk2;
+ uint64_t nbk3;
+ uint64_t nbw1;
+ uint64_t nbw3;
+ uint64_t nbm1;
+ uint64_t nbm3;
+} ggml_metal_kargs_lightning_indexer;
+
typedef struct {
int32_t n_tokens;
int32_t n_iter;
{
n_fuse = ggml_metal_op_cumsum(ctx, idx);
} break;
+ case GGML_OP_LIGHTNING_INDEXER:
+ {
+ n_fuse = ggml_metal_op_lightning_indexer(ctx, idx);
+ } break;
case GGML_OP_DSV4_HC_COMB:
case GGML_OP_DSV4_HC_PRE:
case GGML_OP_DSV4_HC_POST:
return 1;
}
+int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) {
+ ggml_tensor * op = ctx->node(idx);
+
+ ggml_metal_encoder_t enc = ctx->enc;
+
+ GGML_ASSERT(op->op == GGML_OP_LIGHTNING_INDEXER);
+
+ const ggml_tensor * q = op->src[0];
+ const ggml_tensor * k = op->src[1];
+ const ggml_tensor * w = op->src[2];
+ const ggml_tensor * m = op->src[3];
+
+ GGML_ASSERT(q->type == GGML_TYPE_F32);
+ GGML_ASSERT(k->type == GGML_TYPE_F32 ||
+ k->type == GGML_TYPE_F16 ||
+ k->type == GGML_TYPE_BF16 ||
+ k->type == GGML_TYPE_Q4_0 ||
+ k->type == GGML_TYPE_Q4_1 ||
+ k->type == GGML_TYPE_Q5_0 ||
+ k->type == GGML_TYPE_Q5_1 ||
+ k->type == GGML_TYPE_Q8_0);
+ GGML_ASSERT(w->type == GGML_TYPE_F32);
+ GGML_ASSERT(m->type == GGML_TYPE_F16);
+ GGML_ASSERT(op->type == GGML_TYPE_F32);
+
+ GGML_ASSERT(q->ne[0] == OP_LIGHTNING_INDEXER_DK);
+ GGML_ASSERT(q->ne[1] == OP_LIGHTNING_INDEXER_NH);
+
+ ggml_metal_kargs_lightning_indexer args = {
+ /*.n_kv =*/ (int32_t) k->ne[2],
+ /*.n_batch =*/ (int32_t) q->ne[2],
+ /*.mask_ne3 =*/ (int32_t) m->ne[3],
+ /*.nb1 =*/ op->nb[1],
+ /*.nb3 =*/ op->nb[3],
+ /*.nbq1 =*/ q->nb[1],
+ /*.nbq2 =*/ q->nb[2],
+ /*.nbq3 =*/ q->nb[3],
+ /*.nbk2 =*/ k->nb[2],
+ /*.nbk3 =*/ k->nb[3],
+ /*.nbw1 =*/ w->nb[1],
+ /*.nbw3 =*/ w->nb[3],
+ /*.nbm1 =*/ m->nb[1],
+ /*.nbm3 =*/ m->nb[3],
+ };
+
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(q), 1);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(k), 2);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(w), 3);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(m), 4);
+ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 5);
+
+ const int nsg = OP_LIGHTNING_INDEXER_NSG;
+ const int nkptg = OP_LIGHTNING_INDEXER_NKPSG*nsg;
+ const int nbptg = OP_LIGHTNING_INDEXER_NBPTG;
+
+ auto pipeline = ggml_metal_library_get_pipeline_lightning_indexer(ctx->lib, op);
+ ggml_metal_encoder_set_pipeline(enc, pipeline);
+ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0);
+ ggml_metal_encoder_dispatch_threadgroups(enc,
+ (k->ne[2] + nkptg - 1)/nkptg,
+ (q->ne[2] + nbptg - 1)/nbptg,
+ q->ne[3], 32, nsg, 1);
+
+ return 1;
+}
+
int ggml_metal_op_dsv4_hc(ggml_metal_op_t ctx, int idx) {
ggml_tensor * op = ctx->node(idx);
int ggml_metal_op_get_rows (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_set_rows (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_diag (ggml_metal_op_t ctx, int idx);
+int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_dsv4_hc (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_soft_max (ggml_metal_op_t ctx, int idx);
int ggml_metal_op_ssm_conv (ggml_metal_op_t ctx, int idx);
template [[host_name("kernel_count_equal_i32")]] kernel kernel_count_equal_t kernel_count_equal<int32_t>;
+template<
+ typename kd4x4_t,
+ short nl_k,
+ void (*deq_k)(device const kd4x4_t *, short, thread half4x4 &)>
+kernel void kernel_lightning_indexer(
+ constant ggml_metal_kargs_lightning_indexer & args,
+ device const char * q,
+ device const char * k,
+ device const char * w,
+ device const char * m,
+ device char * dst,
+ uint3 tgpig[[threadgroup_position_in_grid]],
+ ushort tiitg[[thread_index_in_threadgroup]],
+ ushort tiisg[[thread_index_in_simdgroup]],
+ ushort sgitg[[simdgroup_index_in_threadgroup]]) {
+ constexpr short DK = OP_LIGHTNING_INDEXER_DK;
+ constexpr short NH = OP_LIGHTNING_INDEXER_NH;
+ constexpr short NHPTG = OP_LIGHTNING_INDEXER_NHPTG;
+ constexpr short NKPSG = OP_LIGHTNING_INDEXER_NKPSG;
+ constexpr short NSG = OP_LIGHTNING_INDEXER_NSG;
+ constexpr short NBPTG = OP_LIGHTNING_INDEXER_NBPTG;
+
+ constexpr short DK4 = DK/4;
+ constexpr short DK8 = DK/8;
+ constexpr short DK16 = DK/16;
+
+ constexpr short NK = NKPSG*NSG; // keys per threadgroup
+ constexpr short NTG = 32*NSG; // threads per threadgroup
+
+ const int i_stream = tgpig.z;
+ const int i_kv_0 = tgpig.x*NK; // first key of this threadgroup
+ const int i_kv = i_kv_0 + sgitg*NKPSG; // first key of this simdgroup
+
+ threadgroup half4x4 sk4x4[NK*DK16];
+ threadgroup half * sk = (threadgroup half *) sk4x4;
+
+ for (short i = tiitg; i < NK*DK16; i += NTG) {
+ const short ik = i/DK16;
+ const short i16 = i%DK16;
+
+ half4x4 tmp;
+
+ if (i_kv_0 + ik < args.n_kv) {
+ device const kd4x4_t * kr = (device const kd4x4_t *) (k + (i_kv_0 + ik)*args.nbk2 + i_stream*args.nbk3);
+
+ deq_k(kr + i16/nl_k, i16%nl_k, tmp);
+ } else {
+ FOR_UNROLL (short j = 0; j < 4; ++j) {
+ tmp[j] = half4(0.0h);
+ }
+ }
+
+ sk4x4[i] = tmp;
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ // K tile of this simdgroup, transposed to [DK, NKPSG]
+ simdgroup_half8x8 mk[DK8];
+
+ FOR_UNROLL (short i = 0; i < DK8; ++i) {
+ simdgroup_load(mk[i], sk + sgitg*NKPSG*DK + 8*i, DK, 0, true);
+ }
+
+ threadgroup half4 sq4[NHPTG*DK4];
+ threadgroup half * sq = (threadgroup half *) sq4;
+
+ threadgroup float sw [NHPTG];
+ threadgroup float sqk[NSG*NHPTG*NKPSG];
+
+ const int i_batch_0 = tgpig.y*NBPTG;
+ const int n_batch = min((int) NBPTG, args.n_batch - i_batch_0);
+
+ for (short ib = 0; ib < n_batch; ++ib) {
+ const int i_batch = i_batch_0 + ib;
+
+ device const char * pq = q + i_batch*args.nbq2 + i_stream*args.nbq3;
+ device const char * pw = w + i_batch*args.nbw1 + i_stream*args.nbw3;
+
+ float score = 0.0f;
+
+ FOR_UNROLL (short i_head = 0; i_head < NH; i_head += NHPTG) {
+ // stage the Q tile [DK, NHPTG] and the (prescaled) head weights
+ for (short i = tiitg; i < NHPTG*DK4; i += NTG) {
+ const short ih = i/DK4;
+ const short i4 = i%DK4;
+
+ device const float4 * q4 = (device const float4 *) (pq + (i_head + ih)*args.nbq1);
+
+ sq4[ih*DK4 + i4] = half4(q4[i4]);
+ }
+
+ if (tiitg < NHPTG) {
+ sw[tiitg] = ((device const float *) pw)[i_head + tiitg];
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+
+ simdgroup_float8x8 mqk = make_filled_simdgroup_matrix<float, 8>(0.0f);
+
+ FOR_UNROLL (short i = 0; i < DK8; ++i) {
+ simdgroup_half8x8 mq;
+
+ simdgroup_load(mq, sq + 8*i, DK, 0, false);
+ simdgroup_multiply_accumulate(mqk, mq, mk[i], mqk);
+ }
+
+ threadgroup float * pqk = sqk + sgitg*NHPTG*NKPSG;
+
+ simdgroup_store(mqk, pqk, NKPSG, 0, false);
+ simdgroup_barrier(mem_flags::mem_threadgroup);
+
+ // one lane per key: ReLU, apply the head weight and accumulate over the head tile
+ if (tiisg < NKPSG) {
+ FOR_UNROLL (short ih = 0; ih < NHPTG; ++ih) {
+ score += max(pqk[ih*NKPSG + tiisg], 0.0f)*sw[ih];
+ }
+ }
+
+ threadgroup_barrier(mem_flags::mem_threadgroup);
+ }
+
+ if (tiisg < NKPSG) {
+ const int ik = i_kv + tiisg;
+ if (ik < args.n_kv) {
+ device const half * pm = (device const half *) (m + i_batch*args.nbm1 + (i_stream % args.mask_ne3)*args.nbm3);
+ device float * pd = (device float *) (dst + i_batch*args.nb1 + i_stream*args.nb3);
+
+ pd[ik] = score + (float) pm[ik];
+ }
+ }
+ }
+}
+
+typedef decltype(kernel_lightning_indexer<half4x4, 1, dequantize_f16>) kernel_lightning_indexer_t;
+
+template [[host_name("kernel_lightning_indexer_f32")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<float4x4, 1, dequantize_f32>;
+template [[host_name("kernel_lightning_indexer_f16")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<half4x4, 1, dequantize_f16>;
+
+#if defined(GGML_METAL_HAS_BF16)
+template [[host_name("kernel_lightning_indexer_bf16")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<bfloat4x4, 1, dequantize_bf16>;
+#endif
+
+template [[host_name("kernel_lightning_indexer_q4_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q4_0, 2, dequantize_q4_0>;
+template [[host_name("kernel_lightning_indexer_q4_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q4_1, 2, dequantize_q4_1>;
+template [[host_name("kernel_lightning_indexer_q5_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q5_0, 2, dequantize_q5_0>;
+template [[host_name("kernel_lightning_indexer_q5_1")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q5_1, 2, dequantize_q5_1>;
+template [[host_name("kernel_lightning_indexer_q8_0")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer<block_q8_0, 2, dequantize_q8_0>;
+
kernel void kernel_dsv4_hc_comb_f32(
constant ggml_metal_kargs_dsv4_hc_comb & args,
device const char * mixes,
*(device float *) (dst + i0*args.nb_d0 + idst*args.nb_d1 + it*args.nb_d2) = result[idst];
}
}
-
}
}
+ for (int kv : { 1, 7, 8, 63, 64, 65 }) {
+ for (ggml_type type_K : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_BF16, GGML_TYPE_Q8_0, GGML_TYPE_Q5_1, GGML_TYPE_Q5_0, GGML_TYPE_Q4_1, GGML_TYPE_Q4_0}) {
+ test_cases.emplace_back(new test_lightning_indexer(128, 64, kv, 32, 4, 1, type_K));
+ }
+ }
+
return test_cases;
}
#ifdef _MSC_VER
overview / pkg / ggml / sources / llama.cpp / commitdiff