metal: implement DSv4 Lightning Indexer (#25893)
* metal: implement F16 Lightning Indexer - Implement GGML_OP_LIGHTNING_INDEXER for 128-dimensional, 64-head inputs with F32 queries and weights plus F16 keys and masks. - Add tiled and tail kernels and test KV lengths around 8- and 64-element boundaries. llama-bench (--mmap 1, -fa 1, -p 512, -n 128; d=0/10k/20k/30k): Before: - pp512: 153.73 ± 0.87 t/s - tg128: 8.91 ± 0.04 t/s - pp512 @ d10000: 73.90 ± 0.39 t/s - tg128 @ d10000: 8.66 ± 0.03 t/s - pp512 @ d20000: 45.83 ± 0.18 t/s - tg128 @ d20000: 8.26 ± 0.03 t/s - pp512 @ d30000: 33.40 ± 0.21 t/s - tg128 @ d30000: 7.94 ± 0.01 t/s After: - pp512: 155.19 ± 0.91 t/s - tg128: 8.95 ± 0.04 t/s - pp512 @ d10000: 86.95 ± 0.69 t/s - tg128 @ d10000: 9.00 ± 0.05 t/s - pp512 @ d20000: 62.01 ± 0.45 t/s - tg128 @ d20000: 8.68 ± 0.04 t/s - pp512 @ d30000: 49.18 ± 0.33 t/s - tg128 @ d30000: 8.60 ± 0.02 t/s Assisted-by: Codex * metal: stage Lightning Indexer K tiles - Stage and dequantize K in F16 threadgroup memory before simdgroup matrix loads. - Zero-fill partial tiles and guard stores so all KV segments use the same numerical path. - Support F32, F16, BF16, Q4_0, Q4_1, Q5_0, Q5_1, and Q8_0 K caches. llama-bench (--mmap 1, -fa on, -p 512, -n 128; d=0/10k/20k): - pp512: 160.38 +/- 1.01 t/s - tg128: 9.08 +/- 0.03 t/s - pp512 @ d10000: 88.37 +/- 0.46 t/s - tg128 @ d10000: 9.07 +/- 0.04 t/s - pp512 @ d20000: 62.53 +/- 0.46 t/s - tg128 @ d20000: 8.84 +/- 0.03 t/s Assisted-by: Codex * dedup Lightning Indexer constants, fix flaky test * cont : fix whitespace --------- Co-authored-by: forforever73 <690105611@qq.com> Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
This commit is contained in:
co-authored by
forforever73
Georgi Gerganov
parent
221f0f6356
commit
1464c62d88
@@ -2,6 +2,7 @@
|
||||
|
||||
#import "ggml-impl.h"
|
||||
#import "ggml-backend-impl.h"
|
||||
#import "ggml-metal-impl.h"
|
||||
|
||||
#include <Foundation/Foundation.h>
|
||||
|
||||
@@ -1308,6 +1309,36 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
|
||||
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 &&
|
||||
|
||||
Reference in New Issue
Block a user