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>

ggml/src/ggml-metal/ggml-metal-ops.cpp

@@ -320,6 +320,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) {
             {
                 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:
@@ -1307,6 +1311,72 @@ int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) {
     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);