diff --git a/crates/skippy-ffi/src/lib.rs b/crates/skippy-ffi/src/lib.rs index d3a6d9317..b15a9ea21 100644 --- a/crates/skippy-ffi/src/lib.rs +++ b/crates/skippy-ffi/src/lib.rs @@ -1,6 +1,6 @@ pub const ABI_VERSION_MAJOR: u32 = 0; pub const ABI_VERSION_MINOR: u32 = 1; -pub const ABI_VERSION_PATCH: u32 = 31; +pub const ABI_VERSION_PATCH: u32 = 32; pub const FEATURE_BACKEND_DEVICES: u64 = 1 << 23; pub const FEATURE_RUNTIME_EVENTS: u64 = 1 << 24; pub const FEATURE_NATIVE_MTP_N1: u64 = 1 << 25; @@ -219,10 +219,18 @@ pub struct RuntimeConfig { pub flash_attn_type: i32, pub load_mode: LoadMode, pub disable_repack: bool, + pub use_mmap_prefetch: bool, + pub use_mmap_buffer: bool, pub filter_tensors_on_load: bool, pub include_embeddings: bool, pub include_output: bool, pub selected_backend_device: *const c_char, + pub glm_dsa_policy_profile: i32, + pub glm_dsa_policy_flags: u32, + pub glm_dsa_short_prefill_max_tokens: i32, + pub glm_dsa_direct_sparse_decode_max_top_k: i32, + pub glm_dsa_dense_sparse_mask_max_bytes: u64, + pub glm_dsa_compact_flash_min_kv: i32, } #[repr(C)] diff --git a/crates/skippy-runtime/src/config.rs b/crates/skippy-runtime/src/config.rs index d0bca3e5c..3d50cc8d8 100644 --- a/crates/skippy-runtime/src/config.rs +++ b/crates/skippy-runtime/src/config.rs @@ -132,10 +132,18 @@ impl RuntimeConfig { flash_attn_type: self.flash_attn_type as i32, load_mode: self.load_mode, disable_repack: false, + use_mmap_prefetch: false, + use_mmap_buffer: false, filter_tensors_on_load: self.filter_tensors_on_load, include_embeddings: self.include_embeddings, include_output: self.include_output, selected_backend_device: selected_backend_device_ptr, + glm_dsa_policy_profile: 0, + glm_dsa_policy_flags: 0, + glm_dsa_short_prefill_max_tokens: 0, + glm_dsa_direct_sparse_decode_max_top_k: 0, + glm_dsa_dense_sparse_mask_max_bytes: 0, + glm_dsa_compact_flash_min_kv: 0, }, _selected_backend_device: selected_backend_device, }) @@ -311,6 +319,22 @@ mod tests { Ok(()) } + #[test] + fn runtime_config_raw_defaults_glm_dsa_controls_to_disabled() -> anyhow::Result<()> { + let raw = RuntimeConfig::default().as_raw()?.raw; + + assert!(!raw.use_mmap_prefetch); + assert!(!raw.use_mmap_buffer); + assert_eq!(raw.glm_dsa_policy_profile, 0); + assert_eq!(raw.glm_dsa_policy_flags, 0); + assert_eq!(raw.glm_dsa_short_prefill_max_tokens, 0); + assert_eq!(raw.glm_dsa_direct_sparse_decode_max_top_k, 0); + assert_eq!(raw.glm_dsa_dense_sparse_mask_max_bytes, 0); + assert_eq!(raw.glm_dsa_compact_flash_min_kv, 0); + assert!(raw.selected_backend_device.is_null()); + Ok(()) + } + #[test] fn runtime_config_raw_uses_smaller_batch_for_unified_kv_defaults() -> anyhow::Result<()> { let config = RuntimeConfig { diff --git a/third_party/llama.cpp/patches/0022-ggml-add-GLM-DSA-sparse-execution-primitives.patch b/third_party/llama.cpp/patches/0022-ggml-add-GLM-DSA-sparse-execution-primitives.patch new file mode 100644 index 000000000..13c30015b --- /dev/null +++ b/third_party/llama.cpp/patches/0022-ggml-add-GLM-DSA-sparse-execution-primitives.patch @@ -0,0 +1,31484 @@ +From dfe419c41613f4cd640998b391899a145abdfc10 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Tue, 14 Jul 2026 09:15:57 +1000 +Subject: [PATCH 19/24] ggml: add GLM-DSA sparse execution primitives + +--- + ggml/include/ggml-rpc.h | 2 +- + ggml/include/ggml.h | 57 + + ggml/src/ggml-backend.cpp | 84 + + ggml/src/ggml-cpu/ggml-cpu.c | 91 +- + ggml/src/ggml-cpu/ops.cpp | 500 +- + ggml/src/ggml-cpu/ops.h | 6 + + ggml/src/ggml-metal/ggml-metal-common.cpp | 7 +- + ggml/src/ggml-metal/ggml-metal-context.m | 232 +- + ggml/src/ggml-metal/ggml-metal-device.cpp | 3334 ++++- + ggml/src/ggml-metal/ggml-metal-device.h | 83 +- + ggml/src/ggml-metal/ggml-metal-device.m | 152 +- + ggml/src/ggml-metal/ggml-metal-impl.h | 331 +- + ggml/src/ggml-metal/ggml-metal-ops.cpp | 13555 +++++++++++++++----- + ggml/src/ggml-metal/ggml-metal-ops.h | 13 + + ggml/src/ggml-metal/ggml-metal.cpp | 5 + + ggml/src/ggml-metal/ggml-metal.metal | 9996 ++++++++++++--- + ggml/src/ggml.c | 235 +- + src/llama-hparams.h | 4 + + 18 files changed, 23411 insertions(+), 5276 deletions(-) + +diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h +index 16ca3394..82d27410 100644 +--- a/ggml/include/ggml-rpc.h ++++ b/ggml/include/ggml-rpc.h +@@ -11,7 +11,7 @@ extern "C" { + #define RPC_PROTO_PATCH_VERSION 3 + + #ifdef __cplusplus +-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); ++static_assert(GGML_OP_COUNT == 107, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); + #endif + + #define GGML_RPC_MAX_SERVERS 16 +diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h +index 35f0c44e..7ef0ec88 100644 +--- a/ggml/include/ggml.h ++++ b/ggml/include/ggml.h +@@ -574,6 +574,11 @@ extern "C" { + GGML_OP_DSV4_HC_COMB, + GGML_OP_DSV4_HC_PRE, + GGML_OP_DSV4_HC_POST, ++ GGML_OP_DSA_SPARSE_MASK, ++ GGML_OP_DSA_SPARSE_ATTN, ++ GGML_OP_DSA_TOP1_ATTN, ++ GGML_OP_MOE_ROUTE_WEIGHTS, ++ GGML_OP_MOE_WEIGHTED_SUM, + + GGML_OP_UNARY, + +@@ -590,6 +595,9 @@ extern "C" { + + GGML_OP_GLU, + ++ // selected expert matmuls followed by a route-weighted reduction ++ GGML_OP_MOE_MUL_MAT_ID, ++ + GGML_OP_COUNT, + }; + +@@ -1673,6 +1681,11 @@ extern "C" { + struct ggml_tensor * a, // data + struct ggml_tensor * b); // row indices + ++ GGML_API struct ggml_tensor * ggml_get_rows_typed( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, // data ++ struct ggml_tensor * b); // row indices ++ + GGML_API struct ggml_tensor * ggml_get_rows_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // gradients of ggml_get_rows result +@@ -2639,6 +2652,50 @@ extern "C" { + struct ggml_tensor * residual, + struct ggml_tensor * post, + struct ggml_tensor * comb); ++ GGML_API struct ggml_tensor * ggml_dsa_sparse_mask( ++ struct ggml_context * ctx, ++ struct ggml_tensor * kq_mask, ++ struct ggml_tensor * top_k); ++ ++ GGML_API struct ggml_tensor * ggml_dsa_sparse_attn( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * kq_mask_rows, ++ struct ggml_tensor * top_k, ++ float scale); ++ ++ GGML_API struct ggml_tensor * ggml_dsa_top1_attn( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * v, ++ struct ggml_tensor * top_k); ++ ++ GGML_API struct ggml_tensor * ggml_moe_weighted_sum( ++ struct ggml_context * ctx, ++ struct ggml_tensor * experts, ++ struct ggml_tensor * weights); ++ ++ // experts: [n_ff, n_embd, n_expert] ++ // input: [n_ff, n_expert_used, n_tokens] ++ // ids: [n_expert_used, n_tokens] ++ // weights: [1, n_expert_used, n_tokens] ++ // result: [n_embd, n_tokens] ++ GGML_API struct ggml_tensor * ggml_moe_mul_mat_id( ++ struct ggml_context * ctx, ++ struct ggml_tensor * experts, ++ struct ggml_tensor * input, ++ struct ggml_tensor * ids, ++ struct ggml_tensor * weights); ++ ++ GGML_API struct ggml_tensor * ggml_moe_route_weights( ++ struct ggml_context * ctx, ++ struct ggml_tensor * probs, ++ struct ggml_tensor * ids, ++ bool norm, ++ float clamp_min, ++ float scale); + + // custom operators + +diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp +index 87615921..c0f90e89 100644 +--- a/ggml/src/ggml-backend.cpp ++++ b/ggml/src/ggml-backend.cpp +@@ -982,6 +982,89 @@ static void ggml_backend_sched_print_assignments(ggml_backend_sched_t sched, str + } + } + ++static bool ggml_backend_sched_glm_route_diag_enabled() { ++ const char * log = getenv("GGML_METAL_MOE_DISPATCH_LOG"); ++ if (log && atoi(log) != 0) { ++ return true; ++ } ++ log = getenv("SKIPPY_GLM_DSA_LOG_METAL_DISPATCH"); ++ return log && atoi(log) != 0; ++} ++ ++static int ggml_backend_sched_split_id_for_node(ggml_backend_sched_t sched, int node_idx) { ++ for (int i = 0; i < sched->n_splits; ++i) { ++ const ggml_backend_sched_split * split = &sched->splits[i]; ++ if (node_idx >= split->i_start && node_idx < split->i_end) { ++ return i; ++ } ++ } ++ return -1; ++} ++ ++static void ggml_backend_sched_log_glm_route_splits(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { ++ if (!ggml_backend_sched_glm_route_diag_enabled()) { ++ return; ++ } ++ ++ for (int i = 0; i < graph->n_nodes; ++i) { ++ ggml_tensor * node = graph->nodes[i]; ++ if (node->name[0] == '\0' || strstr(node->name, "ffn_moe_probs") == NULL || ++ node->op != GGML_OP_UNARY || ggml_get_unary_op(node) != GGML_UNARY_OP_SIGMOID) { ++ continue; ++ } ++ ++ char ops[1536] = {}; ++ size_t offset = 0; ++ int first_split = -1; ++ int last_split = -1; ++ int split_count = 0; ++ ++ for (int j = i; j < graph->n_nodes && j < i + 14 && offset < sizeof(ops); ++j) { ++ ggml_tensor * t = graph->nodes[j]; ++ const int split_id = ggml_backend_sched_split_id_for_node(sched, j); ++ if (j == i) { ++ first_split = split_id; ++ last_split = split_id; ++ split_count = split_id >= 0 ? 1 : 0; ++ } else if (split_id != last_split) { ++ last_split = split_id; ++ if (split_id >= 0) { ++ ++split_count; ++ } ++ } ++ ++ const int backend_id = tensor_backend_id(t); ++ const char * backend_name = ++ backend_id >= 0 && backend_id < sched->n_backends ? ++ ggml_backend_name(sched->backends[backend_id]) : ++ ""; ++ const int written = snprintf( ++ ops + offset, ++ sizeof(ops) - offset, ++ "%s%d:%s/%s@split%d/%s", ++ j == i ? "" : ",", ++ j, ++ ggml_op_name(t->op), ++ t->name[0] == '\0' ? "" : t->name, ++ split_id, ++ backend_name); ++ if (written < 0) { ++ break; ++ } ++ offset += (size_t) written; ++ } ++ ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=topk_moe_route_sched tensor=%s graph_uid=%llu split_start=%d split_end=%d split_count=%d candidate=%s grid_x=1 grid_y=1 grid_z=1 threads_x=1\n", ++ node->name, ++ (unsigned long long) graph->uid, ++ first_split, ++ last_split, ++ split_count, ++ ops); ++ } ++} ++ + static bool ggml_backend_sched_buffer_supported(ggml_backend_sched_t sched, struct ggml_tensor * t, int backend_id) { + ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; + ggml_backend_buffer_type_t buft = NULL; +@@ -1378,6 +1461,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra + if (sched->debug) { + ggml_backend_sched_print_assignments(sched, graph); + } ++ ggml_backend_sched_log_glm_route_splits(sched, graph); + + // swap node_backend_ids and leaf _backend_ids with prevs + { +diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c +index 491316f7..eed46cfe 100644 +--- a/ggml/src/ggml-cpu/ggml-cpu.c ++++ b/ggml/src/ggml-cpu/ggml-cpu.c +@@ -1569,6 +1569,9 @@ static void ggml_compute_forward_mul_mat_id( + + if (src1->type != vec_dot_type) { + incr_ptr_aligned(&wdata_cur, ggml_row_size(vec_dot_type, ggml_nelements(src1)), sizeof(int64_t)); ++ if (src1->type == GGML_TYPE_F16) { ++ incr_ptr_aligned(&wdata_cur, nth * ne10 * sizeof(float), sizeof(int64_t)); ++ } + } + + int64_t * matrix_row_counts = // [n_as] +@@ -1583,15 +1586,22 @@ static void ggml_compute_forward_mul_mat_id( + GGML_ASSERT(params->wsize >= (size_t)((char *) wdata_cur - (char *) params->wdata)); + + if (src1->type != vec_dot_type) { +- char * wdata = params->wdata; ++ void * wdata_conv_cur = params->wdata; + + const size_t nbw0 = ggml_type_size(vec_dot_type); + const size_t nbw1 = ggml_row_size(vec_dot_type, ne10); + const size_t nbw2 = nbw1*ne11; + const size_t nbw3 = nbw2*ne12; + +- assert(params->wsize >= ne13*nbw3); +- GGML_ASSERT(src1->type == GGML_TYPE_F32); ++ char * wdata = incr_ptr_aligned(&wdata_conv_cur, ne13*nbw3, sizeof(int64_t)); ++ float * src1_f32_scratch = NULL; ++ if (src1->type == GGML_TYPE_F16) { ++ src1_f32_scratch = incr_ptr_aligned(&wdata_conv_cur, nth * ne10 * sizeof(float), sizeof(int64_t)); ++ } else { ++ GGML_ASSERT(src1->type == GGML_TYPE_F32); ++ } ++ ++ assert(params->wsize >= (size_t)((char *) wdata_conv_cur - (char *) params->wdata)); + + #if 0 + for (int64_t i13 = 0; i13 < ne13; ++i13) { +@@ -1610,9 +1620,17 @@ static void ggml_compute_forward_mul_mat_id( + size_t bs = ggml_blck_size(vec_dot_type); + int64_t ne10_block_start = (ith * ne10/bs) / nth; + int64_t ne10_block_end = ((ith + 1) * ne10/bs) / nth; +- from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10), +- (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), +- (ne10_block_end - ne10_block_start) * bs); ++ const int64_t n = (ne10_block_end - ne10_block_start) * bs; ++ const char * src1_row = (const char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10; ++ const float * src1_f32 = (const float *) src1_row; ++ if (src1_f32_scratch != NULL) { ++ float * src1_f32_mut = src1_f32_scratch + ith * ne10; ++ ggml_fp16_to_fp32_row((const ggml_fp16_t *) src1_row, src1_f32_mut, n); ++ src1_f32 = src1_f32_mut; ++ } ++ from_float(src1_f32, ++ (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), ++ n); + } + } + } +@@ -2076,6 +2094,30 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm + { + ggml_compute_forward_dsv4_hc_post(params, tensor); + } break; ++ case GGML_OP_DSA_SPARSE_MASK: ++ { ++ ggml_compute_forward_dsa_sparse_mask(params, tensor); ++ } break; ++ case GGML_OP_DSA_SPARSE_ATTN: ++ { ++ ggml_compute_forward_dsa_sparse_attn(params, tensor); ++ } break; ++ case GGML_OP_DSA_TOP1_ATTN: ++ { ++ ggml_compute_forward_dsa_top1_attn(params, tensor); ++ } break; ++ case GGML_OP_MOE_ROUTE_WEIGHTS: ++ { ++ ggml_compute_forward_moe_route_weights(params, tensor); ++ } break; ++ case GGML_OP_MOE_WEIGHTED_SUM: ++ { ++ ggml_compute_forward_moe_weighted_sum(params, tensor); ++ } break; ++ case GGML_OP_MOE_MUL_MAT_ID: ++ { ++ ggml_compute_forward_moe_mul_mat_id(params, tensor); ++ } break; + case GGML_OP_MAP_CUSTOM1: + { + ggml_compute_forward_map_custom1(params, tensor); +@@ -2400,6 +2442,12 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { + case GGML_OP_SSM_CONV: + case GGML_OP_SSM_SCAN: + case GGML_OP_LIGHTNING_INDEXER: ++ case GGML_OP_DSA_SPARSE_MASK: ++ case GGML_OP_DSA_SPARSE_ATTN: ++ case GGML_OP_DSA_TOP1_ATTN: ++ case GGML_OP_MOE_ROUTE_WEIGHTS: ++ case GGML_OP_MOE_WEIGHTED_SUM: ++ case GGML_OP_MOE_MUL_MAT_ID: + { + n_tasks = n_threads; + } break; +@@ -2864,6 +2912,9 @@ struct ggml_cplan ggml_graph_plan( + // src1 + if (src1->type != vec_dot_type) { + cur += ggml_row_size(vec_dot_type, ggml_nelements(src1)) + sizeof(int64_t); ++ if (src1->type == GGML_TYPE_F16) { ++ cur += n_threads * src1->ne[0] * sizeof(float) + sizeof(int64_t); ++ } + } + // matrix_row_counts + cur += n_as * sizeof(int64_t) + sizeof(int64_t); +@@ -2889,7 +2940,7 @@ struct ggml_cplan ggml_graph_plan( + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + { +- cur = ggml_type_size(GGML_TYPE_F32) * node->ne[0] * n_tasks; ++ cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; + } break; + case GGML_OP_CONV_TRANSPOSE_1D: + { +@@ -2988,16 +3039,30 @@ struct ggml_cplan ggml_graph_plan( + const int64_t per_thread = S_v + (K > 1 ? S_v * S_v : 0); + cur = per_thread * sizeof(float) * n_tasks; + } break; ++ case GGML_OP_LIGHTNING_INDEXER: ++ { ++ const int64_t n_embd = node->src[1]->ne[0]; ++ cur = sizeof(float) * (n_embd + CACHE_LINE_SIZE_F32) * n_tasks; ++ } break; ++ case GGML_OP_DSA_SPARSE_ATTN: ++ { ++ const int64_t dk = node->src[0]->ne[0]; ++ const int64_t dv = node->src[2]->ne[0]; ++ const int64_t n_top_k = node->src[4]->ne[0]; ++ cur = sizeof(float) * (dk + dv + n_top_k + CACHE_LINE_SIZE_F32) * n_tasks; ++ } break; ++ case GGML_OP_DSA_TOP1_ATTN: ++ { ++ cur = 0; ++ } break; ++ case GGML_OP_MOE_MUL_MAT_ID: ++ { ++ cur = sizeof(float)*node->src[0]->ne[0]*n_tasks; ++ } break; + case GGML_OP_COUNT: + { + GGML_ABORT("fatal error"); + } +- case GGML_OP_LIGHTNING_INDEXER: +- { +- // temp buffer for dequantizing lightning indexer keys +- const int64_t ne10 = node->src[1]->ne[0]; +- cur += sizeof(float)*ne10*n_tasks; +- } break; + default: + break; + } +diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp +index 42ec809c..61237ed0 100644 +--- a/ggml/src/ggml-cpu/ops.cpp ++++ b/ggml/src/ggml-cpu/ops.cpp +@@ -4926,9 +4926,16 @@ static void ggml_compute_forward_get_rows_f16( + + GGML_ASSERT(i01 >= 0 && i01 < ne01); + +- ggml_cpu_fp16_to_fp32( +- (const ggml_fp16_t*) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), +- (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); ++ if (dst->type == GGML_TYPE_F16) { ++ memcpy( ++ (char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3, ++ (char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03, ++ nc*sizeof(ggml_fp16_t)); ++ } else { ++ ggml_cpu_fp16_to_fp32( ++ (const ggml_fp16_t*) ((char *) src0->data + i01*nb01 + i11*nb02 + i12*nb03), ++ (float *) ((char *) dst->data + i10*nb1 + i11*nb2 + i12*nb3), nc); ++ } + } + } + +@@ -6092,6 +6099,7 @@ static void ggml_compute_forward_rope_flt( + dst_data[1] = src[1]; + } + } ++ + } //attn-heads + } + } +@@ -11229,6 +11237,492 @@ void ggml_compute_forward_dsv4_hc_post( + } + } + ++// ggml_compute_forward_dsa_sparse_mask ++ ++void ggml_compute_forward_dsa_sparse_mask( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; // kq_mask rows: [1, n_kv, n_batch, n_stream] ++ const ggml_tensor * src1 = dst->src[1]; // top_k ++ ++ GGML_ASSERT(dst->type == src0->type); ++ GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16); ++ GGML_ASSERT(src1->type == GGML_TYPE_I32); ++ GGML_ASSERT(dst->ne[0] == src0->ne[0]); ++ GGML_ASSERT(dst->ne[1] == src0->ne[1]); ++ GGML_ASSERT(dst->ne[2] == src0->ne[2]); ++ GGML_ASSERT(dst->ne[3] == src0->ne[3]); ++ GGML_ASSERT(src0->ne[0] == 1); ++ GGML_ASSERT(src1->ne[1] == src0->ne[2]); ++ GGML_ASSERT(src0->ne[3] % src1->ne[2] == 0); ++ GGML_ASSERT(src1->ne[3] == 1); ++ ++ const int64_t n_kv = src0->ne[1]; ++ const int64_t n_batch = src0->ne[2]; ++ const int64_t n_stream = src0->ne[3]; ++ const int64_t n_top_k = src1->ne[0]; ++ ++ const int ith = params->ith; ++ const int nth = params->nth; ++ ++ const int64_t total = n_batch * n_stream; ++ const int64_t dr = (total + nth - 1) / nth; ++ const int64_t i0 = dr * ith; ++ const int64_t i1 = MIN(i0 + dr, total); ++ ++ for (int64_t i = i0; i < i1; ++i) { ++ const int64_t i_batch = i % n_batch; ++ const int64_t i_stream = i / n_batch; ++ ++ for (int64_t i_kv = 0; i_kv < n_kv; ++i_kv) { ++ char * dst_data = (char *) dst->data + i_kv * dst->nb[1] + i_batch * dst->nb[2] + i_stream * dst->nb[3]; ++ ++ if (dst->type == GGML_TYPE_F32) { ++ *(float *) dst_data = -INFINITY; ++ } else { ++ *(ggml_fp16_t *) dst_data = GGML_CPU_FP32_TO_FP16(-INFINITY); ++ } ++ } ++ ++ const int64_t i12 = i_stream % src1->ne[2]; ++ for (int64_t i_top = 0; i_top < n_top_k; ++i_top) { ++ const char * top_k_data = ++ (const char *) src1->data + i_top * src1->nb[0] + i_batch * src1->nb[1] + i12 * src1->nb[2]; ++ const int32_t i_kv = *(const int32_t *) top_k_data; ++ GGML_ASSERT(i_kv >= 0 && i_kv < n_kv); ++ ++ const char * src = (const char *) src0->data + i_kv * src0->nb[1] + i_batch * src0->nb[2] + i_stream * src0->nb[3]; ++ char * dst_data = (char *) dst->data + i_kv * dst->nb[1] + i_batch * dst->nb[2] + i_stream * dst->nb[3]; ++ ++ if (dst->type == GGML_TYPE_F32) { ++ *(float *) dst_data = *(const float *) src; ++ } else { ++ *(ggml_fp16_t *) dst_data = *(const ggml_fp16_t *) src; ++ } ++ } ++ } ++} ++ ++// ggml_compute_forward_dsa_sparse_attn ++ ++static float ggml_dsa_sparse_attn_mask_value(const ggml_tensor * mask, int64_t i_kv, int64_t i_batch, int64_t i_stream) { ++ const char * data = (const char *) mask->data + i_kv * mask->nb[1] + i_batch * mask->nb[2] + i_stream * mask->nb[3]; ++ return mask->type == GGML_TYPE_F32 ? *(const float *) data : GGML_CPU_FP16_TO_FP32(*(const ggml_fp16_t *) data); ++} ++ ++void ggml_compute_forward_dsa_sparse_attn( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; // q: [Dk, n_batch, n_head, n_stream] ++ const ggml_tensor * src1 = dst->src[1]; // k: [Dk, n_kv, n_kv_head, n_stream] ++ const ggml_tensor * src2 = dst->src[2]; // v: [Dv, n_kv, n_v_head, n_stream] ++ const ggml_tensor * src3 = dst->src[3]; // kq mask rows: [1, n_kv, n_batch, n_stream] ++ const ggml_tensor * src4 = dst->src[4]; // top_k: [n_top_k, n_batch, n_top_stream, 1] ++ ++ const float scale = ggml_get_op_params_f32(dst, 0); ++ ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ GGML_ASSERT(src0->type == GGML_TYPE_F32); ++ GGML_ASSERT(src3->type == GGML_TYPE_F32 || src3->type == GGML_TYPE_F16); ++ GGML_ASSERT(src4->type == GGML_TYPE_I32); ++ ++ GGML_ASSERT(dst->ne[0] == src2->ne[0]); ++ GGML_ASSERT(dst->ne[1] == src0->ne[1]); ++ GGML_ASSERT(dst->ne[2] == src0->ne[2]); ++ GGML_ASSERT(dst->ne[3] == src0->ne[3]); ++ ++ const int64_t dk = src0->ne[0]; ++ const int64_t dv = src2->ne[0]; ++ const int64_t n_batch = src0->ne[1]; ++ const int64_t n_head = src0->ne[2]; ++ const int64_t n_stream = src0->ne[3]; ++ const int64_t n_kv = src1->ne[1]; ++ const int64_t n_kv_head = src1->ne[2]; ++ const int64_t n_v_head = src2->ne[2]; ++ const int64_t n_top_k = src4->ne[0]; ++ const int64_t n_top_stream = src4->ne[2]; ++ ++ GGML_ASSERT(src1->ne[0] == dk); ++ GGML_ASSERT(src2->ne[1] == n_kv); ++ GGML_ASSERT(src3->ne[0] == 1); ++ GGML_ASSERT(src3->ne[1] == n_kv); ++ GGML_ASSERT(src3->ne[2] == n_batch); ++ GGML_ASSERT(src3->ne[3] == n_stream); ++ GGML_ASSERT(src4->ne[1] == n_batch); ++ GGML_ASSERT(n_stream % n_top_stream == 0); ++ GGML_ASSERT(src4->ne[3] == 1); ++ GGML_ASSERT(n_head % n_kv_head == 0); ++ GGML_ASSERT(n_head % n_v_head == 0); ++ ++ ggml_to_float_t const k_to_float = ggml_get_type_traits(src1->type)->to_float; ++ ggml_to_float_t const v_to_float = ggml_get_type_traits(src2->type)->to_float; ++ GGML_ASSERT((src1->type == GGML_TYPE_F32 || k_to_float) && "dsa sparse attn: unsupported K type"); ++ GGML_ASSERT((src2->type == GGML_TYPE_F32 || v_to_float) && "dsa sparse attn: unsupported V type"); ++ ++ const int ith = params->ith; ++ const int nth = params->nth; ++ ++ float * k_row_f32 = (float *) params->wdata + ith * (dk + dv + n_top_k + CACHE_LINE_SIZE_F32); ++ float * v_row_f32 = k_row_f32 + dk; ++ float * scores = v_row_f32 + dv; ++ ++ const int64_t total = n_stream * n_head * n_batch; ++ const int64_t dr = (total + nth - 1) / nth; ++ const int64_t i0 = dr * ith; ++ const int64_t i1 = MIN(i0 + dr, total); ++ ++ for (int64_t i = i0; i < i1; ++i) { ++ const int64_t i_batch = i % n_batch; ++ const int64_t i_head = (i / n_batch) % n_head; ++ const int64_t i_stream = i / (n_batch * n_head); ++ ++ const int64_t i_kv_head = i_head / (n_head / n_kv_head); ++ const int64_t i_v_head = i_head / (n_head / n_v_head); ++ const int64_t i_top_stream = i_stream % n_top_stream; ++ ++ const float * q_row = (const float *) ((const char *) src0->data + ++ i_batch * src0->nb[1] + i_head * src0->nb[2] + i_stream * src0->nb[3]); ++ ++ float max_score = -INFINITY; ++ int64_t active_top_end = 0; ++ for (int64_t i_top = 0; i_top < n_top_k; ++i_top) { ++ const char * top_k_data = (const char *) src4->data + ++ i_top * src4->nb[0] + i_batch * src4->nb[1] + i_top_stream * src4->nb[2]; ++ const int32_t i_kv = *(const int32_t *) top_k_data; ++ GGML_ASSERT(i_kv >= 0 && i_kv < n_kv); ++ ++ const float mask = ggml_dsa_sparse_attn_mask_value(src3, i_kv, i_batch, i_stream); ++ if (!isfinite(mask)) { ++ scores[i_top] = -INFINITY; ++ continue; ++ } ++ active_top_end = i_top + 1; ++ ++ const char * k_row = (const char *) src1->data + ++ i_kv * src1->nb[1] + i_kv_head * src1->nb[2] + i_stream * src1->nb[3]; ++ const float * k_row_data; ++ if (k_to_float) { ++ k_to_float(k_row, k_row_f32, dk); ++ k_row_data = k_row_f32; ++ } else { ++ k_row_data = (const float *) k_row; ++ } ++ ++ float qk = 0.0f; ++ ggml_vec_dot_f32(dk, &qk, 0, q_row, 0, k_row_data, 0, 1); ++ scores[i_top] = qk * scale + mask; ++ max_score = MAX(max_score, scores[i_top]); ++ } ++ ++ if (!isfinite(max_score)) { ++ float * dst_row = (float *) ((char *) dst->data + ++ i_batch * dst->nb[1] + i_head * dst->nb[2] + i_stream * dst->nb[3]); ++ memset(dst_row, 0, dv * sizeof(float)); ++ continue; ++ } ++ ++ float sum = 0.0f; ++ for (int64_t i_top = 0; i_top < active_top_end; ++i_top) { ++ scores[i_top] = expf(scores[i_top] - max_score); ++ sum += scores[i_top]; ++ } ++ ++ float * dst_row = (float *) ((char *) dst->data + ++ i_batch * dst->nb[1] + i_head * dst->nb[2] + i_stream * dst->nb[3]); ++ memset(dst_row, 0, dv * sizeof(float)); ++ ++ if (sum == 0.0f || !isfinite(sum)) { ++ continue; ++ } ++ ++ for (int64_t i_top = 0; i_top < active_top_end; ++i_top) { ++ const char * top_k_data = (const char *) src4->data + ++ i_top * src4->nb[0] + i_batch * src4->nb[1] + i_top_stream * src4->nb[2]; ++ const int32_t i_kv = *(const int32_t *) top_k_data; ++ ++ const float p = scores[i_top] / sum; ++ if (p == 0.0f) { ++ continue; ++ } ++ ++ const char * v_row = (const char *) src2->data + ++ i_kv * src2->nb[1] + i_v_head * src2->nb[2] + i_stream * src2->nb[3]; ++ const float * v_row_data; ++ if (v_to_float) { ++ v_to_float(v_row, v_row_f32, dv); ++ v_row_data = v_row_f32; ++ } else { ++ v_row_data = (const float *) v_row; ++ } ++ ++ ggml_vec_mad_f32(dv, dst_row, v_row_data, p); ++ } ++ } ++} ++ ++// ggml_compute_forward_dsa_top1_attn ++ ++void ggml_compute_forward_dsa_top1_attn( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * src0 = dst->src[0]; // q: [Dk, n_batch, n_head, n_stream] ++ const ggml_tensor * src1 = dst->src[1]; // v: [Dv, n_kv, n_v_head, n_stream] ++ const ggml_tensor * src2 = dst->src[2]; // top_k: [1, n_batch, n_top_stream, 1] ++ ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ GGML_ASSERT(src0->type == GGML_TYPE_F32); ++ GGML_ASSERT(src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16); ++ GGML_ASSERT(src2->type == GGML_TYPE_I32); ++ ++ GGML_ASSERT(dst->ne[0] == src1->ne[0]); ++ GGML_ASSERT(dst->ne[1] == src0->ne[1]); ++ GGML_ASSERT(dst->ne[2] == src0->ne[2]); ++ GGML_ASSERT(dst->ne[3] == src0->ne[3]); ++ ++ const int64_t dv = src1->ne[0]; ++ const int64_t n_batch = src0->ne[1]; ++ const int64_t n_head = src0->ne[2]; ++ const int64_t n_stream = src0->ne[3]; ++ const int64_t n_kv = src1->ne[1]; ++ const int64_t n_v_head = src1->ne[2]; ++ const int64_t n_top_stream = src2->ne[2]; ++ ++ GGML_ASSERT(src0->ne[3] == src1->ne[3]); ++ GGML_ASSERT(src2->ne[0] == 1); ++ GGML_ASSERT(src2->ne[1] == n_batch); ++ GGML_ASSERT(n_stream % n_top_stream == 0); ++ GGML_ASSERT(src2->ne[3] == 1); ++ GGML_ASSERT(n_head % n_v_head == 0); ++ ++ ggml_to_float_t const v_to_float = ggml_get_type_traits(src1->type)->to_float; ++ ++ const int ith = params->ith; ++ const int nth = params->nth; ++ ++ const int64_t total = n_stream * n_head * n_batch; ++ const int64_t dr = (total + nth - 1) / nth; ++ const int64_t i0 = dr * ith; ++ const int64_t i1 = MIN(i0 + dr, total); ++ ++ for (int64_t i = i0; i < i1; ++i) { ++ const int64_t i_batch = i % n_batch; ++ const int64_t i_head = (i / n_batch) % n_head; ++ const int64_t i_stream = i / (n_batch * n_head); ++ const int64_t i_v_head = i_head / (n_head / n_v_head); ++ const int64_t i_top_stream = i_stream % n_top_stream; ++ ++ const char * top_k_data = (const char *) src2->data + ++ i_batch * src2->nb[1] + i_top_stream * src2->nb[2]; ++ const int32_t i_kv = *(const int32_t *) top_k_data; ++ GGML_ASSERT(i_kv >= 0 && i_kv < n_kv); ++ ++ const char * v_row = (const char *) src1->data + ++ i_kv * src1->nb[1] + i_v_head * src1->nb[2] + i_stream * src1->nb[3]; ++ float * dst_row = (float *) ((char *) dst->data + ++ i_batch * dst->nb[1] + i_head * dst->nb[2] + i_stream * dst->nb[3]); ++ ++ if (src1->type == GGML_TYPE_F32) { ++ memcpy(dst_row, v_row, dv * sizeof(float)); ++ } else { ++ GGML_ASSERT(v_to_float); ++ v_to_float(v_row, dst_row, dv); ++ } ++ } ++} ++ ++// ggml_compute_forward_moe_route_weights ++ ++void ggml_compute_forward_moe_route_weights( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * probs = dst->src[0]; // [1, n_expert, n_tokens, 1] ++ const ggml_tensor * ids = dst->src[1]; // [n_expert_used, n_tokens, 1, 1] ++ ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ GGML_ASSERT(probs->type == GGML_TYPE_F32); ++ GGML_ASSERT(ids->type == GGML_TYPE_I32); ++ GGML_ASSERT(probs->ne[0] == 1); ++ GGML_ASSERT(dst->ne[0] == 1); ++ GGML_ASSERT(dst->ne[1] == ids->ne[0]); ++ GGML_ASSERT(dst->ne[2] == ids->ne[1]); ++ GGML_ASSERT(probs->ne[2] == ids->ne[1]); ++ ++ const int64_t n_expert = probs->ne[1]; ++ const int64_t n_expert_used = ids->ne[0]; ++ const int64_t n_tokens = ids->ne[1]; ++ const bool norm = ggml_get_op_params_i32(dst, 2) != 0; ++ const float clamp_min = ggml_get_op_params_f32(dst, 0); ++ const float scale = ggml_get_op_params_f32(dst, 1); ++ ++ const int ith = params->ith; ++ const int nth = params->nth; ++ ++ const int64_t dr = (n_tokens + nth - 1) / nth; ++ const int64_t t0 = dr * ith; ++ const int64_t t1 = MIN(t0 + dr, n_tokens); ++ ++ for (int64_t token = t0; token < t1; ++token) { ++ float sum = 0.0f; ++ ++ if (norm) { ++ for (int64_t expert_idx = 0; expert_idx < n_expert_used; ++expert_idx) { ++ const int32_t expert = *(const int32_t *) ((const char *) ids->data + ++ expert_idx * ids->nb[0] + token * ids->nb[1]); ++ GGML_ASSERT(expert >= 0 && expert < n_expert); ++ sum += *(const float *) ((const char *) probs->data + ++ expert * probs->nb[1] + token * probs->nb[2]); ++ } ++ sum = MAX(sum, clamp_min); ++ } ++ ++ for (int64_t expert_idx = 0; expert_idx < n_expert_used; ++expert_idx) { ++ const int32_t expert = *(const int32_t *) ((const char *) ids->data + ++ expert_idx * ids->nb[0] + token * ids->nb[1]); ++ GGML_ASSERT(expert >= 0 && expert < n_expert); ++ ++ float weight = *(const float *) ((const char *) probs->data + ++ expert * probs->nb[1] + token * probs->nb[2]); ++ if (norm) { ++ weight /= sum; ++ } ++ weight *= scale; ++ ++ *(float *) ((char *) dst->data + expert_idx * dst->nb[1] + token * dst->nb[2]) = weight; ++ } ++ } ++} ++ ++// ggml_compute_forward_moe_weighted_sum ++ ++void ggml_compute_forward_moe_weighted_sum( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * experts = dst->src[0]; // [n_embd, n_expert_used, n_tokens, 1] ++ const ggml_tensor * weights = dst->src[1]; // [1, n_expert_used, n_tokens, 1] ++ ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ GGML_ASSERT(experts->type == GGML_TYPE_F32); ++ GGML_ASSERT(weights->type == GGML_TYPE_F32); ++ GGML_ASSERT(weights->ne[0] == 1); ++ GGML_ASSERT(dst->ne[0] == experts->ne[0]); ++ GGML_ASSERT(dst->ne[1] == experts->ne[2]); ++ GGML_ASSERT(experts->ne[1] == weights->ne[1]); ++ GGML_ASSERT(experts->ne[2] == weights->ne[2]); ++ ++ const int64_t n_embd = experts->ne[0]; ++ const int64_t n_expert_used = experts->ne[1]; ++ const int64_t n_tokens = experts->ne[2]; ++ const bool already_weighted = ggml_get_op_params_i32(dst, 0) != 0; ++ ++ const int ith = params->ith; ++ const int nth = params->nth; ++ ++ const int64_t total = n_embd * n_tokens; ++ const int64_t dr = (total + nth - 1) / nth; ++ const int64_t i0 = dr * ith; ++ const int64_t i1 = MIN(i0 + dr, total); ++ ++ for (int64_t i = i0; i < i1; ++i) { ++ const int64_t i_embd = i % n_embd; ++ const int64_t token = i / n_embd; ++ ++ float acc = 0.0f; ++ for (int64_t expert = 0; expert < n_expert_used; ++expert) { ++ const float value = *(const float *) ((const char *) experts->data + ++ i_embd * experts->nb[0] + expert * experts->nb[1] + token * experts->nb[2]); ++ const float weight = already_weighted ? 1.0f : *(const float *) ((const char *) weights->data + ++ expert * weights->nb[1] + token * weights->nb[2]); ++ acc += value * weight; ++ } ++ ++ *(float *) ((char *) dst->data + i_embd * dst->nb[0] + token * dst->nb[1]) = acc; ++ } ++} ++ ++// ggml_compute_forward_moe_mul_mat_id ++ ++void ggml_compute_forward_moe_mul_mat_id( ++ const ggml_compute_params * params, ++ ggml_tensor * dst) { ++ const ggml_tensor * experts = dst->src[0]; // [n_ff, n_embd, n_expert, 1] ++ const ggml_tensor * input = dst->src[1]; // [n_ff, n_expert_used, n_tokens, 1] ++ const ggml_tensor * ids = dst->src[2]; // [n_expert_used, n_tokens, 1, 1] ++ const ggml_tensor * weights = dst->src[3]; // [1, n_expert_used, n_tokens, 1] ++ ++ GGML_ASSERT(dst->type == GGML_TYPE_F32); ++ GGML_ASSERT(input->type == GGML_TYPE_F32 || input->type == GGML_TYPE_F16); ++ GGML_ASSERT(ids->type == GGML_TYPE_I32); ++ GGML_ASSERT(weights->type == GGML_TYPE_F32); ++ ++ const int64_t n_ff = experts->ne[0]; ++ const int64_t n_embd = experts->ne[1]; ++ const int64_t n_expert = experts->ne[2]; ++ const int64_t n_expert_used = ids->ne[0]; ++ const int64_t n_tokens = ids->ne[1]; ++ ++ GGML_ASSERT(input->ne[0] == n_ff); ++ GGML_ASSERT(input->ne[1] == n_expert_used); ++ GGML_ASSERT(input->ne[2] == n_tokens); ++ GGML_ASSERT(weights->ne[0] == 1); ++ GGML_ASSERT(weights->ne[1] == n_expert_used); ++ GGML_ASSERT(weights->ne[2] == n_tokens); ++ GGML_ASSERT(dst->ne[0] == n_embd); ++ GGML_ASSERT(dst->ne[1] == n_tokens); ++ ++ const int ith = params->ith; ++ const int nth = params->nth; ++ const int64_t total = n_embd*n_tokens; ++ const int64_t i0 = (total*ith)/nth; ++ const int64_t i1 = (total*(ith + 1))/nth; ++ ++ float * row_scratch = (float *) params->wdata + ith*n_ff; ++ const ggml_to_float_t to_float = ggml_get_type_traits(experts->type)->to_float; ++ ++ for (int64_t i = i0; i < i1; ++i) { ++ const int64_t i_embd = i % n_embd; ++ const int64_t token = i / n_embd; ++ float acc = 0.0f; ++ ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ const int32_t expert = *(const int32_t *) ((const char *) ids->data + ++ slot*ids->nb[0] + token*ids->nb[1]); ++ GGML_ASSERT(expert >= 0 && expert < n_expert); ++ ++ const char * expert_row = (const char *) experts->data + ++ i_embd*experts->nb[1] + expert*experts->nb[2]; ++ const float * row; ++ if (experts->type == GGML_TYPE_F32) { ++ row = (const float *) expert_row; ++ } else { ++ GGML_ASSERT(to_float != nullptr); ++ to_float(expert_row, row_scratch, n_ff); ++ row = row_scratch; ++ } ++ ++ const char * input_row = (const char *) input->data + ++ slot*input->nb[1] + token*input->nb[2]; ++ float dot = 0.0f; ++ if (input->type == GGML_TYPE_F32) { ++ const float * values = (const float *) input_row; ++ for (int64_t col = 0; col < n_ff; ++col) { ++ dot += row[col]*values[col]; ++ } ++ } else { ++ const ggml_fp16_t * values = (const ggml_fp16_t *) input_row; ++ for (int64_t col = 0; col < n_ff; ++col) { ++ dot += row[col]*GGML_FP16_TO_FP32(values[col]); ++ } ++ } ++ ++ const float weight = *(const float *) ((const char *) weights->data + ++ slot*weights->nb[1] + token*weights->nb[2]); ++ acc += dot*weight; ++ } ++ ++ *(float *) ((char *) dst->data + i_embd*dst->nb[0] + token*dst->nb[1]) = acc; ++ } ++} ++ + // ggml_compute_forward_rwkv_wkv7 + + static void ggml_compute_forward_rwkv_wkv7_f32( +diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h +index 4c1642a6..31a8740f 100644 +--- a/ggml/src/ggml-cpu/ops.h ++++ b/ggml/src/ggml-cpu/ops.h +@@ -109,6 +109,12 @@ void ggml_compute_forward_lightning_indexer(const struct ggml_compute_params * p + void ggml_compute_forward_dsv4_hc_comb(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_dsv4_hc_pre(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_dsv4_hc_post(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_dsa_sparse_mask(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_dsa_sparse_attn(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_dsa_top1_attn(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_moe_route_weights(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_moe_weighted_sum(const struct ggml_compute_params * params, struct ggml_tensor * dst); ++void ggml_compute_forward_moe_mul_mat_id(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_map_custom1(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_map_custom2(const struct ggml_compute_params * params, struct ggml_tensor * dst); + void ggml_compute_forward_map_custom3(const struct ggml_compute_params * params, struct ggml_tensor * dst); +diff --git a/ggml/src/ggml-metal/ggml-metal-common.cpp b/ggml/src/ggml-metal/ggml-metal-common.cpp +index 2eb9820b..30069090 100644 +--- a/ggml/src/ggml-metal/ggml-metal-common.cpp ++++ b/ggml/src/ggml-metal/ggml-metal-common.cpp +@@ -3,6 +3,8 @@ + #include "ggml-impl.h" + #include "ggml-backend-impl.h" + ++#include ++#include + #include + + // represents a memory range (i.e. an interval from a starting address p0 to an ending address p1 in a given buffer pb) +@@ -206,6 +208,10 @@ struct node_info { + } + }; + ++// TODO(GLM-DSA Metal fused-route milestone): keep route packing disabled until ++// it is represented by a backend op or the encoder can safely consume the ++// packed subgraph. ++ + static std::vector ggml_metal_graph_optimize_reorder(const std::vector & nodes) { + // helper to add node src and dst ranges + const auto & h_add = [](ggml_mem_ranges_t mrs, const node_info & node) { +@@ -418,7 +424,6 @@ void ggml_graph_optimize(ggml_cgraph * gf) { + break; + } + } +- + // add the fused tensors into the node info so we can unfuse them later + for (int k = 1; k < f; k++) { + ++i; +diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m +index 32d97cd5..083949ce 100644 +--- a/ggml/src/ggml-metal/ggml-metal-context.m ++++ b/ggml/src/ggml-metal/ggml-metal-context.m +@@ -19,6 +19,114 @@ + // max number of MTLCommandBuffer used to submit a graph for processing + #define GGML_METAL_MAX_COMMAND_BUFFERS 8 + ++// Each command-buffer lane owns a separate private resource so fused kernels ++// can exchange intermediates without synchronizing the shared GGML arena. ++#define GGML_METAL_FUSION_SCRATCH_SIZE (1024*1024) ++ ++static int ggml_metal_skippy_sched_debug(void) { ++ const char * value = getenv("SKIPPY_GLM_DSA_METAL_SCHED_DEBUG"); ++ return value == NULL ? 0 : atoi(value); ++} ++ ++static int ggml_metal_skippy_sched_debug_limit(void) { ++ const char * value = getenv("SKIPPY_GLM_DSA_METAL_SCHED_DEBUG_LIMIT"); ++ return value == NULL ? 8 : atoi(value); ++} ++ ++static int ggml_metal_skippy_cb_timing(void) { ++ const char * value = getenv("SKIPPY_GLM_DSA_METAL_CB_TIMING"); ++ if (value != NULL && atoi(value) != 0) { ++ return 1; ++ } ++ value = getenv("GGML_METAL_CB_TIMING"); ++ return value == NULL ? 0 : atoi(value) != 0; ++} ++ ++static int ggml_metal_skippy_cb_timing_limit(void) { ++ const char * value = getenv("SKIPPY_GLM_DSA_METAL_CB_TIMING_LIMIT"); ++ if (value != NULL) { ++ return atoi(value); ++ } ++ value = getenv("GGML_METAL_CB_TIMING_LIMIT"); ++ return value == NULL ? 16 : atoi(value); ++} ++ ++static int ggml_metal_host_timing(void) { ++ const char * value = getenv("SKIPPY_GLM_DSA_METAL_HOST_TIMING"); ++ if (value != NULL && atoi(value) != 0) { ++ return 1; ++ } ++ value = getenv("GGML_METAL_HOST_TIMING"); ++ return value == NULL ? 0 : atoi(value) != 0; ++} ++ ++static int ggml_metal_host_timing_limit(void) { ++ const char * value = getenv("SKIPPY_GLM_DSA_METAL_HOST_TIMING_LIMIT"); ++ if (value != NULL) { ++ return atoi(value); ++ } ++ value = getenv("GGML_METAL_HOST_TIMING_LIMIT"); ++ return value == NULL ? 16 : atoi(value); ++} ++ ++static void ggml_metal_attach_cb_timing( ++ id cmd_buf, ++ int graph_seq, ++ int cb_idx, ++ int idx_start, ++ int idx_end, ++ int n_nodes) { ++ const int64_t cpu_start_us = ggml_time_us(); ++ [cmd_buf addCompletedHandler:^(id cb) { ++ const int64_t cpu_end_us = ggml_time_us(); ++ const double gpu_start_s = cb.GPUStartTime; ++ const double gpu_end_s = cb.GPUEndTime; ++ const double gpu_ms = gpu_end_s > gpu_start_s ? 1000.0*(gpu_end_s - gpu_start_s) : -1.0; ++ GGML_LOG_INFO( ++ "skippy: metal_cb_timing graph=%d cb=%d node_start=%d node_end=%d nodes=%d total_nodes=%d status=%lu cpu_ms=%.3f gpu_ms=%.3f gpu_start=%.9f gpu_end=%.9f\n", ++ graph_seq, ++ cb_idx, ++ idx_start, ++ idx_end, ++ idx_end - idx_start, ++ n_nodes, ++ (unsigned long) cb.status, ++ (double) (cpu_end_us - cpu_start_us) / 1000.0, ++ gpu_ms, ++ gpu_start_s, ++ gpu_end_s); ++ }]; ++} ++ ++static bool ggml_metal_topk_moe_route_fusion_enabled(void) { ++ const char * value = getenv("GGML_METAL_ENABLE_TOPK_MOE_ROUTE_FUSION"); ++ if (value != NULL) { ++ return atoi(value) != 0; ++ } ++ value = getenv("SKIPPY_GLM_DSA_ENABLE_METAL_TOPK_MOE_FUSION"); ++ if (value != NULL) { ++ return atoi(value) != 0; ++ } ++ value = getenv("GGML_METAL_DISABLE_TOPK_MOE_ROUTE_FUSION"); ++ if (value != NULL && atoi(value) != 0) { ++ return false; ++ } ++ value = getenv("SKIPPY_GLM_DSA_DISABLE_METAL_TOPK_MOE_FUSION"); ++ if (value != NULL && atoi(value) != 0) { ++ return false; ++ } ++ return true; ++} ++ ++static bool ggml_metal_topk_moe_route_single_cb_enabled(void) { ++ const char * value = getenv("GGML_METAL_TOPK_MOE_ROUTE_SINGLE_CB"); ++ if (value != NULL) { ++ return atoi(value) != 0; ++ } ++ value = getenv("SKIPPY_GLM_DSA_TOPK_MOE_ROUTE_SINGLE_CB"); ++ return value != NULL && atoi(value) != 0; ++} ++ + struct ggml_metal_command_buffer { + id obj; + }; +@@ -54,6 +162,7 @@ + + // command buffer state + int n_cb; // number of extra threads used to submit the command buffers ++ int n_cb_active; // number of extra command buffers used by the current graph + int n_nodes_0; // number of nodes submitted by the main thread + int n_nodes_1; // remaining number of nodes submitted by the n_cb threads + int n_nodes_per_cb; +@@ -65,6 +174,7 @@ + + // n_cb command buffers + 1 used by the main thread + struct ggml_metal_command_buffer cmd_bufs[GGML_METAL_MAX_COMMAND_BUFFERS + 1]; ++ ggml_metal_buffer_t fusion_scratch[GGML_METAL_MAX_COMMAND_BUFFERS + 1]; + + // extra command buffers for things like getting, setting and copying tensors + NSMutableArray * cmd_bufs_ext; +@@ -176,7 +286,14 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) { + for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { + res->cmd_bufs[i].obj = nil; + } +- ++ for (int i = 0; i <= GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { ++ res->fusion_scratch[i] = ggml_metal_buffer_init( ++ res->dev, GGML_METAL_FUSION_SCRATCH_SIZE, false); ++ if (res->fusion_scratch[i] == NULL) { ++ GGML_LOG_WARN("%s: failed to allocate fusion scratch lane %d; fused private-resource paths will fall back\n", ++ __func__, i); ++ } ++ } + res->cmd_bufs_ext = [[NSMutableArray alloc] init]; + + res->cmd_buf_last = nil; +@@ -194,6 +311,11 @@ void ggml_metal_free(ggml_metal_t ctx) { + [ctx->cmd_bufs[i].obj release]; + } + } ++ for (int i = 0; i <= GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { ++ if (ctx->fusion_scratch[i]) { ++ ggml_metal_buffer_free(ctx->fusion_scratch[i]); ++ } ++ } + + for (int i = 0; i < (int) ctx->cmd_bufs_ext.count; ++i) { + if (ctx->cmd_bufs_ext[i]) { +@@ -245,7 +367,7 @@ void ggml_metal_synchronize(ggml_metal_t ctx) { + + // check status of all command buffers + { +- const int n_cb = ctx->n_cb; ++ const int n_cb = ctx->n_cb_active; + + for (int cb_idx = 0; cb_idx <= n_cb; ++cb_idx) { + id cmd_buf = ctx->cmd_bufs[cb_idx].obj; +@@ -441,11 +563,11 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * + return GGML_STATUS_FAILED; + } + +- // number of nodes encoded by the main thread (empirically determined) +- const int n_main = MAX(64, 0.1*gf->n_nodes); ++ const bool topk_moe_route_fusion = ggml_metal_topk_moe_route_fusion_enabled(); + +- // number of threads in addition to the main thread +- const int n_cb = ctx->n_cb; ++ // number of nodes encoded by the main thread (empirically determined) ++ const bool topk_moe_route_single_cb = topk_moe_route_fusion && ggml_metal_topk_moe_route_single_cb_enabled(); ++ const int n_main = topk_moe_route_single_cb ? gf->n_nodes : MAX(64, 0.1*gf->n_nodes); + + // keep the memory wired + ggml_metal_device_rsets_keep_alive(ctx->dev); +@@ -458,12 +580,69 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * + // tests on M1 Pro and M2 Ultra using LLaMA models, show that optimal values for n_cb are 1 or 2 + + @autoreleasepool { ++ const int64_t host_start_us = ggml_time_us(); + ctx->gf = gf; + + ctx->n_nodes_0 = MIN(n_main, gf->n_nodes); + ctx->n_nodes_1 = gf->n_nodes - ctx->n_nodes_0; + +- ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb; ++ ctx->n_cb_active = ctx->n_nodes_1 > 0 ? ctx->n_cb : 0; ++ const int n_cb = ctx->n_cb_active; ++ ++ ctx->n_nodes_per_cb = n_cb > 0 ? (ctx->n_nodes_1 + n_cb - 1) / n_cb : 0; ++ ++ static int graph_seq = 0; ++ graph_seq++; ++ ++ const int host_timing = ggml_metal_host_timing(); ++ const int host_timing_limit = ggml_metal_host_timing_limit(); ++ const bool host_timing_should_log = ++ host_timing && (host_timing_limit < 0 || graph_seq <= host_timing_limit); ++ ++ const int skippy_sched_debug = ggml_metal_skippy_sched_debug(); ++ const int skippy_cb_timing = ggml_metal_skippy_cb_timing(); ++ const int skippy_cb_timing_limit = ggml_metal_skippy_cb_timing_limit(); ++ const bool skippy_cb_timing_should_log = ++ skippy_cb_timing && ++ (skippy_cb_timing_limit < 0 || graph_seq <= skippy_cb_timing_limit); ++ if (skippy_sched_debug > 0) { ++ const int skippy_sched_limit = ggml_metal_skippy_sched_debug_limit(); ++ const bool skippy_sched_should_log = skippy_sched_limit < 0 || graph_seq <= skippy_sched_limit; ++ if (!skippy_sched_should_log) { ++ goto skippy_sched_debug_done; ++ } ++ ++ GGML_LOG_INFO( ++ "skippy: metal_sched graph=%d nodes=%d n_main=%d n_cb=%d n_cb_configured=%d n_nodes_0=%d n_nodes_1=%d n_nodes_per_cb=%d use_concurrency=%d use_fusion=%d use_graph_optimize=%d topk_moe_route_fusion=%d topk_moe_route_single_cb=%d\n", ++ graph_seq, ++ gf->n_nodes, ++ n_main, ++ n_cb, ++ ctx->n_cb, ++ ctx->n_nodes_0, ++ ctx->n_nodes_1, ++ ctx->n_nodes_per_cb, ++ ctx->use_concurrency ? 1 : 0, ++ ctx->use_fusion ? 1 : 0, ++ ctx->use_graph_optimize ? 1 : 0, ++ topk_moe_route_fusion ? 1 : 0, ++ topk_moe_route_single_cb ? 1 : 0); ++ ++ if (skippy_sched_debug > 1) { ++ for (int i = 0; i < gf->n_nodes; ++i) { ++ const struct ggml_tensor * node = gf->nodes[i]; ++ GGML_LOG_INFO( ++ "skippy: metal_sched_node graph=%d idx=%d op=%s name=%s\n", ++ graph_seq, ++ i, ++ ggml_op_name(node->op), ++ node->name); ++ } ++ } ++ ++skippy_sched_debug_done: ++ ; ++ } + + if (ctx->capture_compute >= 0) { + ctx->capture_compute--; +@@ -508,6 +687,7 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * + + // the main thread commits the first few commands immediately + // cmd_buf[n_cb] ++ const int64_t main_encode_start_us = host_timing_should_log ? ggml_time_us() : 0; + { + id cmd_buf = [queue commandBufferWithUnretainedReferences]; + [cmd_buf retain]; +@@ -517,16 +697,22 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * + } + ctx->cmd_bufs[n_cb].obj = cmd_buf; + ++ if (skippy_cb_timing_should_log) { ++ ggml_metal_attach_cb_timing(cmd_buf, graph_seq, n_cb, 0, ctx->n_nodes_0, gf->n_nodes); ++ } ++ + [cmd_buf enqueue]; + + ctx->encode_async(n_cb); + } ++ const int64_t main_encode_end_us = host_timing_should_log ? ggml_time_us() : 0; + + // remember the command buffer for the next iteration + ctx->cmd_buf_last = ctx->cmd_bufs[n_cb].obj; + + // prepare the rest of the command buffers asynchronously (optional) + // cmd_buf[0.. n_cb) ++ const int64_t secondary_setup_start_us = host_timing_should_log ? ggml_time_us() : 0; + for (int cb_idx = 0; cb_idx < n_cb; ++cb_idx) { + id cmd_buf = [queue commandBufferWithUnretainedReferences]; + [cmd_buf retain]; +@@ -536,6 +722,12 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * + } + ctx->cmd_bufs[cb_idx].obj = cmd_buf; + ++ if (skippy_cb_timing_should_log) { ++ const int idx_start = ctx->n_nodes_0 + ((cb_idx + 0) * ctx->n_nodes_per_cb); ++ const int idx_end = ctx->n_nodes_0 + MIN((cb_idx == n_cb - 1) ? ctx->n_nodes_1 : (cb_idx + 1) * ctx->n_nodes_per_cb, ctx->n_nodes_1); ++ ggml_metal_attach_cb_timing(cmd_buf, graph_seq, cb_idx, idx_start, idx_end, gf->n_nodes); ++ } ++ + // always enqueue the first two command buffers + // enqueue all of the command buffers if we don't need to abort + if (cb_idx < 2 || ctx->abort_callback == NULL) { +@@ -546,8 +738,25 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * + ctx->cmd_buf_last = cmd_buf; + } + } ++ const int64_t secondary_setup_end_us = host_timing_should_log ? ggml_time_us() : 0; + ++ const int64_t parallel_encode_start_us = host_timing_should_log ? ggml_time_us() : 0; + dispatch_apply(n_cb, ctx->d_queue, ctx->encode_async); ++ const int64_t parallel_encode_end_us = host_timing_should_log ? ggml_time_us() : 0; ++ ++ if (host_timing_should_log) { ++ GGML_LOG_INFO( ++ "skippy: metal_host_timing graph=%d nodes=%d n_main=%d n_cb=%d pre_encode_us=%lld main_encode_us=%lld secondary_setup_us=%lld parallel_encode_us=%lld submit_us=%lld\n", ++ graph_seq, ++ gf->n_nodes, ++ ctx->n_nodes_0, ++ n_cb, ++ (long long) (main_encode_start_us - host_start_us), ++ (long long) (main_encode_end_us - main_encode_start_us), ++ (long long) (secondary_setup_end_us - secondary_setup_start_us), ++ (long long) (parallel_encode_end_us - parallel_encode_start_us), ++ (long long) (parallel_encode_end_us - host_start_us)); ++ } + + // for debugging: block until graph is computed + //[ctx->cmd_buf_last waitUntilCompleted]; +@@ -661,6 +870,11 @@ ggml_metal_event_t ggml_metal_get_ev_cpy(ggml_metal_t ctx) { + } + + void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) { ++ const char * profile_n_cb = getenv("SKIPPY_GLM_DSA_METAL_PROFILE_N_CB"); ++ if (profile_n_cb != NULL && profile_n_cb[0] != '\0') { ++ n_cb = atoi(profile_n_cb); ++ } ++ + if (ctx->n_cb != n_cb) { + ctx->n_cb = MIN(n_cb, GGML_METAL_MAX_COMMAND_BUFFERS); + +@@ -675,7 +889,7 @@ void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) { + + ctx->encode_async = Block_copy(^(size_t iter) { + const int cb_idx = iter; +- const int n_cb_l = ctx->n_cb; ++ const int n_cb_l = ctx->n_cb_active; + + const int n_nodes_0 = ctx->n_nodes_0; + const int n_nodes_1 = ctx->n_nodes_1; +@@ -695,6 +909,8 @@ void ggml_metal_set_n_cb(ggml_metal_t ctx, int n_cb) { + ggml_metal_op_t ctx_op = ggml_metal_op_init( + ctx->dev, + cmd_buf, ++ ggml_metal_buffer_get_base_id(ctx->fusion_scratch[cb_idx]), ++ GGML_METAL_FUSION_SCRATCH_SIZE, + ctx->gf, + idx_start, + idx_end, +diff --git a/ggml/src/ggml-metal/ggml-metal-device.cpp b/ggml/src/ggml-metal/ggml-metal-device.cpp +index 270c1411..088089aa 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.cpp ++++ b/ggml/src/ggml-metal/ggml-metal-device.cpp +@@ -5,6 +5,7 @@ + #include "ggml-impl.h" + + #include ++#include + #include + #include + #include +@@ -17,6 +18,207 @@ struct ggml_metal_device_deleter { + + typedef std::unique_ptr ggml_metal_device_ptr; + ++static int ggml_metal_glm_dsa_mul_mv_nsg_requested(const char * type_env, int fallback) { ++ const char * value = getenv(type_env); ++ if (value == nullptr || value[0] == '\0') { ++ value = getenv("LLAMA_GLM_DSA_MUL_MV_NSG"); ++ } ++ if (value == nullptr || value[0] == '\0') { ++ return fallback; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ case 8: ++ return requested; ++ default: ++ return fallback; ++ } ++} ++ ++static int ggml_metal_glm_dsa_mul_mv_shape_policy_mask() { ++ const char * value = getenv("LLAMA_GLM_DSA_EXPERIMENTAL_MUL_MV_SHAPE_POLICY"); ++ return value != nullptr && value[0] != '\0' ? atoi(value) : 0; ++} ++ ++static bool ggml_metal_glm_dsa_q8_0_attention_shape(const ggml_tensor * op) { ++ const ggml_tensor * weights = op->src[0]; ++ return weights != nullptr && ++ ((weights->ne[0] == 6144 && (weights->ne[1] == 576 || weights->ne[1] == 2048)) || ++ (weights->ne[0] == 2048 && weights->ne[1] == 16384)); ++} ++ ++static bool ggml_metal_glm_dsa_q3_k_attention_output_shape(const ggml_tensor * op) { ++ const ggml_tensor * weights = op->src[0]; ++ return weights != nullptr && weights->ne[0] == 16384 && weights->ne[1] == 6144; ++} ++ ++static bool ggml_metal_glm_dsa_q4_k_gate_up_shape(const ggml_tensor * op) { ++ const ggml_tensor * weights = op->src[0]; ++ return weights != nullptr && weights->ne[0] == 6144 && weights->ne[1] == 2048; ++} ++ ++static bool ggml_metal_glm_dsa_q8_0_mul_mv_row_parallel_enabled() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q8_0_MUL_MV_ROW_PARALLEL"); ++ return value != nullptr && value[0] != '\0' && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q8_0_mul_mv_vector_dot_enabled() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q8_0_MUL_MV_VECTOR_DOT"); ++ return value != nullptr && value[0] != '\0' && atoi(value) != 0; ++} ++ ++static int ggml_metal_glm_dsa_q8_0_mul_mv_nr0_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q8_0_MUL_MV_NR0"); ++ if (value == nullptr || value[0] == '\0') { ++ return N_R0_Q8_0; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ case 8: ++ return requested; ++ default: ++ return N_R0_Q8_0; ++ } ++} ++ ++static int ggml_metal_glm_dsa_q3_k_mul_mv_id_nsg_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_NSG"); ++ if (value == nullptr || value[0] == '\0') { ++ return N_SG_Q3_K; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ return requested; ++ default: ++ return N_SG_Q3_K; ++ } ++} ++ ++static int ggml_metal_glm_dsa_q2_k_mul_mv_id_nsg_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q2_K_MUL_MV_ID_NSG"); ++ if (value == nullptr || value[0] == '\0') { ++ return N_SG_Q2_K; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ return requested; ++ default: ++ return N_SG_Q2_K; ++ } ++} ++ ++static int ggml_metal_glm_dsa_q2_k_mul_mv_id_nr0_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q2_K_MUL_MV_ID_NR0"); ++ if (value == nullptr || value[0] == '\0') { ++ return N_R0_Q2_K; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ case 8: ++ return requested; ++ default: ++ return N_R0_Q2_K; ++ } ++} ++ ++static bool ggml_metal_glm_dsa_q2_k_mul_mv_id_nr0_overridden() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q2_K_MUL_MV_ID_NR0"); ++ return value != nullptr && value[0] != '\0'; ++} ++ ++static bool ggml_metal_glm_dsa_q2_k_mul_mv_id_glm_down_enabled() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q2_K_MUL_MV_ID_GLM_DOWN"); ++ if (value != nullptr && value[0] != '\0') { ++ return atoi(value) != 0; ++ } ++ ++ const char * disabled = getenv("LLAMA_GLM_DSA_DISABLE_Q2_K_MUL_MV_ID_GLM_DOWN"); ++ return disabled == nullptr || disabled[0] == '\0' || atoi(disabled) == 0; ++} ++ ++static int ggml_metal_glm_dsa_q2_gate_up_swiglu_nsg_requested(int fallback) { ++ const char * value = getenv("LLAMA_GLM_DSA_Q2_GATE_UP_SWIGLU_NSG"); ++ if (value == nullptr || value[0] == '\0') { ++ return fallback; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ return requested; ++ default: ++ return fallback; ++ } ++} ++ ++static int ggml_metal_glm_dsa_q2_gate_up_swiglu_nr0_requested(int fallback) { ++ const char * value = getenv("LLAMA_GLM_DSA_Q2_GATE_UP_SWIGLU_NR0"); ++ if (value == nullptr || value[0] == '\0') { ++ return fallback; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ case 8: ++ return requested; ++ default: ++ return fallback; ++ } ++} ++ ++static int ggml_metal_glm_dsa_q3_k_mul_mv_id_nr0_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_NR0"); ++ if (value == nullptr || value[0] == '\0') { ++ return N_R0_Q3_K; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ case 8: ++ return requested; ++ default: ++ return N_R0_Q3_K; ++ } ++} ++ ++static bool ggml_metal_glm_dsa_q3_k_mul_mv_id_glm_down_enabled() { ++ const char * value = getenv("LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_GLM_DOWN"); ++ if (value != nullptr && value[0] != '\0') { ++ return atoi(value) != 0; ++ } ++ ++ const char * disabled = getenv("LLAMA_GLM_DSA_DISABLE_Q3_K_MUL_MV_ID_GLM_DOWN"); ++ return disabled == nullptr || disabled[0] == '\0' || atoi(disabled) == 0; ++} ++ + ggml_metal_device_t ggml_metal_device_get(int device) { + static std::vector devs; + +@@ -145,11 +347,17 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_2d(ggml_met + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows(ggml_metal_library_t lib, ggml_type tsrc) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows(ggml_metal_library_t lib, ggml_type tsrc, ggml_type tdst, bool use_f16_vec4) { + char base[256]; + char name[256]; + +- snprintf(base, 256, "kernel_get_rows_%s", ggml_type_name(tsrc)); ++ if (use_f16_vec4) { ++ snprintf(base, 256, "kernel_get_rows_f16_f16_vec4"); ++ } else if (tsrc == GGML_TYPE_F16 && tdst == GGML_TYPE_F16) { ++ snprintf(base, 256, "kernel_get_rows_f16_f16"); ++ } else { ++ snprintf(base, 256, "kernel_get_rows_%s", ggml_type_name(tsrc)); ++ } + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); +@@ -160,15 +368,20 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows(ggml_me + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, const ggml_tensor * op) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows_packed_f16(ggml_metal_library_t lib) { ++ const char * base = "kernel_get_rows_f16_f16_packed_rows"; ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, base, nullptr); ++ } ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, ggml_type tidx, ggml_type tdst) { + char base[256]; + char name[256]; + +- const auto tsrc = op->src[0]->type; +- const auto tidx = op->src[1]->type; +- const auto tdst = op->type; +- +- snprintf(base, 256, "kernel_set_rows_%s_%s_%s", ggml_type_name(tsrc), ggml_type_name(tidx), ggml_type_name(tdst)); ++ snprintf(base, 256, "kernel_set_rows_%s_%s", ggml_type_name(tdst), ggml_type_name(tidx)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); +@@ -179,31 +392,40 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_me + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag(ggml_metal_library_t lib, const ggml_tensor * op) { +- char base[256]; +- char name[256]; ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_mask_fill(ggml_metal_library_t lib) { ++ const char * base = "kernel_dsa_sparse_mask_fill"; ++ const char * name = "kernel_dsa_sparse_mask_fill"; + +- const int n = op->src[0]->ne[0]; ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } + +- snprintf(base, 256, "kernel_diag_%s", ggml_type_name(op->src[0]->type)); +- snprintf(name, 256, "%s_n=%d", base, n); ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_mask_set(ggml_metal_library_t lib) { ++ const char * base = "kernel_dsa_sparse_mask_set"; ++ const char * name = "kernel_dsa_sparse_mask_set"; + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + +- res.nsg = 1; +- res.smem = 0; +- + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat(ggml_metal_library_t lib, ggml_type tsrc) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_attn(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_DSA_SPARSE_ATTN); ++ + char base[256]; + char name[256]; + +- snprintf(base, 256, "kernel_repeat_%s", ggml_type_name(tsrc)); ++ snprintf(base, 256, "kernel_dsa_sparse_attn_%s_%s_%s", ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->src[3]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); +@@ -214,11 +436,16 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat(ggml_meta + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat(ggml_metal_library_t lib, ggml_type tsrc) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_attn_cached_topk(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_DSA_SPARSE_ATTN); ++ + char base[256]; + char name[256]; + +- snprintf(base, 256, "kernel_concat_%s", ggml_type_name(tsrc)); ++ snprintf(base, 256, "kernel_dsa_sparse_attn_cached_topk_%s_%s_%s", ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->src[3]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); +@@ -229,217 +456,183 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat(ggml_meta + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary(ggml_metal_library_t lib, const ggml_tensor * op) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_attn_decode_grouped(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_DSA_SPARSE_ATTN); ++ + char base[256]; + char name[256]; + +- int op_num = -1; ++ snprintf(base, 256, "kernel_dsa_sparse_attn_decode_grouped_%s_%s_%s", ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->src[3]->type)); ++ snprintf(name, 256, "%s", base); + +- switch (op->op) { +- case GGML_OP_SCALE: op_num = OP_UNARY_NUM_SCALE; break; +- case GGML_OP_FILL: op_num = OP_UNARY_NUM_FILL; break; +- case GGML_OP_CLAMP: op_num = OP_UNARY_NUM_CLAMP; break; +- case GGML_OP_SQR: op_num = OP_UNARY_NUM_SQR; break; +- case GGML_OP_SQRT: op_num = OP_UNARY_NUM_SQRT; break; +- case GGML_OP_SIN: op_num = OP_UNARY_NUM_SIN; break; +- case GGML_OP_COS: op_num = OP_UNARY_NUM_COS; break; +- case GGML_OP_LOG: op_num = OP_UNARY_NUM_LOG; break; +- case GGML_OP_LEAKY_RELU: op_num = OP_UNARY_NUM_LEAKY_RELU; break; +- case GGML_OP_UNARY: +- switch (ggml_get_unary_op(op)) { +- case GGML_UNARY_OP_TANH: op_num = OP_UNARY_NUM_TANH; break; +- case GGML_UNARY_OP_RELU: op_num = OP_UNARY_NUM_RELU; break; +- case GGML_UNARY_OP_SIGMOID: op_num = OP_UNARY_NUM_SIGMOID; break; +- case GGML_UNARY_OP_GELU: op_num = OP_UNARY_NUM_GELU; break; +- case GGML_UNARY_OP_GELU_ERF: op_num = OP_UNARY_NUM_GELU_ERF; break; +- case GGML_UNARY_OP_GELU_QUICK: op_num = OP_UNARY_NUM_GELU_QUICK; break; +- case GGML_UNARY_OP_SILU: op_num = OP_UNARY_NUM_SILU; break; +- case GGML_UNARY_OP_ELU: op_num = OP_UNARY_NUM_ELU; break; +- case GGML_UNARY_OP_NEG: op_num = OP_UNARY_NUM_NEG; break; +- case GGML_UNARY_OP_ABS: op_num = OP_UNARY_NUM_ABS; break; +- case GGML_UNARY_OP_SGN: op_num = OP_UNARY_NUM_SGN; break; +- case GGML_UNARY_OP_STEP: op_num = OP_UNARY_NUM_STEP; break; +- case GGML_UNARY_OP_HARDSWISH: op_num = OP_UNARY_NUM_HARDSWISH; break; +- case GGML_UNARY_OP_HARDSIGMOID: op_num = OP_UNARY_NUM_HARDSIGMOID; break; +- case GGML_UNARY_OP_EXP: op_num = OP_UNARY_NUM_EXP; break; +- case GGML_UNARY_OP_SOFTPLUS: op_num = OP_UNARY_NUM_SOFTPLUS; break; +- case GGML_UNARY_OP_EXPM1: op_num = OP_UNARY_NUM_EXPM1; break; +- case GGML_UNARY_OP_FLOOR: op_num = OP_UNARY_NUM_FLOOR; break; +- case GGML_UNARY_OP_CEIL: op_num = OP_UNARY_NUM_CEIL; break; +- case GGML_UNARY_OP_ROUND: op_num = OP_UNARY_NUM_ROUND; break; +- case GGML_UNARY_OP_TRUNC: op_num = OP_UNARY_NUM_TRUNC; break; +- case GGML_UNARY_OP_XIELU: op_num = OP_UNARY_NUM_XIELU; break; +- default: GGML_ABORT("fatal error"); +- } break; +- default: GGML_ABORT("fatal error"); +- }; ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } + +- const char * t0_str = ggml_type_name(op->src[0]->type); +- const char * t_str = ggml_type_name(op->type); ++ return res; ++} + +- const bool is_c4 = op->src[0]->ne[0] % 4 == 0; +- const bool is_cnt = ggml_is_contiguous(op->src[0]) && ggml_nelements(op) < 32768; ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_selected_row_flash_vec( ++ ggml_metal_library_t lib, ++ int32_t nsg, ++ int32_t nwg) { ++ const char * base = "kernel_selected_row_flash_vec_f16_dk576_dv512"; + +- snprintf(base, 256, "kernel_unary_%s_%s%s", t0_str, t_str, is_c4 ? "_4" : ""); +- snprintf(name, 256, "%s_op=%d_cnt=%d", base, op_num, is_cnt); ++ char name[256]; ++ snprintf(name, sizeof(name), "%s_nsg=%d_nwg=%d", base, nsg, nwg); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + ggml_metal_cv_t cv = ggml_metal_cv_init(); + +- ggml_metal_cv_set_int16(cv, op_num, FC_UNARY + 0); +- ggml_metal_cv_set_bool (cv, is_cnt, FC_UNARY + 1); ++ ggml_metal_cv_set_int32(cv, nsg, FC_FLASH_ATTN_EXT_VEC + 22); ++ ggml_metal_cv_set_int32(cv, nwg, FC_FLASH_ATTN_EXT_VEC + 23); + + res = ggml_metal_library_compile_pipeline(lib, base, name, cv); + + ggml_metal_cv_free(cv); + } + +- res.c4 = is_c4; +- res.cnt = is_cnt; +- + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_ASSERT(ggml_is_contiguous_1(op->src[0])); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_selected_row_flash_pair( ++ ggml_metal_library_t lib, ++ int32_t nwg) { ++ const char * base = "kernel_selected_row_flash_pair_f16_dk576_dv512"; + +- char base[256]; + char name[256]; +- +- const char * op_str = "undefined"; +- switch (op->op) { +- case GGML_OP_GLU: +- switch (ggml_get_glu_op(op)) { +- case GGML_GLU_OP_REGLU: op_str = "reglu"; break; +- case GGML_GLU_OP_GEGLU: op_str = "geglu"; break; +- case GGML_GLU_OP_SWIGLU: op_str = "swiglu"; break; +- case GGML_GLU_OP_SWIGLU_OAI: op_str = "swiglu_oai"; break; +- case GGML_GLU_OP_GEGLU_ERF: op_str = "geglu_erf"; break; +- case GGML_GLU_OP_GEGLU_QUICK: op_str = "geglu_quick"; break; +- default: GGML_ABORT("fatal error"); +- } break; +- default: GGML_ABORT("fatal error"); +- }; +- +- snprintf(base, 256, "kernel_%s_%s", op_str, ggml_type_name(op->src[0]->type)); +- snprintf(name, 256, "%s", base); ++ snprintf(name, sizeof(name), "%s_nwg=%d", base, nwg); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { +- res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int32(cv, nwg, FC_FLASH_ATTN_EXT_VEC + 23); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); + } + + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum(ggml_metal_library_t lib, const ggml_tensor * op) { +- assert(op->op == GGML_OP_SUM); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_selected_row_flash_tiled( ++ ggml_metal_library_t lib, ++ int32_t nwg) { ++ const char * base = "kernel_selected_row_flash_tiled_f16_dk576_dv512"; + +- char base[256]; + char name[256]; +- +- snprintf(base, 256, "kernel_op_sum_%s", ggml_type_name(op->src[0]->type)); +- snprintf(name, 256, "%s", base); ++ snprintf(name, sizeof(name), "%s_nwg=%d", base, nwg); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { +- res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int32(cv, nwg, FC_FLASH_ATTN_EXT_VEC + 23); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); + } + + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_multihead_flash( ++ ggml_metal_library_t lib, ++ int32_t nwg) { ++ const char * base = "kernel_glm_compact_multihead_flash_f16_dk576_dv512"; + +- char base[256]; + char name[256]; +- +- int op_num = -1; +- +- switch (op->op) { +- case GGML_OP_SUM_ROWS: op_num = OP_SUM_ROWS_NUM_SUM_ROWS; break; +- case GGML_OP_MEAN: op_num = OP_SUM_ROWS_NUM_MEAN; break; +- default: GGML_ABORT("fatal error"); +- }; +- +- const char * t0_str = ggml_type_name(op->src[0]->type); +- const char * t_str = ggml_type_name(op->type); +- +- const bool is_c4 = op->src[0]->ne[0] % 4 == 0; +- +- snprintf(base, 256, "kernel_sum_rows_%s_%s%s", t0_str, t_str, is_c4 ? "_4" : ""); +- snprintf(name, 256, "%s_op=%d", base, op_num); ++ snprintf(name, sizeof(name), "%s_nwg=%d", base, nwg); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + ggml_metal_cv_t cv = ggml_metal_cv_init(); + +- ggml_metal_cv_set_int16(cv, op_num, FC_SUM_ROWS + 0); ++ ggml_metal_cv_set_int32(cv, nwg, FC_FLASH_ATTN_EXT_VEC + 23); + + res = ggml_metal_library_compile_pipeline(lib, base, name, cv); + + ggml_metal_cv_free(cv); + } + +- res.smem = 32*sizeof(float); ++ return res; ++} + +- if (is_c4) { +- res.smem *= 4; ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_qk_scores( ++ ggml_metal_library_t lib) { ++ const char * name = "kernel_glm_compact_qk_scores_f16_dk576"; ++ 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); + } +- +- res.c4 = is_c4; +- + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_blk(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_ASSERT(op->op == GGML_OP_CUMSUM); +- +- char base[256]; +- char name[256]; +- +- snprintf(base, 256, "kernel_cumsum_blk_%s", ggml_type_name(op->src[0]->type)); +- snprintf(name, 256, "%s", base); +- ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_scores_v( ++ ggml_metal_library_t lib) { ++ const char * name = "kernel_glm_compact_scores_v_f16_dv512_nwg4"; + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { +- res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); + } +- + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_ASSERT(op->op == GGML_OP_CUMSUM); +- +- char base[256]; +- char name[256]; ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_softmax_prefix( ++ ggml_metal_library_t lib) { ++ const char * name = "kernel_glm_compact_softmax_prefix_f32_nwg4"; ++ 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; ++} + +- snprintf(base, 256, "kernel_cumsum_add_%s", ggml_type_name(op->src[0]->type)); +- snprintf(name, 256, "%s", base); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_chunk_v( ++ ggml_metal_library_t lib) { ++ const char * name = "kernel_glm_compact_chunk_v_f16_dv512"; ++ 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_glm_compact_chunk_fold( ++ ggml_metal_library_t lib) { ++ const char * name = "kernel_glm_compact_chunk_fold_f32_dv512_nwg4"; + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { +- res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr); + } ++ return res; ++} + ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_probs_v_sequential( ++ ggml_metal_library_t lib) { ++ const char * name = "kernel_glm_compact_probs_v_sequential_f16_dv512_nwg4"; ++ 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_tri(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_ASSERT(op->op == GGML_OP_TRI); +- GGML_ASSERT(op->src[0]->nb[0] == ggml_type_size(op->src[0]->type)); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_top1_attn(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_DSA_TOP1_ATTN); + + char base[256]; + char name[256]; + +- const char * op_str = "tri"; +- const int ttype = op->op_params[0]; +- +- snprintf(base, 256, "kernel_%s_%s_%d", op_str, ggml_type_name(op->src[0]->type), ttype); +- ++ snprintf(base, 256, "kernel_dsa_top1_attn_%s", ggml_type_name(op->src[1]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); +@@ -450,50 +643,46 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri(ggml_metal_l + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_ASSERT(!op->src[1] || op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32); +- ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag(ggml_metal_library_t lib, const ggml_tensor * op) { + char base[256]; + char name[256]; + +- const char * suffix = ""; +- +- if (op->src[0]->ne[0] % 4 == 0) { +- suffix = "_4"; +- } +- +- const ggml_type tsrc1 = op->src[1] ? op->src[1]->type : GGML_TYPE_F32; ++ const int n = op->src[0]->ne[0]; + +- snprintf(base, 256, "kernel_soft_max_%s%s", ggml_type_name(tsrc1), suffix); +- snprintf(name, 256, "%s", base); ++ snprintf(base, 256, "kernel_diag_%s", ggml_type_name(op->src[0]->type)); ++ snprintf(name, 256, "%s_n=%d", base, n); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + +- res.smem = 32*sizeof(float); ++ res.nsg = 1; ++ res.smem = 0; + + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); +- +- GGML_ASSERT(ggml_is_contiguous(op->src[0])); +- GGML_ASSERT(ggml_is_contiguous(op->src[1])); +- ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat(ggml_metal_library_t lib, ggml_type tsrc) { + char base[256]; + char name[256]; + +- const char * suffix = ""; ++ snprintf(base, 256, "kernel_repeat_%s", ggml_type_name(tsrc)); ++ snprintf(name, 256, "%s", base); + +- if (op->src[1]->ne[0] % 4 == 0) { +- suffix = "_4"; ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + +- snprintf(base, 256, "kernel_ssm_conv_%s_%s%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type), suffix); ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat(ggml_metal_library_t lib, ggml_type tsrc) { ++ char base[256]; ++ char name[256]; ++ ++ snprintf(base, 256, "kernel_concat_%s", ggml_type_name(tsrc)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); +@@ -504,92 +693,117 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv(ggml_me + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched(ggml_metal_library_t lib, const ggml_tensor * op, int ssm_conv_bs) { +- GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); +- +- GGML_ASSERT(ggml_is_contiguous(op->src[0])); +- GGML_ASSERT(ggml_is_contiguous(op->src[1])); +- ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary(ggml_metal_library_t lib, const ggml_tensor * op) { + char base[256]; + char name[256]; + +- const char * suffix = ""; +- if (op->src[1]->ne[0] % 4 == 0) { +- suffix = "_4"; +- } ++ int op_num = -1; + +- snprintf(base, 256, "kernel_ssm_conv_%s_%s_batched%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type), suffix); +- snprintf(name, 256, "%s_ssm_conv_bs=%d", base, ssm_conv_bs); ++ switch (op->op) { ++ case GGML_OP_SCALE: op_num = OP_UNARY_NUM_SCALE; break; ++ case GGML_OP_FILL: op_num = OP_UNARY_NUM_FILL; break; ++ case GGML_OP_CLAMP: op_num = OP_UNARY_NUM_CLAMP; break; ++ case GGML_OP_SQR: op_num = OP_UNARY_NUM_SQR; break; ++ case GGML_OP_SQRT: op_num = OP_UNARY_NUM_SQRT; break; ++ case GGML_OP_SIN: op_num = OP_UNARY_NUM_SIN; break; ++ case GGML_OP_COS: op_num = OP_UNARY_NUM_COS; break; ++ case GGML_OP_LOG: op_num = OP_UNARY_NUM_LOG; break; ++ case GGML_OP_LEAKY_RELU: op_num = OP_UNARY_NUM_LEAKY_RELU; break; ++ case GGML_OP_UNARY: ++ switch (ggml_get_unary_op(op)) { ++ case GGML_UNARY_OP_TANH: op_num = OP_UNARY_NUM_TANH; break; ++ case GGML_UNARY_OP_RELU: op_num = OP_UNARY_NUM_RELU; break; ++ case GGML_UNARY_OP_SIGMOID: op_num = OP_UNARY_NUM_SIGMOID; break; ++ case GGML_UNARY_OP_GELU: op_num = OP_UNARY_NUM_GELU; break; ++ case GGML_UNARY_OP_GELU_ERF: op_num = OP_UNARY_NUM_GELU_ERF; break; ++ case GGML_UNARY_OP_GELU_QUICK: op_num = OP_UNARY_NUM_GELU_QUICK; break; ++ case GGML_UNARY_OP_SILU: op_num = OP_UNARY_NUM_SILU; break; ++ case GGML_UNARY_OP_ELU: op_num = OP_UNARY_NUM_ELU; break; ++ case GGML_UNARY_OP_NEG: op_num = OP_UNARY_NUM_NEG; break; ++ case GGML_UNARY_OP_ABS: op_num = OP_UNARY_NUM_ABS; break; ++ case GGML_UNARY_OP_SGN: op_num = OP_UNARY_NUM_SGN; break; ++ case GGML_UNARY_OP_STEP: op_num = OP_UNARY_NUM_STEP; break; ++ case GGML_UNARY_OP_HARDSWISH: op_num = OP_UNARY_NUM_HARDSWISH; break; ++ case GGML_UNARY_OP_HARDSIGMOID: op_num = OP_UNARY_NUM_HARDSIGMOID; break; ++ case GGML_UNARY_OP_EXP: op_num = OP_UNARY_NUM_EXP; break; ++ case GGML_UNARY_OP_SOFTPLUS: op_num = OP_UNARY_NUM_SOFTPLUS; break; ++ case GGML_UNARY_OP_EXPM1: op_num = OP_UNARY_NUM_EXPM1; break; ++ case GGML_UNARY_OP_FLOOR: op_num = OP_UNARY_NUM_FLOOR; break; ++ case GGML_UNARY_OP_CEIL: op_num = OP_UNARY_NUM_CEIL; break; ++ case GGML_UNARY_OP_ROUND: op_num = OP_UNARY_NUM_ROUND; break; ++ case GGML_UNARY_OP_TRUNC: op_num = OP_UNARY_NUM_TRUNC; break; ++ case GGML_UNARY_OP_XIELU: op_num = OP_UNARY_NUM_XIELU; break; ++ default: GGML_ABORT("fatal error"); ++ } break; ++ default: GGML_ABORT("fatal error"); ++ }; ++ ++ const char * t0_str = ggml_type_name(op->src[0]->type); ++ const char * t_str = ggml_type_name(op->type); ++ ++ const bool is_c4 = op->src[0]->ne[0] % 4 == 0; ++ const bool is_cnt = ggml_is_contiguous(op->src[0]) && ggml_nelements(op) < 32768; ++ ++ snprintf(base, 256, "kernel_unary_%s_%s%s", t0_str, t_str, is_c4 ? "_4" : ""); ++ snprintf(name, 256, "%s_op=%d_cnt=%d", base, op_num, is_cnt); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + ggml_metal_cv_t cv = ggml_metal_cv_init(); + +- ggml_metal_cv_set_int16(cv, ssm_conv_bs, FC_SSM_CONV + 0); ++ ggml_metal_cv_set_int16(cv, op_num, FC_UNARY + 0); ++ ggml_metal_cv_set_bool (cv, is_cnt, FC_UNARY + 1); + + res = ggml_metal_library_compile_pipeline(lib, base, name, cv); + + ggml_metal_cv_free(cv); + } + ++ res.c4 = is_c4; ++ res.cnt = is_cnt; ++ + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(ggml_is_contiguous_1(op->src[0])); + + char base[256]; + char name[256]; + +- const int nsg = (ne00 + 31)/32; ++ const char * op_str = "undefined"; ++ switch (op->op) { ++ case GGML_OP_GLU: ++ switch (ggml_get_glu_op(op)) { ++ case GGML_GLU_OP_REGLU: op_str = "reglu"; break; ++ case GGML_GLU_OP_GEGLU: op_str = "geglu"; break; ++ case GGML_GLU_OP_SWIGLU: op_str = "swiglu"; break; ++ case GGML_GLU_OP_SWIGLU_OAI: op_str = "swiglu_oai"; break; ++ case GGML_GLU_OP_GEGLU_ERF: op_str = "geglu_erf"; break; ++ case GGML_GLU_OP_GEGLU_QUICK: op_str = "geglu_quick"; break; ++ default: GGML_ABORT("fatal error"); ++ } break; ++ default: GGML_ABORT("fatal error"); ++ }; + +- snprintf(base, 256, "kernel_ssm_scan_%s", ggml_type_name(op->src[0]->type)); +- snprintf(name, 256, "%s_nsg=%d", base, nsg); ++ snprintf(base, 256, "kernel_%s_%s", op_str, ggml_type_name(op->src[0]->type)); ++ snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + +- // Shared memory layout: +- // - sgptg * NW floats for partial sums (nsg * 32) +- // - sgptg floats for shared_x_dt (nsg) +- // - sgptg floats for shared_dA (nsg) +- // Total: nsg * (32 + 2) floats +- res.smem = (32 + 2)*sizeof(float)*nsg; +- + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_library_t lib, const ggml_tensor * op) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_SUM); ++ + char base[256]; + char name[256]; + +- const int64_t C = op->ne[0]; +- const int64_t H = op->src[0]->ne[1]; +- +- switch (op->op) { +- case GGML_OP_RWKV_WKV6: +- { +- GGML_ASSERT(op->src[5]->type == GGML_TYPE_F32); +- GGML_ASSERT(C % H == 0); +- GGML_ASSERT(C / H == 64); +- +- snprintf(base, 256, "kernel_rwkv_wkv6_%s", ggml_type_name(op->src[0]->type)); +- } break; +- case GGML_OP_RWKV_WKV7: +- { +- GGML_ASSERT(op->src[6]->type == GGML_TYPE_F32); +- GGML_ASSERT(C % H == 0); +- GGML_ASSERT(C / H == 64); +- +- snprintf(base, 256, "kernel_rwkv_wkv7_%s", ggml_type_name(op->src[0]->type)); +- } break; +- default: +- GGML_ABORT("fatal error"); +- } +- ++ snprintf(base, 256, "kernel_op_sum_%s", ggml_type_name(op->src[0]->type)); + snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); +@@ -600,562 +814,2266 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_ + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net(ggml_metal_library_t lib, const ggml_tensor * op) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ + char base[256]; + char name[256]; + +- // v is src[2], dimensions: S_v = ne[0], H = ne[1] +- const int ne20 = op->src[2]->ne[0]; // S_v +- const int ne21 = op->src[2]->ne[1]; // H +- const int ne30 = op->src[3]->ne[0]; // G +- // state is src[5], 4D [S_v, S_v, H_v, n_seqs] (s0 only); K is op param 0. +- const int K = ggml_get_op_params_i32(op, 0); ++ int op_num = -1; + +- const int nsg = op->src[2]->ne[0]/32; ++ switch (op->op) { ++ case GGML_OP_SUM_ROWS: op_num = OP_SUM_ROWS_NUM_SUM_ROWS; break; ++ case GGML_OP_MEAN: op_num = OP_SUM_ROWS_NUM_MEAN; break; ++ default: GGML_ABORT("fatal error"); ++ }; + +- GGML_ASSERT(op->src[5]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->ne[0] == ne20 * ne21); +- GGML_ASSERT(ne20 % 32 == 0); ++ const char * t0_str = ggml_type_name(op->src[0]->type); ++ const char * t_str = ggml_type_name(op->type); + +- snprintf(base, 256, "kernel_gated_delta_net_%s_%d", ggml_type_name(op->src[0]->type), nsg); +- snprintf(name, 256, "%s_ne20=%d_ne30=%d_K=%d", base, ne20, ne30, K); ++ const bool is_c4 = op->src[0]->ne[0] % 4 == 0; ++ ++ snprintf(base, 256, "kernel_sum_rows_%s_%s%s", t0_str, t_str, is_c4 ? "_4" : ""); ++ snprintf(name, 256, "%s_op=%d", base, op_num); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + ggml_metal_cv_t cv = ggml_metal_cv_init(); + +- ggml_metal_cv_set_int16(cv, ne20, FC_GATED_DELTA_NET + 0); +- ggml_metal_cv_set_int16(cv, ne30, FC_GATED_DELTA_NET + 1); +- ggml_metal_cv_set_int16(cv, K, FC_GATED_DELTA_NET + 2); ++ ggml_metal_cv_set_int16(cv, op_num, FC_SUM_ROWS + 0); + + res = ggml_metal_library_compile_pipeline(lib, base, name, cv); + + ggml_metal_cv_free(cv); + } + +- res.nsg = nsg; ++ res.smem = 32*sizeof(float); ++ ++ if (is_c4) { ++ res.smem *= 4; ++ } ++ ++ res.c4 = is_c4; + + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri(ggml_metal_library_t lib, const ggml_tensor * op) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_blk(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_CUMSUM); ++ + char base[256]; + char name[256]; + +- const int nsg = 8; +- const int n = op->src[1]->ne[1]; +- const int k = op->src[1]->ne[0]; +- +- snprintf(base, 256, "kernel_solve_tri_%s", ggml_type_name(op->src[0]->type)); +- snprintf(name, 256, "%s_nsg=%d_n=%d_k=%d", base, nsg, n, k); ++ snprintf(base, 256, "kernel_cumsum_blk_%s", ggml_type_name(op->src[0]->type)); ++ snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { +- ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } + +- ggml_metal_cv_set_int16(cv, nsg, FC_SOLVE_TRI + 0); +- ggml_metal_cv_set_int16(cv, n, FC_SOLVE_TRI + 1); +- ggml_metal_cv_set_int16(cv, k, FC_SOLVE_TRI + 2); ++ return res; ++} + +- res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_CUMSUM); + +- ggml_metal_cv_free(cv); +- } ++ char base[256]; ++ char name[256]; + +- res.nsg = nsg; +- res.smem = GGML_PAD(GGML_PAD(n, 32)*nsg*sizeof(float), 16); ++ snprintf(base, 256, "kernel_cumsum_add_%s", ggml_type_name(op->src[0]->type)); ++ snprintf(name, 256, "%s", base); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } + + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext(ggml_metal_library_t lib, const ggml_tensor * op, int nsg, int nxpsg, int r1ptg) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_TRI); ++ GGML_ASSERT(op->src[0]->nb[0] == ggml_type_size(op->src[0]->type)); ++ + char base[256]; + char name[256]; + +- const ggml_type tsrc0 = op->src[0]->type; +- const ggml_type tsrc1 = op->src[1]->type; +- const int ne12 = op->src[1]->ne[2]; +- const int r2 = ne12 / op->src[0]->ne[2]; +- const int r3 = op->src[1]->ne[3] / op->src[0]->ne[3]; ++ const char * op_str = "tri"; ++ const int ttype = op->op_params[0]; + +- GGML_ASSERT(ne12 <= INT16_MAX && r2 <= INT16_MAX && r3 <= INT16_MAX); ++ snprintf(base, 256, "kernel_%s_%s_%d", op_str, ggml_type_name(op->src[0]->type), ttype); + +- snprintf(base, 256, "kernel_mul_mv_ext_%s_%s_r1_%d", ggml_type_name(tsrc0), ggml_type_name(tsrc1), r1ptg); +- snprintf(name, 256, "%s_nsg=%d_nxpsg=%d_ne12=%d_r2=%d_r3=%d", base, nsg, nxpsg, ne12, r2, r3); ++ snprintf(name, 256, "%s", base); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { +- ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } + +- ggml_metal_cv_set_int16(cv, nsg, FC_MUL_MV + 0); +- ggml_metal_cv_set_int16(cv, nxpsg, FC_MUL_MV + 1); +- ggml_metal_cv_set_int16(cv, (int16_t) ne12, FC_MUL_MV + 2); +- ggml_metal_cv_set_int16(cv, (int16_t) r2, FC_MUL_MV + 3); +- ggml_metal_cv_set_int16(cv, (int16_t) r3, FC_MUL_MV + 4); ++ return res; ++} + +- res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(!op->src[1] || op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32); + +- ggml_metal_cv_free(cv); ++ char base[256]; ++ char name[256]; ++ ++ const char * suffix = ""; ++ ++ if (op->src[0]->ne[0] % 4 == 0) { ++ suffix = "_4"; ++ } ++ ++ const ggml_type tsrc1 = op->src[1] ? op->src[1]->type : GGML_TYPE_F32; ++ ++ snprintf(base, 256, "kernel_soft_max_%s%s", ggml_type_name(tsrc1), suffix); ++ snprintf(name, 256, "%s", base); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + ++ res.smem = 32*sizeof(float); ++ + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_metal_library_t lib, const ggml_tensor * op) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ ++ GGML_ASSERT(ggml_is_contiguous(op->src[0])); ++ GGML_ASSERT(ggml_is_contiguous(op->src[1])); ++ + char base[256]; + char name[256]; + +- const ggml_type tsrc0 = op->src[0]->type; +- const ggml_type tsrc1 = op->src[1]->type; ++ const char * suffix = ""; + +- const bool bc_inp = op->src[0]->ne[0] % 32 != 0; ++ if (op->src[1]->ne[0] % 4 == 0) { ++ suffix = "_4"; ++ } + +- constexpr int NRA = SZ_SIMDGROUP * N_MM_BLOCK_Y * N_MM_SIMD_GROUP_Y; +- constexpr int NRB = SZ_SIMDGROUP * N_MM_BLOCK_X * N_MM_SIMD_GROUP_X; ++ snprintf(base, 256, "kernel_ssm_conv_%s_%s%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type), suffix); ++ snprintf(name, 256, "%s", base); + +- const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(lib))->has_tensor; ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } + +- const bool bc_out = has_tensor +- ? (op->ne[0] % NRA != 0 || op->ne[1] % NRB != 0) +- : (op->ne[0] % 64 != 0 || op->ne[1] % 32 != 0); ++ return res; ++} + +- GGML_ASSERT(op->src[1]->ne[2] <= INT16_MAX && op->src[1]->ne[3] <= INT16_MAX); +- const int16_t ne12 = (int16_t) op->src[1]->ne[2]; +- const int16_t ne13 = (int16_t) op->src[1]->ne[3]; +- const int16_t r2 = (int16_t) (ne12 / op->src[0]->ne[2]); +- const int16_t r3 = (int16_t) (ne13 / op->src[0]->ne[3]); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched(ggml_metal_library_t lib, const ggml_tensor * op, int ssm_conv_bs) { ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); + +- snprintf(base, 256, "kernel_mul_mm_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1)); +- snprintf(name, 256, "%s_bci=%d_bco=%d_ne12=%d_ne13=%d_r2=%d_r3=%d", +- base, bc_inp, bc_out, ne12, ne13, r2, r3); ++ GGML_ASSERT(ggml_is_contiguous(op->src[0])); ++ GGML_ASSERT(ggml_is_contiguous(op->src[1])); ++ ++ char base[256]; ++ char name[256]; ++ ++ const char * suffix = ""; ++ if (op->src[1]->ne[0] % 4 == 0) { ++ suffix = "_4"; ++ } ++ ++ snprintf(base, 256, "kernel_ssm_conv_%s_%s_batched%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type), suffix); ++ snprintf(name, 256, "%s_ssm_conv_bs=%d", base, ssm_conv_bs); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + ggml_metal_cv_t cv = ggml_metal_cv_init(); + +- ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0); +- ggml_metal_cv_set_bool(cv, bc_out, FC_MUL_MM + 1); +- ggml_metal_cv_set_int16(cv, ne12, FC_MUL_MM + 2); +- ggml_metal_cv_set_int16(cv, ne13, FC_MUL_MM + 3); +- ggml_metal_cv_set_int16(cv, r2, FC_MUL_MM + 4); +- ggml_metal_cv_set_int16(cv, r3, FC_MUL_MM + 5); ++ ggml_metal_cv_set_int16(cv, ssm_conv_bs, FC_SSM_CONV + 0); + + res = ggml_metal_library_compile_pipeline(lib, base, name, cv); + + ggml_metal_cv_free(cv); + } + +- if (has_tensor) { +- res.nr0 = NRA; +- res.nr1 = NRB; +- +- const size_t smem_a = NRA * N_MM_NK_TOTAL * sizeof(ggml_fp16_t); +- res.smem = smem_a; +- } else { +- res.nr0 = 64; +- res.nr1 = 32; +- +- res.smem = bc_out ? 8192 : (4096 + 2048); +- } +- +- res.nsg = N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y; +- + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_metal_library_t lib, const ggml_tensor * op) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op) { + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + + char base[256]; + char name[256]; + +- int nsg = 0; // number of simdgroups +- int nr0 = 0; // number of src0 rows per simdgroup +- int nr1 = 1; // number of src1 rows per threadgroup ++ const int nsg = (ne00 + 31)/32; + +- size_t smem = 0; // shared memory ++ snprintf(base, 256, "kernel_ssm_scan_%s", ggml_type_name(op->src[0]->type)); ++ snprintf(name, 256, "%s_nsg=%d", base, nsg); + +- const ggml_type tsrc0 = op->src[0]->type; +- const ggml_type tsrc1 = 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, base, name, nullptr); ++ } + +- const char * suffix = ""; ++ // Shared memory layout: ++ // - sgptg * NW floats for partial sums (nsg * 32) ++ // - sgptg floats for shared_x_dt (nsg) ++ // - sgptg floats for shared_dA (nsg) ++ // Total: nsg * (32 + 2) floats ++ res.smem = (32 + 2)*sizeof(float)*nsg; + +- // use custom matrix x vector kernel +- switch (tsrc0) { +- case GGML_TYPE_F32: +- case GGML_TYPE_F16: +- case GGML_TYPE_BF16: +- { +- if (ne00 < 32) { +- nsg = 1; +- nr0 = 32; +- nr1 = 1; +- suffix = "_short"; +- } else { +- nsg = std::min(4, (ne00 + 127) / 128); +- nr0 = 2; +- nr1 = 1; +- smem = 32*sizeof(float)*nr0; +- suffix = ne00 % 4 == 0 ? "_4" : ""; +- } +- } break; +- case GGML_TYPE_Q1_0: +- { +- nsg = N_SG_Q1_0; +- nr0 = N_R0_Q1_0; +- } break; +- case GGML_TYPE_Q2_0: +- { +- nsg = N_SG_Q2_0; +- nr0 = N_R0_Q2_0; +- } break; +- case GGML_TYPE_Q4_0: +- { +- nsg = N_SG_Q4_0; +- nr0 = N_R0_Q4_0; +- } break; +- case GGML_TYPE_Q4_1: +- { +- nsg = N_SG_Q4_1; +- nr0 = N_R0_Q4_1; +- } break; +- case GGML_TYPE_Q5_0: +- { +- nsg = N_SG_Q5_0; +- nr0 = N_R0_Q5_0; +- } break; +- case GGML_TYPE_Q5_1: +- { +- nsg = N_SG_Q5_1; +- nr0 = N_R0_Q5_1; +- } break; +- case GGML_TYPE_Q8_0: +- { +- nsg = N_SG_Q8_0; +- nr0 = N_R0_Q8_0; +- smem = 32*sizeof(float)*N_R0_Q8_0; +- } break; +- case GGML_TYPE_MXFP4: +- { +- nsg = N_SG_MXFP4; +- nr0 = N_R0_MXFP4; +- smem = 32*sizeof(float); +- } break; +- case GGML_TYPE_Q2_K: ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_library_t lib, const ggml_tensor * op) { ++ char base[256]; ++ char name[256]; ++ ++ const int64_t C = op->ne[0]; ++ const int64_t H = op->src[0]->ne[1]; ++ ++ switch (op->op) { ++ case GGML_OP_RWKV_WKV6: + { +- nsg = N_SG_Q2_K; +- nr0 = N_R0_Q2_K; ++ GGML_ASSERT(op->src[5]->type == GGML_TYPE_F32); ++ GGML_ASSERT(C % H == 0); ++ GGML_ASSERT(C / H == 64); ++ ++ snprintf(base, 256, "kernel_rwkv_wkv6_%s", ggml_type_name(op->src[0]->type)); + } break; +- case GGML_TYPE_Q3_K: ++ case GGML_OP_RWKV_WKV7: + { +- nsg = N_SG_Q3_K; +- nr0 = N_R0_Q3_K; +- } break; +- case GGML_TYPE_Q4_K: +- { +- nsg = N_SG_Q4_K; +- nr0 = N_R0_Q4_K; +- } break; +- case GGML_TYPE_Q5_K: +- { +- nsg = N_SG_Q5_K; +- nr0 = N_R0_Q5_K; +- } break; +- case GGML_TYPE_Q6_K: +- { +- nsg = N_SG_Q6_K; +- nr0 = N_R0_Q6_K; +- } break; +- case GGML_TYPE_IQ2_XXS: +- { +- nsg = N_SG_IQ2_XXS; +- nr0 = N_R0_IQ2_XXS; +- smem = 256*8+128; +- } break; +- case GGML_TYPE_IQ2_XS: +- { +- nsg = N_SG_IQ2_XS; +- nr0 = N_R0_IQ2_XS; +- smem = 512*8+128; +- } break; +- case GGML_TYPE_IQ3_XXS: +- { +- nsg = N_SG_IQ3_XXS; +- nr0 = N_R0_IQ3_XXS; +- smem = 256*4+128; +- } break; +- case GGML_TYPE_IQ3_S: +- { +- nsg = N_SG_IQ3_S; +- nr0 = N_R0_IQ3_S; +- smem = 512*4; +- } break; +- case GGML_TYPE_IQ2_S: +- { +- nsg = N_SG_IQ2_S; +- nr0 = N_R0_IQ2_S; +- } break; +- case GGML_TYPE_IQ1_S: +- { +- nsg = N_SG_IQ1_S; +- nr0 = N_R0_IQ1_S; +- } break; +- case GGML_TYPE_IQ1_M: +- { +- nsg = N_SG_IQ1_M; +- nr0 = N_R0_IQ1_M; +- } break; +- case GGML_TYPE_IQ4_NL: +- { +- nsg = N_SG_IQ4_NL; +- nr0 = N_R0_IQ4_NL; +- smem = 32*sizeof(float); +- } break; +- case GGML_TYPE_IQ4_XS: +- { +- nsg = N_SG_IQ4_XS; +- nr0 = N_R0_IQ4_XS; +- smem = 32*sizeof(float); ++ GGML_ASSERT(op->src[6]->type == GGML_TYPE_F32); ++ GGML_ASSERT(C % H == 0); ++ GGML_ASSERT(C / H == 64); ++ ++ snprintf(base, 256, "kernel_rwkv_wkv7_%s", ggml_type_name(op->src[0]->type)); + } break; + default: +- { +- GGML_LOG_ERROR("Asserting on type %d\n", (int) tsrc0); +- GGML_ABORT("not implemented"); +- } +- }; ++ GGML_ABORT("fatal error"); ++ } + +- GGML_ASSERT(ne12 <= INT16_MAX && ne13 <= INT16_MAX); +- const int16_t r2 = (int16_t) (ne12 / ne02); +- const int16_t r3 = (int16_t) (ne13 / ne03); ++ snprintf(name, 256, "%s", base); + +- snprintf(base, 256, "kernel_mul_mv_%s_%s%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1), suffix); +- snprintf(name, 256, "%s_nsg=%d_ne12=%d_r2=%d_r3=%d", base, nsg, ne12, r2, r3); ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net(ggml_metal_library_t lib, const ggml_tensor * op) { ++ char base[256]; ++ char name[256]; ++ ++ // v is src[2], dimensions: S_v = ne[0], H = ne[1] ++ const int ne20 = op->src[2]->ne[0]; // S_v ++ const int ne21 = op->src[2]->ne[1]; // H ++ const int ne30 = op->src[3]->ne[0]; // G ++ // state is src[5], 4D [S_v, S_v, H_v, n_seqs] (s0 only); K is op param 0. ++ const int K = ggml_get_op_params_i32(op, 0); ++ ++ const int nsg = op->src[2]->ne[0]/32; ++ ++ GGML_ASSERT(op->src[5]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->ne[0] == ne20 * ne21); ++ GGML_ASSERT(ne20 % 32 == 0); ++ ++ snprintf(base, 256, "kernel_gated_delta_net_%s_%d", ggml_type_name(op->src[0]->type), nsg); ++ snprintf(name, 256, "%s_ne20=%d_ne30=%d_K=%d", base, ne20, ne30, K); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int16(cv, ne20, FC_GATED_DELTA_NET + 0); ++ ggml_metal_cv_set_int16(cv, ne30, FC_GATED_DELTA_NET + 1); ++ ggml_metal_cv_set_int16(cv, K, FC_GATED_DELTA_NET + 2); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ res.nsg = nsg; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer( ++ ggml_metal_library_t lib, ++ const ggml_tensor * op, ++ bool parallel, ++ bool staged_q) { ++ assert(op->op == GGML_OP_LIGHTNING_INDEXER); ++ assert(!parallel || !staged_q); ++ ++ char base[256]; ++ char name[256]; ++ ++ snprintf(base, 256, "kernel_lightning_indexer_%s%s", ++ parallel ? "parallel_" : staged_q ? "staged_q_" : "", ++ ggml_type_name(op->src[1]->type)); ++ snprintf(name, 256, "%s", base); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri(ggml_metal_library_t lib, const ggml_tensor * op) { ++ char base[256]; ++ char name[256]; ++ ++ const int nsg = 8; ++ const int n = op->src[1]->ne[1]; ++ const int k = op->src[1]->ne[0]; ++ ++ snprintf(base, 256, "kernel_solve_tri_%s", ggml_type_name(op->src[0]->type)); ++ snprintf(name, 256, "%s_nsg=%d_n=%d_k=%d", base, nsg, n, k); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int16(cv, nsg, FC_SOLVE_TRI + 0); ++ ggml_metal_cv_set_int16(cv, n, FC_SOLVE_TRI + 1); ++ ggml_metal_cv_set_int16(cv, k, FC_SOLVE_TRI + 2); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ res.nsg = nsg; ++ res.smem = GGML_PAD(GGML_PAD(n, 32)*nsg*sizeof(float), 16); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext(ggml_metal_library_t lib, const ggml_tensor * op, int nsg, int nxpsg, int r1ptg) { ++ char base[256]; ++ char name[256]; ++ ++ const ggml_type tsrc0 = op->src[0]->type; ++ const ggml_type tsrc1 = op->src[1]->type; ++ const int ne12 = op->src[1]->ne[2]; ++ const int r2 = ne12 / op->src[0]->ne[2]; ++ const int r3 = op->src[1]->ne[3] / op->src[0]->ne[3]; ++ ++ GGML_ASSERT(ne12 <= INT16_MAX && r2 <= INT16_MAX && r3 <= INT16_MAX); ++ ++ snprintf(base, 256, "kernel_mul_mv_ext_%s_%s_r1_%d", ggml_type_name(tsrc0), ggml_type_name(tsrc1), r1ptg); ++ snprintf(name, 256, "%s_nsg=%d_nxpsg=%d_ne12=%d_r2=%d_r3=%d", base, nsg, nxpsg, ne12, r2, r3); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + ggml_metal_cv_t cv = ggml_metal_cv_init(); + + ggml_metal_cv_set_int16(cv, nsg, FC_MUL_MV + 0); ++ ggml_metal_cv_set_int16(cv, nxpsg, FC_MUL_MV + 1); + ggml_metal_cv_set_int16(cv, (int16_t) ne12, FC_MUL_MV + 2); +- ggml_metal_cv_set_int16(cv, r2, FC_MUL_MV + 3); +- ggml_metal_cv_set_int16(cv, r3, FC_MUL_MV + 4); ++ ggml_metal_cv_set_int16(cv, (int16_t) r2, FC_MUL_MV + 3); ++ ggml_metal_cv_set_int16(cv, (int16_t) r3, FC_MUL_MV + 4); + + res = ggml_metal_library_compile_pipeline(lib, base, name, cv); + + ggml_metal_cv_free(cv); + } + +- res.nr0 = nr0; +- res.nr1 = nr1; +- res.nsg = nsg; +- res.smem = smem; ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_metal_library_t lib, const ggml_tensor * op) { ++ char base[256]; ++ char name[256]; ++ ++ const ggml_type tsrc0 = op->src[0]->type; ++ const ggml_type tsrc1 = op->src[1]->type; ++ ++ const bool bc_inp = op->src[0]->ne[0] % 32 != 0; ++ ++ constexpr int NRA = SZ_SIMDGROUP * N_MM_BLOCK_Y * N_MM_SIMD_GROUP_Y; ++ constexpr int NRB = SZ_SIMDGROUP * N_MM_BLOCK_X * N_MM_SIMD_GROUP_X; ++ ++ const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(lib))->has_tensor; ++ ++ const bool bc_out = has_tensor ++ ? (op->ne[0] % NRA != 0 || op->ne[1] % NRB != 0) ++ : (op->ne[0] % 64 != 0 || op->ne[1] % 32 != 0); ++ ++ GGML_ASSERT(op->src[1]->ne[2] <= INT16_MAX && op->src[1]->ne[3] <= INT16_MAX); ++ const int16_t ne12 = (int16_t) op->src[1]->ne[2]; ++ const int16_t ne13 = (int16_t) op->src[1]->ne[3]; ++ const int16_t r2 = (int16_t) (ne12 / op->src[0]->ne[2]); ++ const int16_t r3 = (int16_t) (ne13 / op->src[0]->ne[3]); ++ ++ snprintf(base, 256, "kernel_mul_mm_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1)); ++ snprintf(name, 256, "%s_bci=%d_bco=%d_ne12=%d_ne13=%d_r2=%d_r3=%d", ++ base, bc_inp, bc_out, ne12, ne13, r2, r3); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0); ++ ggml_metal_cv_set_bool(cv, bc_out, FC_MUL_MM + 1); ++ ggml_metal_cv_set_int16(cv, ne12, FC_MUL_MM + 2); ++ ggml_metal_cv_set_int16(cv, ne13, FC_MUL_MM + 3); ++ ggml_metal_cv_set_int16(cv, r2, FC_MUL_MM + 4); ++ ggml_metal_cv_set_int16(cv, r3, FC_MUL_MM + 5); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ if (has_tensor) { ++ res.nr0 = NRA; ++ res.nr1 = NRB; ++ ++ const size_t smem_a = NRA * N_MM_NK_TOTAL * sizeof(ggml_fp16_t); ++ res.smem = smem_a; ++ } else { ++ res.nr0 = 64; ++ res.nr1 = 32; ++ ++ res.smem = bc_out ? 8192 : (4096 + 2048); ++ } ++ ++ res.nsg = N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y; + + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(ggml_metal_library_t lib, int ne02, int ne20) { ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ ++ char base[256]; ++ char name[256]; ++ ++ int nsg = 0; // number of simdgroups ++ int nr0 = 0; // number of src0 rows per simdgroup ++ int nr1 = 1; // number of src1 rows per threadgroup ++ ++ size_t smem = 0; // shared memory ++ ++ const ggml_type tsrc0 = op->src[0]->type; ++ const ggml_type tsrc1 = op->src[1]->type; ++ ++ const char * suffix = ""; ++ ++ // use custom matrix x vector kernel ++ switch (tsrc0) { ++ case GGML_TYPE_F32: ++ case GGML_TYPE_F16: ++ case GGML_TYPE_BF16: ++ { ++ if (ne00 < 32) { ++ nsg = 1; ++ nr0 = 32; ++ nr1 = 1; ++ suffix = "_short"; ++ } else { ++ nsg = std::min(4, (ne00 + 127) / 128); ++ nr0 = 2; ++ nr1 = 1; ++ smem = 32*sizeof(float)*nr0; ++ suffix = ne00 % 4 == 0 ? "_4" : ""; ++ } ++ } break; ++ case GGML_TYPE_Q1_0: ++ { ++ nsg = N_SG_Q1_0; ++ nr0 = N_R0_Q1_0; ++ } break; ++ case GGML_TYPE_TQ2_0: ++ { ++ nsg = N_SG_TQ2_0; ++ nr0 = N_R0_TQ2_0; ++ } break; ++ case GGML_TYPE_Q4_0: ++ { ++ nsg = ggml_metal_glm_dsa_mul_mv_nsg_requested( ++ "LLAMA_GLM_DSA_MUL_MV_Q4_0_NSG", N_SG_Q4_0); ++ nr0 = N_R0_Q4_0; ++ } break; ++ case GGML_TYPE_Q4_1: ++ { ++ nsg = N_SG_Q4_1; ++ nr0 = N_R0_Q4_1; ++ } break; ++ case GGML_TYPE_Q5_0: ++ { ++ nsg = N_SG_Q5_0; ++ nr0 = N_R0_Q5_0; ++ } break; ++ case GGML_TYPE_Q5_1: ++ { ++ nsg = N_SG_Q5_1; ++ nr0 = N_R0_Q5_1; ++ } break; ++ case GGML_TYPE_Q8_0: ++ { ++ const bool row_parallel = ggml_metal_glm_dsa_q8_0_mul_mv_row_parallel_enabled(); ++ const bool vector_dot = !row_parallel && ggml_metal_glm_dsa_q8_0_mul_mv_vector_dot_enabled(); ++ const int requested_nr0 = ggml_metal_glm_dsa_q8_0_mul_mv_nr0_requested(); ++ const int default_nsg = (ggml_metal_glm_dsa_mul_mv_shape_policy_mask() & 1) != 0 && ++ ggml_metal_glm_dsa_q8_0_attention_shape(op) ? 2 : N_SG_Q8_0; ++ nsg = ggml_metal_glm_dsa_mul_mv_nsg_requested( ++ "LLAMA_GLM_DSA_MUL_MV_Q8_0_NSG", row_parallel ? 2 : default_nsg); ++ nr0 = row_parallel ? 4*nsg : requested_nr0; ++ smem = row_parallel ? 0 : 32*sizeof(float)*requested_nr0; ++ suffix = row_parallel ? "_row_parallel" : ++ (requested_nr0 == 1 ? "_r1" : ++ requested_nr0 == 4 ? "_r4" : ++ requested_nr0 == 8 ? "_r8" : ++ vector_dot ? "_vector_dot" : ""); ++ } break; ++ case GGML_TYPE_MXFP4: ++ { ++ nsg = N_SG_MXFP4; ++ nr0 = N_R0_MXFP4; ++ smem = 32*sizeof(float); ++ } break; ++ case GGML_TYPE_Q2_K: ++ { ++ nsg = ggml_metal_glm_dsa_mul_mv_nsg_requested( ++ "LLAMA_GLM_DSA_MUL_MV_Q2_K_NSG", N_SG_Q2_K); ++ nr0 = N_R0_Q2_K; ++ } break; ++ case GGML_TYPE_Q3_K: ++ { ++ const int default_nsg = (ggml_metal_glm_dsa_mul_mv_shape_policy_mask() & 2) != 0 && ++ ggml_metal_glm_dsa_q3_k_attention_output_shape(op) ? 8 : N_SG_Q3_K; ++ nsg = ggml_metal_glm_dsa_mul_mv_nsg_requested( ++ "LLAMA_GLM_DSA_MUL_MV_Q3_K_NSG", default_nsg); ++ nr0 = N_R0_Q3_K; ++ } break; ++ case GGML_TYPE_Q4_K: ++ { ++ const int default_nsg = (ggml_metal_glm_dsa_mul_mv_shape_policy_mask() & 4) != 0 && ++ ggml_metal_glm_dsa_q4_k_gate_up_shape(op) ? 1 : N_SG_Q4_K; ++ nsg = ggml_metal_glm_dsa_mul_mv_nsg_requested( ++ "LLAMA_GLM_DSA_MUL_MV_Q4_K_NSG", default_nsg); ++ nr0 = N_R0_Q4_K; ++ } break; ++ case GGML_TYPE_Q5_K: ++ { ++ nsg = N_SG_Q5_K; ++ nr0 = N_R0_Q5_K; ++ } break; ++ case GGML_TYPE_Q6_K: ++ { ++ nsg = N_SG_Q6_K; ++ nr0 = N_R0_Q6_K; ++ } break; ++ case GGML_TYPE_IQ2_XXS: ++ { ++ nsg = N_SG_IQ2_XXS; ++ nr0 = N_R0_IQ2_XXS; ++ smem = 256*8+128; ++ } break; ++ case GGML_TYPE_IQ2_XS: ++ { ++ nsg = N_SG_IQ2_XS; ++ nr0 = N_R0_IQ2_XS; ++ smem = 512*8+128; ++ } break; ++ case GGML_TYPE_IQ3_XXS: ++ { ++ nsg = N_SG_IQ3_XXS; ++ nr0 = N_R0_IQ3_XXS; ++ smem = 256*4+128; ++ } break; ++ case GGML_TYPE_IQ3_S: ++ { ++ nsg = N_SG_IQ3_S; ++ nr0 = N_R0_IQ3_S; ++ smem = 512*4; ++ } break; ++ case GGML_TYPE_IQ2_S: ++ { ++ nsg = N_SG_IQ2_S; ++ nr0 = N_R0_IQ2_S; ++ } break; ++ case GGML_TYPE_IQ1_S: ++ { ++ nsg = N_SG_IQ1_S; ++ nr0 = N_R0_IQ1_S; ++ } break; ++ case GGML_TYPE_IQ1_M: ++ { ++ nsg = N_SG_IQ1_M; ++ nr0 = N_R0_IQ1_M; ++ } break; ++ case GGML_TYPE_IQ4_NL: ++ { ++ nsg = N_SG_IQ4_NL; ++ nr0 = N_R0_IQ4_NL; ++ smem = 32*sizeof(float); ++ } break; ++ case GGML_TYPE_IQ4_XS: ++ { ++ nsg = N_SG_IQ4_XS; ++ nr0 = N_R0_IQ4_XS; ++ smem = 32*sizeof(float); ++ } break; ++ default: ++ { ++ GGML_LOG_ERROR("Asserting on type %d\n", (int) tsrc0); ++ GGML_ABORT("not implemented"); ++ } ++ }; ++ ++ GGML_ASSERT(ne12 <= INT16_MAX && ne13 <= INT16_MAX); ++ const int16_t r2 = (int16_t) (ne12 / ne02); ++ const int16_t r3 = (int16_t) (ne13 / ne03); ++ ++ snprintf(base, 256, "kernel_mul_mv_%s_%s%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1), suffix); ++ snprintf(name, 256, "%s_nsg=%d_ne12=%d_r2=%d_r3=%d", base, nsg, ne12, r2, r3); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int16(cv, nsg, FC_MUL_MV + 0); ++ ggml_metal_cv_set_int16(cv, (int16_t) ne12, FC_MUL_MV + 2); ++ ggml_metal_cv_set_int16(cv, r2, FC_MUL_MV + 3); ++ ggml_metal_cv_set_int16(cv, r3, FC_MUL_MV + 4); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = nr1; ++ res.nsg = nsg; ++ res.smem = smem; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(ggml_metal_library_t lib, int ne02, int ne20) { ++ char base[256]; ++ char name[256]; ++ ++ snprintf(base, 256, "kernel_mul_mm_id_map0_ne20_%d", ne20); ++ snprintf(name, 256, "%s_ne02=%d", base, ne02); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.smem = (size_t) ne02*ne20*sizeof(uint16_t); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_metal_library_t lib, const ggml_tensor * op) { ++ char base[256]; ++ char name[256]; ++ ++ const ggml_type tsrc0 = op->src[0]->type; ++ const ggml_type tsrc1 = op->src[1]->type; ++ ++ const bool bc_inp = op->src[0]->ne[0] % 32 != 0; ++ ++ snprintf(base, 256, "kernel_mul_mm_id_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1)); ++ snprintf(name, 256, "%s_bci=%d", base, bc_inp); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ res.smem = 8192; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ ++ char base[256]; ++ char name[256]; ++ ++ int nsg = 0; // number of simdgroups ++ int nr0 = 0; // number of src0 rows per simdgroup ++ int nr1 = 1; // number of src1 rows per threadgroup ++ ++ size_t smem = 0; // shared memory ++ ++ const ggml_type tsrc0 = op->src[0]->type; ++ const ggml_type tsrc1 = op->src[1]->type; ++ ++ const char * suffix = ""; ++ ++ // use custom matrix x vector kernel ++ switch (tsrc0) { ++ case GGML_TYPE_F32: ++ case GGML_TYPE_F16: ++ case GGML_TYPE_BF16: ++ { ++ nsg = std::min(4, (ne00 + 127) / 128); ++ nr0 = 2; ++ nr1 = 1; ++ smem = 32*sizeof(float)*nr0; ++ suffix = ne00 % 4 == 0 ? "_4" : ""; ++ } break; ++ case GGML_TYPE_Q1_0: ++ { ++ nsg = N_SG_Q1_0; ++ nr0 = N_R0_Q1_0; ++ } break; ++ case GGML_TYPE_TQ2_0: ++ { ++ nsg = N_SG_TQ2_0; ++ nr0 = N_R0_TQ2_0; ++ } break; ++ case GGML_TYPE_Q4_0: ++ { ++ nsg = N_SG_Q4_0; ++ nr0 = N_R0_Q4_0; ++ } break; ++ case GGML_TYPE_Q4_1: ++ { ++ nsg = N_SG_Q4_1; ++ nr0 = N_R0_Q4_1; ++ } break; ++ case GGML_TYPE_Q5_0: ++ { ++ nsg = N_SG_Q5_0; ++ nr0 = N_R0_Q5_0; ++ } break; ++ case GGML_TYPE_Q5_1: ++ { ++ nsg = N_SG_Q5_1; ++ nr0 = N_R0_Q5_1; ++ } break; ++ case GGML_TYPE_Q8_0: ++ { ++ nsg = N_SG_Q8_0; ++ nr0 = N_R0_Q8_0; ++ smem = 32*sizeof(float)*N_R0_Q8_0; ++ } break; ++ case GGML_TYPE_MXFP4: ++ { ++ nsg = N_SG_MXFP4; ++ nr0 = N_R0_MXFP4; ++ smem = 32*sizeof(float); ++ } break; ++ case GGML_TYPE_Q2_K: ++ { ++ nsg = ggml_metal_glm_dsa_q2_k_mul_mv_id_nsg_requested(); ++ nr0 = ggml_metal_glm_dsa_q2_k_mul_mv_id_nr0_requested(); ++ const bool is_glm_gate_up = ++ op->src[0]->ne[0] == 6144 && ++ op->src[0]->ne[1] == 2048 && ++ op->src[2]->ne[0] == 8 && ++ op->src[2]->ne[1] == 1; ++ const bool is_glm_down = ++ op->src[0]->ne[0] == 2048 && ++ op->src[0]->ne[1] == 6144 && ++ op->src[2]->ne[0] == 8 && ++ op->src[2]->ne[1] == 1; ++ if (is_glm_gate_up && !ggml_metal_glm_dsa_q2_k_mul_mv_id_nr0_overridden()) { ++ nr0 = 8; ++ } ++ if (is_glm_down && !ggml_metal_glm_dsa_q2_k_mul_mv_id_nr0_overridden()) { ++ nr0 = 8; ++ } ++ if (is_glm_down && ++ nr0 == N_R0_Q2_K && ++ ggml_metal_glm_dsa_q2_k_mul_mv_id_glm_down_enabled()) { ++ suffix = "_glm_down"; ++ } else if (nr0 != N_R0_Q2_K) { ++ snprintf(name, sizeof(name), "_r%d", nr0); ++ suffix = name; ++ } ++ } break; ++ case GGML_TYPE_Q3_K: ++ { ++ nsg = ggml_metal_glm_dsa_q3_k_mul_mv_id_nsg_requested(); ++ nr0 = ggml_metal_glm_dsa_q3_k_mul_mv_id_nr0_requested(); ++ const bool is_glm_down = ++ op->src[0]->ne[0] == 2048 && ++ op->src[0]->ne[1] == 6144 && ++ op->src[2]->ne[0] == 8 && ++ op->src[2]->ne[1] == 1; ++ if (is_glm_down && ++ nr0 == N_R0_Q3_K && ++ ggml_metal_glm_dsa_q3_k_mul_mv_id_glm_down_enabled()) { ++ suffix = "_glm_down"; ++ } else if (nr0 != N_R0_Q3_K) { ++ snprintf(name, sizeof(name), "_r%d", nr0); ++ suffix = name; ++ } ++ } break; ++ case GGML_TYPE_Q4_K: ++ { ++ nsg = N_SG_Q4_K; ++ nr0 = N_R0_Q4_K; ++ } break; ++ case GGML_TYPE_Q5_K: ++ { ++ nsg = N_SG_Q5_K; ++ nr0 = N_R0_Q5_K; ++ } break; ++ case GGML_TYPE_Q6_K: ++ { ++ nsg = N_SG_Q6_K; ++ nr0 = N_R0_Q6_K; ++ } break; ++ case GGML_TYPE_IQ2_XXS: ++ { ++ nsg = N_SG_IQ2_XXS; ++ nr0 = N_R0_IQ2_XXS; ++ smem = 256*8+128; ++ } break; ++ case GGML_TYPE_IQ2_XS: ++ { ++ nsg = N_SG_IQ2_XS; ++ nr0 = N_R0_IQ2_XS; ++ smem = 512*8+128; ++ } break; ++ case GGML_TYPE_IQ3_XXS: ++ { ++ nsg = N_SG_IQ3_XXS; ++ nr0 = N_R0_IQ3_XXS; ++ smem = 256*4+128; ++ } break; ++ case GGML_TYPE_IQ3_S: ++ { ++ nsg = N_SG_IQ3_S; ++ nr0 = N_R0_IQ3_S; ++ smem = 512*4; ++ } break; ++ case GGML_TYPE_IQ2_S: ++ { ++ nsg = N_SG_IQ2_S; ++ nr0 = N_R0_IQ2_S; ++ } break; ++ case GGML_TYPE_IQ1_S: ++ { ++ nsg = N_SG_IQ1_S; ++ nr0 = N_R0_IQ1_S; ++ } break; ++ case GGML_TYPE_IQ1_M: ++ { ++ nsg = N_SG_IQ1_M; ++ nr0 = N_R0_IQ1_M; ++ } break; ++ case GGML_TYPE_IQ4_NL: ++ { ++ nsg = N_SG_IQ4_NL; ++ nr0 = N_R0_IQ4_NL; ++ smem = 32*sizeof(float); ++ } break; ++ case GGML_TYPE_IQ4_XS: ++ { ++ nsg = N_SG_IQ4_XS; ++ nr0 = N_R0_IQ4_XS; ++ smem = 32*sizeof(float); ++ } break; ++ default: ++ { ++ GGML_LOG_ERROR("Asserting on type %d\n", (int)op->src[2]->type); ++ GGML_ABORT("not implemented"); ++ } ++ }; ++ ++ snprintf(base, 256, "kernel_mul_mv_id_%s_%s%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1), suffix); ++ snprintf(name, 256, "%s_nsg=%d", base, nsg); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int16(cv, nsg, FC_MUL_MV + 0); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 2); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 3); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 4); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = nr1; ++ res.nsg = nsg; ++ res.smem = smem; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32 || op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ const bool glm_decode_shape = ++ op->src[0]->src[0]->ne[0] >= 2048 && ++ op->src[0]->src[0]->ne[1] >= 1024 && ++ op->src[0]->src[2]->ne[0] == 8; ++ const int default_nsg = glm_decode_shape ? 2 : 1; ++ const int default_nr0 = glm_decode_shape ? 4 : 1; ++ const int nsg = ggml_metal_glm_dsa_q2_gate_up_swiglu_nsg_requested(default_nsg); ++ const int nr0 = ggml_metal_glm_dsa_q2_gate_up_swiglu_nr0_requested(default_nr0); ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int16(cv, nsg, FC_MUL_MV + 0); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 2); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 3); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 4); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = 0; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_vecscale(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ const bool glm_decode_shape = ++ op->src[0]->src[0]->ne[0] >= 2048 && ++ op->src[0]->src[0]->ne[1] >= 1024 && ++ op->src[0]->src[2]->ne[0] == 8; ++ const int default_nsg = glm_decode_shape ? 2 : 1; ++ const int default_nr0 = glm_decode_shape ? 4 : 1; ++ const int nsg = ggml_metal_glm_dsa_q2_gate_up_swiglu_nsg_requested(default_nsg); ++ const int nr0 = ggml_metal_glm_dsa_q2_gate_up_swiglu_nr0_requested(default_nr0); ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_vecscale"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int16(cv, nsg, FC_MUL_MV + 0); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 2); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 3); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 4); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = 0; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32 || op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 1; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = 0; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_share_y(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ constexpr int shared_y_floats = 32*32; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_share_y"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = (nsg*nr0 + shared_y_floats)*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_vecscale(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_vecscale"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_q8_act(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_q8_act"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual_prequant_q8(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[3] != nullptr); ++ GGML_ASSERT(op->src[0]->src[3]->type == GGML_TYPE_Q8_0); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 1; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot1_dual_prequant_q8"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = 0; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual_inblock_q2(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 1; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot1_dual_inblock_q2"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = 0; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_half_y(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_half_y"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_rowtile(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_rowtile"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_r12(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU || op->op == GGML_OP_MUL); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 12; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_r12"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_r16(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU || op->op == GGML_OP_MUL); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 16; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_r16"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot2(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32 || op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 4; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot2"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot8(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32 || op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot8"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++static ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_dual( ++ ggml_metal_library_t lib, ++ const ggml_tensor * op, ++ const char * kernel_name, ++ int nsg, ++ int nr0 = 8) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32 || op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "%s", kernel_name); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = 0; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual(ggml_metal_library_t lib, const ggml_tensor * op) { ++ const char * scan_all = getenv("GGML_METAL_EXPERIMENTAL_Q2_WEIGHT_SCAN"); ++ const char * scan_gate_up = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_WEIGHT_SCAN"); ++ const bool scan = (scan_all && atoi(scan_all) != 0) || (scan_gate_up && atoi(scan_gate_up) != 0); ++ return ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_dual( ++ lib, ++ op, ++ scan ? ++ "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot1_dual_scan" : ++ "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot1_dual", ++ 1); ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot2_dual(ggml_metal_library_t lib, const ggml_tensor * op) { ++ return ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_dual( ++ lib, op, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot2_dual", 2); ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual(ggml_metal_library_t lib, const ggml_tensor * op) { ++ return ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_dual( ++ lib, op, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot4_dual", 4); ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual_r12(ggml_metal_library_t lib, const ggml_tensor * op) { ++ return ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_dual( ++ lib, op, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot4_dual_r12", 4, 12); ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual_r16(ggml_metal_library_t lib, const ggml_tensor * op) { ++ return ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_dual( ++ lib, op, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot4_dual_r16", 4, 16); ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot8_split(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32 || op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 16; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot8_split"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_route_q2_gate_up_swiglu_pair_sg_slot8(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_GLU || op->op == GGML_OP_MUL); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_glm_moe_route_q2_gate_up_swiglu_pair_sg_slot8"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_weights_q2_gate_up_swiglu_pair_sg_slot1(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_GLU); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 1; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_glm_moe_weights_q2_gate_up_swiglu_pair_sg_slot1"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_f16(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_CPY); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_GLU); ++ GGML_ASSERT(op->src[0]->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[0]->src[1]->type == GGML_TYPE_F32 || op->src[0]->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->type == GGML_TYPE_F16); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_f16"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_weighted_reduce(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_MOE_WEIGHTED_SUM || op->op == GGML_OP_MOE_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0] != nullptr); ++ ++ const ggml_tensor * experts = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[0] : op->src[0]->src[0]; ++ const ggml_tensor * input = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[1] : op->src[0]->src[1]; ++ const ggml_tensor * weights = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[3] : op->src[1]; ++ if (op->op == GGML_OP_MOE_WEIGHTED_SUM) { ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ } ++ ++ GGML_ASSERT(experts->type == GGML_TYPE_Q2_K || experts->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(input->type == GGML_TYPE_F32 || input->type == GGML_TYPE_F16); ++ GGML_ASSERT(weights->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ const ggml_type src0_type = experts->type; ++ const int nsg = src0_type == GGML_TYPE_Q3_K ? ++ ggml_metal_glm_dsa_q3_k_mul_mv_id_nsg_requested() : ++ ggml_metal_glm_dsa_q2_k_mul_mv_id_nsg_requested(); ++ const int nr0 = src0_type == GGML_TYPE_Q3_K ? N_R0_Q3_K : 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_%s_weighted_reduce", ggml_type_name(src0_type)); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ ggml_metal_cv_t cv = ggml_metal_cv_init(); ++ ++ ggml_metal_cv_set_int16(cv, nsg, FC_MUL_MV + 0); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 2); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 3); ++ ggml_metal_cv_set_int16(cv, 1, FC_MUL_MV + 4); ++ ++ res = ggml_metal_library_compile_pipeline(lib, base, name, cv); ++ ++ ggml_metal_cv_free(cv); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = 0; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_f16(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_f16"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_shifted(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_shifted"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_f16_shifted(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_f16_shifted"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_vecscale(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ const char * scan_all = getenv("GGML_METAL_EXPERIMENTAL_Q2_WEIGHT_SCAN"); ++ const char * scan_down = getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_WEIGHT_SCAN"); ++ const bool scan = (scan_all && atoi(scan_all) != 0) || (scan_down && atoi(scan_down) != 0); ++ snprintf( ++ base, ++ 256, ++ "%s", ++ scan ? ++ "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_scan" : ++ "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_vecscale"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r4_nb8(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 4; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r4_nb8"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r4_nb8_vecscale(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 4; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r4_nb8_vecscale"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r16_nb8(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 16; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r16_nb8"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = N_R0_Q3_K; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r8"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[0]->ne[0] == 2048); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r8_nb8"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8_w0(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_MOE_WEIGHTED_SUM || op->op == GGML_OP_MOE_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0] != nullptr); ++ ++ const ggml_tensor * experts = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[0] : op->src[0]->src[0]; ++ const ggml_tensor * input = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[1] : op->src[0]->src[1]; ++ const ggml_tensor * ids = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[2] : op->src[0]->src[2]; ++ const ggml_tensor * weights = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[3] : op->src[1]; ++ if (op->op == GGML_OP_MOE_WEIGHTED_SUM) { ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(ggml_get_op_params_i32(op, 0) == 0); ++ } ++ ++ GGML_ASSERT(experts->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(experts->ne[0] == 2048); ++ GGML_ASSERT(input->type == GGML_TYPE_F32); ++ GGML_ASSERT(ids->ne[0] == 8); ++ GGML_ASSERT(weights->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r8_nb8_w0"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8_w0_f16(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[0]->ne[0] == 2048); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F16); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_get_op_params_i32(op, 0) == 0); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r8_nb8_w0_f16"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r6_nb8_w0(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[0]->ne[0] == 2048); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_get_op_params_i32(op, 0) == 0); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 6; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r6_nb8_w0"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r10_nb8_w0(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[0]->ne[0] == 2048); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_get_op_params_i32(op, 0) == 0); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 10; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r10_nb8_w0"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_glm52_w0(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_MOE_WEIGHTED_SUM || op->op == GGML_OP_MOE_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0] != nullptr); ++ ++ const ggml_tensor * experts = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[0] : op->src[0]->src[0]; ++ const ggml_tensor * input = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[1] : op->src[0]->src[1]; ++ const ggml_tensor * ids = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[2] : op->src[0]->src[2]; ++ const ggml_tensor * weights = op->op == GGML_OP_MOE_MUL_MAT_ID ? op->src[3] : op->src[1]; ++ if (op->op == GGML_OP_MOE_WEIGHTED_SUM) { ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(ggml_get_op_params_i32(op, 0) == 0); ++ } ++ ++ GGML_ASSERT(experts->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(experts->ne[0] == 2048); ++ GGML_ASSERT(experts->ne[1] == 6144); ++ GGML_ASSERT(input->type == GGML_TYPE_F32); ++ GGML_ASSERT(ids->ne[0] == 8); ++ GGML_ASSERT(ids->ne[1] == 1); ++ GGML_ASSERT(weights->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_glm52_w0"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r12_nb8_w0(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[0]->ne[0] == 2048); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ GGML_ASSERT(ggml_get_op_params_i32(op, 0) == 0); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 12; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r12_nb8_w0"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r16(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 16; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r16"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_split2(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 16; ++ constexpr int nr0 = N_R0_Q3_K; ++ ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_split2"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); ++ } ++ ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_swiglu_q3_down_weighted(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; ++ + char base[256]; + char name[256]; +- +- snprintf(base, 256, "kernel_mul_mm_id_map0_ne20_%d", ne20); +- snprintf(name, 256, "%s_ne02=%d", base, ne02); ++ snprintf(base, 256, "kernel_glm_moe_swiglu_q3_K_down_weighted"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { + res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + +- res.smem = (size_t) ne02*ne20*sizeof(uint16_t); ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); + + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_metal_library_t lib, const ggml_tensor * op) { +- char base[256]; +- char name[256]; +- +- const ggml_type tsrc0 = op->src[0]->type; +- const ggml_type tsrc1 = op->src[1]->type; ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_swiglu_q2_down_weighted(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(op->src[0]->src[2]->ne[0] == 8); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); + +- const bool bc_inp = op->src[0]->ne[0] % 32 != 0; ++ constexpr int nsg = 8; ++ constexpr int nr0 = 8; + +- snprintf(base, 256, "kernel_mul_mm_id_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1)); +- snprintf(name, 256, "%s_bci=%d", base, bc_inp); ++ char base[256]; ++ char name[256]; ++ snprintf(base, 256, "kernel_glm_moe_swiglu_q2_K_down_weighted"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { +- ggml_metal_cv_t cv = ggml_metal_cv_init(); +- +- ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0); +- +- res = ggml_metal_library_compile_pipeline(lib, base, name, cv); +- +- ggml_metal_cv_free(cv); ++ res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr); + } + +- res.smem = 8192; ++ res.nr0 = nr0; ++ res.nr1 = 1; ++ res.nsg = nsg; ++ res.smem = nsg*nr0*sizeof(float); + + return res; + } + +-ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_metal_library_t lib, const ggml_tensor * op) { +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_accum_atomic(ggml_metal_library_t lib, const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MOE_WEIGHTED_SUM); ++ GGML_ASSERT(op->src[0] != nullptr); ++ GGML_ASSERT(op->src[0]->op == GGML_OP_MUL_MAT_ID); ++ GGML_ASSERT(op->src[0]->src[0]->type == GGML_TYPE_Q3_K); ++ GGML_ASSERT(op->src[0]->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ const int nsg = ggml_metal_glm_dsa_q3_k_mul_mv_id_nsg_requested(); ++ const int nr0 = N_R0_Q3_K; + + char base[256]; + char name[256]; +- +- int nsg = 0; // number of simdgroups +- int nr0 = 0; // number of src0 rows per simdgroup +- int nr1 = 1; // number of src1 rows per threadgroup +- +- size_t smem = 0; // shared memory +- +- const ggml_type tsrc0 = op->src[0]->type; +- const ggml_type tsrc1 = op->src[1]->type; +- +- const char * suffix = ""; +- +- // use custom matrix x vector kernel +- switch (tsrc0) { +- case GGML_TYPE_F32: +- case GGML_TYPE_F16: +- case GGML_TYPE_BF16: +- { +- nsg = std::min(4, (ne00 + 127) / 128); +- nr0 = 2; +- nr1 = 1; +- smem = 32*sizeof(float)*nr0; +- suffix = ne00 % 4 == 0 ? "_4" : ""; +- } break; +- case GGML_TYPE_Q1_0: +- { +- nsg = N_SG_Q1_0; +- nr0 = N_R0_Q1_0; +- } break; +- case GGML_TYPE_Q2_0: +- { +- nsg = N_SG_Q2_0; +- nr0 = N_R0_Q2_0; +- } break; +- case GGML_TYPE_Q4_0: +- { +- nsg = N_SG_Q4_0; +- nr0 = N_R0_Q4_0; +- } break; +- case GGML_TYPE_Q4_1: +- { +- nsg = N_SG_Q4_1; +- nr0 = N_R0_Q4_1; +- } break; +- case GGML_TYPE_Q5_0: +- { +- nsg = N_SG_Q5_0; +- nr0 = N_R0_Q5_0; +- } break; +- case GGML_TYPE_Q5_1: +- { +- nsg = N_SG_Q5_1; +- nr0 = N_R0_Q5_1; +- } break; +- case GGML_TYPE_Q8_0: +- { +- nsg = N_SG_Q8_0; +- nr0 = N_R0_Q8_0; +- smem = 32*sizeof(float)*N_R0_Q8_0; +- } break; +- case GGML_TYPE_MXFP4: +- { +- nsg = N_SG_MXFP4; +- nr0 = N_R0_MXFP4; +- smem = 32*sizeof(float); +- } break; +- case GGML_TYPE_Q2_K: +- { +- nsg = N_SG_Q2_K; +- nr0 = N_R0_Q2_K; +- } break; +- case GGML_TYPE_Q3_K: +- { +- nsg = N_SG_Q3_K; +- nr0 = N_R0_Q3_K; +- } break; +- case GGML_TYPE_Q4_K: +- { +- nsg = N_SG_Q4_K; +- nr0 = N_R0_Q4_K; +- } break; +- case GGML_TYPE_Q5_K: +- { +- nsg = N_SG_Q5_K; +- nr0 = N_R0_Q5_K; +- } break; +- case GGML_TYPE_Q6_K: +- { +- nsg = N_SG_Q6_K; +- nr0 = N_R0_Q6_K; +- } break; +- case GGML_TYPE_IQ2_XXS: +- { +- nsg = N_SG_IQ2_XXS; +- nr0 = N_R0_IQ2_XXS; +- smem = 256*8+128; +- } break; +- case GGML_TYPE_IQ2_XS: +- { +- nsg = N_SG_IQ2_XS; +- nr0 = N_R0_IQ2_XS; +- smem = 512*8+128; +- } break; +- case GGML_TYPE_IQ3_XXS: +- { +- nsg = N_SG_IQ3_XXS; +- nr0 = N_R0_IQ3_XXS; +- smem = 256*4+128; +- } break; +- case GGML_TYPE_IQ3_S: +- { +- nsg = N_SG_IQ3_S; +- nr0 = N_R0_IQ3_S; +- smem = 512*4; +- } break; +- case GGML_TYPE_IQ2_S: +- { +- nsg = N_SG_IQ2_S; +- nr0 = N_R0_IQ2_S; +- } break; +- case GGML_TYPE_IQ1_S: +- { +- nsg = N_SG_IQ1_S; +- nr0 = N_R0_IQ1_S; +- } break; +- case GGML_TYPE_IQ1_M: +- { +- nsg = N_SG_IQ1_M; +- nr0 = N_R0_IQ1_M; +- } break; +- case GGML_TYPE_IQ4_NL: +- { +- nsg = N_SG_IQ4_NL; +- nr0 = N_R0_IQ4_NL; +- smem = 32*sizeof(float); +- } break; +- case GGML_TYPE_IQ4_XS: +- { +- nsg = N_SG_IQ4_XS; +- nr0 = N_R0_IQ4_XS; +- smem = 32*sizeof(float); +- } break; +- default: +- { +- GGML_LOG_ERROR("Asserting on type %d\n", (int)op->src[2]->type); +- GGML_ABORT("not implemented"); +- } +- }; +- +- snprintf(base, 256, "kernel_mul_mv_id_%s_%s%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1), suffix); +- snprintf(name, 256, "%s_nsg=%d", base, nsg); ++ snprintf(base, 256, "kernel_mul_mv_id_q3_K_weighted_accum_atomic"); ++ snprintf(name, 256, "%s_nsg=%d_nr0=%d", base, nsg, nr0); + + ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name); + if (!res.pipeline) { +@@ -1172,9 +3090,32 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_m + } + + res.nr0 = nr0; +- res.nr1 = nr1; ++ res.nr1 = 1; + res.nsg = nsg; +- res.smem = smem; ++ res.smem = 0; ++ ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_zero_f32(ggml_metal_library_t lib) { ++ const char * base = "kernel_zero_f32"; ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, base, nullptr); ++ } ++ return res; ++} ++ ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_q2_selected_weight_scan(ggml_metal_library_t lib) { ++ const char * base = "kernel_glm_moe_q2_selected_weight_scan"; ++ ++ ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base); ++ if (!res.pipeline) { ++ res = ggml_metal_library_compile_pipeline(lib, base, base, nullptr); ++ } ++ ++ res.nsg = 8; ++ res.smem = 8*sizeof(float); + + return res; + } +@@ -1306,6 +3247,89 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml + return res; + } + ++ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe_route(ggml_metal_library_t lib) { ++ const char * name = "kernel_topk_moe_route_f32_i32"; ++ ++ 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_topk_moe_route_sg_reduce(ggml_metal_library_t lib) { ++ const char * name = "kernel_topk_moe_route_f32_i32_sg_reduce"; ++ ++ 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_topk_moe_route_glm_256_8_sg32(ggml_metal_library_t lib) { ++ const char * name = "kernel_topk_moe_route_glm_256_8_sg32"; ++ ++ 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_moe_route_weights(ggml_metal_library_t lib) { ++ const char * name = "kernel_moe_route_weights_f32_i32"; ++ ++ 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_moe_weighted_sum(ggml_metal_library_t lib) { ++ const char * name = "kernel_moe_weighted_sum_f32"; ++ ++ 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_moe_weighted_sum_x4(ggml_metal_library_t lib) { ++ const char * name = "kernel_moe_weighted_sum_f32x4"; ++ ++ 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_glu_weighted(ggml_metal_library_t lib, const ggml_tensor * op) { ++ GGML_ASSERT(op->op == GGML_OP_MUL); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ const bool use_c4 = op->ne[0] % 4 == 0; ++ const char * name = use_c4 ? "kernel_swiglu_weighted_f32_4" : "kernel_swiglu_weighted_f32"; ++ ++ 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); ++ } ++ ++ res.c4 = use_c4; ++ ++ return res; ++} ++ + ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_pad( + ggml_metal_library_t lib, + const struct ggml_tensor * op, +@@ -1523,7 +3547,7 @@ ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_v + const ggml_tensor * op, + int32_t dv, + int32_t nwg) { +- assert(op->op == GGML_OP_FLASH_ATTN_EXT); ++ assert(op->op == GGML_OP_FLASH_ATTN_EXT || op->op == GGML_OP_DSA_SPARSE_ATTN); + + char base[256]; + char name[256]; +diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h +index b36fa811..03567502 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.h ++++ b/ggml/src/ggml-metal/ggml-metal-device.h +@@ -88,7 +88,7 @@ void ggml_metal_encoder_set_threadgroup_memory_size(ggml_metal_encoder_t encoder + void ggml_metal_encoder_dispatch_threadgroups(ggml_metal_encoder_t encoder, int tg0, int tg1, int tg2, int tptg0, int tptg1, int tptg2); + + void ggml_metal_encoder_memory_barrier(ggml_metal_encoder_t encoder); +- ++void ggml_metal_encoder_memory_barrier_buffer(ggml_metal_encoder_t encoder, struct ggml_metal_buffer_id buffer); + void ggml_metal_encoder_end_encoding(ggml_metal_encoder_t encoder); + + // +@@ -111,8 +111,25 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_base + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cpy (ggml_metal_library_t lib, enum ggml_type tsrc, enum ggml_type tdst); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_1d (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_2d (ggml_metal_library_t lib, const struct ggml_tensor * op, enum ggml_op_pool op_pool); +-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows (ggml_metal_library_t lib, enum ggml_type tsrc); +-struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows (ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows (ggml_metal_library_t lib, enum ggml_type tsrc, enum ggml_type tdst, bool use_f16_vec4); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows_packed_f16(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows (ggml_metal_library_t lib, enum ggml_type tidx, enum ggml_type tdst); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_mask_fill(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_mask_set (ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_attn (ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_attn_cached_topk(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_sparse_attn_decode_grouped(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsa_top1_attn (ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_selected_row_flash_vec(ggml_metal_library_t lib, int32_t nsg, int32_t nwg); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_selected_row_flash_pair(ggml_metal_library_t lib, int32_t nwg); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_selected_row_flash_tiled(ggml_metal_library_t lib, int32_t nwg); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_multihead_flash(ggml_metal_library_t lib, int32_t nwg); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_qk_scores(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_scores_v(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_softmax_prefix(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_chunk_v(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_chunk_fold(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_compact_probs_v_sequential(ggml_metal_library_t lib); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat (ggml_metal_library_t lib, enum ggml_type tsrc); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat (ggml_metal_library_t lib, enum ggml_type tsrc); +@@ -129,6 +146,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_ + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net (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, bool parallel, bool staged_q); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext (ggml_metal_library_t lib, const struct ggml_tensor * op, int nsg, int nxpsg, int r1ptg); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm (ggml_metal_library_t lib, const struct ggml_tensor * op); +@@ -136,11 +154,66 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0 (ggml_metal_library_t lib, int ne02, int ne20); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id (ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_vecscale(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_share_y(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_vecscale(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_q8_act(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_half_y(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_rowtile(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_r12(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_r16(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot2(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual_prequant_q8(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual_inblock_q2(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot2_dual(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual_r12(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual_r16(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot8(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot8_split(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_f16(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_route_q2_gate_up_swiglu_pair_sg_slot8(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_weights_q2_gate_up_swiglu_pair_sg_slot1(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_weighted_reduce(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r4_nb8(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r4_nb8_vecscale(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_f16(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_shifted(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_f16_shifted(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_vecscale(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r16_nb8(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8_w0(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8_w0_f16(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r6_nb8_w0(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r10_nb8_w0(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_glm52_w0(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r12_nb8_w0(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r16(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_split2(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_swiglu_q3_down_weighted(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_swiglu_q2_down_weighted(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glm_moe_q2_selected_weight_scan(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_accum_atomic(ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_zero_f32 (ggml_metal_library_t lib); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge (ggml_metal_library_t lib, const struct ggml_tensor * op); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe_route (ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe_route_sg_reduce(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe_route_glm_256_8_sg32(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_route_weights (ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_weighted_sum (ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_weighted_sum_x4(ggml_metal_library_t lib); ++struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu_weighted (ggml_metal_library_t lib, const struct ggml_tensor * op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin (ggml_metal_library_t lib, const struct ggml_tensor * op, int32_t n_fuse ); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one (ggml_metal_library_t lib, enum ggml_op op); + struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_l2_norm (ggml_metal_library_t lib, const struct ggml_tensor * op); +@@ -309,6 +382,10 @@ void ggml_metal_buffer_free (ggml_metal_buffer_t buf); + void * ggml_metal_buffer_get_base (ggml_metal_buffer_t buf); + bool ggml_metal_buffer_is_shared(ggml_metal_buffer_t buf); + ++// Returns the first device resource at offset zero. This is used for ++// backend-owned scratch that is not represented by a host ggml_tensor. ++struct ggml_metal_buffer_id ggml_metal_buffer_get_base_id(ggml_metal_buffer_t buf); ++ + void ggml_metal_buffer_memset_tensor(ggml_metal_buffer_t buf, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); + void ggml_metal_buffer_set_tensor (ggml_metal_buffer_t buf, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size); + void ggml_metal_buffer_get_tensor (ggml_metal_buffer_t buf, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size); +diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m +index 80e47f2c..8becf966 100644 +--- a/ggml/src/ggml-metal/ggml-metal-device.m ++++ b/ggml/src/ggml-metal/ggml-metal-device.m +@@ -2,7 +2,6 @@ + + #import "ggml-impl.h" + #import "ggml-backend-impl.h" +- + #include + + #include +@@ -513,6 +512,15 @@ void ggml_metal_encoder_memory_barrier(ggml_metal_encoder_t encoder) { + [encoder->obj memoryBarrierWithScope:MTLBarrierScopeBuffers]; + } + ++void ggml_metal_encoder_memory_barrier_buffer(ggml_metal_encoder_t encoder, struct ggml_metal_buffer_id buffer) { ++ if (buffer.metal == nil) { ++ return; ++ } ++ ++ id resource = buffer.metal; ++ [encoder->obj memoryBarrierWithResources:&resource count:1]; ++} ++ + void ggml_metal_encoder_end_encoding(ggml_metal_encoder_t encoder) { + [encoder->obj endEncoding]; + } +@@ -1272,6 +1280,137 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te + return true; + case GGML_OP_GATED_DELTA_NET: + return has_simdgroup_reduction && op->src[2]->ne[0] % 32 == 0; ++ case GGML_OP_LIGHTNING_INDEXER: ++ { ++ const bool k_scalar = ++ (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && ++ op->src[1]->nb[0] == ggml_type_size(op->src[1]->type); ++ const bool k_quant = ++ (op->src[1]->type == GGML_TYPE_Q4_0 || ++ op->src[1]->type == GGML_TYPE_Q8_0 || ++ op->src[1]->type == GGML_TYPE_Q2_K || ++ op->src[1]->type == GGML_TYPE_Q3_K || ++ op->src[1]->type == GGML_TYPE_Q4_K || ++ op->src[1]->type == GGML_TYPE_Q5_K || ++ op->src[1]->type == GGML_TYPE_Q6_K) && ++ op->src[1]->ne[0] % ggml_blck_size(op->src[1]->type) == 0; ++ return op->type == GGML_TYPE_F32 && ++ op->src[0]->type == GGML_TYPE_F32 && ++ op->src[2]->type == GGML_TYPE_F32 && ++ (k_scalar || k_quant) && ++ op->nb[0] == sizeof(float) && ++ op->src[0]->nb[0] == sizeof(float) && ++ op->src[2]->nb[0] == sizeof(float) && ++ op->src[0]->ne[0] == op->src[1]->ne[0] && ++ op->src[0]->ne[1] == op->src[2]->ne[0] && ++ op->src[1]->ne[1] == 1 && ++ op->src[0]->ne[2] == op->src[2]->ne[1] && ++ op->src[2]->ne[2] == 1 && ++ op->src[0]->ne[3] == op->src[1]->ne[3] && ++ op->src[1]->ne[3] == op->src[2]->ne[3]; ++ } ++ case GGML_OP_DSA_SPARSE_MASK: ++ return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16) && ++ op->src[0]->type == op->type && ++ op->src[1]->type == GGML_TYPE_I32 && ++ op->src[0]->ne[0] == 1 && ++ op->src[0]->ne[2] == op->src[1]->ne[1] && ++ op->src[0]->ne[3] % op->src[1]->ne[2] == 0 && ++ op->src[1]->ne[3] == 1 && ++ op->ne[0] == op->src[0]->ne[0] && ++ op->ne[1] == op->src[0]->ne[1] && ++ op->ne[2] == op->src[0]->ne[2] && ++ op->ne[3] == op->src[0]->ne[3]; ++ case GGML_OP_DSA_SPARSE_ATTN: ++ return op->type == GGML_TYPE_F32 && ++ op->src[0]->type == GGML_TYPE_F32 && ++ (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && ++ (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && ++ (op->src[3]->type == GGML_TYPE_F32 || op->src[3]->type == GGML_TYPE_F16) && ++ op->src[4]->type == GGML_TYPE_I32 && ++ op->src[0]->nb[0] == sizeof(float) && ++ op->src[1]->nb[0] == ggml_type_size(op->src[1]->type) && ++ op->src[2]->nb[0] == ggml_type_size(op->src[2]->type) && ++ op->src[3]->ne[0] == 1 && ++ op->src[3]->ne[1] == op->src[1]->ne[1] && ++ op->src[3]->ne[2] == op->src[0]->ne[1] && ++ op->src[3]->ne[3] == op->src[0]->ne[3] && ++ op->src[4]->ne[0] <= 4096 && ++ op->src[4]->ne[1] == op->src[0]->ne[1] && ++ op->src[0]->ne[3] % op->src[4]->ne[2] == 0 && ++ op->src[4]->ne[3] == 1 && ++ op->src[0]->ne[0] == op->src[1]->ne[0] && ++ op->src[1]->ne[1] == op->src[2]->ne[1] && ++ op->src[0]->ne[2] % op->src[1]->ne[2] == 0 && ++ op->src[0]->ne[2] % op->src[2]->ne[2] == 0 && ++ op->ne[0] == op->src[2]->ne[0] && ++ op->ne[1] == op->src[0]->ne[1] && ++ op->ne[2] == op->src[0]->ne[2] && ++ op->ne[3] == op->src[0]->ne[3]; ++ case GGML_OP_DSA_TOP1_ATTN: ++ return op->type == GGML_TYPE_F32 && ++ op->src[0]->type == GGML_TYPE_F32 && ++ (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && ++ op->src[2]->type == GGML_TYPE_I32 && ++ op->nb[0] == sizeof(float) && ++ op->src[0]->nb[0] == sizeof(float) && ++ op->src[1]->nb[0] == ggml_type_size(op->src[1]->type) && ++ op->src[2]->ne[0] == 1 && ++ op->src[2]->ne[1] == op->src[0]->ne[1] && ++ op->src[0]->ne[3] == op->src[1]->ne[3] && ++ op->src[0]->ne[3] % op->src[2]->ne[2] == 0 && ++ op->src[2]->ne[3] == 1 && ++ op->src[0]->ne[2] % op->src[1]->ne[2] == 0 && ++ op->ne[0] == op->src[1]->ne[0] && ++ op->ne[1] == op->src[0]->ne[1] && ++ op->ne[2] == op->src[0]->ne[2] && ++ op->ne[3] == op->src[0]->ne[3]; ++ case GGML_OP_MOE_WEIGHTED_SUM: ++ return op->type == GGML_TYPE_F32 && ++ op->src[0]->type == GGML_TYPE_F32 && ++ op->src[1]->type == GGML_TYPE_F32 && ++ op->nb[0] == sizeof(float) && ++ op->src[0]->nb[0] == sizeof(float) && ++ op->src[1]->nb[0] == sizeof(float) && ++ op->src[1]->ne[0] == 1 && ++ op->src[0]->ne[1] == op->src[1]->ne[1] && ++ op->src[0]->ne[2] == op->src[1]->ne[2] && ++ op->src[0]->ne[3] == 1 && ++ op->src[1]->ne[3] == 1 && ++ op->ne[0] == op->src[0]->ne[0] && ++ op->ne[1] == op->src[0]->ne[2]; ++ case GGML_OP_MOE_MUL_MAT_ID: ++ return has_simdgroup_reduction && ++ op->type == GGML_TYPE_F32 && ++ (op->src[0]->type == GGML_TYPE_Q2_K || op->src[0]->type == GGML_TYPE_Q3_K) && ++ (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && ++ op->src[2]->type == GGML_TYPE_I32 && ++ op->src[3]->type == GGML_TYPE_F32 && ++ op->nb[0] == sizeof(float) && ++ op->src[0]->ne[0] == op->src[1]->ne[0] && ++ op->src[2]->ne[0] == op->src[1]->ne[1] && ++ op->src[2]->ne[1] == op->src[1]->ne[2] && ++ op->src[3]->ne[0] == 1 && ++ op->src[3]->ne[1] == op->src[2]->ne[0] && ++ op->src[3]->ne[2] == op->src[2]->ne[1] && ++ op->ne[0] == op->src[0]->ne[1] && ++ op->ne[1] == op->src[1]->ne[2]; ++ case GGML_OP_MOE_ROUTE_WEIGHTS: ++ return op->type == GGML_TYPE_F32 && ++ op->src[0]->type == GGML_TYPE_F32 && ++ op->src[1]->type == GGML_TYPE_I32 && ++ op->nb[0] == sizeof(float) && ++ op->src[0]->nb[0] == sizeof(float) && ++ op->src[1]->nb[0] == sizeof(int32_t) && ++ op->src[0]->ne[0] == 1 && ++ op->src[1]->ne[0] <= 16 && ++ op->src[1]->ne[1] == op->src[0]->ne[2] && ++ op->src[1]->ne[2] == 1 && ++ op->src[1]->ne[3] == 1 && ++ op->src[0]->ne[3] == 1 && ++ op->ne[0] == 1 && ++ op->ne[1] == op->src[1]->ne[0] && ++ op->ne[2] == op->src[1]->ne[1]; + case GGML_OP_SOLVE_TRI: + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: +@@ -1411,6 +1550,7 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te + + // pointers to global device + ggml_metal_device_t dev; ++ + }; + + static void ggml_metal_log_allocated_size(id device, size_t size_aligned) { +@@ -1515,7 +1655,6 @@ ggml_metal_buffer_t ggml_metal_buffer_init(ggml_metal_device_t dev, size_t size, + ggml_metal_buffer_t res = calloc(1, sizeof(struct ggml_metal_buffer)); + + res->dev = dev; +- + const size_t size_page = sysconf(_SC_PAGESIZE); + + size_t size_aligned = size; +@@ -1585,7 +1724,6 @@ ggml_metal_buffer_t ggml_metal_buffer_map(ggml_metal_device_t dev, void * ptr, s + ggml_metal_buffer_t res = calloc(1, sizeof(struct ggml_metal_buffer)); + + res->dev = dev; +- + res->all_data = ptr; + res->all_size = size; + +@@ -1913,3 +2051,11 @@ struct ggml_metal_buffer_id ggml_metal_buffer_get_id(ggml_metal_buffer_t buf, co + + return res; + } ++ ++struct ggml_metal_buffer_id ggml_metal_buffer_get_base_id(ggml_metal_buffer_t buf) { ++ if (buf == NULL || buf->n_buffers <= 0) { ++ return (struct ggml_metal_buffer_id) { nil, 0 }; ++ } ++ ++ return (struct ggml_metal_buffer_id) { buf->buffers[0].metal, 0 }; ++} +diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h +index 330278d0..f7112ee5 100644 +--- a/ggml/src/ggml-metal/ggml-metal-impl.h ++++ b/ggml/src/ggml-metal/ggml-metal-impl.h +@@ -23,6 +23,8 @@ + + #define N_R0_Q1_0 8 + #define N_SG_Q1_0 2 ++#define N_R0_TQ2_0 8 ++#define N_SG_TQ2_0 4 + + #define N_R0_Q2_0 8 + #define N_SG_Q2_0 2 +@@ -48,7 +50,7 @@ + #define N_R0_Q2_K 4 + #define N_SG_Q2_K 2 + +-#define N_R0_Q3_K 2 ++#define N_R0_Q3_K 4 + #define N_SG_Q3_K 2 + + #define N_R0_Q4_K 2 +@@ -549,6 +551,129 @@ typedef struct { + int32_t nr0; + } ggml_metal_kargs_mul_mv_id; + ++typedef struct { ++ int32_t nei0; ++ int32_t nei1; ++ uint64_t nbi1; ++ int32_t ne00; ++ int32_t ne01; ++ int32_t ne02; ++ uint64_t nb00; ++ uint64_t nb01; ++ uint64_t nb02; ++ int32_t ne10; ++ int32_t ne11; ++ int32_t ne12; ++ int32_t ne13; ++ uint64_t nb10; ++ uint64_t nb11; ++ uint64_t nb12; ++ int32_t ne0; ++ int32_t ne1; ++ uint64_t nb1; ++ int32_t nr0; ++ uint64_t weights_nb1; ++ uint64_t weights_nb2; ++ int32_t weighted; ++ int32_t _pad0; ++} ggml_metal_kargs_mul_mv_id_gate_up_swiglu; ++ ++typedef struct { ++ int32_t ne00; ++ uint64_t nb01; ++ int32_t ne10; ++ uint64_t nb11; ++ int32_t ne0; ++ int32_t ne1; ++ uint64_t nb1; ++ int32_t i00; ++ int32_t i10; ++ uint64_t weights_nb1; ++ uint64_t weights_nb2; ++ uint64_t dst_nb1; ++} ggml_metal_kargs_glu_weighted; ++ ++typedef struct { ++ int32_t nei0; ++ int32_t nei1; ++ uint64_t nbi1; ++ int32_t ne00; ++ int32_t ne01; ++ int32_t ne02; ++ uint64_t nb00; ++ uint64_t nb01; ++ uint64_t nb02; ++ int32_t ne10; ++ int32_t ne11; ++ int32_t ne12; ++ uint64_t nb10; ++ uint64_t nb11; ++ uint64_t nb12; ++ uint64_t weights_nb1; ++ uint64_t weights_nb2; ++ int32_t ne0; ++ int32_t ne1; ++ uint64_t nb0; ++ uint64_t nb1; ++ int32_t nr0; ++ int32_t already_weighted; ++} ggml_metal_kargs_mul_mv_id_weighted_reduce; ++ ++typedef struct { ++ uint64_t weights_nb1; ++ uint64_t weights_nb2; ++ uint64_t dst_nb0; ++ uint64_t dst_nb1; ++ int32_t already_weighted; ++ int32_t _pad0; ++} ggml_metal_kargs_mul_mv_id_weighted_reduce_extra; ++ ++typedef struct { ++ uint64_t gate_nb1; ++ uint64_t gate_nb2; ++ uint64_t up_nb1; ++ uint64_t up_nb2; ++} ggml_metal_kargs_glm_moe_swiglu_q3_down; ++ ++typedef struct { ++ uint64_t gate_expert_bytes; ++ uint64_t gate_expert_stride; ++ uint64_t up_expert_bytes; ++ uint64_t up_expert_stride; ++ uint64_t down_expert_bytes; ++ uint64_t down_expert_stride; ++ uint64_t ids_nb1; ++ int32_t top_k; ++ int32_t n_tokens; ++ int32_t n_experts; ++ int32_t chunks_per_expert; ++ int32_t storage_block_bytes; ++} ggml_metal_kargs_glm_moe_q2_weight_scan; ++ ++typedef struct { ++ int32_t n_embd; ++ int32_t n_ff; ++ int32_t n_out; ++ int32_t routed_gate_groups; ++ uint64_t shared_gate_nb1; ++ uint64_t shared_up_nb1; ++ uint64_t shared_down_nb1; ++} ggml_metal_kargs_glm_moe_two_phase; ++ ++typedef struct { ++ uint64_t q_b_nb1; ++ uint64_t wk_b_nb1; ++ uint64_t wk_b_nb2; ++ int32_t q_rank; ++ int32_t q_head_dim; ++ int32_t q_nope_dim; ++ int32_t q_abs_dim; ++ int32_t rope_dim; ++ int32_t n_head; ++ int32_t rope_mode; ++ int32_t _pad0; ++} ggml_metal_kargs_glm_absorbed_q; ++ + // NORM + // RMS_NORM + typedef struct { +@@ -935,6 +1060,49 @@ typedef struct { + uint64_t nb3; + } ggml_metal_kargs_gated_delta_net; + ++typedef struct { ++ int32_t ne00; ++ int32_t ne01; ++ int32_t ne02; ++ int32_t ne03; ++ uint64_t nb00; ++ uint64_t nb01; ++ uint64_t nb02; ++ uint64_t nb03; ++ int32_t ne10; ++ int32_t ne11; ++ int32_t ne12; ++ int32_t ne13; ++ uint64_t nb10; ++ uint64_t nb11; ++ uint64_t nb12; ++ uint64_t nb13; ++ int32_t ne20; ++ int32_t ne21; ++ int32_t ne22; ++ int32_t ne23; ++ uint64_t nb20; ++ uint64_t nb21; ++ uint64_t nb22; ++ uint64_t nb23; ++ int32_t ne30; ++ int32_t ne31; ++ int32_t ne32; ++ int32_t ne33; ++ uint64_t nb30; ++ uint64_t nb31; ++ uint64_t nb32; ++ uint64_t nb33; ++ int32_t ne0; ++ int32_t ne1; ++ int32_t ne2; ++ int32_t ne3; ++ uint64_t nb0; ++ uint64_t nb1; ++ uint64_t nb2; ++ uint64_t nb3; ++} ggml_metal_kargs_lightning_indexer; ++ + typedef struct { + int32_t ne00; + int32_t ne01; +@@ -993,6 +1161,116 @@ typedef struct { + uint64_t nb3; + } ggml_metal_kargs_set_rows; + ++typedef struct { ++ int32_t n_kv; ++ int32_t n_batch; ++ int32_t n_stream; ++ int32_t n_top_k; ++ int32_t n_top_stream; ++ int32_t elem_size; ++ int32_t _pad1; ++ int32_t _pad2; ++ uint64_t nb01; ++ uint64_t nb02; ++ uint64_t nb03; ++ uint64_t nb10; ++ uint64_t nb11; ++ uint64_t nb12; ++ uint64_t nb0; ++ uint64_t nb1; ++ uint64_t nb2; ++ uint64_t nb3; ++} ggml_metal_kargs_dsa_sparse_mask; ++ ++typedef struct { ++ int32_t ne00; ++ int32_t ne01; ++ int32_t ne02; ++ int32_t ne03; ++ uint64_t nb00; ++ uint64_t nb01; ++ uint64_t nb02; ++ uint64_t nb03; ++ int32_t ne10; ++ int32_t ne11; ++ int32_t ne12; ++ int32_t ne13; ++ uint64_t nb10; ++ uint64_t nb11; ++ uint64_t nb12; ++ uint64_t nb13; ++ int32_t ne20; ++ int32_t ne21; ++ int32_t ne22; ++ int32_t ne23; ++ uint64_t nb20; ++ uint64_t nb21; ++ uint64_t nb22; ++ uint64_t nb23; ++ int32_t ne30; ++ int32_t ne31; ++ int32_t ne32; ++ int32_t ne33; ++ uint64_t nb30; ++ uint64_t nb31; ++ uint64_t nb32; ++ uint64_t nb33; ++ int32_t ne40; ++ int32_t ne41; ++ int32_t ne42; ++ int32_t ne43; ++ uint64_t nb40; ++ uint64_t nb41; ++ uint64_t nb42; ++ uint64_t nb43; ++ int32_t ne0; ++ int32_t ne1; ++ int32_t ne2; ++ int32_t ne3; ++ uint64_t nb0; ++ uint64_t nb1; ++ uint64_t nb2; ++ uint64_t nb3; ++ float scale; ++} ggml_metal_kargs_dsa_sparse_attn; ++ ++typedef ggml_metal_kargs_dsa_sparse_attn ggml_metal_kargs_selected_row_flash; ++ ++typedef struct { ++ int32_t ne00; ++ int32_t ne01; ++ int32_t ne02; ++ int32_t ne03; ++ uint64_t nb00; ++ uint64_t nb01; ++ uint64_t nb02; ++ uint64_t nb03; ++ int32_t ne10; ++ int32_t ne11; ++ int32_t ne12; ++ int32_t ne13; ++ uint64_t nb10; ++ uint64_t nb11; ++ uint64_t nb12; ++ uint64_t nb13; ++ int32_t ne20; ++ int32_t ne21; ++ int32_t ne22; ++ int32_t ne23; ++ uint64_t nb20; ++ uint64_t nb21; ++ uint64_t nb22; ++ uint64_t nb23; ++ int32_t ne0; ++ int32_t ne1; ++ int32_t ne2; ++ int32_t ne3; ++ uint64_t nb0; ++ uint64_t nb1; ++ uint64_t nb2; ++ uint64_t nb3; ++} ggml_metal_kargs_dsa_top1_attn; ++ + typedef struct { + int32_t ne00; + int32_t ne01; +@@ -1157,6 +1435,57 @@ typedef struct { + int32_t len; + } ggml_metal_kargs_argsort_merge; + ++typedef struct { ++ int32_t n_expert; ++ int32_t n_tokens; ++ int32_t top_k; ++ int32_t has_bias; ++ int32_t norm; ++ int32_t _pad0; ++ int32_t _pad1; ++ int32_t _pad2; ++ float scale; ++ float clamp_min; ++ uint64_t logits_nb0; ++ uint64_t logits_nb1; ++ uint64_t bias_nb0; ++ uint64_t ids_nb0; ++ uint64_t ids_nb1; ++ uint64_t weights_nb1; ++ uint64_t weights_nb2; ++} ggml_metal_kargs_topk_moe_route; ++ ++typedef struct { ++ int32_t n_expert; ++ int32_t n_tokens; ++ int32_t n_expert_used; ++ int32_t norm; ++ float clamp_min; ++ float scale; ++ int32_t _pad0; ++ int32_t _pad1; ++ uint64_t probs_nb1; ++ uint64_t probs_nb2; ++ uint64_t ids_nb0; ++ uint64_t ids_nb1; ++ uint64_t dst_nb1; ++ uint64_t dst_nb2; ++} ggml_metal_kargs_moe_route_weights; ++ ++typedef struct { ++ int32_t n_embd; ++ int32_t n_tokens; ++ int32_t n_expert_used; ++ int32_t already_weighted; ++ uint64_t experts_nb0; ++ uint64_t experts_nb1; ++ uint64_t experts_nb2; ++ uint64_t weights_nb1; ++ uint64_t weights_nb2; ++ uint64_t dst_nb0; ++ uint64_t dst_nb1; ++} ggml_metal_kargs_moe_weighted_sum; ++ + typedef struct { + int64_t ne0; + float start; +diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp +index c716f118..36378f8e 100644 +--- a/ggml/src/ggml-metal/ggml-metal-ops.cpp ++++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp +@@ -10,8 +10,10 @@ + + #include + #include ++#include + #include + #include ++#include + + static ggml_metal_buffer_id ggml_metal_get_buffer_id(const ggml_tensor * t) { + if (!t) { +@@ -25,266 +27,5445 @@ static ggml_metal_buffer_id ggml_metal_get_buffer_id(const ggml_tensor * t) { + return ggml_metal_buffer_get_id(ctx, t); + } + +-struct ggml_metal_op { +- ggml_metal_op( +- ggml_metal_device_t dev, +- ggml_metal_cmd_buf_t cmd_buf, +- ggml_cgraph * gf, +- int idx_start, +- int idx_end, +- bool use_fusion, +- bool use_concurrency, +- bool use_capture, +- int debug_graph, +- int debug_fusion) { +- this->dev = dev; +- this->lib = ggml_metal_device_get_library(dev); +- this->enc = ggml_metal_encoder_init(cmd_buf, use_concurrency); +- this->mem_ranges = ggml_mem_ranges_init(debug_graph); +- this->idx_start = idx_start; +- this->idx_end = idx_end; +- this->use_fusion = use_fusion; +- this->use_concurrency = use_concurrency; +- this->use_capture = use_capture; +- this->debug_graph = debug_graph; +- this->debug_fusion = debug_fusion; +- this->gf = gf; ++static bool ggml_metal_lightning_indexer_parallel_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_PARALLEL_LIGHTNING_INDEXER"); ++ return value && atoi(value) != 0; ++} + +- idxs.reserve(gf->n_nodes); ++static bool ggml_metal_lightning_indexer_staged_q_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_EXPERIMENTAL_LIGHTNING_INDEXER_STAGED_Q"); ++ return value && atoi(value) != 0; ++} + +- // filter empty nodes +- // TODO: this can be removed when the allocator starts filtering them earlier +- // https://github.com/ggml-org/llama.cpp/pull/16130#issuecomment-3327905830 +- for (int i = idx_start; i < idx_end; i++) { +- if (!ggml_op_is_empty(gf->nodes[i]->op) && !ggml_is_empty(gf->nodes[i])) { +- idxs.push_back(i); +- } +- } ++static int ggml_metal_lightning_indexer_parallel_threads_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_PARALLEL_LIGHTNING_INDEXER_THREADS"); ++ if (value == nullptr || value[0] == '\0') { ++ return 64; + } + +- ~ggml_metal_op() { +- ggml_metal_encoder_end_encoding(this->enc); +- ggml_metal_encoder_free(this->enc); +- ggml_mem_ranges_free(this->mem_ranges); ++ const int requested = atoi(value); ++ switch (requested) { ++ case 32: ++ case 64: ++ case 128: ++ case 256: ++ case 512: ++ case 1024: ++ return requested; ++ default: ++ return 64; + } ++} + +- int n_nodes() const { +- return idxs.size(); ++static bool ggml_metal_glm_dsa_dispatch_log_enabled() { ++ const char * value = getenv("GGML_METAL_MOE_DISPATCH_LOG"); ++ if (value && atoi(value) != 0) { ++ return true; + } ++ value = getenv("SKIPPY_GLM_DSA_LOG_METAL_DISPATCH"); ++ return value && atoi(value) != 0; ++} + +- ggml_tensor * node(int i) const { +- assert(i >= 0 && i < (int) idxs.size()); +- return ggml_graph_node(gf, idxs[i]); ++static bool ggml_metal_glm_dsa_topk_moe_fusion_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_TOPK_MOE_ROUTE_FUSION"); ++ if (value) { ++ return atoi(value) != 0; + } ++ value = getenv("SKIPPY_GLM_DSA_ENABLE_METAL_TOPK_MOE_FUSION"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ value = getenv("GGML_METAL_DISABLE_TOPK_MOE_ROUTE_FUSION"); ++ if (value && atoi(value) != 0) { ++ return false; ++ } ++ value = getenv("SKIPPY_GLM_DSA_DISABLE_METAL_TOPK_MOE_FUSION"); ++ if (value && atoi(value) != 0) { ++ return false; ++ } ++ return true; ++} + +- bool can_fuse(int i0, const ggml_op * ops, int n_ops) const { +- assert(use_fusion); +- assert(i0 >= 0 && i0 < n_nodes()); ++static bool ggml_metal_glm_dsa_topk_moe_route_sg_reduce_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_TOPK_MOE_ROUTE_SG_REDUCE"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_TOPK_MOE_ROUTE_SG_REDUCE"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ value = getenv("GGML_METAL_DISABLE_TOPK_MOE_ROUTE_SG_REDUCE"); ++ if (value && atoi(value) != 0) { ++ return false; ++ } ++ value = getenv("SKIPPY_GLM_DSA_DISABLE_TOPK_MOE_ROUTE_SG_REDUCE"); ++ if (value && atoi(value) != 0) { ++ return false; ++ } ++ return true; ++} + +- if (i0 + n_ops > n_nodes()) { +- return false; +- } ++static bool ggml_metal_glm_dsa_topk_moe_route_glm_256_8_sg32_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_TOPK_MOE_ROUTE_SG32"); ++ return value && atoi(value) != 0; ++} + +- return ggml_can_fuse_ext(gf, idxs.data() + i0, ops, n_ops); ++static bool ggml_metal_glm_dsa_moe_decode_motif_reference_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_GLM_MOE_DECODE_MOTIF_REFERENCE"); ++ if (value) { ++ return atoi(value) != 0; + } ++ value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_MOE_DECODE_MOTIF_REFERENCE"); ++ return value && atoi(value) != 0; ++} + +- ggml_metal_device_t dev; +- ggml_metal_library_t lib; +- ggml_metal_encoder_t enc; +- ggml_mem_ranges_t mem_ranges; ++static bool ggml_metal_glm_dsa_moe_private_scratch_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_PRIVATE_SCRATCH"); ++ return value && atoi(value) != 0; ++} + +- bool use_fusion; +- bool use_concurrency; +- bool use_capture; ++static bool ggml_metal_glm_dsa_moe_two_phase_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_TWO_PHASE"); ++ return value && atoi(value) != 0; ++} + +- int debug_graph; +- int debug_fusion; ++static bool ggml_metal_glm_dsa_moe_dual_lane_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_DUAL_LANE"); ++ return value && atoi(value) != 0; ++} + +-private: +- ggml_cgraph * gf; ++static bool ggml_metal_glm_dsa_moe_dual_lane_gate_slot4_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_DUAL_LANE_GATE_SLOT4"); ++ return value && atoi(value) != 0; ++} + +- int idx_start; +- int idx_end; ++static bool ggml_metal_glm_dsa_absorbed_qkv_phases_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_ABSORBED_QKV_PHASES"); ++ return value && atoi(value) != 0; ++} + +- // non-empty node indices +- std::vector idxs; +-}; ++static bool ggml_metal_glm_dsa_moe_route_anchor_bypass_enabled() { ++ const char * disabled = getenv("GGML_METAL_DISABLE_GLM_MOE_ROUTE_ANCHOR_BYPASS"); ++ if (disabled) { ++ return atoi(disabled) == 0; ++ } ++ const char * value = getenv("GGML_METAL_ENABLE_GLM_MOE_ROUTE_ANCHOR_BYPASS"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ return true; ++} + +-ggml_metal_op_t ggml_metal_op_init( +- ggml_metal_device_t dev, +- ggml_metal_cmd_buf_t cmd_buf, +- ggml_cgraph * gf, +- int idx_start, +- int idx_end, +- bool use_fusion, +- bool use_concurrency, +- bool use_capture, +- int debug_graph, +- int debug_fusion) { +- ggml_metal_op_t res = new ggml_metal_op( +- dev, +- cmd_buf, +- gf, +- idx_start, +- idx_end, +- use_fusion, +- use_concurrency, +- use_capture, +- debug_graph, +- debug_fusion); ++static bool ggml_metal_glm_dsa_moe_swiglu_q3_down_fusion_enabled() { ++ const char * disabled = getenv("GGML_METAL_DISABLE_GLM_MOE_SWIGLU_Q3_DOWN_FUSION"); ++ if (disabled) { ++ return atoi(disabled) == 0; ++ } ++ const char * value = getenv("GGML_METAL_ENABLE_GLM_MOE_SWIGLU_Q3_DOWN_FUSION"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_MOE_SWIGLU_Q3_DOWN_FUSION"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ return true; ++} + +- return res; ++static bool ggml_metal_glm_dsa_moe_swiglu_q2_down_fusion_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_SWIGLU_Q2_DOWN_FUSION"); ++ return value && atoi(value) != 0; + } + +-void ggml_metal_op_free(ggml_metal_op_t ctx) { +- delete ctx; ++static bool ggml_metal_glm_dsa_moe_decode_skip_internal_barriers_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_DECODE_SKIP_INTERNAL_BARRIERS"); ++ return value && atoi(value) != 0; + } + +-int ggml_metal_op_n_nodes(ggml_metal_op_t ctx) { +- return ctx->n_nodes(); ++static bool ggml_metal_glm_dsa_moe_decode_skip_route_barrier_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_DECODE_SKIP_ROUTE_BARRIER"); ++ return value && atoi(value) != 0; + } + +-static bool ggml_metal_op_concurrency_reset(ggml_metal_op_t ctx) { +- if (!ctx->mem_ranges) { +- return true; +- } ++static bool ggml_metal_glm_dsa_moe_decode_skip_gate_up_barrier_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_DECODE_SKIP_GATE_UP_BARRIER"); ++ return value && atoi(value) != 0; ++} + +- ggml_metal_encoder_memory_barrier(ctx->enc); ++static bool ggml_metal_glm_dsa_moe_decode_scoped_barriers_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_DECODE_SCOPED_BARRIERS"); ++ return value && atoi(value) != 0; ++} + +- ggml_mem_ranges_reset(ctx->mem_ranges); ++static bool ggml_metal_glm_dsa_moe_route_gate_up_fusion_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_ROUTE_GATE_UP_FUSION"); ++ return value && atoi(value) != 0; ++} + +- return true; ++static bool ggml_metal_glm_dsa_moe_sort_route_ids_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_SORT_ROUTE_IDS"); ++ return value && atoi(value) != 0; + } + +-static bool ggml_metal_op_concurrency_check(ggml_metal_op_t ctx, const ggml_tensor * node) { +- if (!ctx->mem_ranges) { +- return false; ++static int ggml_metal_glm_dsa_moe_max_active_experts() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_MAX_ACTIVE_EXPERTS"); ++ if (value == nullptr) { ++ return 0; + } + +- return ggml_mem_ranges_check(ctx->mem_ranges, node); ++ const int requested = atoi(value); ++ return requested >= 1 && requested <= 8 ? requested : 0; + } + +-static bool ggml_metal_op_concurrency_add(ggml_metal_op_t ctx, const ggml_tensor * node) { +- if (!ctx->mem_ranges) { +- return true; ++static int ggml_metal_glm_dsa_moe_q2_weight_roofline_chunks() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_Q2_WEIGHT_ROOFLINE_CHUNKS"); ++ if (value == nullptr) { ++ return 0; + } + +- return ggml_mem_ranges_add(ctx->mem_ranges, node); ++ const int requested = atoi(value); ++ return requested >= 1 && requested <= 64 ? requested : 0; + } + +-static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { +- struct ggml_tensor * node = ctx->node(idx); +- +- //GGML_LOG_INFO("%s: encoding node %3d, op = %8s\n", __func__, idx, ggml_op_name(node->op)); ++static bool ggml_metal_glm_dsa_moe_q2_weight_roofline_bypass_route_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_Q2_WEIGHT_ROOFLINE_BYPASS_ROUTE"); ++ return value && atoi(value) != 0; ++} + +- if (ggml_is_empty(node)) { +- return 1; ++static int ggml_metal_glm_dsa_moe_q2_weight_roofline_block_bytes() { ++ const int q2_k_block_bytes = int(ggml_type_size(GGML_TYPE_Q2_K)); ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_Q2_WEIGHT_ROOFLINE_BLOCK_BYTES"); ++ if (value == nullptr) { ++ return q2_k_block_bytes; + } + +- switch (node->op) { +- case GGML_OP_NONE: +- case GGML_OP_RESHAPE: +- case GGML_OP_VIEW: +- case GGML_OP_TRANSPOSE: +- case GGML_OP_PERMUTE: +- { +- // noop -> next node +- if (ctx->debug_graph > 0) { +- GGML_LOG_DEBUG("%s: node[%5d] - %-12s %s\n", __func__, idx, ggml_op_name(node->op), "(noop)"); +- } +- } return 1; +- default: +- { +- } break; +- } ++ const int requested = atoi(value); ++ return requested >= q2_k_block_bytes && requested <= 256 && requested % 4 == 0 ? ++ requested : q2_k_block_bytes; ++} + +- if (!ggml_metal_device_supports_op(ctx->dev, node)) { +- GGML_LOG_ERROR("%s: error: unsupported op '%s'\n", __func__, ggml_op_desc(node)); +- GGML_ABORT("unsupported op"); ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_Q2_GATE_UP_SWIGLU_FUSION"); ++ if (value) { ++ return atoi(value) != 0; + } ++ value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_GATE_UP_SWIGLU"); ++ return value && atoi(value) != 0; ++} + +- if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { +- return 1; ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_default_enabled() { ++ const char * disabled = getenv("GGML_METAL_DISABLE_Q2_GATE_UP_SWIGLU_FUSION"); ++ if (disabled) { ++ return atoi(disabled) == 0; + } ++ return false; ++} + +- int n_fuse = 1; ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_vecscale_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_VECSCALE"); ++ return value && atoi(value) != 0; ++} + +- // check if the current node can run concurrently with other nodes before it +- // the condition is that: +- // - the current node cannot write to any previous src or dst ranges +- // - the current node cannot read from any previous dst ranges +- // +- // if the condition is not satisfied, we put a memory barrier and clear all ranges +- // otherwise, we add the new ranges to the encoding context and process the node concurrently +- // +- { +- const bool is_concurrent = ggml_metal_op_concurrency_check(ctx, node); ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG"); ++ return value && atoi(value) != 0; ++} + +- if (!is_concurrent) { +- ggml_metal_op_concurrency_reset(ctx); +- } ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot8_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SLOT8"); ++ return value && atoi(value) != 0; ++} + +- if (ctx->debug_graph > 0) { +- GGML_LOG_DEBUG("%s: node[%5d] - %-12s %-12s %s\n", __func__, idx, ggml_op_name(node->op), ggml_get_name(node), is_concurrent ? "(concurrent)" : ""); +- } +- if (ctx->debug_graph > 1) { +- GGML_TENSOR_LOCALS( int64_t, ne0, node->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, node->src[0], nb); +- GGML_TENSOR_LOCALS( int64_t, ne1, node->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, node->src[1], nb); +- GGML_TENSOR_LOCALS( int64_t, ne2, node->src[2], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb2, node->src[2], nb); +- GGML_TENSOR_LOCALS( int64_t, ne3, node->src[3], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb3, node->src[3], nb); +- GGML_TENSOR_LOCALS( int64_t, ne, node, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, node, nb); ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot2_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SLOT2"); ++ return value && atoi(value) != 0; ++} + +- if (node->src[0]) { +- GGML_LOG_DEBUG("%s: src0 - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(node->src[0]->type), ne00, ne01, ne02, ne03, nb00, nb01, nb02, nb03, +- ggml_is_contiguous(node->src[0]), node->src[0]->name); +- } +- if (node->src[1]) { +- GGML_LOG_DEBUG("%s: src1 - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(node->src[1]->type), ne10, ne11, ne12, ne13, nb10, nb11, nb12, nb13, +- ggml_is_contiguous(node->src[1]), node->src[1]->name); +- } +- if (node->src[2]) { +- GGML_LOG_DEBUG("%s: src2 - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(node->src[2]->type), ne20, ne21, ne22, ne23, nb20, nb21, nb22, nb23, +- ggml_is_contiguous(node->src[2]), node->src[2]->name); +- } +- if (node->src[3]) { +- GGML_LOG_DEBUG("%s: src3 - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(node->src[3]->type), ne30, ne31, ne32, ne33, nb30, nb31, nb32, nb33, +- ggml_is_contiguous(node->src[3]), node->src[3]->name); +- } +- if (node) { +- GGML_LOG_DEBUG("%s: node - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], 1, %s\n", __func__, ggml_type_name(node->type), ne0, ne1, ne2, ne3, nb0, nb1, nb2, nb3, +- node->name); +- } +- } +- } ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SLOT4_DUAL"); ++ return value && atoi(value) != 0; ++} + +- switch (node->op) { +- case GGML_OP_CONCAT: +- { +- n_fuse = ggml_metal_op_concat(ctx, idx); +- } break; +- case GGML_OP_ADD: +- case GGML_OP_SUB: +- case GGML_OP_MUL: +- case GGML_OP_DIV: +- { +- n_fuse = ggml_metal_op_bin(ctx, idx); +- } break; +- case GGML_OP_ADD_ID: +- { +- n_fuse = ggml_metal_op_add_id(ctx, idx); +- } break; +- case GGML_OP_REPEAT: +- { +- n_fuse = ggml_metal_op_repeat(ctx, idx); +- } break; +- case GGML_OP_ACC: +- { +- n_fuse = ggml_metal_op_acc(ctx, idx); ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_r12_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SLOT4_DUAL_R12"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_r16_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SLOT4_DUAL_R16"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot1_dual_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SLOT1_DUAL"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot2_dual_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SLOT2_DUAL"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot8_split_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SLOT8_SPLIT"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_share_y_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_SHARE_Y"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_vecscale_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_VECSCALE"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_q8_act_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_Q8_ACT"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_prequant_q8_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_PREQUANT_Q8"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_inblock_repack_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_INBLOCK_REPACK"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_half_y_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_HALF_Y"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_rowtile_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_ROWTILE"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_r16_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_R16"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_r12_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_R12"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_any_variant_enabled() { ++ return ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot8_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot2_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_r12_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_r16_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot1_dual_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot2_dual_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot8_split_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_share_y_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_vecscale_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_q8_act_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_prequant_q8_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_inblock_repack_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_half_y_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_rowtile_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_r16_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_r12_enabled(); ++} ++ ++static bool ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_default_enabled() { ++ const char * disabled = getenv("GGML_METAL_DISABLE_Q2_GATE_UP_SWIGLU_PAIR_SG"); ++ if (disabled) { ++ return atoi(disabled) == 0; ++ } ++ return true; ++} ++ ++static bool ggml_metal_glm_dsa_weighted_swiglu_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_WEIGHTED_SWIGLU"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_weighted_reduce_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_DOWN_WEIGHTED_REDUCE_DIRECT"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SLOT_PARALLEL_REDUCE"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_default_enabled() { ++ const char * disabled = getenv("GGML_METAL_DISABLE_Q2_DOWN_SLOT_PARALLEL_REDUCE"); ++ if (disabled) { ++ return atoi(disabled) == 0; ++ } ++ return true; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r16_w1_enabled() { ++ const char * disabled = getenv("GGML_METAL_DISABLE_Q2_DOWN_SLOT_PARALLEL_REDUCE"); ++ if (disabled) { ++ return atoi(disabled) == 0; ++ } ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SLOT_PARALLEL_REDUCE_R16_W1"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ return true; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r16_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SLOT_PARALLEL_REDUCE_R16"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r4_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SLOT_PARALLEL_REDUCE_R4"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_f16_act_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_F16_ACT"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_shift_high_bits_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SHIFT_HIGH_BITS"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_vec_scale_enabled() { ++ const char * disabled = getenv("GGML_METAL_DISABLE_Q2_DOWN_VEC_SCALE"); ++ if (disabled) { ++ return atoi(disabled) == 0; ++ } ++ const char * value = getenv("GGML_METAL_ENABLE_Q2_DOWN_VEC_SCALE"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ value = getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_VEC_SCALE"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ return true; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_weighted_reduce_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_Q3_DOWN_WEIGHTED_REDUCE_FUSION"); ++ if (value) { ++ return atoi(value) != 0; ++ } ++ value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_WEIGHTED_REDUCE_DIRECT"); ++ return value && atoi(value) != 0; ++} ++ ++static const char * ggml_metal_tensor_name(const ggml_tensor * tensor); ++ ++static bool ggml_metal_glm_dsa_q3_down_weighted_reduce_tensor_selected( ++ const ggml_tensor * tensor) { ++ const char * selected = getenv("GGML_METAL_Q3_DOWN_WEIGHTED_REDUCE_TENSOR"); ++ return selected == nullptr || selected[0] == '\0' || ++ strcmp(selected, ggml_metal_tensor_name(tensor)) == 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8_W0"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r6_nb8_w0_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R6_NB8_W0"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r10_nb8_w0_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R10_NB8_W0"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_glm52_w0_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_GLM52_W0"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_f16_act_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_Q3_DOWN_F16_ACT"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_default_enabled() { ++ return false; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w1_default_enabled() { ++ return false; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r12_nb8_w0_enabled() { ++ const char * value = getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R12_NB8_W0"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r16_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE_R16"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_slot_split2_reduce_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_SPLIT2_REDUCE"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q3_down_atomic_accum_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_ATOMIC_ACCUM"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_weighted_reduce_reference_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_DOWN_WEIGHTED_REDUCE_REFERENCE"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_weighted_reduce_noop_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_DOWN_WEIGHTED_REDUCE_NOOP"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_shared_expert_noop_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_SHARED_EXPERT_NOOP"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_routed_expert_noop_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_ROUTED_EXPERT_NOOP"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_indexer_projection_noop_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_INDEXER_PROJECTION_NOOP"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_attention_projection_noop_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_ATTN_PROJECTION_NOOP"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_lightning_indexer_noop_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_LIGHTNING_INDEXER_NOOP"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_selected_row_flash_noop_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_SELECTED_ROW_FLASH_NOOP"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_compact_multihead_flash_enabled(const ggml_tensor * op) { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_MULTIHEAD_FLASH"); ++ if (value == nullptr || atoi(value) == 0) { ++ return false; ++ } ++ ++ const char * tensor = getenv("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_MULTIHEAD_FLASH_TENSOR"); ++ return tensor == nullptr || tensor[0] == '\0' || strstr(ggml_metal_tensor_name(op), tensor) != nullptr; ++} ++ ++static bool ggml_metal_glm_compact_split_exact_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_SPLIT_EXACT"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_compact_sequential_v_diagnostic_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_SEQUENTIAL_V_DIAGNOSTIC"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_compact_legacy_scores_v_diagnostic_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_LEGACY_SCORES_V_DIAGNOSTIC"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_compact_dump_scores_diagnostic_enabled() { ++ const char * value = getenv("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_DUMP_SCORES_DIAGNOSTIC"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_moe_route_weights_slot0_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_MOE_ROUTE_WEIGHTS_SLOT0"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_q2_down_weighted_reduce_stock_slot0_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_DOWN_WEIGHTED_REDUCE_STOCK_SLOT0"); ++ return value && atoi(value) != 0; ++} ++ ++static bool ggml_metal_glm_dsa_selected_row_flash_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_SELECTED_ROW_FLASH"); ++ return value == nullptr || value[0] == '\0' || atoi(value) != 0; ++} ++ ++static int ggml_metal_glm_dsa_sparse_attn_threads_requested() { ++ const char * value = getenv("SKIPPY_GLM_DSA_SPARSE_ATTN_THREADS"); ++ if (value == nullptr || value[0] == '\0') { ++ return 256; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 32: ++ case 64: ++ case 128: ++ case 256: ++ return requested; ++ default: ++ return 256; ++ } ++} ++ ++static int ggml_metal_glm_dsa_sparse_attn_threads_for_shape(int requested, int n_batch, int n_top_k) { ++ // Large-top-k prefill shapes are sensitive to the 256-thread sparse-attn ++ // kernel on Apple Metal and can leave rows unwritten. Keep decode and ++ // small-top-k prefill on the requested path, but cap large prefill rows to ++ // the shape that passes backend parity. ++ if (n_batch > 1 && n_top_k >= 64) { ++ return std::min(requested, 32); ++ } ++ if (n_top_k > 512) { ++ return std::min(requested, 128); ++ } ++ return requested; ++} ++ ++static bool ggml_metal_glm_dsa_sparse_attn_cache_topk_enabled() { ++ const char * value = getenv("SKIPPY_GLM_DSA_SPARSE_ATTN_CACHE_TOPK"); ++ return value && atoi(value) != 0; ++} ++ ++static int ggml_metal_glm_dsa_sparse_attn_decode_group_heads_requested() { ++ const char * value = getenv("SKIPPY_GLM_DSA_SPARSE_ATTN_DECODE_GROUP_HEADS"); ++ if (value == nullptr || value[0] == '\0') { ++ return 1; ++ } ++ ++ const int requested = atoi(value); ++ switch (requested) { ++ case 2: ++ case 4: ++ return requested; ++ default: ++ return 1; ++ } ++} ++ ++static int ggml_metal_glm_dsa_mul_mm_id_min_tokens_requested() { ++ const char * value = getenv("SKIPPY_GLM_DSA_MUL_MM_ID_MIN_TOKENS"); ++ if (value == nullptr || value[0] == '\0') { ++ return 32; ++ } ++ ++ const int requested = atoi(value); ++ return requested > 0 ? requested : 32; ++} ++ ++static const char * ggml_metal_tensor_name(const ggml_tensor * tensor) { ++ return tensor != nullptr && tensor->name[0] != '\0' ? tensor->name : ""; ++} ++ ++static void ggml_metal_log_topk_moe_route_encode_candidate(ggml_metal_op_t ctx, int idx, const char * reason); ++ ++struct ggml_metal_topk_moe_route_fusion { ++ ggml_tensor * logits = nullptr; ++ ggml_tensor * ids = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * bias = nullptr; ++ ggml_tensor * clamp = nullptr; ++ ggml_tensor * scale = nullptr; ++ int n_fuse = 0; ++}; ++ ++struct ggml_metal_mul_mv_id_gate_up_swiglu_fusion { ++ ggml_tensor * up = nullptr; ++ ggml_tensor * gate = nullptr; ++ ggml_tensor * glu = nullptr; ++ ggml_tensor * cast = nullptr; ++ ggml_tensor * weighted = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * q8 = nullptr; ++ int n_fuse = 0; ++}; ++ ++struct ggml_metal_weighted_swiglu_fusion { ++ ggml_tensor * glu = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * weighted = nullptr; ++ int n_fuse = 0; ++}; ++ ++struct ggml_metal_mul_mv_id_weighted_reduce_fusion { ++ ggml_tensor * down = nullptr; ++ ggml_tensor * shared_gate = nullptr; ++ ggml_tensor * shared_up = nullptr; ++ ggml_tensor * weighted_sum = nullptr; ++ int shared_gate_offset = -1; ++ int shared_up_offset = -1; ++ int weighted_sum_offset = -1; ++ int n_fuse = 0; ++}; ++ ++struct ggml_metal_glm_moe_private_bindings; ++ ++int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_repeat(ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_unary(ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_moe_weighted_sum(ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_moe_mul_mat_id(ggml_metal_op_t ctx, int idx); ++static int ggml_metal_op_weighted_swiglu(ggml_metal_op_t ctx, int idx); ++static int ggml_metal_op_mul_mv_id_gate_up_swiglu( ++ ggml_metal_op_t ctx, ++ int idx, ++ ggml_tensor * src1_override = nullptr); ++static int ggml_metal_op_glm_moe_route_gate_up_swiglu( ++ ggml_metal_op_t ctx, ++ const struct ggml_metal_glm_moe_decode_motif_reference & motif, ++ ggml_tensor * src1_override, ++ const ggml_metal_glm_moe_private_bindings * private_bindings = nullptr); ++static int ggml_metal_encode_moe_mul_mat_id( ++ ggml_metal_op_t ctx, ++ ggml_tensor * op, ++ const ggml_metal_glm_moe_private_bindings * private_bindings); ++static int ggml_metal_op_mul_mv_id_weighted_reduce( ++ ggml_metal_op_t ctx, ++ int idx, ++ bool subgraph_owned = false); ++ ++struct ggml_metal_op { ++ ggml_metal_op( ++ ggml_metal_device_t dev, ++ ggml_metal_cmd_buf_t cmd_buf, ++ ggml_metal_buffer_id fusion_scratch, ++ size_t fusion_scratch_size, ++ ggml_cgraph * gf, ++ int idx_start, ++ int idx_end, ++ bool use_fusion, ++ bool use_concurrency, ++ bool use_capture, ++ int debug_graph, ++ int debug_fusion) { ++ this->dev = dev; ++ this->lib = ggml_metal_device_get_library(dev); ++ this->enc = ggml_metal_encoder_init(cmd_buf, use_concurrency); ++ this->mem_ranges = ggml_mem_ranges_init(debug_graph); ++ this->fusion_scratch = fusion_scratch; ++ this->fusion_scratch_size = fusion_scratch_size; ++ this->idx_start = idx_start; ++ this->idx_end = idx_end; ++ this->use_fusion = use_fusion; ++ this->use_concurrency = use_concurrency; ++ this->use_capture = use_capture; ++ this->debug_graph = debug_graph; ++ this->debug_fusion = debug_fusion; ++ this->gf = gf; ++ ++ idxs.reserve(gf->n_nodes); ++ ++ // filter empty nodes ++ // TODO: this can be removed when the allocator starts filtering them earlier ++ // https://github.com/ggml-org/llama.cpp/pull/16130#issuecomment-3327905830 ++ for (int i = idx_start; i < idx_end; i++) { ++ if (!ggml_op_is_empty(gf->nodes[i]->op) && !ggml_is_empty(gf->nodes[i])) { ++ idxs.push_back(i); ++ } ++ } ++ ++ } ++ ++ ~ggml_metal_op() { ++ ggml_metal_encoder_end_encoding(this->enc); ++ ggml_metal_encoder_free(this->enc); ++ ggml_mem_ranges_free(this->mem_ranges); ++ } ++ ++ int n_nodes() const { ++ return idxs.size(); ++ } ++ ++ ggml_tensor * node(int i) const { ++ assert(i >= 0 && i < (int) idxs.size()); ++ return ggml_graph_node(gf, idxs[i]); ++ } ++ ++ int graph_index(int i) const { ++ assert(i >= 0 && i < (int) idxs.size()); ++ return idxs[i]; ++ } ++ ++ int graph_node_count() const { ++ return gf->n_nodes; ++ } ++ ++ uint64_t graph_uid() const { ++ return gf->uid; ++ } ++ ++ int split_start() const { ++ return idx_start; ++ } ++ ++ int split_end() const { ++ return idx_end; ++ } ++ ++ ggml_tensor * graph_node(int i) const { ++ assert(i >= 0 && i < gf->n_nodes); ++ return ggml_graph_node(gf, i); ++ } ++ ++ int32_t graph_node_use_count(int i) const { ++ assert(i >= 0 && i < gf->n_nodes); ++ return ggml_node_get_use_count(gf, i); ++ } ++ ++ int filtered_count_for_graph_span(int i0, int graph_count) const { ++ const int graph_start = graph_index(i0); ++ const int graph_end = graph_start + graph_count; ++ int count = 0; ++ for (int i = i0; i < (int) idxs.size() && idxs[i] < graph_end; ++i) { ++ ++count; ++ } ++ return count; ++ } ++ ++ bool can_fuse_graph_subgraph(int graph_i0, const ggml_op * ops, int n_ops, const int * outputs, int n_outputs) const { ++ assert(use_fusion); ++ return ggml_can_fuse_subgraph(gf, graph_i0, n_ops, ops, outputs, n_outputs); ++ } ++ ++ bool can_fuse(int i0, const ggml_op * ops, int n_ops) const { ++ assert(use_fusion); ++ assert(i0 >= 0 && i0 < n_nodes()); ++ ++ if (i0 + n_ops > n_nodes()) { ++ return false; ++ } ++ ++ return ggml_can_fuse_ext(gf, idxs.data() + i0, ops, n_ops); ++ } ++ ++ bool can_fuse_subgraph(int i0, const ggml_op * ops, int n_ops, const int * output_offsets, int n_outputs) const { ++ assert(use_fusion); ++ assert(i0 >= 0 && i0 < n_nodes()); ++ ++ if (i0 + n_ops > n_nodes()) { ++ return false; ++ } ++ ++ int outputs[4]; ++ GGML_ASSERT(n_outputs <= 4); ++ for (int i = 0; i < n_outputs; ++i) { ++ outputs[i] = idxs[i0 + output_offsets[i]]; ++ } ++ ++ return ggml_can_fuse_subgraph_ext(gf, idxs.data() + i0, n_ops, ops, outputs, n_outputs); ++ } ++ ++ bool can_fuse_filtered_subgraph( ++ int i0, ++ const ggml_op * ops, ++ int n_ops, ++ const int * output_offsets, ++ int n_outputs) const { ++ assert(use_fusion); ++ assert(i0 >= 0 && i0 < n_nodes()); ++ if (i0 + n_ops > n_nodes()) { ++ return false; ++ } ++ ++ const int graph_start = graph_index(i0); ++ const int graph_end = graph_index(i0 + n_ops - 1) + 1; ++ const auto is_output = [&](int offset) { ++ for (int i = 0; i < n_outputs; ++i) { ++ if (output_offsets[i] == offset) { ++ return true; ++ } ++ } ++ return false; ++ }; ++ const auto is_inside_graph_span = [&](const ggml_tensor * tensor) { ++ for (int graph_idx = graph_start; graph_idx < graph_end; ++graph_idx) { ++ if (graph_node(graph_idx) == tensor) { ++ return true; ++ } ++ } ++ return false; ++ }; ++ ++ for (int rel = 0; rel < n_ops; ++rel) { ++ const int graph_idx = graph_index(i0 + rel); ++ const ggml_tensor * node = graph_node(graph_idx); ++ if (node->op != ops[rel] || (node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { ++ return false; ++ } ++ if (is_output(rel)) { ++ continue; ++ } ++ if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { ++ return false; ++ } ++ ++ int span_uses = 0; ++ for (int consumer_idx = graph_start; consumer_idx < graph_end; ++consumer_idx) { ++ const ggml_tensor * consumer = graph_node(consumer_idx); ++ for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { ++ span_uses += consumer->src[src_idx] == node ? 1 : 0; ++ } ++ } ++ if (span_uses != graph_node_use_count(graph_idx)) { ++ return false; ++ } ++ ++ for (const ggml_tensor * view_src = node->view_src; ++ view_src != nullptr; ++ view_src = view_src->view_src) { ++ if (!is_inside_graph_span(view_src)) { ++ return false; ++ } ++ } ++ } ++ return true; ++ } ++ ++ void set_fused_range_outputs(int output0, int output1 = -1, int output2 = -1) { ++ fused_range_output_count = 0; ++ const int outputs[] = { output0, output1, output2 }; ++ for (int output : outputs) { ++ if (output >= 0) { ++ fused_range_outputs[fused_range_output_count++] = output; ++ } ++ } ++ } ++ ++ bool has_fused_range_outputs() const { ++ return fused_range_output_count > 0; ++ } ++ ++ bool tracks_fused_range_output(int offset) const { ++ for (int i = 0; i < fused_range_output_count; ++i) { ++ if (fused_range_outputs[i] == offset) { ++ return true; ++ } ++ } ++ return false; ++ } ++ ++ void clear_fused_range_outputs() { ++ fused_range_output_count = 0; ++ } ++ ++ ggml_metal_device_t dev; ++ ggml_metal_library_t lib; ++ ggml_metal_encoder_t enc; ++ ggml_mem_ranges_t mem_ranges; ++ ggml_metal_buffer_id fusion_scratch; ++ size_t fusion_scratch_size; ++ ++ bool use_fusion; ++ bool use_concurrency; ++ bool use_capture; ++ ++ int debug_graph; ++ int debug_fusion; ++ ++ int fused_range_outputs[3] = { -1, -1, -1 }; ++ int fused_range_output_count = 0; ++ ++private: ++ ggml_cgraph * gf; ++ ++ int idx_start; ++ int idx_end; ++ ++ // non-empty node indices ++ std::vector idxs; ++}; ++ ++static void ggml_metal_log_selected_row_flash_candidate(ggml_metal_op_t ctx, int idx) { ++ if (!ggml_metal_glm_dsa_dispatch_log_enabled() || !ggml_metal_glm_dsa_selected_row_flash_enabled()) { ++ return; ++ } ++ ++ const ggml_tensor * get_rows = ctx->node(idx); ++ if (get_rows == nullptr || get_rows->op != GGML_OP_GET_ROWS || ++ get_rows->name[0] == '\0' || std::strstr(get_rows->name, "dsa_compact_") == nullptr) { ++ return; ++ } ++ ++ const int graph_idx = ctx->graph_index(idx); ++ const int use_count = ctx->graph_node_use_count(graph_idx); ++ const int graph_end = std::min(ctx->split_end(), ctx->graph_node_count()); ++ ++ const ggml_tensor * consumer = nullptr; ++ int consumer_graph_idx = -1; ++ int consumer_src_slot = -1; ++ int consumer_count = 0; ++ const ggml_tensor * view = nullptr; ++ const ggml_tensor * flash = nullptr; ++ int flash_graph_idx = -1; ++ for (int graph_i = graph_idx + 1; graph_i < graph_end; ++graph_i) { ++ const ggml_tensor * candidate = ctx->graph_node(graph_i); ++ if (candidate->op == GGML_OP_VIEW && candidate->src[0] == get_rows && view == nullptr) { ++ view = candidate; ++ } ++ if (candidate->op == GGML_OP_FLASH_ATTN_EXT && ++ candidate->src[1] == get_rows && ++ candidate->src[2] != nullptr && ++ (candidate->src[2] == get_rows || ++ (candidate->src[2]->op == GGML_OP_VIEW && candidate->src[2]->src[0] == get_rows))) { ++ flash = candidate; ++ flash_graph_idx = graph_i; ++ } ++ for (int src_i = 0; src_i < GGML_MAX_SRC; ++src_i) { ++ if (candidate->src[src_i] == get_rows) { ++ ++consumer_count; ++ if (consumer == nullptr) { ++ consumer = candidate; ++ consumer_graph_idx = graph_i; ++ consumer_src_slot = src_i; ++ } ++ } ++ } ++ } ++ ++ const bool is_flash = consumer != nullptr && consumer->op == GGML_OP_FLASH_ATTN_EXT; ++ const bool is_k_or_v = is_flash && (consumer_src_slot == 1 || consumer_src_slot == 2); ++ const bool packed_kv_view = flash != nullptr && flash->src[1] == get_rows && ++ flash->src[2] != nullptr && flash->src[2]->op == GGML_OP_VIEW && flash->src[2]->src[0] == get_rows; ++ const char * reason = packed_kv_view ? "accepted_packed_kv_view" : ++ consumer == nullptr ? "consumer_not_found" : ++ !is_flash ? "consumer_not_flash_attn_ext" : ++ !is_k_or_v ? "consumer_not_kv" : ++ consumer_count != use_count ? "consumer_count_mismatch" : ++ "accepted_graph_pattern"; ++ ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=selected_row_flash_candidate tensor=%s reason=%s graph_uid=%llu split_start=%d split_end=%d filtered_idx=%d filtered_nodes=%d graph_idx=%d use_count=%d get_rows_uses=%d consumer_count=%d consumer_graph_idx=%d consumer_op=%s consumer_tensor=%s consumer_src_slot=%d next_tensor=%s flash_graph_idx=%d generic=0 view=%d src_type=%s rows_type=%s dst_type=%s rows=%lld width=%lld consumer_kv=%lld consumer_heads=%lld consumer_stream=%lld kv=%lld heads=%lld stream=%lld grid_x=1 grid_y=1 grid_z=1 threads_x=1\n", ++ ggml_metal_tensor_name(get_rows), ++ reason, ++ (unsigned long long) ctx->graph_uid(), ++ ctx->split_start(), ++ ctx->split_end(), ++ idx, ++ ctx->n_nodes(), ++ graph_idx, ++ use_count, ++ use_count, ++ consumer_count, ++ consumer_graph_idx, ++ consumer ? ggml_op_name(consumer->op) : "none", ++ ggml_metal_tensor_name(consumer), ++ consumer_src_slot, ++ ggml_metal_tensor_name(flash), ++ flash_graph_idx, ++ view != nullptr, ++ get_rows->src[0] ? ggml_type_name(get_rows->src[0]->type) : "none", ++ get_rows->src[1] ? ggml_type_name(get_rows->src[1]->type) : "none", ++ ggml_type_name(get_rows->type), ++ (long long) get_rows->ne[1], ++ (long long) get_rows->ne[0], ++ consumer ? (long long) consumer->ne[1] : -1LL, ++ consumer ? (long long) consumer->ne[2] : -1LL, ++ consumer ? (long long) consumer->ne[3] : -1LL, ++ flash && flash->src[1] ? (long long) flash->src[1]->ne[1] : 0LL, ++ flash && flash->src[0] ? (long long) flash->src[0]->ne[2] : 0LL, ++ flash && flash->src[0] ? (long long) flash->src[0]->ne[3] : 0LL); ++} ++ ++static bool ggml_metal_selected_row_flash_shape_ok(const ggml_tensor * flash) { ++ if (flash == nullptr || flash->op != GGML_OP_FLASH_ATTN_EXT || ++ !ggml_metal_glm_dsa_selected_row_flash_enabled()) { ++ return false; ++ } ++ ++ ggml_tensor * rows = flash->src[1]; ++ ggml_tensor * view = flash->src[2]; ++ if (flash->src[0] == nullptr || rows == nullptr || view == nullptr || ++ rows->op != GGML_OP_GET_ROWS || view->op != GGML_OP_VIEW || view->src[0] != rows || ++ rows->src[0] == nullptr || rows->src[1] == nullptr) { ++ return false; ++ } ++ ++ const float max_bias = ggml_get_op_params_f32(flash, 1); ++ const float logit_softcap = ggml_get_op_params_f32(flash, 2); ++ if (flash->src[3] != nullptr || flash->src[4] != nullptr || max_bias != 0.0f || logit_softcap != 0.0f) { ++ return false; ++ } ++ ++ return flash->src[0]->type == GGML_TYPE_F32 && ++ rows->src[0]->type == GGML_TYPE_F16 && ++ rows->src[1]->type == GGML_TYPE_I32 && ++ rows->type == GGML_TYPE_F16 && ++ view->type == GGML_TYPE_F16 && ++ flash->type == GGML_TYPE_F32 && ++ rows->ne[0] == rows->src[0]->ne[0] && ++ view->ne[0] > 0 && ++ view->ne[0] <= rows->src[0]->ne[0] && ++ rows->ne[1] == rows->src[1]->ne[0] && ++ rows->src[1]->ne[0] <= 4096 && ++ flash->src[0]->ne[0] == rows->src[0]->ne[0] && ++ flash->ne[0] == view->ne[0]; ++} ++ ++static ggml_tensor * ggml_metal_find_selected_row_flash_consumer(ggml_metal_op_t ctx, int idx, ggml_tensor * rows) { ++ if (!ctx->use_fusion || rows == nullptr || !ggml_metal_glm_dsa_selected_row_flash_enabled()) { ++ return nullptr; ++ } ++ ++ const int graph_idx = ctx->graph_index(idx); ++ const int graph_end = std::min(ctx->split_end(), ctx->graph_node_count()); ++ for (int graph_i = graph_idx + 1; graph_i < graph_end; ++graph_i) { ++ ggml_tensor * candidate = ctx->graph_node(graph_i); ++ if (candidate->op == GGML_OP_FLASH_ATTN_EXT && ++ candidate->src[1] == rows && ++ candidate->src[2] != nullptr && ++ candidate->src[2]->op == GGML_OP_VIEW && ++ candidate->src[2]->src[0] == rows && ++ ggml_metal_selected_row_flash_shape_ok(candidate)) { ++ return candidate; ++ } ++ } ++ ++ return nullptr; ++} ++ ++static bool ggml_metal_selected_row_flash_can_defer_compact_k_rows(ggml_metal_op_t ctx, int idx, const ggml_tensor * rows) { ++ if (rows == nullptr || rows->op != GGML_OP_GET_ROWS || !ggml_metal_glm_dsa_selected_row_flash_enabled()) { ++ return false; ++ } ++ ++ if (rows->name[0] == '\0' || std::strstr(rows->name, "dsa_compact_k_topk_rows") == nullptr) { ++ return false; ++ } ++ ++ if (rows->src[0] == nullptr || rows->src[1] == nullptr) { ++ return false; ++ } ++ ++ const int graph_idx = ctx->graph_index(idx); ++ const int32_t use_count = ctx->graph_node_use_count(graph_idx); ++ ++ return use_count == 2 && ++ rows->src[0]->type == GGML_TYPE_F16 && ++ rows->src[1]->type == GGML_TYPE_I32 && ++ rows->type == GGML_TYPE_F16 && ++ rows->ne[0] == rows->src[0]->ne[0] && ++ rows->ne[1] == rows->src[1]->ne[0] && ++ rows->src[1]->ne[0] <= 4096; ++} ++ ++static bool ggml_metal_selected_row_flash_vec_shape_ok(const ggml_tensor * op) { ++ if (!ggml_metal_selected_row_flash_shape_ok(op)) { ++ return false; ++ } ++ ++ const ggml_tensor * q = op->src[0]; ++ const ggml_tensor * rows = op->src[1]; ++ const ggml_tensor * view = op->src[2]; ++ ++ return q->ne[0] == 576 && ++ view->ne[0] == 512 && ++ q->ne[1] == 1 && ++ rows->src[0]->ne[2] == 1 && ++ view->ne[2] == 1 && ++ rows->src[0]->ne[3] == 1 && ++ view->ne[3] == 1; ++} ++ ++ggml_metal_op_t ggml_metal_op_init( ++ ggml_metal_device_t dev, ++ ggml_metal_cmd_buf_t cmd_buf, ++ ggml_metal_buffer_id fusion_scratch, ++ size_t fusion_scratch_size, ++ ggml_cgraph * gf, ++ int idx_start, ++ int idx_end, ++ bool use_fusion, ++ bool use_concurrency, ++ bool use_capture, ++ int debug_graph, ++ int debug_fusion) { ++ ggml_metal_op_t res = new ggml_metal_op( ++ dev, ++ cmd_buf, ++ fusion_scratch, ++ fusion_scratch_size, ++ gf, ++ idx_start, ++ idx_end, ++ use_fusion, ++ use_concurrency, ++ use_capture, ++ debug_graph, ++ debug_fusion); ++ ++ return res; ++} ++ ++void ggml_metal_op_free(ggml_metal_op_t ctx) { ++ delete ctx; ++} ++ ++int ggml_metal_op_n_nodes(ggml_metal_op_t ctx) { ++ return ctx->n_nodes(); ++} ++ ++static bool ggml_metal_op_concurrency_reset(ggml_metal_op_t ctx) { ++ if (!ctx->mem_ranges) { ++ return true; ++ } ++ ++ ggml_metal_encoder_memory_barrier(ctx->enc); ++ ++ ggml_mem_ranges_reset(ctx->mem_ranges); ++ ++ return true; ++} ++ ++static void ggml_metal_op_internal_phase_barrier(ggml_metal_op_t ctx) { ++ ggml_metal_encoder_memory_barrier(ctx->enc); ++ if (ctx->mem_ranges) { ++ ggml_mem_ranges_reset(ctx->mem_ranges); ++ } ++} ++ ++static bool ggml_metal_op_concurrency_reset_tensors( ++ ggml_metal_op_t ctx, ++ const ggml_tensor * tensor0, ++ const ggml_tensor * tensor1, ++ const ggml_tensor * tensor2) { ++ if (!ctx->mem_ranges) { ++ return true; ++ } ++ ++ ggml_metal_encoder_memory_barrier_buffer(ctx->enc, ggml_metal_get_buffer_id(tensor0)); ++ ggml_metal_encoder_memory_barrier_buffer(ctx->enc, ggml_metal_get_buffer_id(tensor1)); ++ ggml_metal_encoder_memory_barrier_buffer(ctx->enc, ggml_metal_get_buffer_id(tensor2)); ++ ggml_mem_ranges_reset(ctx->mem_ranges); ++ ++ return true; ++} ++ ++static bool ggml_metal_op_concurrency_check(ggml_metal_op_t ctx, const ggml_tensor * node) { ++ if (!ctx->mem_ranges) { ++ return false; ++ } ++ ++ return ggml_mem_ranges_check(ctx->mem_ranges, node); ++} ++ ++static bool ggml_metal_op_concurrency_add(ggml_metal_op_t ctx, const ggml_tensor * node) { ++ if (!ctx->mem_ranges) { ++ return true; ++ } ++ ++ return ggml_mem_ranges_add(ctx->mem_ranges, node); ++} ++ ++static bool ggml_metal_tensor_name_contains(const ggml_tensor * tensor, const char * needle) { ++ return tensor != nullptr && std::strstr(ggml_metal_tensor_name(tensor), needle) != nullptr; ++} ++ ++static bool ggml_metal_tensor_is_shared_expert_gate(const ggml_tensor * tensor) { ++ return ggml_metal_tensor_name_contains(tensor, "ffn_shexp_gate") || ++ ggml_metal_tensor_name_contains(tensor, "ffn_gate"); ++} ++ ++static bool ggml_metal_tensor_is_shared_expert_up(const ggml_tensor * tensor) { ++ return ggml_metal_tensor_name_contains(tensor, "ffn_shexp_up") || ++ ggml_metal_tensor_name_contains(tensor, "ffn_up"); ++} ++ ++static bool ggml_metal_tensor_is_glm_dsa_qk_moe_expert(const ggml_tensor * tensor) { ++ return tensor != nullptr && (tensor->type == GGML_TYPE_Q2_K || tensor->type == GGML_TYPE_Q3_K); ++} ++ ++static void ggml_metal_log_glm_dsa_moe_motif_candidate(ggml_metal_op_t ctx, int idx) { ++ if (!ggml_metal_glm_dsa_dispatch_log_enabled() || idx >= ctx->n_nodes()) { ++ return; ++ } ++ ++ int route_weights_idx = -1; ++ int gate_idx = -1; ++ int up_idx = -1; ++ int glu_idx = -1; ++ int down_idx = -1; ++ int shared_gate_idx = -1; ++ int shared_up_idx = -1; ++ int weighted_sum_idx = -1; ++ ++ for (int j = idx; j < ctx->n_nodes() && j < idx + 18; ++j) { ++ ggml_tensor * node = ctx->node(j); ++ if (route_weights_idx < 0 && ++ node->op == GGML_OP_MOE_ROUTE_WEIGHTS && ++ ggml_metal_tensor_name_contains(node, "ffn_moe_route_weights")) { ++ route_weights_idx = j; ++ } else if (gate_idx < 0 && ++ node->op == GGML_OP_MUL_MAT_ID && ++ ggml_metal_tensor_name_contains(node, "ffn_moe_gate")) { ++ gate_idx = j; ++ } else if (up_idx < 0 && ++ node->op == GGML_OP_MUL_MAT_ID && ++ ggml_metal_tensor_name_contains(node, "ffn_moe_up")) { ++ up_idx = j; ++ } else if (glu_idx < 0 && ++ node->op == GGML_OP_GLU && ++ ggml_metal_tensor_name_contains(node, "ffn_moe_swiglu")) { ++ glu_idx = j; ++ } else if (down_idx < 0 && ++ node->op == GGML_OP_MUL_MAT_ID && ++ ggml_metal_tensor_name_contains(node, "ffn_moe_down")) { ++ down_idx = j; ++ } else if (down_idx >= 0 && shared_gate_idx < 0 && ++ node->op == GGML_OP_MUL_MAT && ++ ggml_metal_tensor_is_shared_expert_gate(node)) { ++ shared_gate_idx = j; ++ } else if (down_idx >= 0 && shared_up_idx < 0 && ++ node->op == GGML_OP_MUL_MAT && ++ ggml_metal_tensor_is_shared_expert_up(node)) { ++ shared_up_idx = j; ++ } else if (weighted_sum_idx < 0 && ++ node->op == GGML_OP_MOE_WEIGHTED_SUM && ++ ggml_metal_tensor_name_contains(node, "ffn_moe_out")) { ++ weighted_sum_idx = j; ++ break; ++ } ++ } ++ ++ if (route_weights_idx < 0 || gate_idx < 0 || up_idx < 0 || glu_idx < 0 || down_idx < 0 || weighted_sum_idx < 0) { ++ return; ++ } ++ ++ ggml_tensor * route_weights = ctx->node(route_weights_idx); ++ ggml_tensor * gate = ctx->node(gate_idx); ++ ggml_tensor * up = ctx->node(up_idx); ++ ggml_tensor * down = ctx->node(down_idx); ++ ggml_tensor * shared_gate = shared_gate_idx >= 0 ? ctx->node(shared_gate_idx) : nullptr; ++ ggml_tensor * shared_up = shared_up_idx >= 0 ? ctx->node(shared_up_idx) : nullptr; ++ ggml_tensor * weighted_sum = ctx->node(weighted_sum_idx); ++ ++ const bool natural_order = ++ route_weights_idx < gate_idx && ++ gate_idx < up_idx && ++ up_idx < glu_idx && ++ glu_idx < down_idx && ++ down_idx < weighted_sum_idx; ++ ++ const bool backend_candidate = ++ natural_order && ++ ggml_metal_tensor_is_glm_dsa_qk_moe_expert(gate->src[0]) && ++ ggml_metal_tensor_is_glm_dsa_qk_moe_expert(up->src[0]) && ++ ggml_metal_tensor_is_glm_dsa_qk_moe_expert(down->src[0]) && ++ gate->src[1] != nullptr && gate->src[1]->type == GGML_TYPE_F32 && ++ up->src[1] != nullptr && up->src[1]->type == GGML_TYPE_F32 && ++ down->src[1] != nullptr && down->src[1]->type == GGML_TYPE_F32 && ++ gate->src[2] != nullptr && gate->src[2]->type == GGML_TYPE_I32 && ++ up->src[2] != nullptr && up->src[2]->type == GGML_TYPE_I32 && ++ down->src[2] != nullptr && down->src[2]->type == GGML_TYPE_I32 && ++ gate->type == GGML_TYPE_F32 && ++ up->type == GGML_TYPE_F32 && ++ down->type == GGML_TYPE_F32 && ++ weighted_sum->type == GGML_TYPE_F32; ++ ++ bool subgraph_fusable = false; ++ if (ctx->use_fusion && down_idx + 3 < ctx->n_nodes() && down_idx + 3 == weighted_sum_idx && ++ shared_gate_idx == down_idx + 1 && shared_up_idx == down_idx + 2) { ++ const ggml_op ops[] = { ++ GGML_OP_MUL_MAT_ID, ++ GGML_OP_MUL_MAT, ++ GGML_OP_MUL_MAT, ++ GGML_OP_MOE_WEIGHTED_SUM, ++ }; ++ const int outputs[] = { 1, 2, 3 }; ++ subgraph_fusable = ctx->can_fuse_subgraph(down_idx, ops, 4, outputs, 3); ++ } ++ ++ const int filtered_gap = weighted_sum_idx - down_idx; ++ const int graph_gap = ctx->graph_index(weighted_sum_idx) - ctx->graph_index(down_idx); ++ ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=glm_dsa_moe_motif_candidate tensor=%s route_weights=%s shared_gate=%s shared_up=%s weighted_sum=%s reason=full_motif natural_order=%d backend_candidate=%d subgraph_fusable=%d motif_nodes=%d fusion_outputs=%d filtered_gap=%d graph_gap=%d weighted_sum_gap=%d weighted_sum_graph_gap=%d src0_type=%s src1_type=%s ids_type=%s dst_type=%s experts=%lld used_experts=%lld tokens=%lld grid_x=1 grid_y=1 grid_z=1 threads_x=1\n", ++ ggml_metal_tensor_name(down), ++ ggml_metal_tensor_name(route_weights), ++ ggml_metal_tensor_name(shared_gate), ++ ggml_metal_tensor_name(shared_up), ++ ggml_metal_tensor_name(weighted_sum), ++ natural_order ? 1 : 0, ++ backend_candidate ? 1 : 0, ++ subgraph_fusable ? 1 : 0, ++ 4, ++ 3, ++ filtered_gap, ++ graph_gap, ++ filtered_gap, ++ graph_gap, ++ ggml_type_name(down->src[0]->type), ++ ggml_type_name(down->src[1]->type), ++ ggml_type_name(down->src[2]->type), ++ ggml_type_name(down->type), ++ (long long) down->src[0]->ne[2], ++ (long long) down->src[2]->ne[0], ++ (long long) down->src[2]->ne[1]); ++} ++ ++static void ggml_metal_log_topk_moe_route_encode_candidate(ggml_metal_op_t ctx, int idx, const char * reason) { ++ if (!ggml_metal_glm_dsa_dispatch_log_enabled() || idx >= ctx->n_nodes()) { ++ return; ++ } ++ ++ const ggml_tensor * node = ctx->node(idx); ++ if (node->name[0] == '\0' || std::strstr(node->name, "ffn_moe_probs") == nullptr) { ++ return; ++ } ++ ++ char ops[512] = {}; ++ size_t offset = 0; ++ for (int j = idx; j < ctx->n_nodes() && j < idx + 14 && offset < sizeof(ops); ++j) { ++ const ggml_tensor * t = ctx->node(j); ++ const int written = snprintf( ++ ops + offset, ++ sizeof(ops) - offset, ++ "%s%s/%s", ++ j == idx ? "" : ",", ++ ggml_op_name(t->op), ++ ggml_metal_tensor_name(t)); ++ if (written < 0) { ++ break; ++ } ++ offset += (size_t) written; ++ } ++ ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=topk_moe_route_encode tensor=%s candidate=%s reason=%s graph_uid=%llu split_start=%d split_end=%d filtered_nodes=%d graph_nodes=%d graph_idx=%d grid_x=1 grid_y=1 grid_z=1 threads_x=1\n", ++ ggml_metal_tensor_name(node), ++ ops, ++ reason, ++ (unsigned long long) ctx->graph_uid(), ++ ctx->split_start(), ++ ctx->split_end(), ++ ctx->n_nodes(), ++ ctx->graph_node_count(), ++ ctx->graph_index(idx)); ++ ggml_metal_log_glm_dsa_moe_motif_candidate(ctx, idx); ++} ++ ++static bool ggml_metal_match_topk_moe_route_fusion( ++ ggml_metal_op_t ctx, ++ int idx, ++ ggml_metal_topk_moe_route_fusion & fusion) { ++ if (!ctx->use_fusion || !ggml_metal_glm_dsa_topk_moe_fusion_enabled()) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "disabled_or_short"); ++ return false; ++ } ++ ++ const int graph_idx = ctx->graph_index(idx); ++ if (graph_idx + 11 > ctx->graph_node_count()) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "disabled_or_short"); ++ return false; ++ } ++ ++ ggml_tensor * sigmoid = ctx->graph_node(graph_idx); ++ if (sigmoid->op != GGML_OP_UNARY || ggml_get_unary_op(sigmoid) != GGML_UNARY_OP_SIGMOID || ++ sigmoid->src[0] == nullptr || sigmoid->src[0]->type != GGML_TYPE_F32) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "not_sigmoid"); ++ return false; ++ } ++ ++ if (idx + 4 <= ctx->n_nodes()) { ++ ggml_tensor * local_sigmoid = ctx->node(idx); ++ ggml_tensor * biased_probs = ctx->node(idx + 1); ++ ggml_tensor * ids = ctx->node(idx + 2); ++ ggml_tensor * weights = ctx->node(idx + 3); ++ ++ const bool compact_sequence = ++ local_sigmoid == sigmoid && ++ biased_probs->op == GGML_OP_ADD && biased_probs->src[0] == sigmoid && ++ biased_probs->src[1] != nullptr && biased_probs->src[1]->type == GGML_TYPE_F32 && ++ ids->op == GGML_OP_TOP_K && ids->src[0] == biased_probs && ids->type == GGML_TYPE_I32 && ++ weights->op == GGML_OP_MOE_ROUTE_WEIGHTS && weights->src[0] != nullptr && weights->src[1] == ids && ++ weights->src[0]->op == GGML_OP_RESHAPE && weights->src[0]->src[0] == sigmoid; ++ ++ if (compact_sequence) { ++ const bool compact_shape = ++ ids->ne[0] > 0 && ids->ne[0] <= 16 && ids->ne[1] == sigmoid->ne[1] && ++ sigmoid->ne[0] > 0 && sigmoid->ne[0] <= 256 && sigmoid->ne[1] > 0 && ++ weights->type == GGML_TYPE_F32 && weights->ne[0] == 1 && ++ weights->ne[1] == ids->ne[0] && weights->ne[2] == ids->ne[1]; ++ if (!compact_shape) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "shape"); ++ return false; ++ } ++ ++ const bool compact_safe = ++ (sigmoid->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && ++ (biased_probs->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && ++ ctx->graph_node_use_count(ctx->graph_index(idx + 1)) == 1; ++ if (!compact_safe) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "can_fuse"); ++ return false; ++ } ++ ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "fused"); ++ fusion.logits = sigmoid->src[0]; ++ fusion.ids = ids; ++ fusion.weights = weights; ++ fusion.bias = biased_probs->src[1]; ++ fusion.clamp = weights; ++ fusion.scale = nullptr; ++ fusion.n_fuse = 4; ++ return true; ++ } ++ ++ } ++ ++ if (graph_idx + 6 <= ctx->graph_node_count()) { ++ ggml_tensor * reshape_probs = ctx->graph_node(graph_idx + 1); ++ ggml_tensor * biased_probs = ctx->graph_node(graph_idx + 2); ++ ggml_tensor * argsort = ctx->graph_node(graph_idx + 3); ++ ggml_tensor * ids = ctx->graph_node(graph_idx + 4); ++ ggml_tensor * weights = ctx->graph_node(graph_idx + 5); ++ ++ const bool argsort_route_weights_graph_sequence = ++ reshape_probs->op == GGML_OP_RESHAPE && reshape_probs->src[0] == sigmoid && ++ biased_probs->op == GGML_OP_ADD && biased_probs->src[0] == sigmoid && ++ biased_probs->src[1] != nullptr && biased_probs->src[1]->type == GGML_TYPE_F32 && ++ argsort->op == GGML_OP_ARGSORT && argsort->src[0] == biased_probs && ++ ids->op == GGML_OP_VIEW && ids->src[0] == argsort && ids->type == GGML_TYPE_I32 && ++ weights->op == GGML_OP_MOE_ROUTE_WEIGHTS && weights->src[0] == reshape_probs && weights->src[1] == ids; ++ ++ if (argsort_route_weights_graph_sequence) { ++ const bool argsort_route_weights_shape = ++ ids->ne[0] > 0 && ids->ne[0] <= 16 && ids->ne[1] == sigmoid->ne[1] && ++ sigmoid->ne[0] > 0 && sigmoid->ne[0] <= 256 && sigmoid->ne[1] > 0 && ++ weights->type == GGML_TYPE_F32 && weights->ne[0] == 1 && ++ weights->ne[1] == ids->ne[0] && weights->ne[2] == ids->ne[1]; ++ if (!argsort_route_weights_shape) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "shape"); ++ return false; ++ } ++ ++ const ggml_op ops[] = { ++ GGML_OP_UNARY, ++ GGML_OP_RESHAPE, ++ GGML_OP_ADD, ++ GGML_OP_ARGSORT, ++ GGML_OP_VIEW, ++ GGML_OP_MOE_ROUTE_WEIGHTS, ++ }; ++ const int outputs[] = { graph_idx + 4, graph_idx + 5 }; ++ if (!ctx->can_fuse_graph_subgraph(graph_idx, ops, 6, outputs, 2)) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "can_fuse"); ++ return false; ++ } ++ ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "fused"); ++ fusion.logits = sigmoid->src[0]; ++ fusion.ids = ids; ++ fusion.weights = weights; ++ fusion.bias = biased_probs->src[1]; ++ fusion.clamp = weights; ++ fusion.scale = nullptr; ++ fusion.n_fuse = ctx->filtered_count_for_graph_span(idx, 6); ++ return true; ++ } ++ } ++ ++ ggml_tensor * reshape_probs = ctx->graph_node(graph_idx + 1); ++ ggml_tensor * biased_probs = ctx->graph_node(graph_idx + 2); ++ ggml_tensor * argsort = ctx->graph_node(graph_idx + 3); ++ ggml_tensor * ids = ctx->graph_node(graph_idx + 4); ++ ggml_tensor * get_rows = ctx->graph_node(graph_idx + 5); ++ ggml_tensor * weights_2d = ctx->graph_node(graph_idx + 6); ++ ggml_tensor * weights_sum = ctx->graph_node(graph_idx + 7); ++ ggml_tensor * clamp = ctx->graph_node(graph_idx + 8); ++ ggml_tensor * div = ctx->graph_node(graph_idx + 9); ++ ggml_tensor * weights_3d = ctx->graph_node(graph_idx + 10); ++ ++ if (reshape_probs->op != GGML_OP_RESHAPE || reshape_probs->src[0] != sigmoid || ++ biased_probs->op != GGML_OP_ADD || biased_probs->src[0] != sigmoid || ++ biased_probs->src[1] == nullptr || biased_probs->src[1]->type != GGML_TYPE_F32 || ++ argsort->op != GGML_OP_ARGSORT || argsort->src[0] != biased_probs || ++ ids->op != GGML_OP_VIEW || ids->src[0] != argsort || ids->type != GGML_TYPE_I32 || ++ get_rows->op != GGML_OP_GET_ROWS || get_rows->src[0] != reshape_probs || get_rows->src[1] != ids || ++ weights_2d->op != GGML_OP_RESHAPE || weights_2d->src[0] != get_rows || ++ weights_sum->op != GGML_OP_SUM_ROWS || weights_sum->src[0] != weights_2d || ++ clamp->op != GGML_OP_CLAMP || clamp->src[0] != weights_sum || ++ div->op != GGML_OP_DIV || div->src[0] != weights_2d || div->src[1] != clamp || ++ weights_3d->op != GGML_OP_RESHAPE || weights_3d->src[0] != div) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "shape_or_sequence"); ++ return false; ++ } ++ ++ ggml_tensor * weights = weights_3d; ++ int graph_n_fuse = 11; ++ ggml_tensor * scale = nullptr; ++ if (graph_idx + 11 < ctx->graph_node_count()) { ++ ggml_tensor * maybe_scale = ctx->graph_node(graph_idx + 11); ++ if (maybe_scale->op == GGML_OP_SCALE && maybe_scale->src[0] == weights_3d) { ++ weights = maybe_scale; ++ scale = maybe_scale; ++ graph_n_fuse = 12; ++ } ++ } ++ ++ if (ids->ne[0] <= 0 || ids->ne[0] > 16 || ids->ne[1] != sigmoid->ne[1] || ++ sigmoid->ne[0] > 256 || sigmoid->ne[1] <= 0 || ++ weights->type != GGML_TYPE_F32 || weights->ne[0] != 1 || weights->ne[1] != ids->ne[0] || ++ weights->ne[2] != ids->ne[1]) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "shape"); ++ return false; ++ } ++ ++ const ggml_op ops_with_scale[] = { ++ GGML_OP_UNARY, ++ GGML_OP_RESHAPE, ++ GGML_OP_ADD, ++ GGML_OP_ARGSORT, ++ GGML_OP_VIEW, ++ GGML_OP_GET_ROWS, ++ GGML_OP_RESHAPE, ++ GGML_OP_SUM_ROWS, ++ GGML_OP_CLAMP, ++ GGML_OP_DIV, ++ GGML_OP_RESHAPE, ++ GGML_OP_SCALE, ++ }; ++ const ggml_op ops_without_scale[] = { ++ GGML_OP_UNARY, ++ GGML_OP_RESHAPE, ++ GGML_OP_ADD, ++ GGML_OP_ARGSORT, ++ GGML_OP_VIEW, ++ GGML_OP_GET_ROWS, ++ GGML_OP_RESHAPE, ++ GGML_OP_SUM_ROWS, ++ GGML_OP_CLAMP, ++ GGML_OP_DIV, ++ GGML_OP_RESHAPE, ++ }; ++ const int outputs[] = { graph_idx + 4, graph_idx + graph_n_fuse - 1 }; ++ if (!ctx->can_fuse_graph_subgraph( ++ graph_idx, ++ scale ? ops_with_scale : ops_without_scale, ++ graph_n_fuse, ++ outputs, ++ 2)) { ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "can_fuse"); ++ return false; ++ } ++ ++ ggml_metal_log_topk_moe_route_encode_candidate(ctx, idx, "fused"); ++ fusion.logits = sigmoid->src[0]; ++ fusion.ids = ids; ++ fusion.weights = weights; ++ fusion.bias = biased_probs->src[1]; ++ fusion.clamp = clamp; ++ fusion.scale = scale; ++ fusion.n_fuse = ctx->filtered_count_for_graph_span(idx, graph_n_fuse); ++ return true; ++} ++ ++static int ggml_metal_op_topk_moe_route_fused(ggml_metal_op_t ctx, int idx) { ++ ggml_metal_topk_moe_route_fusion fusion; ++ if (!ggml_metal_match_topk_moe_route_fusion(ctx, idx, fusion)) { ++ return 0; ++ } ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS(uint64_t, nb_logits, fusion.logits, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb_bias, fusion.bias, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb_ids, fusion.ids, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb_weights,fusion.weights, nb); ++ ++ const bool route_weights_op = fusion.weights->op == GGML_OP_MOE_ROUTE_WEIGHTS; ++ const float scale = fusion.scale ? ggml_get_op_params_f32(fusion.scale, 0) : ++ (route_weights_op ? ggml_get_op_params_f32(fusion.weights, 1) : 1.0f); ++ const float clamp_min = ggml_get_op_params_f32(fusion.clamp, 0); ++ const int32_t norm = route_weights_op ? ggml_get_op_params_i32(fusion.weights, 2) != 0 : 1; ++ ++ ggml_metal_kargs_topk_moe_route args = { ++ /*.n_expert =*/ (int32_t) fusion.logits->ne[0], ++ /*.n_tokens =*/ (int32_t) fusion.logits->ne[1], ++ /*.top_k =*/ (int32_t) fusion.ids->ne[0], ++ /*.has_bias =*/ fusion.bias ? 1 : 0, ++ /*.norm =*/ norm, ++ /*._pad0 =*/ ggml_metal_glm_dsa_moe_sort_route_ids_enabled() ? 1 : 0, ++ /*._pad1 =*/ ggml_metal_glm_dsa_moe_max_active_experts(), ++ /*._pad2 =*/ 0, ++ /*.scale =*/ scale, ++ /*.clamp_min =*/ clamp_min, ++ /*.logits_nb0 =*/ nb_logits0, ++ /*.logits_nb1 =*/ nb_logits1, ++ /*.bias_nb0 =*/ nb_bias0, ++ /*.ids_nb0 =*/ nb_ids0, ++ /*.ids_nb1 =*/ nb_ids1, ++ /*.weights_nb1 =*/ nb_weights1, ++ /*.weights_nb2 =*/ nb_weights2, ++ }; ++ ++ const bool glm_256_8_sg32 = ++ ggml_metal_glm_dsa_topk_moe_route_glm_256_8_sg32_enabled() && ++ fusion.logits->ne[0] == 256 && ++ fusion.ids->ne[0] == 8; ++ const bool sg_reduce = !glm_256_8_sg32 && ggml_metal_glm_dsa_topk_moe_route_sg_reduce_enabled(); ++ auto pipeline = glm_256_8_sg32 ? ++ ggml_metal_library_get_pipeline_topk_moe_route_glm_256_8_sg32(lib) : ++ (sg_reduce ? ++ ggml_metal_library_get_pipeline_topk_moe_route_sg_reduce(lib) : ++ ggml_metal_library_get_pipeline_topk_moe_route(lib)); ++ ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ int ida = 0; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(fusion.logits), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(fusion.bias), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(fusion.ids), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(fusion.weights), ida++); ++ ++ const int nth = glm_256_8_sg32 ? 32 : 256; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=topk_moe_route_fused kernel=%s tensor=%s logits=%s ids=%s weights=%s graph_uid=%llu split_start=%d split_end=%d experts=%d tokens=%d top_k=%d sort_ids=%d fused_nodes=%d scale=%g grid_x=%d grid_y=1 grid_z=1 threads_x=%d\n", ++ glm_256_8_sg32 ? "glm_256_8_sg32" : (sg_reduce ? "simdgroup_reduce" : "parallel_reduce"), ++ ggml_metal_tensor_name(ctx->node(idx)), ++ ggml_metal_tensor_name(fusion.logits), ++ ggml_metal_tensor_name(fusion.ids), ++ ggml_metal_tensor_name(fusion.weights), ++ (unsigned long long) ctx->graph_uid(), ++ ctx->split_start(), ++ ctx->split_end(), ++ args.n_expert, ++ args.n_tokens, ++ args.top_k, ++ args._pad0, ++ fusion.n_fuse, ++ (double) args.scale, ++ args.n_tokens, ++ nth); ++ } ++ ggml_metal_encoder_dispatch_threadgroups(enc, args.n_tokens, 1, 1, nth, 1, 1); ++ ++ return fusion.n_fuse; ++} ++ ++struct ggml_metal_glm_moe_decode_motif_reference { ++ ggml_metal_topk_moe_route_fusion route; ++ ggml_tensor * ids = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * gate = nullptr; ++ ggml_tensor * up = nullptr; ++ ggml_tensor * glu = nullptr; ++ ggml_tensor * weighted_down_input = nullptr; ++ ggml_tensor * down = nullptr; ++ ggml_tensor * out = nullptr; ++ ggml_tensor * shared_gate = nullptr; ++ ggml_tensor * shared_up = nullptr; ++ ggml_tensor * shared_glu = nullptr; ++ ggml_tensor * shared_down = nullptr; ++ ggml_tensor * final_out = nullptr; ++ int shared_gate_offset = -1; ++ int shared_up_offset = -1; ++ int shared_glu_offset = -1; ++ int shared_down_offset = -1; ++ int final_out_offset = -1; ++ int out_offset = -1; ++ ggml_tensor * route_anchor_src = nullptr; ++ int route_n_fuse = 0; ++ int n_fuse = 0; ++ bool has_route_anchor = false; ++ bool has_weighted_down = false; ++ bool has_shared_expert_tail = false; ++ bool has_native_down = false; ++ bool final_only_fusable = false; ++}; ++ ++struct ggml_metal_glm_moe_private_bindings { ++ ggml_metal_buffer_id ids; ++ ggml_metal_buffer_id weights; ++ ggml_metal_buffer_id activation; ++}; ++ ++struct ggml_metal_routed_moe_decode_contract { ++ const char * family = "generic"; ++ const ggml_tensor * ids = nullptr; ++ const ggml_tensor * weights = nullptr; ++ const ggml_tensor * gate = nullptr; ++ const ggml_tensor * up = nullptr; ++ const ggml_tensor * glu = nullptr; ++ const ggml_tensor * down = nullptr; ++ const ggml_tensor * out = nullptr; ++ const ggml_tensor * gate_w = nullptr; ++ const ggml_tensor * up_w = nullptr; ++ const ggml_tensor * down_w = nullptr; ++ const ggml_tensor * cur = nullptr; ++ int route_n_fuse = 0; ++ int n_fuse = 0; ++ bool has_route_anchor = false; ++ bool has_weighted_down = false; ++ bool has_shared_expert_tail = false; ++ long long n_expert = 0; ++ long long top_k = 0; ++ long long n_tokens = 0; ++ long long n_embd = 0; ++ long long n_ff = 0; ++ long long out_embd = 0; ++ uint64_t selected_gate_weight_bytes = 0; ++ uint64_t selected_up_weight_bytes = 0; ++ uint64_t selected_down_weight_bytes = 0; ++ uint64_t selected_weight_bytes = 0; ++ uint64_t unfused_intermediate_bytes = 0; ++ uint64_t current_fused_intermediate_bytes = 0; ++ bool gate_up_pair_sg_shape = false; ++ bool q2_weighted_reduce_shape = false; ++ bool q3_weighted_reduce_shape = false; ++}; ++ ++static bool ggml_metal_make_routed_moe_decode_contract( ++ const char * family, ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ ggml_metal_routed_moe_decode_contract & contract) { ++ const ggml_tensor * gate_w = motif.gate != nullptr ? motif.gate->src[0] : nullptr; ++ const ggml_tensor * up_w = motif.up != nullptr ? motif.up->src[0] : nullptr; ++ const ggml_tensor * down_w = motif.down != nullptr ? motif.down->src[0] : nullptr; ++ const ggml_tensor * cur = motif.gate != nullptr ? motif.gate->src[1] : nullptr; ++ if (gate_w == nullptr || up_w == nullptr || down_w == nullptr || ++ cur == nullptr || motif.ids == nullptr || motif.weights == nullptr || ++ motif.gate == nullptr || motif.up == nullptr || motif.glu == nullptr || ++ motif.down == nullptr || motif.out == nullptr) { ++ return false; ++ } ++ ++ contract.family = family; ++ contract.ids = motif.ids; ++ contract.weights = motif.weights; ++ contract.gate = motif.gate; ++ contract.up = motif.up; ++ contract.glu = motif.glu; ++ contract.down = motif.down; ++ contract.out = motif.out; ++ contract.gate_w = gate_w; ++ contract.up_w = up_w; ++ contract.down_w = down_w; ++ contract.cur = cur; ++ contract.route_n_fuse = motif.route_n_fuse; ++ contract.n_fuse = motif.n_fuse; ++ contract.has_route_anchor = motif.has_route_anchor; ++ contract.has_weighted_down = motif.has_weighted_down; ++ contract.has_shared_expert_tail = motif.has_shared_expert_tail; ++ contract.top_k = motif.ids->ne[0]; ++ contract.n_tokens = motif.ids->ne[1]; ++ contract.n_expert = gate_w->ne[2]; ++ contract.n_embd = cur->ne[0]; ++ contract.n_ff = motif.glu->ne[0]; ++ contract.out_embd = motif.out->ne[0]; ++ contract.selected_gate_weight_bytes = ++ uint64_t(contract.top_k) * uint64_t(contract.n_ff) * ++ ggml_row_size(gate_w->type, contract.n_embd); ++ contract.selected_up_weight_bytes = ++ uint64_t(contract.top_k) * uint64_t(contract.n_ff) * ++ ggml_row_size(up_w->type, contract.n_embd); ++ contract.selected_down_weight_bytes = ++ uint64_t(contract.top_k) * uint64_t(contract.out_embd) * ++ ggml_row_size(down_w->type, contract.n_ff); ++ contract.selected_weight_bytes = ++ contract.selected_gate_weight_bytes + ++ contract.selected_up_weight_bytes + ++ contract.selected_down_weight_bytes; ++ const uint64_t slot_activation_bytes = ++ uint64_t(contract.top_k) * uint64_t(contract.n_tokens) * ++ uint64_t(contract.n_ff) * sizeof(float); ++ const uint64_t out_activation_bytes = ++ uint64_t(contract.n_tokens) * uint64_t(contract.out_embd) * sizeof(float); ++ contract.unfused_intermediate_bytes = ++ 3 * slot_activation_bytes + out_activation_bytes; ++ contract.current_fused_intermediate_bytes = ++ slot_activation_bytes + out_activation_bytes; ++ contract.gate_up_pair_sg_shape = ++ gate_w->type == GGML_TYPE_Q2_K && ++ up_w->type == GGML_TYPE_Q2_K && ++ contract.n_embd == 6144 && ++ contract.n_ff == 2048 && ++ contract.top_k == 8 && ++ contract.n_tokens == 1; ++ contract.q2_weighted_reduce_shape = ++ down_w->type == GGML_TYPE_Q2_K && ++ down_w->ne[0] == contract.n_ff && ++ down_w->ne[1] == contract.out_embd && ++ contract.top_k == 8 && ++ contract.n_tokens == 1; ++ contract.q3_weighted_reduce_shape = ++ down_w->type == GGML_TYPE_Q3_K && ++ down_w->ne[0] == contract.n_ff && ++ down_w->ne[1] == contract.out_embd && ++ contract.top_k == 8 && ++ contract.n_tokens == 1; ++ return true; ++} ++ ++static bool ggml_metal_glm_moe_can_scan_q2_selected_weights( ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ int chunks_per_expert, ++ int storage_block_bytes) { ++ ggml_metal_routed_moe_decode_contract contract; ++ if (chunks_per_expert <= 0 || ++ !ggml_metal_make_routed_moe_decode_contract("glm_dsa", motif, contract)) { ++ return false; ++ } ++ ++ const uint64_t vector_bytes = 4*sizeof(uint32_t); ++ const uint64_t q2_k_block_bytes = ggml_type_size(GGML_TYPE_Q2_K); ++ const auto stored_bytes = [storage_block_bytes, q2_k_block_bytes](uint64_t bytes, ggml_type type) { ++ if (type != GGML_TYPE_Q2_K) { ++ return bytes; ++ } ++ GGML_ASSERT(bytes % q2_k_block_bytes == 0); ++ return bytes/q2_k_block_bytes*uint64_t(storage_block_bytes); ++ }; ++ const uint64_t gate_expert_bytes = stored_bytes( ++ contract.selected_gate_weight_bytes/uint64_t(contract.top_k), contract.gate_w->type); ++ const uint64_t up_expert_bytes = stored_bytes( ++ contract.selected_up_weight_bytes/uint64_t(contract.top_k), contract.up_w->type); ++ const uint64_t down_expert_bytes = stored_bytes( ++ contract.selected_down_weight_bytes/uint64_t(contract.top_k), contract.down_w->type); ++ ++ return contract.gate_up_pair_sg_shape && ++ (contract.q2_weighted_reduce_shape || contract.q3_weighted_reduce_shape) && ++ storage_block_bytes >= int(q2_k_block_bytes) && ++ storage_block_bytes <= 256 && ++ storage_block_bytes % 4 == 0 && ++ (storage_block_bytes == int(q2_k_block_bytes) || contract.n_expert >= 3*contract.top_k + 2) && ++ ggml_is_contiguous(contract.gate_w) && ++ ggml_is_contiguous(contract.up_w) && ++ ggml_is_contiguous(contract.down_w) && ++ gate_expert_bytes % vector_bytes == 0 && ++ up_expert_bytes % vector_bytes == 0 && ++ down_expert_bytes % vector_bytes == 0 && ++ ggml_nelements(contract.out) >= contract.top_k*chunks_per_expert; ++} ++ ++static int ggml_metal_op_glm_moe_q2_selected_weight_scan( ++ ggml_metal_op_t ctx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ int chunks_per_expert, ++ int storage_block_bytes) { ++ GGML_ASSERT(ggml_metal_glm_moe_can_scan_q2_selected_weights( ++ motif, chunks_per_expert, storage_block_bytes)); ++ ++ ggml_metal_routed_moe_decode_contract contract; ++ GGML_ASSERT(ggml_metal_make_routed_moe_decode_contract("glm_dsa", motif, contract)); ++ ++ const uint64_t top_k = uint64_t(contract.top_k); ++ const uint64_t q2_k_block_bytes = ggml_type_size(GGML_TYPE_Q2_K); ++ const auto stored_bytes = [storage_block_bytes, q2_k_block_bytes](uint64_t bytes, ggml_type type) { ++ if (type != GGML_TYPE_Q2_K) { ++ return bytes; ++ } ++ GGML_ASSERT(bytes % q2_k_block_bytes == 0); ++ return bytes/q2_k_block_bytes*uint64_t(storage_block_bytes); ++ }; ++ ggml_metal_kargs_glm_moe_q2_weight_scan args = { ++ /*.gate_expert_bytes =*/ stored_bytes(contract.selected_gate_weight_bytes/top_k, contract.gate_w->type), ++ /*.gate_expert_stride =*/ contract.gate_w->nb[2], ++ /*.up_expert_bytes =*/ stored_bytes(contract.selected_up_weight_bytes/top_k, contract.up_w->type), ++ /*.up_expert_stride =*/ contract.up_w->nb[2], ++ /*.down_expert_bytes =*/ stored_bytes(contract.selected_down_weight_bytes/top_k, contract.down_w->type), ++ /*.down_expert_stride =*/ contract.down_w->nb[2], ++ /*.ids_nb1 =*/ contract.ids->nb[1], ++ /*.top_k =*/ int32_t(contract.top_k), ++ /*.n_tokens =*/ int32_t(contract.n_tokens), ++ /*.n_experts =*/ int32_t(contract.n_expert), ++ /*.chunks_per_expert =*/ chunks_per_expert, ++ /*.storage_block_bytes=*/ storage_block_bytes, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_glm_moe_q2_selected_weight_scan(ctx->lib); ++ ggml_metal_encoder_set_pipeline(ctx->enc, pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(contract.gate_w), 1); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(contract.up_w), 2); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(contract.down_w), 3); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(contract.ids), 4); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(contract.out), 5); ++ ggml_metal_encoder_set_threadgroup_memory_size(ctx->enc, pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups( ++ ctx->enc, ++ chunks_per_expert, ++ int(contract.top_k), ++ int(contract.n_tokens), ++ 256, ++ 1, ++ 1); ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_q2_selected_weight_roofline tensor=%s selected_weight_bytes=%llu storage_block_bytes=%d chunks_per_expert=%d threadgroups=%lld threads_per_threadgroup=256\n", ++ ggml_metal_tensor_name(contract.out), ++ (unsigned long long) (stored_bytes(contract.selected_gate_weight_bytes, contract.gate_w->type) + ++ stored_bytes(contract.selected_up_weight_bytes, contract.up_w->type) + ++ stored_bytes(contract.selected_down_weight_bytes, contract.down_w->type)), ++ storage_block_bytes, ++ chunks_per_expert, ++ contract.top_k*contract.n_tokens*chunks_per_expert); ++ } ++ ++ return motif.n_fuse; ++} ++ ++static void ggml_metal_log_routed_moe_decode_motif_contract( ++ ggml_metal_op_t ctx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif) { ++ if (!ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ return; ++ } ++ ++ ggml_metal_routed_moe_decode_contract contract; ++ if (!ggml_metal_make_routed_moe_decode_contract("glm_dsa", motif, contract)) { ++ return; ++ } ++ ++ const double weight_to_current_intermediate = ++ contract.current_fused_intermediate_bytes == 0 ? 0.0 : ++ double(contract.selected_weight_bytes) / double(contract.current_fused_intermediate_bytes); ++ ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=routed_moe_decode_motif_contract family=%s tensor=%s ids=%s weights=%s gate=%s up=%s glu=%s down=%s out=%s graph_uid=%llu route_nodes=%d motif_nodes=%d route_anchor=%d weighted_down=%d shared_expert_tail=%d gate_type=%s up_type=%s down_type=%s cur_type=%s out_type=%s experts=%lld top_k=%lld tokens=%lld embd=%lld ff=%lld out_embd=%lld selected_weight_bytes=%llu selected_gate_weight_bytes=%llu selected_up_weight_bytes=%llu selected_down_weight_bytes=%llu unfused_intermediate_bytes=%llu current_fused_intermediate_bytes=%llu weight_to_current_intermediate=%.2f gate_up_pair_sg_shape=%d q2_weighted_reduce_shape=%d q3_weighted_reduce_shape=%d split_start=%d split_end=%d\n", ++ contract.family, ++ ggml_metal_tensor_name(contract.out), ++ ggml_metal_tensor_name(contract.ids), ++ ggml_metal_tensor_name(contract.weights), ++ ggml_metal_tensor_name(contract.gate), ++ ggml_metal_tensor_name(contract.up), ++ ggml_metal_tensor_name(contract.glu), ++ ggml_metal_tensor_name(contract.down), ++ ggml_metal_tensor_name(contract.out), ++ (unsigned long long) ctx->graph_uid(), ++ contract.route_n_fuse, ++ contract.n_fuse, ++ contract.has_route_anchor ? 1 : 0, ++ contract.has_weighted_down ? 1 : 0, ++ contract.has_shared_expert_tail ? 1 : 0, ++ ggml_type_name(contract.gate_w->type), ++ ggml_type_name(contract.up_w->type), ++ ggml_type_name(contract.down_w->type), ++ ggml_type_name(contract.cur->type), ++ ggml_type_name(contract.out->type), ++ contract.n_expert, ++ contract.top_k, ++ contract.n_tokens, ++ contract.n_embd, ++ contract.n_ff, ++ contract.out_embd, ++ (unsigned long long) contract.selected_weight_bytes, ++ (unsigned long long) contract.selected_gate_weight_bytes, ++ (unsigned long long) contract.selected_up_weight_bytes, ++ (unsigned long long) contract.selected_down_weight_bytes, ++ (unsigned long long) contract.unfused_intermediate_bytes, ++ (unsigned long long) contract.current_fused_intermediate_bytes, ++ weight_to_current_intermediate, ++ contract.gate_up_pair_sg_shape ? 1 : 0, ++ contract.q2_weighted_reduce_shape ? 1 : 0, ++ contract.q3_weighted_reduce_shape ? 1 : 0, ++ ctx->split_start(), ++ ctx->split_end()); ++} ++ ++static bool ggml_metal_match_glm_moe_decode_motif_reference( ++ ggml_metal_op_t ctx, ++ int idx, ++ ggml_metal_glm_moe_decode_motif_reference & motif) { ++ const bool enabled = ++ ggml_metal_glm_dsa_moe_decode_motif_reference_enabled() || ++ ggml_metal_glm_dsa_moe_private_scratch_enabled() || ++ ggml_metal_glm_dsa_moe_two_phase_enabled() || ++ ggml_metal_glm_dsa_moe_dual_lane_enabled(); ++ if (!ctx->use_fusion || !enabled) { ++ return false; ++ } ++ ++ ggml_metal_topk_moe_route_fusion route; ++ ggml_tensor * route_weights = ctx->node(idx); ++ if (route_weights->op == GGML_OP_MOE_ROUTE_WEIGHTS && ++ route_weights->src[0] != nullptr && ++ route_weights->src[1] != nullptr && ++ route_weights->src[0]->type == GGML_TYPE_F32 && ++ route_weights->src[1]->type == GGML_TYPE_I32 && ++ route_weights->type == GGML_TYPE_F32 && ++ route_weights->ne[0] == 1 && ++ route_weights->ne[1] == route_weights->src[1]->ne[0] && ++ route_weights->ne[2] == route_weights->src[1]->ne[1]) { ++ route.ids = route_weights->src[1]; ++ route.weights = route_weights; ++ route.n_fuse = 1; ++ } else if (!ggml_metal_match_topk_moe_route_fusion(ctx, idx, route)) { ++ return false; ++ } ++ ++ int cursor = idx + route.n_fuse; ++ if (cursor >= ctx->n_nodes()) { ++ return false; ++ } ++ ++ bool has_route_anchor = false; ++ if (cursor + 4 < ctx->n_nodes() && ++ ctx->node(cursor + 0)->op == GGML_OP_SUM && ++ ctx->node(cursor + 1)->op == GGML_OP_SCALE && ++ ctx->node(cursor + 2)->op == GGML_OP_REPEAT && ++ ctx->node(cursor + 3)->op == GGML_OP_ADD && ++ ggml_metal_tensor_name_contains(ctx->node(cursor + 0), "ffn_moe_route_anchor_sum") && ++ ggml_metal_tensor_name_contains(ctx->node(cursor + 3), "ffn_moe_cur_route_anchored")) { ++ ggml_tensor * anchor_sum = ctx->node(cursor + 0); ++ ggml_tensor * anchor_zero = ctx->node(cursor + 1); ++ ggml_tensor * anchor_repeat = ctx->node(cursor + 2); ++ ggml_tensor * anchor_cur = ctx->node(cursor + 3); ++ if (anchor_sum->src[0] != route.weights || ++ anchor_zero->src[0] != anchor_sum || ++ anchor_repeat->src[0] != anchor_zero || ++ (anchor_cur->src[0] != anchor_repeat && anchor_cur->src[1] != anchor_repeat)) { ++ return false; ++ } ++ motif.route_anchor_src = anchor_cur->src[0] == anchor_repeat ? anchor_cur->src[1] : anchor_cur->src[0]; ++ has_route_anchor = true; ++ cursor += 4; ++ } ++ ++ if (cursor + 5 > ctx->n_nodes()) { ++ return false; ++ } ++ ++ ggml_tensor * gate = ctx->node(cursor + 0); ++ ggml_tensor * up = ctx->node(cursor + 1); ++ ggml_tensor * glu = ctx->node(cursor + 2); ++ if (gate->op != GGML_OP_MUL_MAT_ID || ++ up->op != GGML_OP_MUL_MAT_ID || ++ glu->op != GGML_OP_GLU || ++ !ggml_metal_tensor_name_contains(gate, "ffn_moe_gate") || ++ !ggml_metal_tensor_name_contains(up, "ffn_moe_up") || ++ !ggml_metal_tensor_name_contains(glu, "ffn_moe_swiglu") || ++ ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || ++ ggml_get_op_params_i32(glu, 1) != 0) { ++ return false; ++ } ++ ++ if (gate->src[1] != up->src[1] || ++ gate->src[2] != route.ids || ++ up->src[2] != route.ids || ++ gate->type != GGML_TYPE_F32 || ++ up->type != GGML_TYPE_F32 || ++ glu->type != GGML_TYPE_F32 || ++ !ggml_are_same_shape(gate, up) || ++ !ggml_are_same_shape(gate, glu)) { ++ return false; ++ } ++ ++ if (glu->src[0] != gate || glu->src[1] != up) { ++ return false; ++ } ++ ++ cursor += 3; ++ ++ bool has_weighted_down = false; ++ ggml_tensor * down_input = glu; ++ if (cursor < ctx->n_nodes() && ++ ctx->node(cursor)->op == GGML_OP_MUL && ++ ggml_metal_tensor_name_contains(ctx->node(cursor), "ffn_moe_down_weighted_input")) { ++ ggml_tensor * weighted = ctx->node(cursor); ++ const bool weighted_sources_ok = ++ (weighted->src[0] == glu && weighted->src[1] == route.weights) || ++ (weighted->src[0] == route.weights && weighted->src[1] == glu); ++ if (!weighted_sources_ok || weighted->type != GGML_TYPE_F32) { ++ return false; ++ } ++ has_weighted_down = true; ++ down_input = weighted; ++ cursor++; ++ } ++ if (cursor < ctx->n_nodes() && ++ ctx->node(cursor)->op == GGML_OP_CPY && ++ ctx->node(cursor)->src[0] == glu && ++ ctx->node(cursor)->type == GGML_TYPE_F16 && ++ ggml_metal_tensor_name_contains(ctx->node(cursor), "ffn_moe_swiglu_f16")) { ++ down_input = ctx->node(cursor); ++ cursor++; ++ } ++ ++ if (ggml_metal_glm_dsa_moe_private_scratch_enabled()) { ++ ggml_tensor * native_down = nullptr; ++ ggml_tensor * shared_gate = nullptr; ++ ggml_tensor * shared_up = nullptr; ++ int native_down_offset = -1; ++ int shared_gate_offset = -1; ++ int shared_up_offset = -1; ++ ++ for (int rel = 0; rel < 4 && cursor + rel < ctx->n_nodes(); ++rel) { ++ ggml_tensor * candidate = ctx->node(cursor + rel); ++ if (candidate->op == GGML_OP_MOE_MUL_MAT_ID) { ++ native_down = candidate; ++ native_down_offset = rel; ++ } else if (candidate->op == GGML_OP_MUL_MAT && ++ ggml_metal_tensor_is_shared_expert_gate(candidate)) { ++ shared_gate = candidate; ++ shared_gate_offset = rel; ++ } else if (candidate->op == GGML_OP_MUL_MAT && ++ ggml_metal_tensor_is_shared_expert_up(candidate)) { ++ shared_up = candidate; ++ shared_up_offset = rel; ++ } ++ } ++ ++ if (native_down == nullptr || shared_gate == nullptr || shared_up == nullptr || ++ native_down->src[0] == nullptr || native_down->src[1] == nullptr || ++ native_down->src[2] == nullptr || native_down->src[3] == nullptr) { ++ return false; ++ } ++ ++ const bool native_shape_ok = ++ !has_weighted_down && ++ native_down->src[0]->type == GGML_TYPE_Q3_K && ++ native_down->src[1] == glu && ++ native_down->src[2] == route.ids && ++ native_down->src[3] == route.weights && ++ native_down->type == GGML_TYPE_F32 && ++ native_down->src[0]->ne[0] == 2048 && ++ native_down->src[0]->ne[1] == 6144 && ++ native_down->src[0]->ne[2] == 256 && ++ native_down->src[1]->type == GGML_TYPE_F32 && ++ native_down->src[1]->ne[0] == 2048 && ++ native_down->src[1]->ne[1] == 8 && ++ native_down->src[1]->ne[2] == 1 && ++ native_down->src[2]->type == GGML_TYPE_I32 && ++ native_down->src[2]->ne[0] == 8 && ++ native_down->src[2]->ne[1] == 1 && ++ native_down->src[3]->type == GGML_TYPE_F32 && ++ native_down->ne[0] == 6144 && ++ native_down->ne[1] == 1 && ++ shared_gate->type == GGML_TYPE_F32 && ++ shared_up->type == GGML_TYPE_F32; ++ if (!native_shape_ok) { ++ return false; ++ } ++ ++ const int native_offset = cursor + native_down_offset - idx; ++ const int native_shared_gate_offset = cursor + shared_gate_offset - idx; ++ const int native_shared_up_offset = cursor + shared_up_offset - idx; ++ const int n_fuse = std::max({ ++ native_offset, ++ native_shared_gate_offset, ++ native_shared_up_offset, ++ }) + 1; ++ ++ std::vector ops; ++ ops.reserve(n_fuse); ++ for (int rel = 0; rel < n_fuse; ++rel) { ++ ops.push_back(ctx->node(idx + rel)->op); ++ } ++ const int outputs[] = { ++ native_offset, ++ native_shared_gate_offset, ++ native_shared_up_offset, ++ }; ++ if (!ctx->can_fuse_subgraph(idx, ops.data(), n_fuse, outputs, 3)) { ++ return false; ++ } ++ ++ motif.route = route; ++ motif.ids = route.ids; ++ motif.weights = route.weights; ++ motif.gate = gate; ++ motif.up = up; ++ motif.glu = glu; ++ motif.down = native_down; ++ motif.out = native_down; ++ motif.shared_gate = shared_gate; ++ motif.shared_up = shared_up; ++ motif.shared_gate_offset = native_shared_gate_offset; ++ motif.shared_up_offset = native_shared_up_offset; ++ motif.out_offset = native_offset; ++ motif.route_anchor_src = has_route_anchor ? motif.route_anchor_src : gate->src[1]; ++ motif.route_n_fuse = route.n_fuse; ++ motif.n_fuse = n_fuse; ++ motif.has_route_anchor = has_route_anchor; ++ motif.has_weighted_down = false; ++ motif.has_shared_expert_tail = true; ++ motif.has_native_down = true; ++ return true; ++ } ++ ++ if (cursor + 2 > ctx->n_nodes()) { ++ return false; ++ } ++ ++ ggml_tensor * down = ctx->node(cursor + 0); ++ ggml_tensor * out = nullptr; ++ ggml_tensor * shared_gate = nullptr; ++ ggml_tensor * shared_up = nullptr; ++ int out_offset = -1; ++ int shared_gate_offset = -1; ++ int shared_up_offset = -1; ++ ++ for (int rel = 1; rel < 4 && cursor + rel < ctx->n_nodes(); ++rel) { ++ ggml_tensor * candidate = ctx->node(cursor + rel); ++ if (candidate->op == GGML_OP_MOE_WEIGHTED_SUM && ++ candidate->src[0] == down && ++ ggml_metal_tensor_name_contains(candidate, "ffn_moe_out")) { ++ out = candidate; ++ out_offset = rel; ++ } else if (candidate->op == GGML_OP_MUL_MAT && ++ ggml_metal_tensor_is_shared_expert_gate(candidate)) { ++ shared_gate = candidate; ++ shared_gate_offset = rel; ++ } else if (candidate->op == GGML_OP_MUL_MAT && ++ ggml_metal_tensor_is_shared_expert_up(candidate)) { ++ shared_up = candidate; ++ shared_up_offset = rel; ++ } ++ } ++ ++ if ((shared_gate == nullptr) != (shared_up == nullptr)) { ++ return false; ++ } ++ const bool has_shared_expert_tail = shared_gate != nullptr; ++ const int tail_n_fuse = std::max({ ++ out_offset, ++ shared_gate_offset, ++ shared_up_offset, ++ }) + 1; ++ if (down->op != GGML_OP_MUL_MAT_ID || ++ out == nullptr || ++ !ggml_metal_tensor_name_contains(down, "ffn_moe_down") || ++ !ggml_metal_tensor_name_contains(out, "ffn_moe_out") || ++ down->src[1] != down_input || ++ down->src[2] != route.ids || ++ out->src[0] != down || ++ out->src[1] != route.weights || ++ down->type != GGML_TYPE_F32 || ++ out->type != GGML_TYPE_F32 || ++ (has_shared_expert_tail && (shared_gate->type != GGML_TYPE_F32 || shared_up->type != GGML_TYPE_F32))) { ++ return false; ++ } ++ ++ const bool expert_types_ok = ++ ggml_metal_tensor_is_glm_dsa_qk_moe_expert(gate->src[0]) && ++ ggml_metal_tensor_is_glm_dsa_qk_moe_expert(up->src[0]) && ++ ggml_metal_tensor_is_glm_dsa_qk_moe_expert(down->src[0]); ++ const bool shape_ok = ++ expert_types_ok && ++ gate->src[0]->ne[2] == down->src[0]->ne[2] && ++ gate->src[2]->ne[0] == down->src[2]->ne[0] && ++ gate->src[2]->ne[1] == down->src[2]->ne[1] && ++ out->ne[0] == down->ne[0] && ++ out->ne[1] == down->ne[2]; ++ if (!shape_ok) { ++ return false; ++ } ++ ++ motif.route = route; ++ motif.ids = route.ids; ++ motif.weights = route.weights; ++ motif.gate = gate; ++ motif.up = up; ++ motif.glu = glu; ++ motif.weighted_down_input = has_weighted_down ? down_input : nullptr; ++ motif.down = down; ++ motif.out = out; ++ motif.shared_gate = shared_gate; ++ motif.shared_up = shared_up; ++ motif.shared_gate_offset = has_shared_expert_tail ? cursor + shared_gate_offset - idx : -1; ++ motif.shared_up_offset = has_shared_expert_tail ? cursor + shared_up_offset - idx : -1; ++ motif.out_offset = cursor + out_offset - idx; ++ if (!has_route_anchor) { ++ motif.route_anchor_src = gate->src[1]; ++ } ++ motif.route_n_fuse = route.n_fuse; ++ motif.n_fuse = cursor + tail_n_fuse - idx; ++ motif.has_route_anchor = has_route_anchor; ++ motif.has_weighted_down = has_weighted_down; ++ motif.has_shared_expert_tail = has_shared_expert_tail; ++ return true; ++} ++ ++static bool ggml_metal_extend_glm_moe_two_phase( ++ ggml_metal_op_t ctx, ++ int idx, ++ ggml_metal_glm_moe_decode_motif_reference & motif) { ++ if ((!ggml_metal_glm_dsa_moe_two_phase_enabled() && ++ !ggml_metal_glm_dsa_moe_dual_lane_enabled()) || ++ !motif.has_shared_expert_tail || motif.has_native_down || ++ motif.shared_gate == nullptr || motif.shared_up == nullptr || ++ motif.route_anchor_src == nullptr || ++ idx + motif.n_fuse + 3 > ctx->n_nodes()) { ++ return false; ++ } ++ ++ const int shared_glu_offset = motif.n_fuse; ++ const int shared_down_offset = shared_glu_offset + 1; ++ const int final_out_offset = shared_down_offset + 1; ++ ggml_tensor * shared_glu = ctx->node(idx + shared_glu_offset); ++ ggml_tensor * shared_down = ctx->node(idx + shared_down_offset); ++ ggml_tensor * final_out = ctx->node(idx + final_out_offset); ++ ++ const bool shared_pair_ok = ++ (shared_glu->src[0] == motif.shared_gate && shared_glu->src[1] == motif.shared_up) || ++ (shared_glu->src[0] == motif.shared_up && shared_glu->src[1] == motif.shared_gate); ++ const bool final_pair_ok = ++ (final_out->src[0] == motif.out && final_out->src[1] == shared_down) || ++ (final_out->src[0] == shared_down && final_out->src[1] == motif.out); ++ const bool shape_ok = ++ motif.gate->src[0]->type == GGML_TYPE_Q2_K && ++ motif.up->src[0]->type == GGML_TYPE_Q2_K && ++ motif.down->src[0]->type == GGML_TYPE_Q3_K && ++ motif.shared_gate->src[0]->type == GGML_TYPE_Q4_K && ++ motif.shared_up->src[0]->type == GGML_TYPE_Q4_K && ++ shared_down->src[0]->type == GGML_TYPE_Q4_K && ++ motif.route_anchor_src->type == GGML_TYPE_F32 && ++ motif.route_anchor_src->ne[0] == 6144 && ++ motif.route_anchor_src->ne[1] == 1 && ++ motif.ids->ne[0] == 8 && motif.ids->ne[1] == 1 && ++ motif.glu->ne[0] == 2048 && motif.glu->ne[1] == 8 && motif.glu->ne[2] == 1 && ++ shared_glu->ne[0] == 2048 && shared_glu->ne[1] == 1 && ++ shared_down->ne[0] == 6144 && shared_down->ne[1] == 1 && ++ final_out->ne[0] == 6144 && final_out->ne[1] == 1; ++ if (shared_glu->op != GGML_OP_GLU || ++ ggml_get_glu_op(shared_glu) != GGML_GLU_OP_SWIGLU || ++ ggml_get_op_params_i32(shared_glu, 1) != 0 || ++ !shared_pair_ok || ++ shared_down->op != GGML_OP_MUL_MAT || ++ shared_down->src[1] != shared_glu || ++ final_out->op != GGML_OP_ADD || ++ !final_pair_ok || ++ !shape_ok) { ++ return false; ++ } ++ ++ motif.shared_glu = shared_glu; ++ motif.shared_down = shared_down; ++ motif.final_out = final_out; ++ motif.shared_glu_offset = shared_glu_offset; ++ motif.shared_down_offset = shared_down_offset; ++ motif.final_out_offset = final_out_offset; ++ motif.n_fuse += 3; ++ ++ std::vector ops; ++ ops.reserve(motif.n_fuse); ++ for (int rel = 0; rel < motif.n_fuse; ++rel) { ++ ops.push_back(ctx->node(idx + rel)->op); ++ } ++ const int final_only_output[] = { motif.final_out_offset }; ++ motif.final_only_fusable = ++ ctx->can_fuse_subgraph(idx, ops.data(), motif.n_fuse, final_only_output, 1) || ++ ctx->can_fuse_filtered_subgraph(idx, ops.data(), motif.n_fuse, final_only_output, 1); ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_dual_lane_contract tensor=%s fused_nodes=%d final_only_fusable=%d routed_use_count=%d final_use_count=%d\n", ++ ggml_metal_tensor_name(motif.final_out), ++ motif.n_fuse, ++ motif.final_only_fusable ? 1 : 0, ++ ctx->graph_node_use_count(ctx->graph_index(idx + motif.out_offset)), ++ ctx->graph_node_use_count(ctx->graph_index(idx + motif.final_out_offset))); ++ } ++ return true; ++} ++ ++static int ggml_metal_op_reference_dispatch_existing_node(ggml_metal_op_t ctx, int idx) { ++ switch (ctx->node(idx)->op) { ++ case GGML_OP_SUM: ++ return ggml_metal_op_sum(ctx, idx); ++ case GGML_OP_SCALE: ++ case GGML_OP_UNARY: ++ case GGML_OP_CLAMP: ++ return ggml_metal_op_unary(ctx, idx); ++ case GGML_OP_REPEAT: ++ return ggml_metal_op_repeat(ctx, idx); ++ case GGML_OP_ADD: ++ case GGML_OP_MUL: ++ case GGML_OP_DIV: ++ case GGML_OP_SUB: ++ return ggml_metal_op_bin(ctx, idx); ++ case GGML_OP_MUL_MAT_ID: ++ return ggml_metal_op_mul_mat_id(ctx, idx); ++ case GGML_OP_MUL_MAT: ++ return ggml_metal_op_mul_mat(ctx, idx); ++ case GGML_OP_GLU: ++ return ggml_metal_op_glu(ctx, idx); ++ case GGML_OP_MOE_WEIGHTED_SUM: ++ return ggml_metal_op_moe_weighted_sum(ctx, idx); ++ default: ++ return 0; ++ } ++} ++ ++static const ggml_tensor * ggml_metal_view_root(const ggml_tensor * tensor) { ++ while (tensor != nullptr && tensor->view_src != nullptr) { ++ tensor = tensor->view_src; ++ } ++ return tensor; ++} ++ ++static bool ggml_metal_pair_is( ++ const ggml_tensor * pair, ++ const ggml_tensor * lhs, ++ const ggml_tensor * rhs) { ++ return pair != nullptr && ++ ((ggml_metal_view_root(pair->src[0]) == ggml_metal_view_root(lhs) && ++ ggml_metal_view_root(pair->src[1]) == ggml_metal_view_root(rhs)) || ++ (ggml_metal_view_root(pair->src[0]) == ggml_metal_view_root(rhs) && ++ ggml_metal_view_root(pair->src[1]) == ggml_metal_view_root(lhs))); ++} ++ ++static int ggml_metal_op_glm_absorbed_q(ggml_metal_op_t ctx, int idx) { ++ constexpr int n_fuse = 5; ++ if (!ctx->use_fusion || !ggml_metal_glm_dsa_absorbed_qkv_phases_enabled() || ++ idx + n_fuse > ctx->n_nodes()) { ++ return 0; ++ } ++ ++ ggml_tensor * q = ctx->node(idx + 0); ++ ggml_tensor * q_abs = ctx->node(idx + 1); ++ ggml_tensor * q_rope = ctx->node(idx + 2); ++ ggml_tensor * kv = ctx->node(idx + 3); ++ ggml_tensor * q_pack = ctx->node(idx + 4); ++ const int rope_mode = ((const int32_t *) q_rope->op_params)[2]; ++ ++ const bool ops_ok = ++ q->op == GGML_OP_MUL_MAT && q_abs->op == GGML_OP_MUL_MAT && ++ q_rope->op == GGML_OP_ROPE && kv->op == GGML_OP_MUL_MAT && ++ q_pack->op == GGML_OP_CONCAT; ++ const bool q_shape_ok = ++ q->type == GGML_TYPE_F32 && ggml_is_contiguous(q) && ++ q->ne[0] == 16384 && q->ne[1] == 1 && ++ q->src[0] != nullptr && q->src[0]->type == GGML_TYPE_Q8_0 && ++ q->src[0]->ne[0] == 2048 && q->src[0]->ne[1] == 16384 && ++ q->src[1] != nullptr && q->src[1]->type == GGML_TYPE_F32 && ++ ggml_is_contiguous(q->src[1]) && q->src[1]->ne[0] == 2048; ++ const bool q_abs_shape_ok = ++ q_abs->type == GGML_TYPE_F32 && ggml_is_contiguous(q_abs) && ++ q_abs->ne[0] == 512 && q_abs->ne[1] == 1 && q_abs->ne[2] == 64 && ++ q_abs->src[0] != nullptr && q_abs->src[0]->type == GGML_TYPE_Q4_0 && ++ q_abs->src[0]->ne[0] == 192 && q_abs->src[0]->ne[1] == 512 && ++ q_abs->src[0]->ne[2] == 64 && q_abs->src[1] != nullptr; ++ const bool rope_shape_ok = ++ q_rope->type == GGML_TYPE_F32 && ggml_is_contiguous(q_rope) && ++ q_rope->ne[0] == 64 && q_rope->ne[1] == 64 && ++ q_rope->src[1] != nullptr && q_rope->src[1]->type == GGML_TYPE_I32 && ++ q_rope->src[1]->ne[0] == 1 && q_rope->src[2] == nullptr && ++ (rope_mode == GGML_ROPE_TYPE_NORMAL || rope_mode == GGML_ROPE_TYPE_NEOX); ++ const bool kv_shape_ok = ++ kv->type == GGML_TYPE_F32 && ggml_is_contiguous(kv) && ++ kv->ne[0] == 576 && kv->ne[1] == 1; ++ const bool pack_shape_ok = ++ q_pack->type == GGML_TYPE_F32 && ggml_is_contiguous(q_pack) && ++ q_pack->ne[0] == 576 && q_pack->ne[1] == 64; ++ const bool dependencies = ++ ggml_metal_view_root(q_abs->src[1]) == q && ++ ggml_metal_view_root(q_rope->src[0]) == q && ++ ggml_metal_view_root(q_pack->src[0]) == q_abs && ++ ggml_metal_view_root(q_pack->src[1]) == q_rope && ++ ggml_metal_tensor_name_contains(q, "q-") && ++ ggml_metal_tensor_name_contains(q->src[0], "attn_q_b.weight") && ++ ggml_metal_tensor_name_contains(q_abs, "q_nope_absorbed") && ++ ggml_metal_tensor_name_contains(q_abs->src[0], "attn_k_b.weight") && ++ ggml_metal_tensor_name_contains(q_rope, "q_pe") && ++ ggml_metal_tensor_name_contains(kv, "kv_cmpr_pe") && ++ ggml_metal_tensor_name_contains(q_pack, "Qcur"); ++ if (!ops_ok || !q_shape_ok || !q_abs_shape_ok || !rope_shape_ok || ++ !kv_shape_ok || !pack_shape_ok || !dependencies) { ++ return 0; ++ } ++ ++ ggml_metal_kargs_glm_absorbed_q fused_args = { ++ /*.q_b_nb1 =*/ q->src[0]->nb[1], ++ /*.wk_b_nb1 =*/ q_abs->src[0]->nb[1], ++ /*.wk_b_nb2 =*/ q_abs->src[0]->nb[2], ++ /*.q_rank =*/ 2048, ++ /*.q_head_dim =*/ 256, ++ /*.q_nope_dim =*/ 192, ++ /*.q_abs_dim =*/ 512, ++ /*.rope_dim =*/ 64, ++ /*.n_head =*/ 64, ++ /*.rope_mode =*/ rope_mode, ++ /*._pad0 =*/ 0, ++ }; ++ ++ const int n_past = ((const int32_t *) q_rope->op_params)[0]; ++ const int n_dims = ((const int32_t *) q_rope->op_params)[1]; ++ const int n_ctx_orig = ((const int32_t *) q_rope->op_params)[4]; ++ float freq_base; ++ float freq_scale; ++ float ext_factor; ++ float attn_factor; ++ float beta_fast; ++ float beta_slow; ++ memcpy(&freq_base, (const int32_t *) q_rope->op_params + 5, sizeof(float)); ++ memcpy(&freq_scale, (const int32_t *) q_rope->op_params + 6, sizeof(float)); ++ memcpy(&ext_factor, (const int32_t *) q_rope->op_params + 7, sizeof(float)); ++ memcpy(&attn_factor, (const int32_t *) q_rope->op_params + 8, sizeof(float)); ++ memcpy(&beta_fast, (const int32_t *) q_rope->op_params + 9, sizeof(float)); ++ memcpy(&beta_slow, (const int32_t *) q_rope->op_params + 10, sizeof(float)); ++ ggml_metal_kargs_rope rope_args = { ++ /*.ne00 =*/ 64, /*.ne01 =*/ 1, /*.ne02 =*/ 1, /*.ne03 =*/ 1, ++ /*.nb00 =*/ sizeof(float), /*.nb01 =*/ 64*sizeof(float), ++ /*.nb02 =*/ 64*sizeof(float), /*.nb03 =*/ 64*sizeof(float), ++ /*.ne0 =*/ 64, /*.ne1 =*/ 1, /*.ne2 =*/ 1, /*.ne3 =*/ 1, ++ /*.nb0 =*/ sizeof(float), /*.nb1 =*/ 64*sizeof(float), ++ /*.nb2 =*/ 64*sizeof(float), /*.nb3 =*/ 64*sizeof(float), ++ /*.n_past =*/ n_past, ++ /*.n_dims =*/ n_dims, ++ /*.n_ctx_orig =*/ n_ctx_orig, ++ /*.freq_base =*/ freq_base, ++ /*.freq_scale =*/ freq_scale, ++ /*.ext_factor =*/ ext_factor, ++ /*.attn_factor =*/ attn_factor, ++ /*.beta_fast =*/ beta_fast, ++ /*.beta_slow =*/ beta_slow, ++ /*.sect_0 =*/ ((const int32_t *) q_rope->op_params)[11], ++ /*.sect_1 =*/ ((const int32_t *) q_rope->op_params)[12], ++ /*.sect_2 =*/ ((const int32_t *) q_rope->op_params)[13], ++ /*.sect_3 =*/ ((const int32_t *) q_rope->op_params)[14], ++ /*.src2 =*/ false, ++ }; ++ ++ const char * kernel = "kernel_glm_absorbed_q_q8_q4"; ++ auto pipeline = ggml_metal_library_get_pipeline(ctx->lib, kernel); ++ if (!pipeline.pipeline) { ++ pipeline = ggml_metal_library_compile_pipeline(ctx->lib, kernel, kernel, nullptr); ++ } ++ int ida = 0; ++ ggml_metal_encoder_set_pipeline(ctx->enc, pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &fused_args, sizeof(fused_args), ida++); ++ ggml_metal_encoder_set_bytes(ctx->enc, &rope_args, sizeof(rope_args), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(q->src[0]), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(q->src[1]), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(q_abs->src[0]), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(q_rope->src[1]), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(q->src[1]), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(q_pack), ida++); ++ ggml_metal_encoder_set_threadgroup_memory_size(ctx->enc, 320*sizeof(float), 0); ++ ggml_metal_encoder_dispatch_threadgroups(ctx->enc, 64, 1, 1, 1024, 1, 1); ++ ++ const int kv_encoded = ggml_metal_op_mul_mat(ctx, idx + 3); ++ GGML_ASSERT(kv_encoded == 1); ++ ctx->set_fused_range_outputs(3, 4); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=glm_absorbed_q_fused kernel=%s tensor=%s fused_nodes=%d dispatch_groups=2 grid_x=64 grid_y=1 grid_z=1 threads_x=1024\n", ++ kernel, ++ ggml_metal_tensor_name(q_pack), ++ n_fuse); ++ } ++ return n_fuse; ++} ++ ++static bool ggml_metal_glm_moe_can_fuse_swiglu_q3_down(const ggml_metal_glm_moe_decode_motif_reference & motif) { ++ if (!ggml_metal_glm_dsa_moe_swiglu_q3_down_fusion_enabled()) { ++ return false; ++ } ++ if (motif.has_weighted_down || ++ motif.gate == nullptr || motif.up == nullptr || ++ motif.glu == nullptr || motif.down == nullptr || motif.out == nullptr || ++ motif.ids == nullptr || motif.weights == nullptr) { ++ return false; ++ } ++ return motif.gate->type == GGML_TYPE_F32 && ++ motif.up->type == GGML_TYPE_F32 && ++ motif.glu->type == GGML_TYPE_F32 && ++ motif.down->type == GGML_TYPE_F32 && ++ motif.out->type == GGML_TYPE_F32 && ++ motif.down->src[0]->type == GGML_TYPE_Q3_K && ++ motif.down->src[1] == motif.glu && ++ motif.down->src[2] == motif.ids && ++ motif.out->src[0] == motif.down && ++ motif.out->src[1] == motif.weights && ++ motif.ids->ne[0] == 8 && ++ motif.ids->ne[1] == 1 && ++ motif.down->src[0]->ne[0] == 2048 && ++ motif.down->src[0]->ne[1] == 6144 && ++ motif.down->src[0]->ne[2] == 256 && ++ motif.out->ne[0] == 6144 && ++ motif.out->ne[1] == 1 && ++ motif.gate->ne[0] == motif.down->src[0]->ne[0] && ++ motif.up->ne[0] == motif.down->src[0]->ne[0] && ++ motif.gate->ne[1] == motif.ids->ne[0] && ++ motif.up->ne[1] == motif.ids->ne[0] && ++ motif.gate->ne[2] == motif.ids->ne[1] && ++ motif.up->ne[2] == motif.ids->ne[1] && ++ motif.glu->ne[0] == motif.gate->ne[0] && ++ motif.glu->ne[1] == motif.gate->ne[1] && ++ motif.glu->ne[2] == motif.gate->ne[2]; ++} ++ ++static int ggml_metal_op_glm_moe_swiglu_q3_down_weighted( ++ ggml_metal_op_t ctx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif) { ++ GGML_ASSERT(ggml_metal_glm_moe_can_fuse_swiglu_q3_down(motif)); ++ ++ ggml_tensor * down = motif.down; ++ ggml_tensor * out = motif.out; ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, down->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, down->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, motif.glu, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, motif.glu, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, motif.ids, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, motif.ids, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, motif.weights, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, out, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, out, nb); ++ ++ auto pipeline = ggml_metal_library_get_pipeline_glm_moe_swiglu_q3_down_weighted(lib, out); ++ const int nr0 = pipeline.nr0; ++ const int nsg = pipeline.nsg; ++ ++ ggml_metal_kargs_mul_mv_id args = { ++ /*.nei0 =*/ ne20, ++ /*.nei1 =*/ ne21, ++ /*.nbi1 =*/ nb21, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nb1 =*/ nb1, ++ /*.nr0 =*/ nr0, ++ }; ++ ggml_metal_kargs_mul_mv_id_weighted_reduce_extra extra = { ++ /*.weights_nb1 =*/ nb31, ++ /*.weights_nb2 =*/ nb32, ++ /*.dst_nb0 =*/ nb0, ++ /*.dst_nb1 =*/ nb1, ++ /*.already_weighted =*/ ggml_get_op_params_i32(out, 0) != 0 ? 1 : 0, ++ /*._pad0 =*/ 0, ++ }; ++ ggml_metal_kargs_glm_moe_swiglu_q3_down swiglu_args = { ++ /*.gate_nb1 =*/ motif.gate->nb[1], ++ /*.gate_nb2 =*/ motif.gate->nb[2], ++ /*.up_nb1 =*/ motif.up->nb[1], ++ /*.up_nb2 =*/ motif.up->nb[2], ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.gate), 2); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.up), 3); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(out), 4); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.ids), 5); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.weights),6); ++ ggml_metal_encoder_set_bytes(enc, &extra, sizeof(extra), 7); ++ ggml_metal_encoder_set_bytes(enc, &swiglu_args, sizeof(swiglu_args), 8); ++ ++ const int grid_x = (ne01 + nr0 - 1)/nr0; ++ const int grid_y = ne21; ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, 1, 32, nsg, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ const ggml_metal_buffer_id bid_gate = ggml_metal_get_buffer_id(motif.gate); ++ const ggml_metal_buffer_id bid_up = ggml_metal_get_buffer_id(motif.up); ++ const ggml_metal_buffer_id bid_out = ggml_metal_get_buffer_id(out); ++ const size_t gate_begin = bid_gate.offs; ++ const size_t gate_end = gate_begin + ggml_nbytes(motif.gate); ++ const size_t up_begin = bid_up.offs; ++ const size_t up_end = up_begin + ggml_nbytes(motif.up); ++ const size_t out_begin = bid_out.offs; ++ const size_t out_end = out_begin + ggml_nbytes(out); ++ const bool gate_out_overlap = ++ bid_gate.metal == bid_out.metal && gate_begin < out_end && out_begin < gate_end; ++ const bool up_out_overlap = ++ bid_up.metal == bid_out.metal && up_begin < out_end && out_begin < up_end; ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_swiglu_q3_down_weighted tensor=%s gate=%s up=%s down=%s weights=%s gate_out_overlap=%d up_out_overlap=%d gate_offs=%llu gate_nbytes=%llu up_offs=%llu up_nbytes=%llu out_offs=%llu out_nbytes=%llu nr0=%d nsg=%d grid_x=%d grid_y=%d threads_x=%d threads_y=%d\n", ++ ggml_metal_tensor_name(out), ++ ggml_metal_tensor_name(motif.gate), ++ ggml_metal_tensor_name(motif.up), ++ ggml_metal_tensor_name(down), ++ ggml_metal_tensor_name(motif.weights), ++ gate_out_overlap ? 1 : 0, ++ up_out_overlap ? 1 : 0, ++ (unsigned long long) gate_begin, ++ (unsigned long long) ggml_nbytes(motif.gate), ++ (unsigned long long) up_begin, ++ (unsigned long long) ggml_nbytes(motif.up), ++ (unsigned long long) out_begin, ++ (unsigned long long) ggml_nbytes(out), ++ nr0, ++ nsg, ++ grid_x, ++ grid_y, ++ 32, ++ nsg); ++ } ++ ++ return 3; ++} ++ ++static bool ggml_metal_glm_moe_can_fuse_swiglu_q2_down(const ggml_metal_glm_moe_decode_motif_reference & motif) { ++ if (!ggml_metal_glm_dsa_moe_swiglu_q2_down_fusion_enabled()) { ++ return false; ++ } ++ if (motif.has_weighted_down || ++ motif.gate == nullptr || motif.up == nullptr || ++ motif.glu == nullptr || motif.down == nullptr || motif.out == nullptr || ++ motif.ids == nullptr || motif.weights == nullptr) { ++ return false; ++ } ++ return motif.gate->type == GGML_TYPE_F32 && ++ motif.up->type == GGML_TYPE_F32 && ++ motif.glu->type == GGML_TYPE_F32 && ++ motif.down->type == GGML_TYPE_F32 && ++ motif.out->type == GGML_TYPE_F32 && ++ motif.down->src[0]->type == GGML_TYPE_Q2_K && ++ motif.down->src[1] == motif.glu && ++ motif.down->src[2] == motif.ids && ++ motif.out->src[0] == motif.down && ++ motif.out->src[1] == motif.weights && ++ ggml_get_op_params_i32(motif.out, 0) == 0 && ++ motif.ids->ne[0] == 8 && ++ motif.ids->ne[1] == 1 && ++ motif.down->src[0]->ne[0] == 2048 && ++ motif.down->src[0]->ne[1] == 6144 && ++ motif.down->src[0]->ne[2] == 256 && ++ motif.out->ne[0] == 6144 && ++ motif.out->ne[1] == 1 && ++ motif.gate->ne[0] == motif.down->src[0]->ne[0] && ++ motif.up->ne[0] == motif.down->src[0]->ne[0] && ++ motif.gate->ne[1] == motif.ids->ne[0] && ++ motif.up->ne[1] == motif.ids->ne[0] && ++ motif.gate->ne[2] == motif.ids->ne[1] && ++ motif.up->ne[2] == motif.ids->ne[1] && ++ motif.glu->ne[0] == motif.gate->ne[0] && ++ motif.glu->ne[1] == motif.gate->ne[1] && ++ motif.glu->ne[2] == motif.gate->ne[2]; ++} ++ ++static int ggml_metal_op_glm_moe_swiglu_q2_down_weighted( ++ ggml_metal_op_t ctx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif) { ++ GGML_ASSERT(ggml_metal_glm_moe_can_fuse_swiglu_q2_down(motif)); ++ ++ ggml_tensor * down = motif.down; ++ ggml_tensor * out = motif.out; ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, down->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, down->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, motif.glu, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, motif.glu, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, motif.ids, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, motif.ids, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, motif.weights, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, out, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, out, nb); ++ ++ auto pipeline = ggml_metal_library_get_pipeline_glm_moe_swiglu_q2_down_weighted(lib, out); ++ const int nr0 = pipeline.nr0; ++ const int nsg = pipeline.nsg; ++ ++ ggml_metal_kargs_mul_mv_id args = { ++ /*.nei0 =*/ ne20, ++ /*.nei1 =*/ ne21, ++ /*.nbi1 =*/ nb21, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nb1 =*/ nb1, ++ /*.nr0 =*/ nr0, ++ }; ++ ggml_metal_kargs_mul_mv_id_weighted_reduce_extra extra = { ++ /*.weights_nb1 =*/ nb31, ++ /*.weights_nb2 =*/ nb32, ++ /*.dst_nb0 =*/ nb0, ++ /*.dst_nb1 =*/ nb1, ++ /*.already_weighted =*/ ggml_get_op_params_i32(out, 0) != 0 ? 1 : 0, ++ /*._pad0 =*/ 0, ++ }; ++ ggml_metal_kargs_glm_moe_swiglu_q3_down swiglu_args = { ++ /*.gate_nb1 =*/ motif.gate->nb[1], ++ /*.gate_nb2 =*/ motif.gate->nb[2], ++ /*.up_nb1 =*/ motif.up->nb[1], ++ /*.up_nb2 =*/ motif.up->nb[2], ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.gate), 2); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.up), 3); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(out), 4); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.ids), 5); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.weights),6); ++ ggml_metal_encoder_set_bytes(enc, &extra, sizeof(extra), 7); ++ ggml_metal_encoder_set_bytes(enc, &swiglu_args, sizeof(swiglu_args), 8); ++ ++ const int grid_x = (ne01 + nr0 - 1)/nr0; ++ const int grid_y = ne21; ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, 1, 32, nsg, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_swiglu_q2_down_weighted tensor=%s gate=%s up=%s down=%s weights=%s nr0=%d nsg=%d grid_x=%d grid_y=%d threads_x=%d threads_y=%d\n", ++ ggml_metal_tensor_name(out), ++ ggml_metal_tensor_name(motif.gate), ++ ggml_metal_tensor_name(motif.up), ++ ggml_metal_tensor_name(down), ++ ggml_metal_tensor_name(motif.weights), ++ nr0, ++ nsg, ++ grid_x, ++ grid_y, ++ 32, ++ nsg); ++ } ++ ++ return 3; ++} ++ ++static bool ggml_metal_glm_moe_route_gate_up_swiglu_shape_ok( ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ const ggml_tensor * src1) { ++ if (motif.route.logits == nullptr || ++ motif.route.ids == nullptr || ++ motif.route.weights == nullptr || ++ motif.route.clamp == nullptr || ++ motif.gate == nullptr || ++ motif.up == nullptr || ++ motif.glu == nullptr || ++ src1 == nullptr) { ++ return false; ++ } ++ ++ const ggml_tensor * dst = motif.weighted_down_input != nullptr ? motif.weighted_down_input : motif.glu; ++ return motif.route.logits->type == GGML_TYPE_F32 && ++ motif.route.logits->ne[0] == 256 && ++ motif.route.logits->ne[1] == 1 && ++ motif.route.ids->type == GGML_TYPE_I32 && ++ motif.route.ids->ne[0] == 8 && ++ motif.route.ids->ne[1] == 1 && ++ motif.route.weights->type == GGML_TYPE_F32 && ++ motif.route.weights->ne[1] == 8 && ++ motif.route.weights->ne[2] == 1 && ++ motif.up->src[0]->type == GGML_TYPE_Q2_K && ++ motif.gate->src[0]->type == GGML_TYPE_Q2_K && ++ motif.up->src[1] == motif.gate->src[1] && ++ motif.up->src[2] == motif.route.ids && ++ motif.gate->src[2] == motif.route.ids && ++ src1->type == GGML_TYPE_F32 && ++ ggml_are_same_shape(src1, motif.up->src[1]) && ++ motif.up->type == GGML_TYPE_F32 && ++ motif.gate->type == GGML_TYPE_F32 && ++ motif.glu->type == GGML_TYPE_F32 && ++ dst->type == GGML_TYPE_F32 && ++ motif.up->src[0]->ne[0] >= 2048 && ++ motif.up->src[0]->ne[1] >= 1024 && ++ motif.up->src[0]->ne[2] == 256 && ++ motif.gate->src[0]->ne[0] >= 2048 && ++ motif.gate->src[0]->ne[1] >= 1024 && ++ motif.gate->src[0]->ne[2] == 256 && ++ motif.gate->src[0]->ne[0] == motif.up->src[0]->ne[0] && ++ motif.gate->src[0]->ne[1] == motif.up->src[0]->ne[1] && ++ motif.glu->ne[0] == motif.up->src[0]->ne[1] && ++ motif.glu->ne[1] == 8 && ++ motif.glu->ne[2] == 1 && ++ dst->ne[0] == motif.glu->ne[0] && ++ dst->ne[1] == motif.glu->ne[1] && ++ dst->ne[2] == motif.glu->ne[2]; ++} ++ ++static bool ggml_metal_glm_moe_can_fuse_route_gate_up_swiglu( ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ const ggml_tensor * src1) { ++ return ggml_metal_glm_dsa_moe_route_gate_up_fusion_enabled() && ++ ggml_metal_glm_moe_route_gate_up_swiglu_shape_ok(motif, src1); ++} ++ ++static int ggml_metal_op_glm_moe_route_gate_up_swiglu( ++ ggml_metal_op_t ctx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ ggml_tensor * src1_override, ++ const ggml_metal_glm_moe_private_bindings * private_bindings) { ++ ggml_tensor * src1 = src1_override != nullptr ? src1_override : motif.up->src[1]; ++ const bool private_path = private_bindings != nullptr; ++ const bool can_fuse = private_path ? ++ ggml_metal_glm_moe_route_gate_up_swiglu_shape_ok(motif, src1) : ++ ggml_metal_glm_moe_can_fuse_route_gate_up_swiglu(motif, src1); ++ if (!can_fuse) { ++ return 0; ++ } ++ ++ ggml_tensor * dst = motif.weighted_down_input != nullptr ? motif.weighted_down_input : motif.glu; ++ ggml_metal_topk_moe_route_fusion route = motif.route; ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS(uint64_t, nb_logits, route.logits, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb_bias, route.bias, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb_ids, route.ids, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb_weights, route.weights, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne0, motif.up->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, motif.up->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, src1, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, src1, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, route.ids, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, route.ids, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, dst, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, dst, nb); ++ ++ const bool route_weights_op = route.weights->op == GGML_OP_MOE_ROUTE_WEIGHTS; ++ const float scale = route.scale ? ggml_get_op_params_f32(route.scale, 0) : ++ (route_weights_op ? ggml_get_op_params_f32(route.weights, 1) : 1.0f); ++ const float clamp_min = ggml_get_op_params_f32(route.clamp, 0); ++ const int32_t norm = route_weights_op ? ggml_get_op_params_i32(route.weights, 2) != 0 : 1; ++ ++ ggml_metal_kargs_topk_moe_route route_args = { ++ /*.n_expert =*/ (int32_t) route.logits->ne[0], ++ /*.n_tokens =*/ (int32_t) route.logits->ne[1], ++ /*.top_k =*/ (int32_t) route.ids->ne[0], ++ /*.has_bias =*/ route.bias ? 1 : 0, ++ /*.norm =*/ norm, ++ /*._pad0 =*/ ggml_metal_glm_dsa_moe_sort_route_ids_enabled() ? 1 : 0, ++ /*._pad1 =*/ ggml_metal_glm_dsa_moe_max_active_experts(), ++ /*._pad2 =*/ 0, ++ /*.scale =*/ scale, ++ /*.clamp_min =*/ clamp_min, ++ /*.logits_nb0 =*/ nb_logits0, ++ /*.logits_nb1 =*/ nb_logits1, ++ /*.bias_nb0 =*/ nb_bias0, ++ /*.ids_nb0 =*/ nb_ids0, ++ /*.ids_nb1 =*/ nb_ids1, ++ /*.weights_nb1 =*/ nb_weights1, ++ /*.weights_nb2 =*/ nb_weights2, ++ }; ++ ggml_metal_kargs_mul_mv_id_gate_up_swiglu gate_args = { ++ /*.nei0 =*/ ne20, ++ /*.nei1 =*/ ne21, ++ /*.nbi1 =*/ nb21, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nb1 =*/ nb1, ++ /*.nr0 =*/ 8, ++ /*.weights_nb1 =*/ nb_weights1, ++ /*.weights_nb2 =*/ nb_weights2, ++ /*.weighted =*/ motif.weighted_down_input != nullptr ? 1 : 0, ++ /*._pad0 =*/ ggml_metal_glm_dsa_moe_max_active_experts(), ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_glm_moe_route_q2_gate_up_swiglu_pair_sg_slot8(lib, dst); ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ const ggml_metal_buffer_id ids_buffer = private_path ? ++ private_bindings->ids : ggml_metal_get_buffer_id(route.ids); ++ const ggml_metal_buffer_id weights_buffer = private_path ? ++ private_bindings->weights : ggml_metal_get_buffer_id(route.weights); ++ const ggml_metal_buffer_id activation_buffer = private_path ? ++ private_bindings->activation : ggml_metal_get_buffer_id(dst); ++ ++ int ida = 0; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &route_args, sizeof(route_args), ida++); ++ ggml_metal_encoder_set_bytes (enc, &gate_args, sizeof(gate_args), ida++); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(route.logits), ida++); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(route.bias ? route.bias : route.logits), ida++); ++ ggml_metal_encoder_set_buffer(enc, ids_buffer, ida++); ++ ggml_metal_encoder_set_buffer(enc, weights_buffer, ida++); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.up->src[0]), ida++); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(motif.gate->src[0]), ida++); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(src1), ida++); ++ ggml_metal_encoder_set_buffer(enc, activation_buffer, ida++); ++ ++ const int nr0 = pipeline.nr0; ++ const int nsg = pipeline.nsg; ++ const int grid_x = (ne01 + nr0 - 1)/nr0; ++ const int grid_z = ne21; ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, 1, grid_z, 32, nsg, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_route_q2_gate_up_swiglu tensor=%s logits=%s ids=%s weights=%s gate=%s up=%s src1=%s src1_override=%d weighted=%d private_scratch=%d experts=%d top_k=%d sort_ids=%d tokens=%d embd=%d ff=%d grid_x=%d grid_z=%d threads_x=32 threads_y=%d\n", ++ ggml_metal_tensor_name(dst), ++ ggml_metal_tensor_name(route.logits), ++ ggml_metal_tensor_name(route.ids), ++ ggml_metal_tensor_name(route.weights), ++ ggml_metal_tensor_name(motif.gate), ++ ggml_metal_tensor_name(motif.up), ++ ggml_metal_tensor_name(src1), ++ src1_override != nullptr ? 1 : 0, ++ motif.weighted_down_input != nullptr ? 1 : 0, ++ private_path ? 1 : 0, ++ route_args.n_expert, ++ route_args.top_k, ++ route_args._pad0, ++ route_args.n_tokens, ++ ne00, ++ ne01, ++ grid_x, ++ grid_z, ++ nsg); ++ } ++ ++ return motif.weighted_down_input != nullptr ? 4 : 3; ++} ++ ++static bool ggml_metal_glm_moe_weights_gate_up_swiglu_shape_ok( ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ const ggml_tensor * src1) { ++ const ggml_tensor * route_weights = motif.route.weights; ++ if (route_weights == nullptr || route_weights->op != GGML_OP_MOE_ROUTE_WEIGHTS || ++ route_weights->src[0] == nullptr || route_weights->src[1] == nullptr || ++ motif.route.ids == nullptr || motif.gate == nullptr || motif.up == nullptr || ++ motif.glu == nullptr || src1 == nullptr) { ++ return false; ++ } ++ ++ return route_weights->src[0]->type == GGML_TYPE_F32 && ++ route_weights->src[1] == motif.route.ids && ++ motif.route.ids->type == GGML_TYPE_I32 && ++ motif.route.ids->ne[0] == 8 && ++ motif.route.ids->ne[1] == 1 && ++ route_weights->type == GGML_TYPE_F32 && ++ route_weights->ne[0] == 1 && ++ route_weights->ne[1] == 8 && ++ route_weights->ne[2] == 1 && ++ motif.up->src[0]->type == GGML_TYPE_Q2_K && ++ motif.gate->src[0]->type == GGML_TYPE_Q2_K && ++ motif.up->src[1] == motif.gate->src[1] && ++ motif.up->src[2] == motif.route.ids && ++ motif.gate->src[2] == motif.route.ids && ++ src1->type == GGML_TYPE_F32 && ++ ggml_are_same_shape(src1, motif.up->src[1]) && ++ motif.up->type == GGML_TYPE_F32 && ++ motif.gate->type == GGML_TYPE_F32 && ++ motif.glu->type == GGML_TYPE_F32 && ++ motif.up->src[0]->ne[0] == 6144 && ++ motif.up->src[0]->ne[1] == 2048 && ++ motif.up->src[0]->ne[2] == 256 && ++ motif.gate->src[0]->ne[0] == 6144 && ++ motif.gate->src[0]->ne[1] == 2048 && ++ motif.gate->src[0]->ne[2] == 256 && ++ motif.glu->ne[0] == 2048 && ++ motif.glu->ne[1] == 8 && ++ motif.glu->ne[2] == 1; ++} ++ ++static int ggml_metal_op_glm_moe_weights_gate_up_swiglu( ++ ggml_metal_op_t ctx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ ggml_tensor * src1, ++ const ggml_metal_glm_moe_private_bindings & bindings) { ++ if (!ggml_metal_glm_moe_weights_gate_up_swiglu_shape_ok(motif, src1)) { ++ return 0; ++ } ++ ++ ggml_tensor * route_weights = motif.route.weights; ++ ggml_tensor * probs = route_weights->src[0]; ++ ggml_tensor * ids = motif.route.ids; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne_route_probs, probs, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb_route_probs, probs, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne_route_ids, ids, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb_route_ids, ids, nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb_route_weights, route_weights, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne0, motif.up->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, motif.up->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, src1, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, src1, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, ids, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, ids, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, motif.glu, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, motif.glu, nb); ++ ++ ggml_metal_kargs_moe_route_weights route_args = { ++ /*.n_expert =*/ ne_route_probs1, ++ /*.n_tokens =*/ ne_route_ids1, ++ /*.n_expert_used =*/ ne_route_ids0, ++ /*.norm =*/ ggml_get_op_params_i32(route_weights, 2), ++ /*.clamp_min =*/ ggml_get_op_params_f32(route_weights, 0), ++ /*.scale =*/ ggml_get_op_params_f32(route_weights, 1), ++ /*._pad0 =*/ ggml_metal_glm_dsa_moe_route_weights_slot0_enabled() ? 1 : 0, ++ /*._pad1 =*/ 0, ++ /*.probs_nb1 =*/ nb_route_probs1, ++ /*.probs_nb2 =*/ nb_route_probs2, ++ /*.ids_nb0 =*/ nb_route_ids0, ++ /*.ids_nb1 =*/ nb_route_ids1, ++ /*.dst_nb1 =*/ nb_route_weights1, ++ /*.dst_nb2 =*/ nb_route_weights2, ++ }; ++ ggml_metal_kargs_mul_mv_id_gate_up_swiglu gate_args = { ++ /*.nei0 =*/ ne20, ++ /*.nei1 =*/ ne21, ++ /*.nbi1 =*/ nb21, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nb1 =*/ nb1, ++ /*.nr0 =*/ 8, ++ /*.weights_nb1 =*/ nb_route_weights1, ++ /*.weights_nb2 =*/ nb_route_weights2, ++ /*.weighted =*/ 0, ++ /*._pad0 =*/ ggml_metal_glm_dsa_moe_max_active_experts(), ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_glm_moe_weights_q2_gate_up_swiglu_pair_sg_slot1( ++ ctx->lib, motif.glu); ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ int ida = 0; ++ ggml_metal_encoder_set_pipeline(ctx->enc, pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &route_args, sizeof(route_args), ida++); ++ ggml_metal_encoder_set_bytes(ctx->enc, &gate_args, sizeof(gate_args), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(probs), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, bindings.ids, ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, bindings.weights, ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.up->src[0]), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.gate->src[0]), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(src1), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, bindings.activation, ida++); ++ ++ const int grid_x = (ne01 + pipeline.nr0 - 1)/pipeline.nr0; ++ ggml_metal_encoder_set_threadgroup_memory_size(ctx->enc, pipeline.smem, 0); ++ const int grid_z = ne21*ne20; ++ ggml_metal_encoder_dispatch_threadgroups( ++ ctx->enc, grid_x, 1, grid_z, 32, pipeline.nsg, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_weights_q2_gate_up_swiglu tensor=%s probs=%s ids=%s weights=%s private_scratch=1 experts=%d top_k=%d tokens=%d embd=%d ff=%d grid_x=%d grid_z=%d threads_x=32 threads_y=%d\n", ++ ggml_metal_tensor_name(motif.glu), ++ ggml_metal_tensor_name(probs), ++ ggml_metal_tensor_name(ids), ++ ggml_metal_tensor_name(route_weights), ++ route_args.n_expert, ++ route_args.n_expert_used, ++ route_args.n_tokens, ++ ne00, ++ ne01, ++ grid_x, ++ grid_z, ++ pipeline.nsg); ++ } ++ ++ return motif.route_n_fuse + 3; ++} ++ ++static size_t ggml_metal_fusion_scratch_align(size_t offset) { ++ constexpr size_t alignment = 256; ++ return (offset + alignment - 1) & ~(alignment - 1); ++} ++ ++static bool ggml_metal_make_glm_moe_private_bindings( ++ ggml_metal_op_t ctx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif, ++ ggml_metal_glm_moe_private_bindings & bindings) { ++ if (ctx->fusion_scratch.metal == nullptr || ctx->fusion_scratch_size == 0 || ++ motif.ids == nullptr || motif.weights == nullptr || motif.glu == nullptr) { ++ return false; ++ } ++ ++ size_t cursor = 0; ++ bindings.ids = ctx->fusion_scratch; ++ bindings.ids.offs += cursor; ++ cursor = ggml_metal_fusion_scratch_align(cursor + ggml_nbytes(motif.ids)); ++ ++ bindings.weights = ctx->fusion_scratch; ++ bindings.weights.offs += cursor; ++ cursor = ggml_metal_fusion_scratch_align(cursor + ggml_nbytes(motif.weights)); ++ ++ bindings.activation = ctx->fusion_scratch; ++ bindings.activation.offs += cursor; ++ cursor = ggml_metal_fusion_scratch_align(cursor + ggml_nbytes(motif.glu)); ++ ++ return cursor <= ctx->fusion_scratch_size; ++} ++ ++static int ggml_metal_op_glm_moe_private_scratch_decode( ++ ggml_metal_op_t ctx, ++ int idx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif) { ++ const bool fused_topk_route = ++ ggml_metal_glm_moe_route_gate_up_swiglu_shape_ok(motif, motif.route_anchor_src); ++ const bool native_route_weights = ++ ggml_metal_glm_moe_weights_gate_up_swiglu_shape_ok(motif, motif.route_anchor_src); ++ if (!ggml_metal_glm_dsa_moe_private_scratch_enabled() || ++ !motif.has_native_down || !motif.has_shared_expert_tail || ++ motif.out_offset < 0 || motif.shared_gate_offset < 0 || motif.shared_up_offset < 0 || ++ motif.route_anchor_src == nullptr || ++ (!fused_topk_route && !native_route_weights)) { ++ return 0; ++ } ++ ++ ggml_metal_glm_moe_private_bindings bindings; ++ if (!ggml_metal_make_glm_moe_private_bindings(ctx, motif, bindings)) { ++ return 0; ++ } ++ ++ if (native_route_weights) { ++ // TOP_K is outside this owned span, so retain its graph allocation as ++ // the exact selected-ID source while weights and activation stay private. ++ bindings.ids = ggml_metal_get_buffer_id(motif.ids); ++ } ++ ++ const int route_gate_up_n = fused_topk_route ? ++ ggml_metal_op_glm_moe_route_gate_up_swiglu( ++ ctx, motif, motif.route_anchor_src, &bindings) : ++ ggml_metal_op_glm_moe_weights_gate_up_swiglu( ++ ctx, motif, motif.route_anchor_src, bindings); ++ if (route_gate_up_n <= 0) { ++ return 0; ++ } ++ ++ const int shared_gate_n = ggml_metal_op_reference_dispatch_existing_node( ++ ctx, idx + motif.shared_gate_offset); ++ const int shared_up_n = ggml_metal_op_reference_dispatch_existing_node( ++ ctx, idx + motif.shared_up_offset); ++ GGML_ASSERT(shared_gate_n == 1); ++ GGML_ASSERT(shared_up_n == 1); ++ ++ // Only the private routed resource is a dependency of the Q3 terminal. ++ // The two shared-expert projections remain eligible to overlap it. ++ ggml_metal_encoder_memory_barrier_buffer(ctx->enc, bindings.activation); ++ ++ const int down_n = ggml_metal_encode_moe_mul_mat_id(ctx, motif.down, &bindings); ++ GGML_ASSERT(down_n == 1); ++ ++ ctx->set_fused_range_outputs( ++ motif.out_offset, ++ motif.shared_gate_offset, ++ motif.shared_up_offset); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_private_scratch_decode tensor=%s fused_nodes=%d route_nodes=%d out_offset=%d shared_gate_offset=%d shared_up_offset=%d scratch_bytes=%llu activation_bytes=%llu dispatch_groups=4 graph_uid=%llu split_start=%d split_end=%d\n", ++ ggml_metal_tensor_name(motif.out), ++ motif.n_fuse, ++ motif.route_n_fuse, ++ motif.out_offset, ++ motif.shared_gate_offset, ++ motif.shared_up_offset, ++ (unsigned long long) ctx->fusion_scratch_size, ++ (unsigned long long) ggml_nbytes(motif.glu), ++ (unsigned long long) ctx->graph_uid(), ++ ctx->split_start(), ++ ctx->split_end()); ++ } ++ ++ return motif.n_fuse; ++} ++ ++static void ggml_metal_encode_glm_moe_shared_mul_mv( ++ ggml_metal_op_t ctx, ++ ggml_tensor * op, ++ ggml_metal_buffer_id src, ++ ggml_metal_buffer_id dst) { ++ GGML_ASSERT(op->op == GGML_OP_MUL_MAT); ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_Q4_K); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->ne[1] == 1); ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ ++ GGML_ASSERT(ne00 == ne10); ++ GGML_ASSERT(ne12 % ne02 == 0); ++ GGML_ASSERT(ne13 % ne03 == 0); ++ ++ const int16_t r2 = ne12/ne02; ++ const int16_t r3 = ne13/ne03; ++ auto pipeline = ggml_metal_library_get_pipeline_mul_mv(ctx->lib, op); ++ const int nr0 = pipeline.nr0; ++ const int nr1 = pipeline.nr1; ++ const int nsg = pipeline.nsg; ++ ++ ggml_metal_kargs_mul_mv args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nr0 =*/ nr0, ++ /*.r2 =*/ r2, ++ /*.r3 =*/ r3, ++ }; ++ ++ ggml_metal_encoder_set_pipeline(ctx->enc, pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer(ctx->enc, src, 2); ++ ggml_metal_encoder_set_buffer(ctx->enc, dst, 3); ++ ggml_metal_encoder_set_threadgroup_memory_size(ctx->enc, pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups( ++ ctx->enc, ++ (ne01 + nr0*nsg - 1)/(nr0*nsg), ++ (ne11 + nr1 - 1)/nr1, ++ ne12*ne13, ++ 32, ++ nsg, ++ 1); ++} ++ ++static int ggml_metal_op_glm_moe_two_phase( ++ ggml_metal_op_t ctx, ++ int idx, ++ const ggml_metal_glm_moe_decode_motif_reference & motif) { ++ const bool dual_lane = ggml_metal_glm_dsa_moe_dual_lane_enabled(); ++ if ((!ggml_metal_glm_dsa_moe_two_phase_enabled() && !dual_lane) || ++ (dual_lane && !motif.final_only_fusable) || ++ motif.shared_glu == nullptr || motif.shared_down == nullptr || motif.final_out == nullptr) { ++ return 0; ++ } ++ ++ ggml_metal_buffer_id shared_gate_scratch = { nullptr, 0 }; ++ ggml_metal_buffer_id shared_up_scratch = { nullptr, 0 }; ++ ggml_metal_buffer_id shared_activation_scratch = { nullptr, 0 }; ++ if (dual_lane) { ++ constexpr size_t shared_projection_bytes = 2048*sizeof(float); ++ size_t scratch_cursor = 0; ++ shared_gate_scratch = ctx->fusion_scratch; ++ shared_gate_scratch.offs += scratch_cursor; ++ scratch_cursor = ggml_metal_fusion_scratch_align(scratch_cursor + shared_projection_bytes); ++ shared_up_scratch = ctx->fusion_scratch; ++ shared_up_scratch.offs += scratch_cursor; ++ scratch_cursor = ggml_metal_fusion_scratch_align(scratch_cursor + shared_projection_bytes); ++ shared_activation_scratch = ctx->fusion_scratch; ++ shared_activation_scratch.offs += scratch_cursor; ++ scratch_cursor = ggml_metal_fusion_scratch_align(scratch_cursor + shared_projection_bytes); ++ if (ctx->fusion_scratch.metal == nullptr || scratch_cursor > ctx->fusion_scratch_size) { ++ return 0; ++ } ++ } ++ ++ const int route_n = motif.route.weights->op == GGML_OP_MOE_ROUTE_WEIGHTS && motif.route.n_fuse == 1 ? ++ ggml_metal_op_moe_route_weights(ctx, idx) : ++ ggml_metal_op_topk_moe_route_fused(ctx, idx); ++ GGML_ASSERT(route_n == motif.route_n_fuse); ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ ggml_tensor * routed_up_w = motif.up->src[0]; ++ ggml_tensor * routed_gate_w = motif.gate->src[0]; ++ ggml_tensor * routed_down_w = motif.down->src[0]; ++ ggml_tensor * shared_gate = dual_lane ? motif.shared_glu->src[0] : motif.shared_gate; ++ ggml_tensor * shared_up = dual_lane ? motif.shared_glu->src[1] : motif.shared_up; ++ ggml_tensor * shared_up_w = shared_up->src[0]; ++ ggml_tensor * shared_gate_w = shared_gate->src[0]; ++ ggml_tensor * shared_down_w = motif.shared_down->src[0]; ++ ggml_tensor * cur = motif.route_anchor_src; ++ ++ ggml_metal_kargs_mul_mv_id_gate_up_swiglu gate_args = { ++ /*.nei0 =*/ (int32_t) motif.ids->ne[0], ++ /*.nei1 =*/ (int32_t) motif.ids->ne[1], ++ /*.nbi1 =*/ motif.ids->nb[1], ++ /*.ne00 =*/ (int32_t) routed_up_w->ne[0], ++ /*.ne01 =*/ (int32_t) routed_up_w->ne[1], ++ /*.ne02 =*/ (int32_t) routed_up_w->ne[2], ++ /*.nb00 =*/ routed_up_w->nb[0], ++ /*.nb01 =*/ routed_up_w->nb[1], ++ /*.nb02 =*/ routed_up_w->nb[2], ++ /*.ne10 =*/ (int32_t) cur->ne[0], ++ /*.ne11 =*/ (int32_t) cur->ne[1], ++ /*.ne12 =*/ (int32_t) cur->ne[2], ++ /*.ne13 =*/ (int32_t) cur->ne[3], ++ /*.nb10 =*/ cur->nb[0], ++ /*.nb11 =*/ cur->nb[1], ++ /*.nb12 =*/ cur->nb[2], ++ /*.ne0 =*/ (int32_t) motif.glu->ne[0], ++ /*.ne1 =*/ (int32_t) motif.glu->ne[1], ++ /*.nb1 =*/ motif.glu->nb[1], ++ /*.nr0 =*/ 8, ++ /*.weights_nb1 =*/ motif.weights->nb[1], ++ /*.weights_nb2 =*/ motif.weights->nb[2], ++ /*.weighted =*/ 0, ++ /*._pad0 =*/ 0, ++ }; ++ ggml_metal_kargs_glm_moe_two_phase phase_args = { ++ /*.n_embd =*/ 6144, ++ /*.n_ff =*/ 2048, ++ /*.n_out =*/ 6144, ++ /*.routed_gate_groups =*/ 2048, ++ /*.shared_gate_nb1 =*/ shared_gate_w->nb[1], ++ /*.shared_up_nb1 =*/ shared_up_w->nb[1], ++ /*.shared_down_nb1 =*/ shared_down_w->nb[1], ++ }; ++ ++ const bool gate_slot4 = dual_lane && ggml_metal_glm_dsa_moe_dual_lane_gate_slot4_enabled(); ++ const char * gate_kernel = gate_slot4 ? ++ "kernel_glm_moe_dual_lane_gate_slot4" : "kernel_glm_moe_two_phase_gate"; ++ auto gate_pipeline = ggml_metal_library_get_pipeline(ctx->lib, gate_kernel); ++ if (!gate_pipeline.pipeline) { ++ gate_pipeline = ggml_metal_library_compile_pipeline(ctx->lib, gate_kernel, gate_kernel, nullptr); ++ } ++ int ida = 0; ++ ggml_metal_encoder_set_pipeline(ctx->enc, gate_pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &gate_args, sizeof(gate_args), ida++); ++ ggml_metal_encoder_set_bytes(ctx->enc, &phase_args, sizeof(phase_args), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(routed_up_w), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(routed_gate_w), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(shared_up_w), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(shared_gate_w), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(cur), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.glu), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.shared_glu), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.ids), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.weights), ida++); ++ ggml_metal_encoder_set_threadgroup_memory_size(ctx->enc, 8*8*sizeof(float), 0); ++ const int32_t gate_groups = gate_slot4 ? ++ ((phase_args.n_ff + gate_args.nr0 - 1)/gate_args.nr0)*((gate_args.nei0 + 3)/4) : ++ (dual_lane ? phase_args.routed_gate_groups : phase_args.routed_gate_groups + phase_args.n_ff); ++ const int32_t gate_nsg = gate_slot4 ? 4 : 1; ++ GGML_ASSERT(32*gate_nsg <= ggml_metal_pipeline_max_theads_per_threadgroup(gate_pipeline)); ++ ggml_metal_encoder_dispatch_threadgroups(ctx->enc, gate_groups, 1, 1, 32, gate_nsg, 1); ++ ++ if (dual_lane) { ++ ggml_metal_encode_glm_moe_shared_mul_mv( ++ ctx, ++ shared_gate, ++ ggml_metal_get_buffer_id(shared_gate->src[1]), ++ shared_gate_scratch); ++ ggml_metal_encode_glm_moe_shared_mul_mv( ++ ctx, ++ shared_up, ++ ggml_metal_get_buffer_id(shared_up->src[1]), ++ shared_up_scratch); ++ } ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ if (dual_lane) { ++ const char * shared_glu_kernel = "kernel_glm_moe_dual_lane_swiglu"; ++ auto shared_glu_pipeline = ggml_metal_library_get_pipeline(ctx->lib, shared_glu_kernel); ++ if (!shared_glu_pipeline.pipeline) { ++ shared_glu_pipeline = ggml_metal_library_compile_pipeline( ++ ctx->lib, shared_glu_kernel, shared_glu_kernel, nullptr); ++ } ++ ida = 0; ++ ggml_metal_encoder_set_pipeline(ctx->enc, shared_glu_pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &phase_args, sizeof(phase_args), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, shared_gate_scratch, ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, shared_up_scratch, ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, shared_activation_scratch, ida++); ++ ggml_metal_encoder_dispatch_threadgroups( ++ ctx->enc, (phase_args.n_ff + 255)/256, 1, 1, 256, 1, 1); ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ ++ ggml_metal_kargs_mul_mv_id down_args = { ++ /*.nei0 =*/ (int32_t) motif.ids->ne[0], ++ /*.nei1 =*/ (int32_t) motif.ids->ne[1], ++ /*.nbi1 =*/ motif.ids->nb[1], ++ /*.ne00 =*/ (int32_t) routed_down_w->ne[0], ++ /*.ne01 =*/ (int32_t) routed_down_w->ne[1], ++ /*.ne02 =*/ (int32_t) routed_down_w->ne[2], ++ /*.nb00 =*/ routed_down_w->nb[0], ++ /*.nb01 =*/ routed_down_w->nb[1], ++ /*.nb02 =*/ routed_down_w->nb[2], ++ /*.ne10 =*/ (int32_t) motif.glu->ne[0], ++ /*.ne11 =*/ (int32_t) motif.glu->ne[1], ++ /*.ne12 =*/ (int32_t) motif.glu->ne[2], ++ /*.ne13 =*/ (int32_t) motif.glu->ne[3], ++ /*.nb10 =*/ motif.glu->nb[0], ++ /*.nb11 =*/ motif.glu->nb[1], ++ /*.nb12 =*/ motif.glu->nb[2], ++ /*.ne0 =*/ (int32_t) motif.final_out->ne[0], ++ /*.ne1 =*/ (int32_t) motif.final_out->ne[1], ++ /*.nb1 =*/ motif.final_out->nb[1], ++ /*.nr0 =*/ 8, ++ }; ++ ggml_metal_kargs_mul_mv_id_weighted_reduce_extra down_extra = { ++ /*.weights_nb1 =*/ motif.weights->nb[1], ++ /*.weights_nb2 =*/ motif.weights->nb[2], ++ /*.dst_nb0 =*/ motif.final_out->nb[0], ++ /*.dst_nb1 =*/ motif.final_out->nb[1], ++ /*.already_weighted =*/ 0, ++ /*._pad0 =*/ 0, ++ }; ++ ++ if (dual_lane) { ++ auto routed_down_pipeline = ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8_w0( ++ ctx->lib, motif.out); ++ ggml_metal_encoder_set_pipeline(ctx->enc, routed_down_pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &down_args, sizeof(down_args), 0); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(routed_down_w), 1); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.glu), 2); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.out), 3); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.ids), 4); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.weights), 5); ++ ggml_metal_encoder_set_bytes(ctx->enc, &down_extra, sizeof(down_extra), 6); ++ ggml_metal_encoder_set_threadgroup_memory_size(ctx->enc, routed_down_pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups( ++ ctx->enc, (phase_args.n_out + 7)/8, 1, 1, 32, 8, 1); ++ ++ ggml_metal_encode_glm_moe_shared_mul_mv( ++ ctx, ++ motif.shared_down, ++ shared_activation_scratch, ++ ggml_metal_get_buffer_id(motif.shared_down)); ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ const char * final_add_kernel = "kernel_glm_moe_dual_lane_add"; ++ auto final_add_pipeline = ggml_metal_library_get_pipeline(ctx->lib, final_add_kernel); ++ if (!final_add_pipeline.pipeline) { ++ final_add_pipeline = ggml_metal_library_compile_pipeline( ++ ctx->lib, final_add_kernel, final_add_kernel, nullptr); ++ } ++ ida = 0; ++ ggml_metal_encoder_set_pipeline(ctx->enc, final_add_pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &phase_args, sizeof(phase_args), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.out), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.shared_down), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.final_out), ida++); ++ ggml_metal_encoder_dispatch_threadgroups( ++ ctx->enc, (phase_args.n_out + 255)/256, 1, 1, 256, 1, 1); ++ ctx->set_fused_range_outputs(motif.final_out_offset); ++ } else { ++ const char * down_kernel = "kernel_glm_moe_two_phase_down"; ++ auto down_pipeline = ggml_metal_library_get_pipeline(ctx->lib, down_kernel); ++ if (!down_pipeline.pipeline) { ++ down_pipeline = ggml_metal_library_compile_pipeline( ++ ctx->lib, down_kernel, down_kernel, nullptr); ++ } ++ ida = 0; ++ ggml_metal_encoder_set_pipeline(ctx->enc, down_pipeline); ++ ggml_metal_encoder_set_bytes(ctx->enc, &down_args, sizeof(down_args), ida++); ++ ggml_metal_encoder_set_bytes(ctx->enc, &down_extra, sizeof(down_extra), ida++); ++ ggml_metal_encoder_set_bytes(ctx->enc, &phase_args, sizeof(phase_args), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(routed_down_w), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.glu), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(shared_down_w), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.shared_glu), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.out), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.final_out), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.ids), ida++); ++ ggml_metal_encoder_set_buffer(ctx->enc, ggml_metal_get_buffer_id(motif.weights), ida++); ++ ggml_metal_encoder_set_threadgroup_memory_size(ctx->enc, 8*8*sizeof(float), 0); ++ ggml_metal_encoder_dispatch_threadgroups( ++ ctx->enc, (phase_args.n_out + 7)/8, 1, 1, 32, 8, 1); ++ ctx->set_fused_range_outputs(motif.out_offset, motif.final_out_offset); ++ } ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=%s tensor=%s fused_nodes=%d route_nodes=%d gate_groups=%d down_groups=%d dispatch_groups=%d final_only=%d\n", ++ dual_lane ? "glm_moe_dual_lane" : "glm_moe_two_phase", ++ ggml_metal_tensor_name(motif.final_out), ++ motif.n_fuse, ++ motif.route_n_fuse, ++ gate_groups, ++ (phase_args.n_out + 7)/8, ++ dual_lane ? 8 : 3, ++ dual_lane ? 1 : 0); ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=%s kernel=%s tensor=%s fused_nodes=%d dispatch_groups=%d grid_x=%d grid_y=1 grid_z=1 threads_x=32\n", ++ dual_lane ? "glm_moe_dual_lane" : "glm_moe_two_phase", ++ dual_lane ? "dual_lane_q2_q3_q4" : "two_phase_q2_q3_q4", ++ ggml_metal_tensor_name(motif.final_out), ++ motif.n_fuse, ++ dual_lane ? 8 : 3, ++ gate_groups); ++ } ++ ++ return motif.n_fuse; ++} ++ ++static int ggml_metal_op_glm_moe_decode_motif_reference(ggml_metal_op_t ctx, int idx) { ++ ggml_metal_glm_moe_decode_motif_reference motif; ++ if (!ggml_metal_match_glm_moe_decode_motif_reference(ctx, idx, motif)) { ++ return 0; ++ } ++ if (ggml_metal_extend_glm_moe_two_phase(ctx, idx, motif)) { ++ const int two_phase_n = ggml_metal_op_glm_moe_two_phase(ctx, idx, motif); ++ if (two_phase_n > 0) { ++ return two_phase_n; ++ } ++ } ++ if (!ggml_metal_glm_dsa_moe_decode_motif_reference_enabled() && ++ !ggml_metal_glm_dsa_moe_private_scratch_enabled()) { ++ return 0; ++ } ++ ggml_metal_log_routed_moe_decode_motif_contract(ctx, motif); ++ ++ const int private_scratch_n = ggml_metal_op_glm_moe_private_scratch_decode(ctx, idx, motif); ++ if (private_scratch_n > 0) { ++ return private_scratch_n; ++ } ++ if (motif.has_native_down) { ++ return 0; ++ } ++ ++ const bool skip_internal_barriers = ggml_metal_glm_dsa_moe_decode_skip_internal_barriers_enabled(); ++ const bool skip_route_barrier = ++ skip_internal_barriers || ggml_metal_glm_dsa_moe_decode_skip_route_barrier_enabled(); ++ const bool skip_gate_up_barrier = ++ skip_internal_barriers || ggml_metal_glm_dsa_moe_decode_skip_gate_up_barrier_enabled(); ++ auto motif_barrier = [&]() { ++ if (!skip_internal_barriers) { ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ }; ++ auto route_barrier = [&]() { ++ if (!skip_route_barrier) { ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ }; ++ auto gate_up_barrier = [&]() { ++ if (!skip_gate_up_barrier) { ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ }; ++ auto route_gate_up_barrier = [&]() { ++ if (skip_gate_up_barrier) { ++ return; ++ } ++ ++ if (ggml_metal_glm_dsa_moe_decode_scoped_barriers_enabled()) { ++ const ggml_tensor * route_gate_up_dst = ++ motif.weighted_down_input != nullptr ? motif.weighted_down_input : motif.glu; ++ ggml_metal_op_concurrency_reset_tensors(ctx, route_gate_up_dst, motif.ids, motif.weights); ++ return; ++ } ++ ++ ggml_metal_op_concurrency_reset(ctx); ++ }; ++ bool route_dispatched = false; ++ int motif_dispatch_groups = 0; ++ int motif_reference_nodes = 0; ++ int motif_route_dispatches = 0; ++ int motif_route_gate_up_dispatches = 0; ++ int motif_gate_up_dispatches = 0; ++ int motif_down_out_dispatches = 0; ++ int motif_fused_tail_dispatches = 0; ++ auto dispatch_route = [&]() { ++ if (!route_dispatched) { ++ const int route_n_fuse = motif.route.weights->op == GGML_OP_MOE_ROUTE_WEIGHTS && ++ motif.route.n_fuse == 1 ? ++ ggml_metal_op_moe_route_weights(ctx, idx) : ++ ggml_metal_op_topk_moe_route_fused(ctx, idx); ++ GGML_ASSERT(route_n_fuse == motif.route_n_fuse); ++ route_dispatched = true; ++ ++motif_dispatch_groups; ++ ++motif_route_dispatches; ++ route_barrier(); ++ } ++ }; ++ ++ const int q2_weight_roofline_chunks = ggml_metal_glm_dsa_moe_q2_weight_roofline_chunks(); ++ const int q2_weight_roofline_block_bytes = ggml_metal_glm_dsa_moe_q2_weight_roofline_block_bytes(); ++ if (ggml_metal_glm_moe_can_scan_q2_selected_weights( ++ motif, q2_weight_roofline_chunks, q2_weight_roofline_block_bytes)) { ++ if (!ggml_metal_glm_dsa_moe_q2_weight_roofline_bypass_route_enabled()) { ++ dispatch_route(); ++ } ++ return ggml_metal_op_glm_moe_q2_selected_weight_scan( ++ ctx, motif, q2_weight_roofline_chunks, q2_weight_roofline_block_bytes); ++ } ++ ++ int rel = motif.route_n_fuse; ++ bool skipped_route_anchor = false; ++ ++ const bool fused_q3_tail_enabled = ggml_metal_glm_moe_can_fuse_swiglu_q3_down(motif); ++ const bool fused_q2_tail_enabled = ggml_metal_glm_moe_can_fuse_swiglu_q2_down(motif); ++ const bool fused_tail_enabled = fused_q3_tail_enabled || fused_q2_tail_enabled; ++ if (fused_tail_enabled) { ++ dispatch_route(); ++ if (motif.has_route_anchor) { ++ for (int i = 0; i < 4; ++i) { ++ const int n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(n == 1); ++ rel += n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ motif_barrier(); ++ } ++ } ++ ++ const int gate_n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(gate_n == 1); ++ rel += gate_n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ ++ const int up_n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(up_n == 1); ++ rel += up_n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ gate_up_barrier(); ++ ++ if (motif.has_shared_expert_tail) { ++ GGML_ASSERT(motif.shared_gate_offset >= 0); ++ GGML_ASSERT(motif.shared_up_offset >= 0); ++ const int shared_gate_n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + motif.shared_gate_offset); ++ const int shared_up_n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + motif.shared_up_offset); ++ GGML_ASSERT(shared_gate_n == 1); ++ GGML_ASSERT(shared_up_n == 1); ++ motif_dispatch_groups += 2; ++ motif_reference_nodes += 2; ++ } ++ ++ const int fused_tail = fused_q2_tail_enabled ? ++ ggml_metal_op_glm_moe_swiglu_q2_down_weighted(ctx, motif) : ++ ggml_metal_op_glm_moe_swiglu_q3_down_weighted(ctx, motif); ++ GGML_ASSERT(fused_tail == 3); ++ ++motif_dispatch_groups; ++ ++motif_fused_tail_dispatches; ++ rel = motif.n_fuse; ++ GGML_ASSERT(rel == motif.n_fuse); ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_decode_motif_reference_tail tensor=%s mode=%s fused_tail=%d\n", ++ ggml_metal_tensor_name(motif.out), ++ fused_q2_tail_enabled ? "swiglu_q2_down_weighted" : "swiglu_q3_down_weighted", ++ fused_tail); ++ } ++ } else { ++ bool consumed_weighted_down = false; ++ bool consumed_route_gate_up = false; ++ ggml_tensor * gate_up_src1_override = nullptr; ++ if (ggml_metal_glm_dsa_moe_route_anchor_bypass_enabled() && ++ motif.has_route_anchor && motif.route_anchor_src != nullptr && motif.route_anchor_src->type == GGML_TYPE_F32 && ++ ggml_are_same_shape(motif.route_anchor_src, motif.gate->src[1])) { ++ gate_up_src1_override = motif.route_anchor_src; ++ rel += 4; ++ skipped_route_anchor = true; ++ } ++ ++ if (!motif.has_route_anchor || skipped_route_anchor) { ++ const int route_gate_up_glu_n = ++ ggml_metal_op_glm_moe_route_gate_up_swiglu(ctx, motif, gate_up_src1_override); ++ if (route_gate_up_glu_n > 0) { ++ rel += route_gate_up_glu_n; ++ consumed_route_gate_up = true; ++ consumed_weighted_down = motif.has_weighted_down && route_gate_up_glu_n == 4; ++ ++motif_dispatch_groups; ++ ++motif_route_gate_up_dispatches; ++ route_gate_up_barrier(); ++ } ++ } ++ ++ if (!consumed_route_gate_up) { ++ dispatch_route(); ++ if (skipped_route_anchor) { ++ rel = motif.route_n_fuse + 4; ++ } ++ } ++ ++ if (!consumed_route_gate_up && motif.has_route_anchor && !skipped_route_anchor) { ++ for (int i = 0; i < 4; ++i) { ++ const int n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(n == 1); ++ rel += n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ motif_barrier(); ++ } ++ } ++ ++ if (!consumed_route_gate_up) { ++ const int gate_up_glu_n = ggml_metal_op_mul_mv_id_gate_up_swiglu( ++ ctx, ++ idx + rel, ++ gate_up_src1_override); ++ if (gate_up_glu_n > 0) { ++ rel += gate_up_glu_n; ++ consumed_weighted_down = motif.has_weighted_down && gate_up_glu_n == 4; ++ ++motif_dispatch_groups; ++ ++motif_gate_up_dispatches; ++ } else { ++ if (skipped_route_anchor) { ++ rel = motif.route_n_fuse; ++ skipped_route_anchor = false; ++ for (int i = 0; i < 4; ++i) { ++ const int n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(n == 1); ++ rel += n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ motif_barrier(); ++ } ++ } ++ const int gate_n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(gate_n == 1); ++ rel += gate_n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ ++ const int up_n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(up_n == 1); ++ rel += up_n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ gate_up_barrier(); ++ ++ const int glu_n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(glu_n == 1); ++ rel += glu_n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ } ++ gate_up_barrier(); ++ } ++ ++ if (motif.has_weighted_down && !consumed_weighted_down) { ++ const int weighted_n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(weighted_n == 1); ++ rel += weighted_n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ motif_barrier(); ++ } ++ ++ const int down_out_n = ggml_metal_op_mul_mv_id_weighted_reduce(ctx, idx + rel, true); ++ if (down_out_n > 0) { ++ rel += down_out_n; ++ ++motif_dispatch_groups; ++ ++motif_down_out_dispatches; ++ } else { ++ while (rel < motif.n_fuse) { ++ const int n = ggml_metal_op_reference_dispatch_existing_node(ctx, idx + rel); ++ GGML_ASSERT(n == 1); ++ rel += n; ++ ++motif_dispatch_groups; ++ ++motif_reference_nodes; ++ motif_barrier(); ++ } ++ } ++ GGML_ASSERT(rel == motif.n_fuse); ++ } ++ ++ if (motif.has_shared_expert_tail) { ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "ggml_metal: moe_dispatch op=glm_moe_decode_motif_reference tensor=%s last_node=%s fused_nodes=%d route_nodes=%d route_anchor=%d route_anchor_bypassed=%d weighted_down=%d shared_expert_tail=%d dispatch_groups=%d reference_nodes=%d route_dispatches=%d route_gate_up_dispatches=%d gate_up_dispatches=%d down_out_dispatches=%d fused_tail_dispatches=%d graph_uid=%llu split_start=%d split_end=%d\n", ++ ggml_metal_tensor_name(motif.out), ++ ggml_metal_tensor_name(ctx->node(idx + motif.n_fuse - 1)), ++ motif.n_fuse, ++ motif.route_n_fuse, ++ motif.has_route_anchor ? 1 : 0, ++ skipped_route_anchor ? 1 : 0, ++ motif.has_weighted_down ? 1 : 0, ++ motif.has_shared_expert_tail ? 1 : 0, ++ motif_dispatch_groups, ++ motif_reference_nodes, ++ motif_route_dispatches, ++ motif_route_gate_up_dispatches, ++ motif_gate_up_dispatches, ++ motif_down_out_dispatches, ++ motif_fused_tail_dispatches, ++ (unsigned long long) ctx->graph_uid(), ++ ctx->split_start(), ++ ctx->split_end()); ++ } ++ ++ return motif.n_fuse; ++} ++ ++static bool ggml_metal_match_mul_mv_id_gate_up_swiglu( ++ ggml_metal_op_t ctx, ++ int idx, ++ ggml_metal_mul_mv_id_gate_up_swiglu_fusion & fusion) { ++ const bool explicit_enabled = ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_vecscale_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_any_variant_enabled(); ++ const bool default_enabled = ggml_metal_glm_dsa_q2_gate_up_swiglu_default_enabled(); ++ const bool enabled = explicit_enabled || default_enabled || ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_default_enabled(); ++ if (!ctx->use_fusion || !enabled || idx + 3 > ctx->n_nodes()) { ++ return false; ++ } ++ ++ ggml_tensor * first = ctx->node(idx); ++ ggml_tensor * second = ctx->node(idx + 1); ++ ggml_tensor * glu = ctx->node(idx + 2); ++ ggml_tensor * next = idx + 3 < ctx->n_nodes() ? ctx->node(idx + 3) : nullptr; ++ if (first->op != GGML_OP_MUL_MAT_ID || second->op != GGML_OP_MUL_MAT_ID || glu->op != GGML_OP_GLU) { ++ return false; ++ } ++ ++ if (ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || ggml_get_op_params_i32(glu, 1) != 0) { ++ return false; ++ } ++ ++ ggml_tensor * gate = glu->src[0]; ++ ggml_tensor * up = glu->src[1]; ++ const bool same_pair = (gate == first && up == second) || (gate == second && up == first); ++ if (!same_pair) { ++ return false; ++ } ++ ++ if (up->src[0] == nullptr || gate->src[0] == nullptr || ++ up->src[1] == nullptr || gate->src[1] == nullptr || ++ up->src[2] == nullptr || gate->src[2] == nullptr) { ++ return false; ++ } ++ ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * cast = nullptr; ++ ggml_tensor * weighted = next; ++ if (next != nullptr && next->op == GGML_OP_CPY && next->src[0] == glu && next->type == GGML_TYPE_F16 && ++ ggml_is_contiguous_1(next)) { ++ cast = next; ++ weighted = nullptr; ++ } ++ if (weighted != nullptr && weighted->op == GGML_OP_MUL && ++ ggml_metal_tensor_name_contains(weighted, "ffn_moe_down_weighted_input")) { ++ if (weighted->src[0] == glu) { ++ weights = weighted->src[1]; ++ } else if (weighted->src[1] == glu) { ++ weights = weighted->src[0]; ++ } ++ } ++ ++ const bool weighted_shape_ok = ++ weights == nullptr || ++ (weighted != nullptr && ++ weighted->type == GGML_TYPE_F32 && ++ weights->type == GGML_TYPE_F32 && ++ weighted->ne[0] == glu->ne[0] && ++ weighted->ne[1] == glu->ne[1] && ++ weighted->ne[2] == glu->ne[2] && ++ weights->ne[0] == 1 && ++ weights->ne[1] == glu->ne[1] && ++ weights->ne[2] == glu->ne[2] && ++ ggml_is_contiguous_1(weighted)); ++ ++ const bool cast_shape_ok = ++ cast == nullptr || ++ (cast->ne[0] == glu->ne[0] && ++ cast->ne[1] == glu->ne[1] && ++ cast->ne[2] == glu->ne[2] && ++ cast->ne[3] == glu->ne[3]); ++ ++ const bool shape_ok = ++ up->src[0]->type == GGML_TYPE_Q2_K && ++ gate->src[0]->type == GGML_TYPE_Q2_K && ++ up->src[1]->type == GGML_TYPE_F32 && ++ gate->src[1] == up->src[1] && ++ gate->src[2] == up->src[2] && ++ up->type == GGML_TYPE_F32 && ++ gate->type == GGML_TYPE_F32 && ++ glu->type == GGML_TYPE_F32 && ++ ggml_are_same_shape(up->src[0], gate->src[0]) && ++ ggml_are_same_stride(up->src[0], gate->src[0]) && ++ ggml_are_same_shape(up, gate) && ++ ggml_are_same_shape(up, glu) && ++ up->src[0]->ne[0] >= 2048 && ++ up->src[0]->ne[1] >= 1024 && ++ up->src[2]->ne[0] == 8 && ++ up->src[2]->ne[1] >= 1 && ++ up->src[2]->ne[1] <= 32 && ++ up->src[0]->ne[0] % ggml_blck_size(up->src[0]->type) == 0 && ++ up->src[0]->ne[1] > 0 && ++ up->src[0]->ne[2] >= up->src[2]->ne[0] && ++ cast_shape_ok && ++ weighted_shape_ok; ++ if (!shape_ok) { ++ return false; ++ } ++ const bool default_pair_sg_shape = ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_default_enabled() && ++ up->src[0]->ne[0] == 6144 && ++ up->src[0]->ne[1] == 2048 && ++ up->src[2]->ne[0] == 8 && ++ up->src[2]->ne[1] == 1; ++ const bool default_plain_shape = ++ default_enabled && ++ up->src[0]->ne[0] == 6144 && ++ up->src[0]->ne[1] == 2048 && ++ up->src[2]->ne[0] == 8 && ++ up->src[2]->ne[1] == 1; ++ if (!explicit_enabled && !default_pair_sg_shape && !default_plain_shape) { ++ return false; ++ } ++ ++ if (weights != nullptr) { ++ const ggml_op ops[] = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU, GGML_OP_MUL }; ++ const int outputs[] = { 3 }; ++ if (!ctx->can_fuse_subgraph(idx, ops, 4, outputs, 1)) { ++ return false; ++ } ++ } else if (cast != nullptr) { ++ const ggml_op ops[] = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU, GGML_OP_CPY }; ++ const int outputs[] = { 3 }; ++ if (!ctx->can_fuse_subgraph(idx, ops, 4, outputs, 1)) { ++ return false; ++ } ++ } else { ++ const ggml_op ops[] = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }; ++ const int outputs[] = { 2 }; ++ if (!ctx->can_fuse_subgraph(idx, ops, 3, outputs, 1)) { ++ return false; ++ } ++ } ++ ++ fusion.up = up; ++ fusion.gate = gate; ++ fusion.glu = glu; ++ fusion.cast = cast; ++ fusion.weighted = weights != nullptr ? weighted : nullptr; ++ fusion.weights = weights; ++ fusion.q8 = up->src[3] != nullptr && ++ up->src[3] == gate->src[3] && ++ up->src[3]->type == GGML_TYPE_Q8_0 && ++ ggml_are_same_shape(up->src[3], up->src[1]) ? up->src[3] : nullptr; ++ fusion.n_fuse = (weights != nullptr || cast != nullptr) ? 4 : 3; ++ return true; ++} ++ ++static int ggml_metal_op_mul_mv_id_gate_up_swiglu( ++ ggml_metal_op_t ctx, ++ int idx, ++ ggml_tensor * src1_override) { ++ ggml_metal_mul_mv_id_gate_up_swiglu_fusion fusion; ++ if (!ggml_metal_match_mul_mv_id_gate_up_swiglu(ctx, idx, fusion)) { ++ return 0; ++ } ++ ++ ggml_tensor * up = fusion.up; ++ ggml_tensor * gate = fusion.gate; ++ ggml_tensor * glu = fusion.glu; ++ ggml_tensor * dst = fusion.weighted != nullptr ? fusion.weighted : (fusion.cast != nullptr ? fusion.cast : fusion.glu); ++ ggml_tensor * weights = fusion.weights; ++ ggml_tensor * q8 = fusion.q8; ++ ggml_tensor * src1 = src1_override != nullptr ? src1_override : up->src[1]; ++ if (src1_override != nullptr && ++ (src1_override->type != GGML_TYPE_F32 || !ggml_are_same_shape(src1_override, up->src[1]))) { ++ return 0; ++ } ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ if (ggml_metal_glm_routed_expert_noop_enabled() && dst->type == GGML_TYPE_F32) { ++ auto pipeline = ggml_metal_library_get_pipeline_zero_f32(lib); ++ const int nth = 256; ++ const int64_t n_tg = (ggml_nelements(dst) + nth - 1)/nth; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(dst), 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, nth, 1, 1); ++ ggml_metal_op_concurrency_reset(ctx); ++ return fusion.n_fuse; ++ } ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, up->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, up->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, src1, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, src1, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, up->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, up->src[2], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, dst, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, dst, nb); ++ const uint64_t weights_nb1 = weights != nullptr ? weights->nb[1] : 0; ++ const uint64_t weights_nb2 = weights != nullptr ? weights->nb[2] : 0; ++ ++ const bool pair_sg_default_shape = ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_default_enabled() && ++ ne00 == 6144 && ++ ne01 == 2048 && ++ ne20 == 8 && ++ ne21 == 1; ++ const bool pair_sg = ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_enabled() || ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_any_variant_enabled() || ++ pair_sg_default_shape; ++ if (dst->type == GGML_TYPE_F16 && !pair_sg) { ++ return 0; ++ } ++ const bool pair_sg_slot8 = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot8_enabled(); ++ const bool pair_sg_slot2 = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot2_enabled(); ++ const bool pair_sg_slot4_dual_default = pair_sg_default_shape && src1_override != nullptr; ++ const bool pair_sg_slot4_dual = pair_sg && dst->type == GGML_TYPE_F32 && ++ (ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_enabled() || ++ pair_sg_slot4_dual_default); ++ const bool pair_sg_slot4_dual_r12 = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_r12_enabled(); ++ const bool pair_sg_slot4_dual_r16 = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot4_dual_r16_enabled(); ++ const bool pair_sg_slot1_dual = pair_sg && dst->type == GGML_TYPE_F32 && ++ (ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot1_dual_enabled() || ++ (pair_sg_default_shape && !pair_sg_slot4_dual_default)); ++ const bool pair_sg_slot2_dual = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot2_dual_enabled(); ++ const bool pair_sg_slot8_split = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_slot8_split_enabled(); ++ const bool pair_sg_share_y = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_share_y_enabled() && ++ ne00 == 6144 && ++ ne01 == 2048 && ++ ne20 == 8 && ++ ne21 == 1 && ++ src1->type == GGML_TYPE_F32; ++ const bool pair_sg_vecscale = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_vecscale_enabled() && ++ ne00 == 6144 && ++ ne01 == 2048 && ++ ne20 == 8 && ++ ne21 == 1 && ++ src1->type == GGML_TYPE_F32; ++ const bool pair_sg_q8_act = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_q8_act_enabled() && ++ ne00 == 6144 && ++ ne01 == 2048 && ++ ne20 == 8 && ++ ne21 == 1 && ++ src1->type == GGML_TYPE_F32; ++ const bool pair_sg_prequant_q8 = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_prequant_q8_enabled() && ++ q8 != nullptr && ++ ne00 == 6144 && ++ ne01 == 2048 && ++ ne20 == 8 && ++ ne21 == 1 && ++ src1->type == GGML_TYPE_F32; ++ const bool pair_sg_inblock_q2 = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_inblock_repack_enabled() && ++ ne00 == 6144 && ++ ne01 == 2048 && ++ ne20 == 8 && ++ ne21 == 1 && ++ src1->type == GGML_TYPE_F32; ++ const bool pair_sg_half_y = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_half_y_enabled() && ++ ne00 == 6144 && ++ ne01 == 2048 && ++ ne20 == 8 && ++ ne21 == 1 && ++ src1->type == GGML_TYPE_F32; ++ const bool pair_sg_rowtile = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_rowtile_enabled() && ++ ne00 == 6144 && ++ ne01 == 2048 && ++ ne20 == 8 && ++ ne21 == 1 && ++ src1->type == GGML_TYPE_F32; ++ const bool pair_sg_r16 = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_r16_enabled(); ++ const bool pair_sg_r12 = pair_sg && dst->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_pair_sg_r12_enabled(); ++ const bool plain_vecscale = !pair_sg && ++ ggml_metal_glm_dsa_q2_gate_up_swiglu_vecscale_enabled() && ++ dst->type == GGML_TYPE_F32 && ++ src1->type == GGML_TYPE_F32; ++ ggml_metal_pipeline_with_params pipeline = {}; ++ pipeline = pair_sg_inblock_q2 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual_inblock_q2(lib, glu) : ++ (pair_sg_prequant_q8 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual_prequant_q8(lib, glu) : ++ (pair_sg ? ++ (pair_sg_r12 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_r12(lib, dst) : ++ (pair_sg_r16 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_r16(lib, dst) : ++ (pair_sg_slot2 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot2(lib, glu) : ++ (pair_sg_slot1_dual ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot1_dual(lib, glu) : ++ (pair_sg_slot2_dual ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot2_dual(lib, glu) : ++ (pair_sg_slot4_dual_r12 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual_r12(lib, glu) : ++ (pair_sg_slot4_dual_r16 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual_r16(lib, glu) : ++ (pair_sg_slot4_dual ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot4_dual(lib, glu) : ++ (pair_sg_slot8_split ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot8_split(lib, glu) : ++ (pair_sg_slot8 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_slot8(lib, glu) : ++ (pair_sg_share_y ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_share_y(lib, glu) : ++ (pair_sg_q8_act ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_q8_act(lib, glu) : ++ (pair_sg_half_y ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_half_y(lib, glu) : ++ (pair_sg_rowtile ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_rowtile(lib, glu) : ++ (pair_sg_vecscale ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_vecscale(lib, glu) : ++ (dst->type == GGML_TYPE_F16 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg_f16(lib, dst) : ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_pair_sg(lib, glu))))))))))))))))) : ++ (plain_vecscale ? ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu_vecscale(lib, glu) : ++ ggml_metal_library_get_pipeline_mul_mv_id_gate_up_swiglu(lib, glu)))); ++ const int nr0 = pipeline.nr0; ++ const int nr1 = pipeline.nr1; ++ const int nsg = pipeline.nsg; ++ ++ ggml_metal_kargs_mul_mv_id_gate_up_swiglu args = { ++ /*.nei0 =*/ ne20, ++ /*.nei1 =*/ ne21, ++ /*.nbi1 =*/ nb21, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nb1 =*/ nb1, ++ /*.nr0 =*/ nr0, ++ /*.weights_nb1 =*/ weights_nb1, ++ /*.weights_nb2 =*/ weights_nb2, ++ /*.weighted =*/ weights != nullptr ? 1 : 0, ++ /*._pad0 =*/ ggml_metal_glm_dsa_moe_max_active_experts(), ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(up->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(gate->src[0]), 2); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(src1), 3); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(dst), 4); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(up->src[2]), 5); ++ ggml_metal_encoder_set_buffer(enc, weights != nullptr ? ++ ggml_metal_get_buffer_id(weights) : ggml_metal_get_buffer_id(up->src[2]), 6); ++ ggml_metal_encoder_set_buffer(enc, q8 != nullptr ? ++ ggml_metal_get_buffer_id(q8) : ggml_metal_get_buffer_id(src1), 7); ++ ++ int slots_per_threadgroup = 1; ++ if (pair_sg) { ++ if (pair_sg_inblock_q2 || pair_sg_prequant_q8) { ++ slots_per_threadgroup = 1; ++ } else if (pair_sg_r12 || pair_sg_r16 || pair_sg_slot4_dual || ++ pair_sg_slot4_dual_r12 || pair_sg_slot4_dual_r16 || ++ pair_sg_share_y || pair_sg_vecscale || pair_sg_q8_act || ++ pair_sg_half_y || pair_sg_rowtile) { ++ slots_per_threadgroup = 4; ++ } else if (pair_sg_slot8 || pair_sg_slot8_split) { ++ slots_per_threadgroup = 8; ++ } else if (pair_sg_slot2 || pair_sg_slot2_dual) { ++ slots_per_threadgroup = 2; ++ } else if (pair_sg_slot1_dual) { ++ slots_per_threadgroup = 1; ++ } else { ++ slots_per_threadgroup = 4; ++ } ++ } ++ const int64_t ne123 = pair_sg ? ++ ((ne20 + slots_per_threadgroup - 1)/slots_per_threadgroup)*ne21 : ++ ne20*ne21; ++ const int grid_x = pair_sg ? (ne01 + nr0 - 1)/nr0 : (ne01 + nr0*nsg - 1)/(nr0*nsg); ++ const int grid_y = (1 + nr1 - 1)/nr1; ++ const int grid_z = ne123; ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, nsg, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ const char * gate_up_op = ++ pair_sg_r12 ? "mul_mv_id_q2_gate_up_swiglu_r12" : ++ pair_sg_r16 ? "mul_mv_id_q2_gate_up_swiglu_r16" : ++ pair_sg_slot2 ? "mul_mv_id_q2_gate_up_swiglu_slot2" : ++ pair_sg_slot1_dual ? "mul_mv_id_q2_gate_up_swiglu_slot1_dual" : ++ pair_sg_slot2_dual ? "mul_mv_id_q2_gate_up_swiglu_slot2_dual" : ++ pair_sg_slot4_dual_r12 ? "mul_mv_id_q2_gate_up_swiglu_slot4_dual_r12" : ++ pair_sg_slot4_dual_r16 ? "mul_mv_id_q2_gate_up_swiglu_slot4_dual_r16" : ++ pair_sg_slot4_dual ? "mul_mv_id_q2_gate_up_swiglu_slot4_dual" : ++ pair_sg_slot8_split ? "mul_mv_id_q2_gate_up_swiglu_slot8_split" : ++ pair_sg_slot8 ? "mul_mv_id_q2_gate_up_swiglu_slot8" : ++ pair_sg_share_y ? "mul_mv_id_q2_gate_up_swiglu_share_y" : ++ pair_sg_inblock_q2 ? "mul_mv_id_q2_gate_up_swiglu_inblock_q2" : ++ pair_sg_q8_act ? "mul_mv_id_q2_gate_up_swiglu_q8_act" : ++ pair_sg_prequant_q8 ? "mul_mv_id_q2_gate_up_swiglu_prequant_q8" : ++ pair_sg_half_y ? "mul_mv_id_q2_gate_up_swiglu_half_y" : ++ pair_sg_rowtile ? "mul_mv_id_q2_gate_up_swiglu_rowtile" : ++ (pair_sg_vecscale || plain_vecscale) ? "mul_mv_id_q2_gate_up_swiglu_vecscale" : ++ pair_sg ? "mul_mv_id_q2_gate_up_swiglu_pair_sg" : ++ "mul_mv_id_q2_gate_up_swiglu"; ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=%s tensor=%s up=%s gate=%s weights=%s src1=%s src1_override=%d src0_type=%s src1_type=%s ids_type=%s dst_type=%s ne00=%d ne01=%d experts=%d used_experts=%d tokens=%d weighted=%d nr0=%d nr1=%d nsg=%d grid_x=%d grid_y=%d grid_z=%d threads_x=%d threads_y=%d fused_nodes=%d\n", ++ gate_up_op, ++ ggml_metal_tensor_name(dst), ++ ggml_metal_tensor_name(up), ++ ggml_metal_tensor_name(gate), ++ ggml_metal_tensor_name(weights), ++ ggml_metal_tensor_name(src1), ++ src1_override != nullptr ? 1 : 0, ++ ggml_type_name(up->src[0]->type), ++ ggml_type_name(src1->type), ++ ggml_type_name(up->src[2]->type), ++ ggml_type_name(dst->type), ++ ne00, ++ ne01, ++ ne02, ++ ne20, ++ ne21, ++ weights != nullptr ? 1 : 0, ++ nr0, ++ nr1, ++ nsg, ++ grid_x, ++ grid_y, ++ grid_z, ++ 32, ++ nsg, ++ fusion.n_fuse); ++ } ++ ++ // The generic concurrency tracker only sees the first node in this fused ++ // span. Publish the fused destination before a later node consumes it. ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ return fusion.n_fuse; ++} ++ ++static bool ggml_metal_match_mul_mv_id_weighted_reduce( ++ ggml_metal_op_t ctx, ++ int idx, ++ ggml_metal_mul_mv_id_weighted_reduce_fusion & fusion, ++ bool subgraph_owned) { ++ const bool enabled = ++ ggml_metal_glm_dsa_q2_down_weighted_reduce_enabled() || ++ ggml_metal_glm_dsa_q2_down_f16_act_enabled() || ++ ggml_metal_glm_dsa_q2_down_vec_scale_enabled() || ++ ggml_metal_glm_dsa_q3_down_weighted_reduce_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r6_nb8_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r10_nb8_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_glm52_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_default_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w1_default_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r12_nb8_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r16_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_split2_reduce_enabled() || ++ ggml_metal_glm_dsa_q3_down_atomic_accum_enabled() || ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_enabled() || ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_default_enabled() || ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r4_enabled() || ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r16_enabled() || ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r16_w1_enabled(); ++ if (!ctx->use_fusion || !enabled || idx + 2 > ctx->n_nodes()) { ++ if (enabled && ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_q2_down_weighted_reduce_reject idx=%d reason=precheck use_fusion=%d n_nodes=%d\n", ++ idx, ++ ctx->use_fusion ? 1 : 0, ++ ctx->n_nodes()); ++ } ++ return false; ++ } ++ ++ ggml_tensor * down = ctx->node(idx); ++ ggml_tensor * weighted_sum = nullptr; ++ ggml_tensor * shared_gate = nullptr; ++ ggml_tensor * shared_up = nullptr; ++ int weighted_sum_offset = -1; ++ int shared_gate_offset = -1; ++ int shared_up_offset = -1; ++ ++ for (int rel = 1; rel < 4 && idx + rel < ctx->n_nodes(); ++rel) { ++ ggml_tensor * candidate = ctx->node(idx + rel); ++ if (candidate->op == GGML_OP_MOE_WEIGHTED_SUM && candidate->src[0] == down && ++ ggml_metal_tensor_name_contains(candidate, "ffn_moe_out")) { ++ weighted_sum = candidate; ++ weighted_sum_offset = rel; ++ } else if (candidate->op == GGML_OP_MUL_MAT && ggml_metal_tensor_is_shared_expert_gate(candidate)) { ++ shared_gate = candidate; ++ shared_gate_offset = rel; ++ } else if (candidate->op == GGML_OP_MUL_MAT && ggml_metal_tensor_is_shared_expert_up(candidate)) { ++ shared_up = candidate; ++ shared_up_offset = rel; ++ } ++ } ++ ++ auto reject = [&](const char * reason) { ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_q2_down_weighted_reduce_reject idx=%d reason=%s node0=%s:%s shared_gate=%s:%s shared_up=%s:%s weighted_sum=%s:%s\n", ++ idx, ++ reason, ++ ggml_metal_tensor_name(down), ++ ggml_op_name(down->op), ++ ggml_metal_tensor_name(shared_gate), ++ shared_gate ? ggml_op_name(shared_gate->op) : "none", ++ ggml_metal_tensor_name(shared_up), ++ shared_up ? ggml_op_name(shared_up->op) : "none", ++ ggml_metal_tensor_name(weighted_sum), ++ weighted_sum ? ggml_op_name(weighted_sum->op) : "none"); ++ } ++ return false; ++ }; ++ ++ if (down->op != GGML_OP_MUL_MAT_ID || weighted_sum == nullptr) { ++ return reject("op_window"); ++ } ++ if ((shared_gate == nullptr) != (shared_up == nullptr)) { ++ return reject("partial_shared_window"); ++ } ++ if (weighted_sum->src[0] != down || weighted_sum->src[1] == nullptr) { ++ return reject("weighted_sum_sources"); ++ } ++ if (!ggml_metal_tensor_name_contains(down, "ffn_moe_down") || ++ !ggml_metal_tensor_name_contains(weighted_sum, "ffn_moe_out")) { ++ return reject("tensor_names"); ++ } ++ if (shared_gate != nullptr && ++ (!ggml_metal_tensor_is_shared_expert_gate(shared_gate) || ++ !ggml_metal_tensor_is_shared_expert_up(shared_up))) { ++ return reject("tensor_names"); ++ } ++ if (down->src[0] == nullptr || down->src[1] == nullptr || down->src[2] == nullptr) { ++ return reject("down_sources"); ++ } ++ ++ const ggml_type down_type = down->src[0]->type; ++ const bool down_type_ok = down_type == GGML_TYPE_Q2_K || down_type == GGML_TYPE_Q3_K; ++ const bool q2_default_shape = ++ down_type == GGML_TYPE_Q2_K && ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_default_enabled() && ++ down->src[0]->ne[0] == 2048 && ++ down->src[2]->ne[0] == 8 && ++ ggml_get_op_params_i32(weighted_sum, 0) == 0; ++ if (down_type == GGML_TYPE_Q2_K && ++ !ggml_metal_glm_dsa_q2_down_weighted_reduce_enabled() && ++ !ggml_metal_glm_dsa_q2_down_f16_act_enabled() && ++ !ggml_metal_glm_dsa_q2_down_vec_scale_enabled() && ++ !ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_enabled() && ++ !q2_default_shape && ++ !ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r4_enabled() && ++ !ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r16_enabled() && ++ !ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r16_w1_enabled()) { ++ return reject("q2_direct_disabled"); ++ } ++ if (down_type == GGML_TYPE_Q2_K && ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_enabled() && ++ down->src[0]->ne[0] != 2048) { ++ return reject("q2_slot_parallel_shape"); ++ } ++ const bool shared_shape_ok = ++ shared_gate == nullptr || ++ (shared_gate->type == GGML_TYPE_F32 && shared_up->type == GGML_TYPE_F32); ++ const bool glm_down_shape = ++ down->src[0]->ne[0] >= 1024 && ++ down->src[0]->ne[1] >= 2048; ++ const bool src1_type_ok = ++ down->src[1]->type == GGML_TYPE_F32 || ++ (down_type == GGML_TYPE_Q2_K && ++ down->src[1]->type == GGML_TYPE_F16 && ++ ggml_metal_glm_dsa_q2_down_f16_act_enabled()) || ++ (down_type == GGML_TYPE_Q3_K && ++ down->src[1]->type == GGML_TYPE_F16 && ++ ggml_metal_glm_dsa_q3_down_f16_act_enabled()); ++ const bool shape_ok = ++ down_type_ok && ++ glm_down_shape && ++ src1_type_ok && ++ down->src[2]->type == GGML_TYPE_I32 && ++ weighted_sum->src[1]->type == GGML_TYPE_F32 && ++ down->type == GGML_TYPE_F32 && ++ shared_shape_ok && ++ weighted_sum->type == GGML_TYPE_F32 && ++ down->src[0]->ne[0] % ggml_blck_size(down_type) == 0 && ++ down->src[0]->ne[0] == down->src[1]->ne[0] && ++ down->src[0]->ne[1] == weighted_sum->ne[0] && ++ down->src[0]->ne[2] >= down->src[2]->ne[0] && ++ down->src[2]->ne[0] == 8 && ++ down->src[1]->ne[1] == 8 && ++ down->src[1]->ne[2] == down->src[2]->ne[1] && ++ weighted_sum->src[1]->ne[0] == 1 && ++ weighted_sum->src[1]->ne[1] == down->src[2]->ne[0] && ++ weighted_sum->src[1]->ne[2] == down->src[2]->ne[1] && ++ weighted_sum->ne[0] == down->ne[0] && ++ weighted_sum->ne[1] == down->ne[2]; ++ if (!shape_ok) { ++ return reject("shape"); ++ } ++ const bool q3_default_w0_shape = ++ down_type == GGML_TYPE_Q3_K && ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_default_enabled() && ++ down->src[0]->ne[0] == 2048 && ++ down->src[2]->ne[1] == 1 && ++ ggml_get_op_params_i32(weighted_sum, 0) == 0; ++ const bool q3_default_w1_shape = ++ down_type == GGML_TYPE_Q3_K && ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w1_default_enabled() && ++ down->src[0]->ne[0] == 2048 && ++ down->src[2]->ne[1] == 1 && ++ ggml_get_op_params_i32(weighted_sum, 0) != 0; ++ if (down_type == GGML_TYPE_Q3_K && ++ !ggml_metal_glm_dsa_q3_down_weighted_reduce_tensor_selected(weighted_sum)) { ++ return reject("q3_tensor"); ++ } ++ if (down_type == GGML_TYPE_Q3_K && ++ !q3_default_w0_shape && ++ !q3_default_w1_shape && ++ !ggml_metal_glm_dsa_q3_down_weighted_reduce_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r6_nb8_w0_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r10_nb8_w0_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_glm52_w0_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r12_nb8_w0_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r16_enabled() && ++ !ggml_metal_glm_dsa_q3_down_slot_split2_reduce_enabled() && ++ !ggml_metal_glm_dsa_q3_down_atomic_accum_enabled()) { ++ return reject("q3_direct_disabled"); ++ } ++ ++ const int n_fuse = std::max({ ++ weighted_sum_offset, ++ shared_gate_offset, ++ shared_up_offset, ++ }) + 1; ++ ggml_op ops[4]; ++ for (int rel = 0; rel < n_fuse; ++rel) { ++ ops[rel] = ctx->node(idx + rel)->op; ++ } ++ int outputs[3]; ++ int n_outputs = 0; ++ outputs[n_outputs++] = weighted_sum_offset; ++ if (shared_gate != nullptr) { ++ outputs[n_outputs++] = shared_gate_offset; ++ outputs[n_outputs++] = shared_up_offset; ++ } ++ if (!subgraph_owned && !ctx->can_fuse_subgraph(idx, ops, n_fuse, outputs, n_outputs)) { ++ return reject("can_fuse_subgraph"); ++ } ++ ++ fusion.down = down; ++ fusion.shared_gate = shared_gate; ++ fusion.shared_up = shared_up; ++ fusion.weighted_sum = weighted_sum; ++ fusion.shared_gate_offset = shared_gate_offset; ++ fusion.shared_up_offset = shared_up_offset; ++ fusion.weighted_sum_offset = weighted_sum_offset; ++ fusion.n_fuse = n_fuse; ++ return true; ++} ++ ++static int ggml_metal_op_mul_mv_id_weighted_reduce( ++ ggml_metal_op_t ctx, ++ int idx, ++ bool subgraph_owned) { ++ ggml_metal_mul_mv_id_weighted_reduce_fusion fusion; ++ if (!ggml_metal_match_mul_mv_id_weighted_reduce(ctx, idx, fusion, subgraph_owned)) { ++ return 0; ++ } ++ if (ggml_metal_glm_dsa_q2_down_weighted_reduce_noop_enabled()) { ++ return 0; ++ } ++ ++ ggml_tensor * down = fusion.down; ++ ggml_tensor * weighted_sum = fusion.weighted_sum; ++ ++ if (ggml_metal_glm_dsa_q2_down_weighted_reduce_reference_enabled()) { ++ for (int rel = 0; rel < fusion.n_fuse; ++rel) { ++ int encoded = 0; ++ if (rel == 0) { ++ encoded = ggml_metal_op_mul_mat_id(ctx, idx); ++ } else if (rel == fusion.shared_gate_offset || rel == fusion.shared_up_offset) { ++ encoded = ggml_metal_op_mul_mat(ctx, idx + rel); ++ } else if (rel == fusion.weighted_sum_offset) { ++ encoded = ggml_metal_op_moe_weighted_sum(ctx, idx + rel); ++ } ++ GGML_ASSERT(encoded == 1); ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=mul_mv_id_q2_down_weighted_reduce_reference tensor=%s down=%s weights=%s src0_type=%s dst_type=%s grid_x=1 grid_y=1 grid_z=1 threads_x=1 fused_nodes=%d\n", ++ ggml_metal_tensor_name(weighted_sum), ++ ggml_metal_tensor_name(down), ++ ggml_metal_tensor_name(weighted_sum->src[1]), ++ ggml_type_name(down->src[0]->type), ++ ggml_type_name(weighted_sum->type), ++ fusion.n_fuse); ++ } ++ return fusion.n_fuse; ++ } ++ ++ auto encode_shared = [&]() { ++ if (fusion.shared_gate == nullptr) { ++ return; ++ } ++ const int shared_gate_fuse = ggml_metal_op_mul_mat(ctx, idx + fusion.shared_gate_offset); ++ const int shared_up_fuse = ggml_metal_op_mul_mat(ctx, idx + fusion.shared_up_offset); ++ GGML_ASSERT(shared_gate_fuse == 1); ++ GGML_ASSERT(shared_up_fuse == 1); ++ ggml_metal_op_concurrency_reset(ctx); ++ }; ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, down->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, down->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, down->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, down->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, down->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, down->src[2], nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, weighted_sum->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ned, down, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nbd, down, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, weighted_sum, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, weighted_sum, nb); ++ ++ const bool preserve_routed_input = fusion.shared_gate != nullptr; ++ const uint64_t scratch_nb0 = ggml_type_size(down->src[1]->type); ++ const uint64_t scratch_nb1 = scratch_nb0*ne10; ++ const uint64_t scratch_nb2 = scratch_nb1*ne11; ++ const uint64_t scratch_nb3 = scratch_nb2*ne12; ++ ggml_metal_buffer_id fusion_scratch = ggml_metal_get_buffer_id(down); ++ fusion_scratch.offs += ggml_nbytes(down); ++ fusion_scratch.offs += ggml_metal_op_mul_mat_id_extra_tpe(down); ++ fusion_scratch.offs += ggml_metal_op_mul_mat_id_extra_ids(down); ++ ggml_metal_buffer_id routed_input = ggml_metal_get_buffer_id(down->src[1]); ++ if (preserve_routed_input) { ++ GGML_ASSERT(ggml_blck_size(down->src[1]->type) == 1); ++ GGML_ASSERT(scratch_nb3*ne13 <= ggml_metal_op_mul_mat_id_extra_src1_scratch(down)); ++ ++ routed_input = fusion_scratch; ++ ++ auto copy_pipeline = ggml_metal_library_get_pipeline_cpy( ++ lib, down->src[1]->type, down->src[1]->type); ++ ggml_metal_kargs_cpy copy_args = { ++ /*.nk0 =*/ ne10, ++ /*.ne00 =*/ ne10, ++ /*.ne01 =*/ ne11, ++ /*.ne02 =*/ ne12, ++ /*.ne03 =*/ ne13, ++ /*.nb00 =*/ nb10, ++ /*.nb01 =*/ nb11, ++ /*.nb02 =*/ nb12, ++ /*.nb03 =*/ nb13, ++ /*.ne0 =*/ ne10, ++ /*.ne1 =*/ ne11, ++ /*.ne2 =*/ ne12, ++ /*.ne3 =*/ ne13, ++ /*.nb0 =*/ scratch_nb0, ++ /*.nb1 =*/ scratch_nb1, ++ /*.nb2 =*/ scratch_nb2, ++ /*.nb3 =*/ scratch_nb3, ++ }; ++ const int nth = std::min( ++ ggml_metal_pipeline_max_theads_per_threadgroup(copy_pipeline), ne10); ++ const int nw0 = (ne10 + nth - 1)/nth; ++ ++ ggml_metal_encoder_set_pipeline(enc, copy_pipeline); ++ ggml_metal_encoder_set_bytes(enc, ©_args, sizeof(copy_args), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[1]), 1); ++ ggml_metal_encoder_set_buffer(enc, routed_input, 2); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nw0*ne11, ne12, ne13, nth, 1, 1); ++ } ++ ++ // The shared experts may reuse the routed activation's allocation after ++ // the original down node. Preserve that activation first, then keep the ++ // copy and both shared matvecs in the same concurrent Metal span. ++ encode_shared(); ++ ++ // The fused kernel consumes selected-expert IDs through the down node even ++ // when it replaces the following weighted reduction. Make the producer's ++ // writes visible before this manually encoded cross-node read. ++ ggml_metal_encoder_memory_barrier_buffer( ++ enc, ggml_metal_get_buffer_id(down->src[2])); ++ ++ const bool use_output_scratch = ++ down->src[0]->type == GGML_TYPE_Q3_K && !preserve_routed_input; ++ const ggml_metal_buffer_id weighted_reduce_dst = use_output_scratch ? ++ fusion_scratch : ggml_metal_get_buffer_id(weighted_sum); ++ ++ if (ggml_metal_glm_routed_expert_noop_enabled() && weighted_sum->type == GGML_TYPE_F32) { ++ auto pipeline = ggml_metal_library_get_pipeline_zero_f32(lib); ++ const int nth = 256; ++ const int64_t n_tg = (ggml_nelements(weighted_sum) + nth - 1)/nth; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_buffer(enc, weighted_reduce_dst, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, nth, 1, 1); ++ ggml_metal_op_concurrency_reset(ctx); ++ return fusion.n_fuse; ++ } ++ ++ auto pipeline = ggml_metal_library_get_pipeline_mul_mv_id_weighted_reduce(lib, weighted_sum); ++ const int nr0 = pipeline.nr0; ++ const int nr1 = pipeline.nr1; ++ const int nsg = pipeline.nsg; ++ ++ ggml_metal_kargs_mul_mv_id args = { ++ /*.nei0 =*/ ne20, ++ /*.nei1 =*/ ne21, ++ /*.nbi1 =*/ nb21, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ preserve_routed_input ? scratch_nb0 : nb10, ++ /*.nb11 =*/ preserve_routed_input ? scratch_nb1 : nb11, ++ /*.nb12 =*/ preserve_routed_input ? scratch_nb2 : nb12, ++ /*.ne0 =*/ ned0, ++ /*.ne1 =*/ ned1, ++ /*.nb1 =*/ nbd1, ++ /*.nr0 =*/ nr0, ++ }; ++ ggml_metal_kargs_mul_mv_id_weighted_reduce_extra extra = { ++ /*.weights_nb1 =*/ nb31, ++ /*.weights_nb2 =*/ nb32, ++ /*.dst_nb0 =*/ use_output_scratch ? sizeof(float) : nb0, ++ /*.dst_nb1 =*/ use_output_scratch ? sizeof(float)*ne0 : nb1, ++ /*.already_weighted =*/ ggml_get_op_params_i32(weighted_sum, 0) != 0 ? 1 : 0, ++ /*._pad0 =*/ ggml_metal_glm_dsa_q2_down_weighted_reduce_stock_slot0_enabled() ? 1 : 0, ++ }; ++ auto publish_weighted_output = [&]() { ++ ggml_metal_op_concurrency_reset(ctx); ++ if (!use_output_scratch) { ++ return; ++ } ++ ++ GGML_ASSERT(sizeof(float)*ggml_nelements(weighted_sum) <= ++ ggml_metal_op_mul_mat_id_extra_src1_scratch(down)); ++ auto copy_pipeline = ggml_metal_library_get_pipeline_cpy( ++ lib, GGML_TYPE_F32, GGML_TYPE_F32); ++ ggml_metal_kargs_cpy copy_args = { ++ /*.nk0 =*/ ne0, ++ /*.ne00 =*/ ne0, ++ /*.ne01 =*/ ne1, ++ /*.ne02 =*/ ne2, ++ /*.ne03 =*/ ne3, ++ /*.nb00 =*/ sizeof(float), ++ /*.nb01 =*/ sizeof(float)*ne0, ++ /*.nb02 =*/ sizeof(float)*ne0*ne1, ++ /*.nb03 =*/ sizeof(float)*ne0*ne1*ne2, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ }; ++ const int nth = std::min( ++ ggml_metal_pipeline_max_theads_per_threadgroup(copy_pipeline), ne0); ++ const int nw0 = (ne0 + nth - 1)/nth; ++ ++ ggml_metal_encoder_set_pipeline(enc, copy_pipeline); ++ ggml_metal_encoder_set_bytes(enc, ©_args, sizeof(copy_args), 0); ++ ggml_metal_encoder_set_buffer(enc, fusion_scratch, 1); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weighted_sum), 2); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nw0*ne1, ne2, ne3, nth, 1, 1); ++ ggml_metal_op_concurrency_reset(ctx); ++ }; ++ const bool q3_default_w0_shape = ++ down->src[0]->type == GGML_TYPE_Q3_K && ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_default_enabled() && ++ ne00 == 2048 && ++ ne21 == 1 && ++ extra.already_weighted == 0; ++ const bool q3_default_w1_shape = ++ down->src[0]->type == GGML_TYPE_Q3_K && ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w1_default_enabled() && ++ ne00 == 2048 && ++ ne21 == 1 && ++ extra.already_weighted != 0; ++ const bool q3_f16_w0_shape = ++ down->src[0]->type == GGML_TYPE_Q3_K && ++ down->src[1]->type == GGML_TYPE_F16 && ++ ggml_metal_glm_dsa_q3_down_f16_act_enabled() && ++ ne00 == 2048 && ++ extra.already_weighted == 0; ++ const bool q2_slot_parallel_requested = ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_enabled(); ++ const bool q2_slot_parallel_r4 = ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r4_enabled(); ++ const bool q2_slot_parallel_r16 = ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r16_enabled(); ++ const bool q2_slot_parallel_r16_w1 = ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_r16_w1_enabled() && ++ ne00 == 2048 && ++ ne20 == 8 && ++ extra.already_weighted != 0; ++ const bool q2_slot_parallel_default_shape = ++ ggml_metal_glm_dsa_q2_down_slot_parallel_reduce_default_enabled() && ++ ne00 == 2048 && ++ ne20 == 8 && ++ extra.already_weighted == 0; ++ const bool q2_f16_act_shape = ++ down->src[0]->type == GGML_TYPE_Q2_K && ++ down->src[1]->type == GGML_TYPE_F16 && ++ ggml_metal_glm_dsa_q2_down_f16_act_enabled() && ++ ne00 == 2048 && ++ ne20 == 8 && ++ extra.already_weighted == 0; ++ const bool q2_shift_high_bits_shape = ++ down->src[0]->type == GGML_TYPE_Q2_K && ++ ggml_metal_glm_dsa_q2_down_shift_high_bits_enabled() && ++ (down->src[1]->type == GGML_TYPE_F32 || down->src[1]->type == GGML_TYPE_F16) && ++ ne00 == 2048 && ++ ne20 == 8 && ++ extra.already_weighted == 0; ++ const bool q2_vec_scale_shape = ++ down->src[0]->type == GGML_TYPE_Q2_K && ++ down->src[1]->type == GGML_TYPE_F32 && ++ ggml_metal_glm_dsa_q2_down_vec_scale_enabled() && ++ ne00 == 2048 && ++ ne20 == 8 && ++ extra.already_weighted == 0; ++ if (down->src[0]->type == GGML_TYPE_Q2_K && ++ (q2_shift_high_bits_shape || q2_f16_act_shape || q2_vec_scale_shape || q2_slot_parallel_requested || q2_slot_parallel_default_shape || q2_slot_parallel_r4 || q2_slot_parallel_r16 || q2_slot_parallel_r16_w1) && ++ ne00 == 2048 && ++ ne20 == 8) { ++ ggml_metal_pipeline_with_params pipeline_slot_parallel; ++ const char * q2_slot_parallel_op = "mul_mv_id_q2_down_slot_parallel_reduce_r8"; ++ if (q2_shift_high_bits_shape && down->src[1]->type == GGML_TYPE_F16) { ++ pipeline_slot_parallel = ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_f16_shifted(lib, weighted_sum); ++ q2_slot_parallel_op = "mul_mv_id_q2_down_slot_parallel_reduce_r8_f16_shifted"; ++ } else if (q2_shift_high_bits_shape) { ++ pipeline_slot_parallel = ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_shifted(lib, weighted_sum); ++ q2_slot_parallel_op = "mul_mv_id_q2_down_slot_parallel_reduce_r8_shifted"; ++ } else if (q2_f16_act_shape) { ++ pipeline_slot_parallel = ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_f16(lib, weighted_sum); ++ q2_slot_parallel_op = "mul_mv_id_q2_down_slot_parallel_reduce_r8_f16"; ++ } else if (q2_slot_parallel_r4 && q2_vec_scale_shape) { ++ pipeline_slot_parallel = ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r4_nb8_vecscale(lib, weighted_sum); ++ q2_slot_parallel_op = "mul_mv_id_q2_down_slot_parallel_reduce_r4_vecscale"; ++ } else if (q2_vec_scale_shape) { ++ pipeline_slot_parallel = ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8_vecscale(lib, weighted_sum); ++ q2_slot_parallel_op = "mul_mv_id_q2_down_slot_parallel_reduce_r8_vecscale"; ++ } else if (q2_slot_parallel_r4) { ++ pipeline_slot_parallel = ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r4_nb8(lib, weighted_sum); ++ q2_slot_parallel_op = "mul_mv_id_q2_down_slot_parallel_reduce_r4"; ++ } else if (q2_slot_parallel_r16 || q2_slot_parallel_r16_w1) { ++ pipeline_slot_parallel = ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r16_nb8(lib, weighted_sum); ++ q2_slot_parallel_op = "mul_mv_id_q2_down_slot_parallel_reduce_r16"; ++ } else { ++ pipeline_slot_parallel = ggml_metal_library_get_pipeline_mul_mv_id_q2_weighted_reduce_slots_sg_r8_nb8(lib, weighted_sum); ++ } ++ const int nr0_slot_parallel = pipeline_slot_parallel.nr0; ++ const int nsg_slot_parallel = pipeline_slot_parallel.nsg; ++ const int grid_x_slot_parallel = (ne01 + nr0_slot_parallel - 1)/nr0_slot_parallel; ++ const int grid_y_slot_parallel = ne21; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline_slot_parallel); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_bytes(enc, &extra, sizeof(extra), 6); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, routed_input, 2); ++ ggml_metal_encoder_set_buffer(enc, weighted_reduce_dst, 3); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[2]), 4); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weighted_sum->src[1]), 5); ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline_slot_parallel.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x_slot_parallel, grid_y_slot_parallel, 1, 32, nsg_slot_parallel, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=%s tensor=%s down=%s src1=%s ids=%s weights=%s already_weighted=%d nr0=%d nsg=%d grid_x=%d grid_y=%d grid_z=1 threads_x=%d threads_y=%d fused_nodes=%d\n", ++ q2_slot_parallel_op, ++ ggml_metal_tensor_name(weighted_sum), ++ ggml_metal_tensor_name(down), ++ ggml_metal_tensor_name(down->src[1]), ++ ggml_metal_tensor_name(down->src[2]), ++ ggml_metal_tensor_name(weighted_sum->src[1]), ++ extra.already_weighted, ++ nr0_slot_parallel, ++ nsg_slot_parallel, ++ grid_x_slot_parallel, ++ grid_y_slot_parallel, ++ 32, ++ nsg_slot_parallel, ++ fusion.n_fuse); ++ } ++ ++ // The fused result replaces both MUL_MAT_ID and MOE_WEIGHTED_SUM, but ++ // the generic tracker only registered the first node in that span. ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ return fusion.n_fuse; ++ } ++ ++ if (down->src[0]->type == GGML_TYPE_Q3_K && ++ (ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r6_nb8_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r10_nb8_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_glm52_w0_enabled() || ++ q3_default_w0_shape || ++ q3_default_w1_shape || ++ q3_f16_w0_shape || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r12_nb8_w0_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r16_enabled() || ++ ggml_metal_glm_dsa_q3_down_slot_split2_reduce_enabled())) { ++ auto pipeline_slot_parallel = ggml_metal_glm_dsa_q3_down_slot_split2_reduce_enabled() ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_split2(lib, weighted_sum) : ++ (ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r16_enabled() ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r16(lib, weighted_sum) : ++ (ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_glm52_w0_enabled() && ne00 == 2048 && ne01 == 6144 && ne20 == 8 && ne21 == 1 && extra.already_weighted == 0 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_glm52_w0(lib, weighted_sum) : ++ (ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r10_nb8_w0_enabled() && ne00 == 2048 && extra.already_weighted == 0 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r10_nb8_w0(lib, weighted_sum) : ++ (ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r12_nb8_w0_enabled() && ne00 == 2048 && extra.already_weighted == 0 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r12_nb8_w0(lib, weighted_sum) : ++ (ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r6_nb8_w0_enabled() && ne00 == 2048 && extra.already_weighted == 0 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r6_nb8_w0(lib, weighted_sum) : ++ (q3_f16_w0_shape ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8_w0_f16(lib, weighted_sum) : ++ ((ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_w0_enabled() || q3_default_w0_shape) && ne00 == 2048 && extra.already_weighted == 0 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8_w0(lib, weighted_sum) : ++ (q3_default_w1_shape && ne00 == 2048 && extra.already_weighted != 0 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8(lib, weighted_sum) : ++ (ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_nb8_enabled() && ne00 == 2048 ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8(lib, weighted_sum) : ++ (ggml_metal_glm_dsa_q3_down_slot_parallel_reduce_r8_enabled() ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8(lib, weighted_sum) : ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg(lib, weighted_sum))))))))))); ++ const int nr0_slot_parallel = pipeline_slot_parallel.nr0; ++ const int nsg_slot_parallel = pipeline_slot_parallel.nsg; ++ const int grid_x_slot_parallel = (ne01 + nr0_slot_parallel - 1)/nr0_slot_parallel; ++ const int grid_y_slot_parallel = ne21; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline_slot_parallel); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_bytes(enc, &extra, sizeof(extra), 6); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, routed_input, 2); ++ ggml_metal_encoder_set_buffer(enc, weighted_reduce_dst, 3); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[2]), 4); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weighted_sum->src[1]), 5); ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline_slot_parallel.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x_slot_parallel, grid_y_slot_parallel, 1, 32, nsg_slot_parallel, 1); ++ ++ publish_weighted_output(); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=mul_mv_id_q3_down_slot_parallel_reduce tensor=%s down=%s src1=%s ids=%s weights=%s already_weighted=%d nr0=%d nsg=%d grid_x=%d grid_y=%d grid_z=1 threads_x=%d threads_y=%d fused_nodes=%d\n", ++ ggml_metal_tensor_name(weighted_sum), ++ ggml_metal_tensor_name(down), ++ ggml_metal_tensor_name(down->src[1]), ++ ggml_metal_tensor_name(down->src[2]), ++ ggml_metal_tensor_name(weighted_sum->src[1]), ++ extra.already_weighted, ++ nr0_slot_parallel, ++ nsg_slot_parallel, ++ grid_x_slot_parallel, ++ grid_y_slot_parallel, ++ 32, ++ nsg_slot_parallel, ++ fusion.n_fuse); ++ } ++ ++ return fusion.n_fuse; ++ } ++ ++ if (down->src[0]->type == GGML_TYPE_Q3_K && ggml_metal_glm_dsa_q3_down_atomic_accum_enabled()) { ++ { ++ auto zero_pipeline = ggml_metal_library_get_pipeline_zero_f32(lib); ++ const int64_t n_zero = ggml_nelements(weighted_sum); ++ const int nth = 256; ++ const int64_t n_tg = (n_zero + nth - 1) / nth; ++ ++ ggml_metal_encoder_set_pipeline(enc, zero_pipeline); ++ ggml_metal_encoder_set_buffer(enc, weighted_reduce_dst, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, nth, 1, 1); ++ } ++ ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ auto pipeline_atomic = ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_accum_atomic(lib, weighted_sum); ++ const int nr0_atomic = pipeline_atomic.nr0; ++ const int nsg_atomic = pipeline_atomic.nsg; ++ const int grid_x_atomic = (ne01 + nr0_atomic*nsg_atomic - 1)/(nr0_atomic*nsg_atomic); ++ const int grid_y_atomic = ne21; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline_atomic); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_kargs_mul_mv_id_weighted_reduce_extra atomic_extra = extra; ++ atomic_extra._pad0 = 0; ++ ggml_metal_encoder_set_bytes(enc, &atomic_extra, sizeof(atomic_extra), 6); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, routed_input, 2); ++ ggml_metal_encoder_set_buffer(enc, weighted_reduce_dst, 3); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[2]), 4); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weighted_sum->src[1]), 5); ++ ++ const int grid_z_atomic = ne20; ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline_atomic.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x_atomic, grid_y_atomic, grid_z_atomic, 32, nsg_atomic, 1); ++ ++ publish_weighted_output(); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=mul_mv_id_q3_down_atomic_accum tensor=%s down=%s src1=%s ids=%s weights=%s already_weighted=%d nr0=%d nsg=%d grid_x=%d grid_y=%d grid_z=%d threads_x=%d threads_y=%d fused_nodes=%d\n", ++ ggml_metal_tensor_name(weighted_sum), ++ ggml_metal_tensor_name(down), ++ ggml_metal_tensor_name(down->src[1]), ++ ggml_metal_tensor_name(down->src[2]), ++ ggml_metal_tensor_name(weighted_sum->src[1]), ++ extra.already_weighted, ++ nr0_atomic, ++ nsg_atomic, ++ grid_x_atomic, ++ grid_y_atomic, ++ grid_z_atomic, ++ 32, ++ nsg_atomic, ++ fusion.n_fuse); ++ } ++ ++ return fusion.n_fuse; ++ } ++ ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_bytes(enc, &extra, sizeof(extra), 6); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, routed_input, 2); ++ ggml_metal_encoder_set_buffer(enc, weighted_reduce_dst, 3); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(down->src[2]), 4); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(weighted_sum->src[1]), 5); ++ ++ const int grid_x = (ne01 + nr0*nsg - 1)/(nr0*nsg); ++ const int grid_y = ne21; ++ const int grid_z = 1; ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, nsg, 1); ++ publish_weighted_output(); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ const ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(down->src[1]); ++ const ggml_metal_buffer_id bid_ids = ggml_metal_get_buffer_id(down->src[2]); ++ const ggml_metal_buffer_id bid_w = ggml_metal_get_buffer_id(weighted_sum->src[1]); ++ const ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(weighted_sum); ++ const size_t src1_begin = bid_src1.offs; ++ const size_t src1_end = src1_begin + ggml_nbytes(down->src[1]); ++ const size_t ids_begin = bid_ids.offs; ++ const size_t ids_end = ids_begin + ggml_nbytes(down->src[2]); ++ const size_t w_begin = bid_w.offs; ++ const size_t w_end = w_begin + ggml_nbytes(weighted_sum->src[1]); ++ const size_t dst_begin = bid_dst.offs; ++ const size_t dst_end = dst_begin + ggml_nbytes(weighted_sum); ++ const bool src1_dst_overlap = ++ bid_src1.metal == bid_dst.metal && src1_begin < dst_end && dst_begin < src1_end; ++ const bool ids_dst_overlap = ++ bid_ids.metal == bid_dst.metal && ids_begin < dst_end && dst_begin < ids_end; ++ const bool weights_dst_overlap = ++ bid_w.metal == bid_dst.metal && w_begin < dst_end && dst_begin < w_end; ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=mul_mv_id_down_weighted_reduce tensor=%s down=%s src1=%s ids=%s weights=%s src0_type=%s src1_type=%s ids_type=%s dst_type=%s ne00=%d ne01=%d experts=%d ne10=%d ne11=%d ne12=%d nb10=%llu nb11=%llu nb12=%llu used_experts=%d tokens=%d ids_nb1=%llu weights_nb1=%llu weights_nb2=%llu already_weighted=%d dst_ne0=%d dst_ne1=%d dst_nb0=%llu dst_nb1=%llu src1_dst_overlap=%d ids_dst_overlap=%d weights_dst_overlap=%d src1_offs=%llu src1_nbytes=%llu ids_offs=%llu ids_nbytes=%llu weights_offs=%llu weights_nbytes=%llu dst_offs=%llu dst_nbytes=%llu same_src1_buffer=%d same_ids_buffer=%d same_weights_buffer=%d nr0=%d nr1=%d nsg=%d grid_x=%d grid_y=%d grid_z=%d threads_x=%d threads_y=%d fused_nodes=%d\n", ++ ggml_metal_tensor_name(weighted_sum), ++ ggml_metal_tensor_name(down), ++ ggml_metal_tensor_name(down->src[1]), ++ ggml_metal_tensor_name(down->src[2]), ++ ggml_metal_tensor_name(weighted_sum->src[1]), ++ ggml_type_name(down->src[0]->type), ++ ggml_type_name(down->src[1]->type), ++ ggml_type_name(down->src[2]->type), ++ ggml_type_name(weighted_sum->type), ++ ne00, ++ ne01, ++ ne02, ++ ne10, ++ ne11, ++ ne12, ++ (unsigned long long) nb10, ++ (unsigned long long) nb11, ++ (unsigned long long) nb12, ++ ne20, ++ ne21, ++ (unsigned long long) nb21, ++ (unsigned long long) nb31, ++ (unsigned long long) nb32, ++ extra.already_weighted, ++ ne0, ++ ne1, ++ (unsigned long long) nb0, ++ (unsigned long long) nb1, ++ src1_dst_overlap ? 1 : 0, ++ ids_dst_overlap ? 1 : 0, ++ weights_dst_overlap ? 1 : 0, ++ (unsigned long long) src1_begin, ++ (unsigned long long) ggml_nbytes(down->src[1]), ++ (unsigned long long) ids_begin, ++ (unsigned long long) ggml_nbytes(down->src[2]), ++ (unsigned long long) w_begin, ++ (unsigned long long) ggml_nbytes(weighted_sum->src[1]), ++ (unsigned long long) dst_begin, ++ (unsigned long long) ggml_nbytes(weighted_sum), ++ bid_src1.metal == bid_dst.metal ? 1 : 0, ++ bid_ids.metal == bid_dst.metal ? 1 : 0, ++ bid_w.metal == bid_dst.metal ? 1 : 0, ++ nr0, ++ nr1, ++ nsg, ++ grid_x, ++ grid_y, ++ grid_z, ++ 32, ++ nsg, ++ fusion.n_fuse); ++ } ++ ++ // Internal barriers clear the scheduler's tracked ranges before this ++ // fused output is encoded, so publish it before later graph nodes read it. ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ return fusion.n_fuse; ++} ++ ++static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { ++ struct ggml_tensor * node = ctx->node(idx); ++ ++ //GGML_LOG_INFO("%s: encoding node %3d, op = %8s\n", __func__, idx, ggml_op_name(node->op)); ++ ++ if (ggml_is_empty(node)) { ++ return 1; ++ } ++ ++ switch (node->op) { ++ case GGML_OP_NONE: ++ case GGML_OP_RESHAPE: ++ case GGML_OP_VIEW: ++ case GGML_OP_TRANSPOSE: ++ case GGML_OP_PERMUTE: ++ { ++ // noop -> next node ++ if (ctx->debug_graph > 0) { ++ GGML_LOG_DEBUG("%s: node[%5d] - %-12s %s\n", __func__, idx, ggml_op_name(node->op), "(noop)"); ++ } ++ } return 1; ++ default: ++ { ++ } break; ++ } ++ ++ if (!ggml_metal_device_supports_op(ctx->dev, node)) { ++ GGML_LOG_ERROR("%s: error: unsupported op '%s'\n", __func__, ggml_op_desc(node)); ++ GGML_ABORT("unsupported op"); ++ } ++ ++ if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { ++ return 1; ++ } ++ ++ int n_fuse = 1; ++ ++ // check if the current node can run concurrently with other nodes before it ++ // the condition is that: ++ // - the current node cannot write to any previous src or dst ranges ++ // - the current node cannot read from any previous dst ranges ++ // ++ // if the condition is not satisfied, we put a memory barrier and clear all ranges ++ // otherwise, we add the new ranges to the encoding context and process the node concurrently ++ // ++ { ++ const bool is_concurrent = ggml_metal_op_concurrency_check(ctx, node); ++ ++ if (!is_concurrent) { ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ ++ if (ctx->debug_graph > 0) { ++ GGML_LOG_DEBUG("%s: node[%5d] - %-12s %-12s %s\n", __func__, idx, ggml_op_name(node->op), ggml_get_name(node), is_concurrent ? "(concurrent)" : ""); ++ } ++ if (ctx->debug_graph > 1) { ++ GGML_TENSOR_LOCALS( int64_t, ne0, node->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, node->src[0], nb); ++ GGML_TENSOR_LOCALS( int64_t, ne1, node->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, node->src[1], nb); ++ GGML_TENSOR_LOCALS( int64_t, ne2, node->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, node->src[2], nb); ++ GGML_TENSOR_LOCALS( int64_t, ne3, node->src[3], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, node->src[3], nb); ++ GGML_TENSOR_LOCALS( int64_t, ne, node, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, node, nb); ++ ++ if (node->src[0]) { ++ GGML_LOG_DEBUG("%s: src0 - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(node->src[0]->type), ne00, ne01, ne02, ne03, nb00, nb01, nb02, nb03, ++ ggml_is_contiguous(node->src[0]), node->src[0]->name); ++ } ++ if (node->src[1]) { ++ GGML_LOG_DEBUG("%s: src1 - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(node->src[1]->type), ne10, ne11, ne12, ne13, nb10, nb11, nb12, nb13, ++ ggml_is_contiguous(node->src[1]), node->src[1]->name); ++ } ++ if (node->src[2]) { ++ GGML_LOG_DEBUG("%s: src2 - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(node->src[2]->type), ne20, ne21, ne22, ne23, nb20, nb21, nb22, nb23, ++ ggml_is_contiguous(node->src[2]), node->src[2]->name); ++ } ++ if (node->src[3]) { ++ GGML_LOG_DEBUG("%s: src3 - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(node->src[3]->type), ne30, ne31, ne32, ne33, nb30, nb31, nb32, nb33, ++ ggml_is_contiguous(node->src[3]), node->src[3]->name); ++ } ++ if (node) { ++ GGML_LOG_DEBUG("%s: node - %4s [%5lld, %5lld, %5lld, %5lld] [%5lld, %5lld, %5lld, %5lld], 1, %s\n", __func__, ggml_type_name(node->type), ne0, ne1, ne2, ne3, nb0, nb1, nb2, nb3, ++ node->name); ++ } ++ } ++ } ++ ++ if (node->op == GGML_OP_UNARY) { ++ n_fuse = ggml_metal_op_glm_moe_decode_motif_reference(ctx, idx); ++ if (n_fuse > 0) { ++ goto done; ++ } ++ n_fuse = ggml_metal_op_topk_moe_route_fused(ctx, idx); ++ if (n_fuse > 0) { ++ goto done; ++ } ++ n_fuse = 1; ++ } ++ ++ switch (node->op) { ++ case GGML_OP_CONCAT: ++ { ++ n_fuse = ggml_metal_op_concat(ctx, idx); ++ } break; ++ case GGML_OP_ADD: ++ { ++ n_fuse = ggml_metal_op_bin(ctx, idx); ++ } break; ++ case GGML_OP_SUB: ++ case GGML_OP_MUL: ++ case GGML_OP_DIV: ++ { ++ n_fuse = ggml_metal_op_bin(ctx, idx); ++ } break; ++ case GGML_OP_ADD_ID: ++ { ++ n_fuse = ggml_metal_op_add_id(ctx, idx); ++ } break; ++ case GGML_OP_REPEAT: ++ { ++ n_fuse = ggml_metal_op_repeat(ctx, idx); ++ } break; ++ case GGML_OP_ACC: ++ { ++ n_fuse = ggml_metal_op_acc(ctx, idx); + } break; + case GGML_OP_SCALE: + case GGML_OP_FILL: +@@ -301,7 +5482,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { + } break; + case GGML_OP_GLU: + { +- n_fuse = ggml_metal_op_glu(ctx, idx); ++ n_fuse = ggml_metal_op_weighted_swiglu(ctx, idx); ++ if (n_fuse == 0) { ++ n_fuse = ggml_metal_op_glu(ctx, idx); ++ } + } break; + case GGML_OP_SUM: + { +@@ -337,17 +5521,57 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { + { + n_fuse = ggml_metal_op_gated_delta_net(ctx, idx); + } break; ++ case GGML_OP_LIGHTNING_INDEXER: ++ { ++ n_fuse = ggml_metal_op_lightning_indexer(ctx, idx); ++ } break; ++ case GGML_OP_DSA_SPARSE_MASK: ++ { ++ n_fuse = ggml_metal_op_dsa_sparse_mask(ctx, idx); ++ } break; ++ case GGML_OP_DSA_SPARSE_ATTN: ++ { ++ n_fuse = ggml_metal_op_dsa_sparse_attn(ctx, idx); ++ } break; ++ case GGML_OP_DSA_TOP1_ATTN: ++ { ++ n_fuse = ggml_metal_op_dsa_top1_attn(ctx, idx); ++ } break; ++ case GGML_OP_MOE_ROUTE_WEIGHTS: ++ { ++ n_fuse = ggml_metal_op_glm_moe_decode_motif_reference(ctx, idx); ++ if (n_fuse == 0) { ++ n_fuse = ggml_metal_op_moe_route_weights(ctx, idx); ++ } ++ } break; ++ case GGML_OP_MOE_WEIGHTED_SUM: ++ { ++ n_fuse = ggml_metal_op_moe_weighted_sum(ctx, idx); ++ } break; ++ case GGML_OP_MOE_MUL_MAT_ID: ++ { ++ n_fuse = ggml_metal_op_moe_mul_mat_id(ctx, idx); ++ } break; + case GGML_OP_SOLVE_TRI: + { + n_fuse = ggml_metal_op_solve_tri(ctx, idx); + } break; + case GGML_OP_MUL_MAT: + { +- n_fuse = ggml_metal_op_mul_mat(ctx, idx); ++ n_fuse = ggml_metal_op_glm_absorbed_q(ctx, idx); ++ if (n_fuse == 0) { ++ n_fuse = ggml_metal_op_mul_mat(ctx, idx); ++ } + } break; + case GGML_OP_MUL_MAT_ID: + { +- n_fuse = ggml_metal_op_mul_mat_id(ctx, idx); ++ n_fuse = ggml_metal_op_mul_mv_id_weighted_reduce(ctx, idx); ++ if (n_fuse == 0) { ++ n_fuse = ggml_metal_op_mul_mv_id_gate_up_swiglu(ctx, idx); ++ } ++ if (n_fuse == 0) { ++ n_fuse = ggml_metal_op_mul_mat_id(ctx, idx); ++ } + } break; + case GGML_OP_GET_ROWS: + { +@@ -482,62 +5706,1421 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { + n_fuse = ggml_metal_op_count_equal(ctx, idx); + } break; + default: +- { +- GGML_LOG_ERROR("%s: error: node %3d, op = %8s not implemented\n", __func__, idx, ggml_op_name(node->op)); +- GGML_ABORT("fatal error"); +- } ++ { ++ GGML_LOG_ERROR("%s: error: node %3d, op = %8s not implemented\n", __func__, idx, ggml_op_name(node->op)); ++ GGML_ABORT("fatal error"); ++ } ++ } ++ ++done: ++ if (ctx->debug_graph > 0) { ++ if (n_fuse > 1) { ++ GGML_LOG_DEBUG("%s: fuse %d ops\n", __func__, n_fuse); ++ } ++ } ++ ++ // A backend-owned fusion can keep internal graph values in a private ++ // resource. In that case only publish the graph-visible outputs to the ++ // ordinary arena dependency tracker. ++ const bool track_selected_outputs = ctx->has_fused_range_outputs(); ++ for (int i = 0; i < n_fuse; ++i) { ++ if (track_selected_outputs && !ctx->tracks_fused_range_output(i)) { ++ continue; ++ } ++ if (!ggml_metal_op_concurrency_add(ctx, ctx->node(idx + i))) { ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ } ++ ctx->clear_fused_range_outputs(); ++ ++ return n_fuse; ++} ++ ++int ggml_metal_op_encode(ggml_metal_op_t ctx, int idx) { ++ if (ctx->use_capture) { ++ ggml_metal_encoder_debug_group_push(ctx->enc, ggml_op_desc(ctx->node(idx))); ++ } ++ ++ int res = ggml_metal_op_encode_impl(ctx, idx); ++ if (idx + res > ctx->n_nodes()) { ++ GGML_ABORT("fusion error: nodes spanning multiple encoders have been fused. this indicates a bug in the fusion logic %s", ++ "https://github.com/ggml-org/llama.cpp/pull/14849"); ++ } ++ ++ if (ctx->use_capture) { ++ ggml_metal_encoder_debug_group_pop(ctx->enc); ++ } ++ ++ return res; ++} ++ ++int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ const int32_t dim = ((const int32_t *) op->op_params)[0]; ++ ++ ggml_metal_kargs_concat args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.dim =*/ dim, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_concat(lib, op->type); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ++ int nth = std::min(256, ne0); ++ ++ // when rows are small, we can batch them together in a single threadgroup ++ int nrptg = 1; ++ if (nth < 256) { ++ nrptg = std::min((256 + nth - 1) / nth, ne1); ++ if (nrptg * nth > 256) { ++ nrptg = 256 / nth; ++ } ++ } ++ ++ const int nw0 = (ne1 + nrptg - 1) / nrptg; ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, nw0, ne2, ne3, nth, nrptg, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_repeat(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ auto pipeline = ggml_metal_library_get_pipeline_repeat(lib, op->type); ++ ++ ggml_metal_kargs_repeat args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ++ const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne0); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_acc(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[1])); ++ ++ const size_t pnb1 = ((const int32_t *) op->op_params)[0]; ++ const size_t pnb2 = ((const int32_t *) op->op_params)[1]; ++ const size_t pnb3 = ((const int32_t *) op->op_params)[2]; ++ const size_t offs = ((const int32_t *) op->op_params)[3]; ++ ++ const bool inplace = (bool) ((const int32_t *) op->op_params)[4]; ++ ++ if (!inplace) { ++ // run a separate kernel to cpy src->dst ++ // not sure how to avoid this ++ // TODO: make a simpler cpy_bytes kernel ++ ++ //const id pipeline = ctx->pipelines[GGML_METAL_PIPELINE_TYPE_CPY_F32_F32].obj; ++ auto pipeline = ggml_metal_library_get_pipeline_cpy(lib, op->src[0]->type, op->type); ++ ++ ggml_metal_kargs_cpy args = { ++ /*.nk0 =*/ ne00, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ++ const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); ++ ++ ggml_metal_op_concurrency_reset(ctx); ++ } ++ ++ ggml_metal_kargs_bin args = { ++ /*.ne00 =*/ ne10, ++ /*.ne01 =*/ ne11, ++ /*.ne02 =*/ ne12, ++ /*.ne03 =*/ ne13, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ pnb1, ++ /*.nb02 =*/ pnb2, ++ /*.nb03 =*/ pnb3, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne0 =*/ ne10, ++ /*.ne1 =*/ ne11, ++ /*.ne2 =*/ ne12, ++ /*.ne3 =*/ ne13, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ pnb1, ++ /*.nb2 =*/ pnb2, ++ /*.nb3 =*/ pnb3, ++ /*.offs =*/ offs, ++ /*.o1 =*/ { 0 }, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_bin_one(lib, GGML_OP_ADD); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ++ const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ ++ int nth = 1; ++ ++ while (2*nth < args.ne0 && nth < nth_max) { ++ nth *= 2; ++ } ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne11, ne12, ne13, nth, 1, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_unary(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ ++ ggml_metal_kargs_unary args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.slope =*/ 0.0, ++ /*.scale =*/ 0.0, ++ /*.bias =*/ 0.0, ++ /*.val =*/ 0.0, ++ /*.min =*/ 0.0, ++ /*.max =*/ 0.0, ++ }; ++ ++ if (op->op == GGML_OP_LEAKY_RELU) { ++ args.slope = ggml_get_op_params_f32(op, 0); ++ } ++ ++ if (op->op == GGML_OP_SCALE) { ++ args.scale = ggml_get_op_params_f32(op, 0); ++ args.bias = ggml_get_op_params_f32(op, 1); ++ } ++ ++ if (op->op == GGML_OP_FILL) { ++ args.val = ggml_get_op_params_f32(op, 0); ++ } ++ ++ if (op->op == GGML_OP_CLAMP) { ++ args.min = ggml_get_op_params_f32(op, 0); ++ args.max = ggml_get_op_params_f32(op, 1); ++ } ++ ++ if (op->op == GGML_OP_UNARY && ggml_get_unary_op(op) == GGML_UNARY_OP_XIELU) { ++ args.slope = ggml_get_op_params_f32(op, 1); // alpha_n ++ args.scale = ggml_get_op_params_f32(op, 2); // alpha_p ++ args.bias = ggml_get_op_params_f32(op, 3); // beta ++ args.val = ggml_get_op_params_f32(op, 4); // eps ++ } ++ ++ auto pipeline = ggml_metal_library_get_pipeline_unary(lib, op); ++ ++ if (pipeline.c4) { ++ args.ne00 = ne00/4; ++ args.ne0 = ne0/4; ++ } ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ ++ if (pipeline.cnt) { ++ const int n = pipeline.c4 ? ggml_nelements(op)/4 : ggml_nelements(op); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, n, 1, 1, 1, 1, 1); ++ } else { ++ const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ const int nth = MIN(args.ne00, nth_max); ++ const int nk0 = (args.ne00 + nth - 1)/nth; ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, nk0*ne01, ne02, ne03, nth, 1, 1); ++ } ++ ++ return 1; ++} ++ ++static bool ggml_metal_match_weighted_swiglu( ++ ggml_metal_op_t ctx, ++ int idx, ++ ggml_metal_weighted_swiglu_fusion & fusion) { ++ if (!ctx->use_fusion || !ggml_metal_glm_dsa_weighted_swiglu_enabled() || idx + 2 > ctx->n_nodes()) { ++ return false; ++ } ++ ++ ggml_tensor * glu = ctx->node(idx); ++ ggml_tensor * weighted = ctx->node(idx + 1); ++ if (glu->op != GGML_OP_GLU || weighted->op != GGML_OP_MUL) { ++ return false; ++ } ++ if (ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || ggml_get_op_params_i32(glu, 1) != 0) { ++ return false; ++ } ++ if (!ggml_metal_tensor_name_contains(glu, "ffn_moe_swiglu") || ++ !ggml_metal_tensor_name_contains(weighted, "ffn_moe_down_weighted_input")) { ++ return false; ++ } ++ ++ ggml_tensor * weights = nullptr; ++ if (weighted->src[0] == glu) { ++ weights = weighted->src[1]; ++ } else if (weighted->src[1] == glu) { ++ weights = weighted->src[0]; ++ } else { ++ return false; ++ } ++ ++ const bool shape_ok = ++ glu->src[0] != nullptr && ++ glu->src[1] != nullptr && ++ weights != nullptr && ++ glu->type == GGML_TYPE_F32 && ++ glu->src[0]->type == GGML_TYPE_F32 && ++ glu->src[1]->type == GGML_TYPE_F32 && ++ weighted->type == GGML_TYPE_F32 && ++ weights->type == GGML_TYPE_F32 && ++ weighted->ne[0] == glu->ne[0] && ++ weighted->ne[1] == glu->ne[1] && ++ weighted->ne[2] == glu->ne[2] && ++ weights->ne[0] == 1 && ++ weights->ne[1] == glu->ne[1] && ++ weights->ne[2] == glu->ne[2] && ++ ggml_is_contiguous_1(glu->src[0]) && ++ ggml_is_contiguous_1(glu->src[1]) && ++ ggml_is_contiguous_1(weighted); ++ if (!shape_ok) { ++ return false; ++ } ++ ++ const ggml_op ops[] = { GGML_OP_GLU, GGML_OP_MUL }; ++ const int outputs[] = { 1 }; ++ if (!ctx->can_fuse_subgraph(idx, ops, 2, outputs, 1)) { ++ return false; ++ } ++ ++ fusion.glu = glu; ++ fusion.weights = weights; ++ fusion.weighted = weighted; ++ fusion.n_fuse = 2; ++ return true; ++} ++ ++static int ggml_metal_op_weighted_swiglu(ggml_metal_op_t ctx, int idx) { ++ ggml_metal_weighted_swiglu_fusion fusion; ++ if (!ggml_metal_match_weighted_swiglu(ctx, idx, fusion)) { ++ return 0; ++ } ++ ++ ggml_tensor * glu = fusion.glu; ++ ggml_tensor * weights = fusion.weights; ++ ggml_tensor * weighted = fusion.weighted; ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, glu->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, glu->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, glu->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, glu->src[1], nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, weights, nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, weighted, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, weighted, nb); ++ ++ GGML_ASSERT(ggml_are_same_shape(glu->src[0], glu->src[1])); ++ ++ auto pipeline = ggml_metal_library_get_pipeline_glu_weighted(lib, weighted); ++ ++ ggml_metal_kargs_glu_weighted args = { ++ /*.ne00 =*/ ne00, ++ /*.nb01 =*/ nb01, ++ /*.ne10 =*/ ne10, ++ /*.nb11 =*/ nb11, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nb1 =*/ nb1, ++ /*.i00 =*/ 0, ++ /*.i10 =*/ 0, ++ /*.weights_nb1 =*/ nb21, ++ /*.weights_nb2 =*/ nb22, ++ /*.dst_nb1 =*/ nb1, ++ }; ++ ++ const int64_t nrows = ggml_nrows(weighted); ++ const int32_t nth = std::min( ++ ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ++ pipeline.c4 ? std::max(1, ne00/4) : ne00/2); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(glu->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(glu->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(weights), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(weighted), 4); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=weighted_swiglu tensor=%s glu=%s weights=%s src0_type=%s weights_type=%s dst_type=%s ne0=%d rows=%lld slots=%d tokens=%d c4=%d grid_x=%lld grid_y=1 grid_z=1 threads_x=%d fused_nodes=%d\n", ++ ggml_metal_tensor_name(weighted), ++ ggml_metal_tensor_name(glu), ++ ggml_metal_tensor_name(weights), ++ ggml_type_name(glu->src[0]->type), ++ ggml_type_name(weights->type), ++ ggml_type_name(weighted->type), ++ ne0, ++ (long long) nrows, ++ ne1, ++ ne2, ++ pipeline.c4 ? 1 : 0, ++ (long long) nrows, ++ nth, ++ fusion.n_fuse); ++ } ++ ++ return fusion.n_fuse; ++} ++ ++int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ if (op->src[1]) { ++ GGML_ASSERT(ggml_are_same_shape(op->src[0], op->src[1])); ++ } ++ ++ auto pipeline = ggml_metal_library_get_pipeline_glu(lib, op); ++ ++ const int32_t swp = ggml_get_op_params_i32(op, 1); ++ const float alpha = ggml_get_op_params_f32(op, 2); ++ const float limit = ggml_get_op_params_f32(op, 3); ++ ++ const int32_t i00 = swp ? ne0 : 0; ++ const int32_t i10 = swp ? 0 : ne0; ++ ++ ggml_metal_kargs_glu args = { ++ /*.ne00 =*/ ne00, ++ /*.nb01 =*/ nb01, ++ /*.ne10 =*/ op->src[1] ? ne10 : ne00, ++ /*.nb11 =*/ op->src[1] ? nb11 : nb01, ++ /*.ne0 =*/ ne0, ++ /*.nb1 =*/ nb1, ++ /*.i00 =*/ op->src[1] ? 0 : i00, ++ /*.i10 =*/ op->src[1] ? 0 : i10, ++ /*.alpha=*/ alpha, ++ /*.limit=*/ limit ++ }; ++ ++ const int64_t nrows = ggml_nrows(op->src[0]); ++ ++ const int32_t nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00/2); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ if (op->src[1]) { ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ } else { ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 2); ++ } ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled() && ++ ggml_metal_tensor_name_contains(op, "ffn_moe_swiglu")) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=glu kernel=%s tensor=%s gate=%s up=%s src0_type=%s src1_type=%s dst_type=%s ne0=%d rows=%lld slots=%d tokens=%d grid_x=%lld grid_y=1 grid_z=1 threads_x=%d\n", ++ ggml_glu_op_name(ggml_get_glu_op(op)), ++ ggml_metal_tensor_name(op), ++ ggml_metal_tensor_name(op->src[0]), ++ ggml_metal_tensor_name(op->src[1]), ++ ggml_type_name(op->src[0]->type), ++ op->src[1] ? ggml_type_name(op->src[1]->type) : ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->type), ++ ne0, ++ (long long) nrows, ++ ne1, ++ ne2, ++ (long long) nrows, ++ nth); ++ } ++ ++ return 1; ++} ++ ++int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ const uint64_t n = (uint64_t) ggml_nelements(op->src[0]); ++ ++ ggml_metal_kargs_sum args = { ++ /*.np =*/ n, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_sum(lib, op); ++ ++ int nth = 32; // SIMD width ++ ++ while (nth < (int) n && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ nth *= 2; ++ } ++ ++ nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ nth = std::min(nth, (int) n); ++ ++ const int nsg = (nth + 31) / 32; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, nsg * sizeof(float), 0); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, nth, 1, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_sum_rows(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ ++ ggml_metal_kargs_sum_rows args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_sum_rows(lib, op); ++ ++ if (pipeline.c4) { ++ args.ne00 = ne00/4; ++ args.ne0 = ne0/4; ++ } ++ ++ int nth = 32; // SIMD width ++ ++ while (nth < args.ne00 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ nth *= 2; ++ } ++ ++ nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ nth = std::min(nth, (int) args.ne00); ++ ++ const size_t smem = pipeline.smem; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_cumsum(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ auto pipeline_blk = ggml_metal_library_get_pipeline_cumsum_blk(lib, op); ++ ++ int nth = 1; ++ while (nth < ne00 && 2*nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_blk)) { ++ nth *= 2; ++ } ++ ++ GGML_ASSERT(ne00 <= nth*nth); ++ ++ const int64_t net0 = (ne00 + nth - 1) / nth; ++ const int64_t net1 = ne01; ++ const int64_t net2 = ne02; ++ const int64_t net3 = ne03; ++ ++ const uint64_t nbt0 = sizeof(float); ++ const uint64_t nbt1 = net0*nbt0; ++ const uint64_t nbt2 = net1*nbt1; ++ const uint64_t nbt3 = net2*nbt2; ++ ++ const size_t smem = GGML_PAD(32*sizeof(float), 16); ++ ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ ++ ggml_metal_buffer_id bid_tmp = bid_dst; ++ bid_tmp.offs += ggml_nbytes(op); ++ ++ { ++ ggml_metal_kargs_cumsum_blk args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.net0 =*/ net0, ++ /*.net1 =*/ net1, ++ /*.net2 =*/ net2, ++ /*.net3 =*/ net3, ++ /*.nbt0 =*/ nbt0, ++ /*.nbt1 =*/ nbt1, ++ /*.nbt2 =*/ nbt2, ++ /*.nbt3 =*/ nbt3, ++ /*.outb =*/ ne00 > nth, ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline_blk); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 3); ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, net0*ne01, ne02, ne03, nth, 1, 1); ++ } ++ ++ if (ne00 > nth) { ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ { ++ ggml_metal_kargs_cumsum_blk args = { ++ /*.ne00 =*/ net0, ++ /*.ne01 =*/ net1, ++ /*.ne02 =*/ net2, ++ /*.ne03 =*/ net3, ++ /*.nb00 =*/ nbt0, ++ /*.nb01 =*/ nbt1, ++ /*.nb02 =*/ nbt2, ++ /*.nb03 =*/ nbt3, ++ /*.net0 =*/ net0, ++ /*.net1 =*/ net1, ++ /*.net2 =*/ net2, ++ /*.net3 =*/ net3, ++ /*.nbt0 =*/ nbt0, ++ /*.nbt1 =*/ nbt1, ++ /*.nbt2 =*/ nbt2, ++ /*.nbt3 =*/ nbt3, ++ /*.outb =*/ false, ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline_blk); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 3); ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, net1, net2, net3, nth, 1, 1); ++ } ++ ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ { ++ auto pipeline_add = ggml_metal_library_get_pipeline_cumsum_add(lib, op); ++ ++ ggml_metal_kargs_cumsum_add args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.net0 =*/ net0, ++ /*.net1 =*/ net1, ++ /*.net2 =*/ net2, ++ /*.net3 =*/ net3, ++ /*.nbt0 =*/ nbt0, ++ /*.nbt1 =*/ nbt1, ++ /*.nbt2 =*/ nbt2, ++ /*.nbt3 =*/ nbt3, ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline_add); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, net0*ne01, ne02, ne03, nth, 1, 1); ++ } ++ } ++ ++ return 1; ++} ++ ++static int32_t ggml_metal_glm_dsa_packed_gather_rows_per_tg( ++ const ggml_metal_device_props * props_dev) { ++ const char * value = getenv("LLAMA_GLM_DSA_EXPERIMENTAL_PACKED_GATHER_ROWS_PER_TG"); ++ if (value == nullptr || value[0] == '\0') { ++ return props_dev->device_id == GGML_METAL_DEVICE_M3_ULTRA ? 16 : 1; ++ } ++ ++ switch (atoi(value)) { ++ case 2: ++ case 4: ++ case 8: ++ case 16: ++ case 32: ++ case 64: ++ return atoi(value); ++ default: ++ return 1; ++ } ++} ++ ++static int32_t ggml_metal_glm_dsa_packed_gather_threads_per_row( ++ const ggml_metal_device_props * props_dev) { ++ const char * value = getenv("LLAMA_GLM_DSA_EXPERIMENTAL_PACKED_GATHER_THREADS_PER_ROW"); ++ if (value == nullptr || value[0] == '\0') { ++ return props_dev->device_id == GGML_METAL_DEVICE_M3_ULTRA ? 32 : 64; ++ } ++ ++ switch (atoi(value)) { ++ case 16: ++ case 32: ++ case 64: ++ return atoi(value); ++ default: ++ return 64; ++ } ++} ++ ++int ggml_metal_op_get_rows(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_log_selected_row_flash_candidate(ctx, idx); ++ if (ggml_metal_find_selected_row_flash_consumer(ctx, idx, op) != nullptr) { ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=selected_row_flash_skip tensor=%s reason=deferred_to_flash grid_x=1 grid_y=1 grid_z=1 threads_x=1\n", ++ ggml_metal_tensor_name(op)); ++ } ++ return 1; ++ } ++ if (ggml_metal_selected_row_flash_can_defer_compact_k_rows(ctx, idx, op)) { ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=selected_row_flash_skip tensor=%s reason=deferred_compact_k_contract grid_x=1 grid_y=1 grid_z=1 threads_x=1\n", ++ ggml_metal_tensor_name(op)); ++ } ++ return 1; ++ } ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ const bool f16_to_f16 = op->src[0]->type == GGML_TYPE_F16 && op->type == GGML_TYPE_F16; ++ ++ ggml_metal_kargs_get_rows args = { ++ /*.ne00t =*/ ggml_is_quantized(op->src[0]->type) ? ne00/16 : ne00, ++ /*.ne00 =*/ ne00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ }; ++ ++ const ggml_metal_buffer_id src0_buffer = ggml_metal_get_buffer_id(op->src[0]); ++ const ggml_metal_buffer_id src1_buffer = ggml_metal_get_buffer_id(op->src[1]); ++ const ggml_metal_buffer_id dst_buffer = ggml_metal_get_buffer_id(op); ++ ++ const bool use_f16_vec4 = f16_to_f16 && ++ args.ne00 >= 4 && ++ args.nb01%8 == 0 && ++ args.nb1%8 == 0 && ++ src0_buffer.offs%8 == 0 && ++ dst_buffer.offs%8 == 0; ++ if (use_f16_vec4) { ++ args.ne00t = (args.ne00 + 3)/4; ++ } ++ ++ const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ const int32_t requested_rows_per_tg = ggml_metal_glm_dsa_packed_gather_rows_per_tg(props_dev); ++ const int32_t requested_threads_per_row = ggml_metal_glm_dsa_packed_gather_threads_per_row(props_dev); ++ const bool use_packed_rows = use_f16_vec4 && ++ requested_rows_per_tg > 1 && ++ requested_rows_per_tg*requested_threads_per_row <= 1024 && ++ args.ne00 == 576 && ++ args.ne00%4 == 0 && ++ args.ne10 >= 256 && ++ ne11 == 1 && ++ ne12 == 1; ++ ++ auto pipeline = use_packed_rows ? ++ ggml_metal_library_get_pipeline_get_rows_packed_f16(lib) : ++ ggml_metal_library_get_pipeline_get_rows(lib, op->src[0]->type, op->type, use_f16_vec4); ++ ++ const int nth = use_packed_rows ? requested_threads_per_row : ++ std::min(args.ne00t, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ const int rows_per_tg = use_packed_rows ? requested_rows_per_tg : 1; ++ GGML_ASSERT(nth*rows_per_tg <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ ++ const int nw0 = (args.ne00t + nth - 1)/nth; ++ const int grid_x = use_packed_rows ? (ne10 + rows_per_tg - 1)/rows_per_tg : nw0*ne10; ++ const int grid_y = ne11; ++ const int grid_z = ne12; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, src0_buffer, 1); ++ ggml_metal_encoder_set_buffer (enc, src1_buffer, 2); ++ ggml_metal_encoder_set_buffer (enc, dst_buffer, 3); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=get_rows kernel=%s tensor=%s src_type=%s top_k_type=%s dst_type=%s rows=%lld rows_per_tg=%d grid_x=%d grid_y=%d grid_z=%d threads_x=%d threads_y=%d\n", ++ use_packed_rows ? "packed_rows" : ++ (use_f16_vec4 ? "typed_vec4" : op->src[0]->type == op->type ? "typed" : "promote"), ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->type), ++ (long long) ne10, ++ rows_per_tg, ++ grid_x, ++ grid_y, ++ grid_z, ++ nth, ++ rows_per_tg); ++ } ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, nth, rows_per_tg, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_set_rows(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ auto pipeline = ggml_metal_library_get_pipeline_set_rows(lib, op->src[1]->type, op->type); ++ ++ const int32_t nk0 = ne0/ggml_blck_size(op->type); ++ ++ int nth = 32; // SIMD width ++ ++ while (nth < nk0 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ nth *= 2; ++ } ++ ++ int nrptg = 1; ++ if (nth > nk0) { ++ nrptg = (nth + nk0 - 1)/nk0; ++ nth = nk0; ++ ++ if (nrptg*nth > ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ nrptg--; ++ } ++ } ++ ++ nth = std::min(nth, nk0); ++ ++ ggml_metal_kargs_set_rows args = { ++ /*.nk0 =*/ nk0, ++ /*.ne01 =*/ ne01, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ }; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nrptg - 1)/nrptg, ne02, ne03, nth, nrptg, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS(int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS(int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ ggml_metal_kargs_diag args = { ++ /*.ne00 =*/ne00, ++ /*.ne01 =*/ne01, ++ /*.ne02 =*/ne02, ++ /*.ne03 =*/ne03, ++ /*.nb00 =*/nb00, ++ /*.nb01 =*/nb01, ++ /*.nb02 =*/nb02, ++ /*.nb03 =*/nb03, ++ /*.ne0 =*/ne0, ++ /*.ne1 =*/ne1, ++ /*.ne2 =*/ne2, ++ /*.ne3 =*/ne3, ++ /*.nb0 =*/nb0, ++ /*.nb1 =*/nb1, ++ /*.nb2 =*/nb2, ++ /*.nb3 =*/nb3, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_diag(lib, op); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 2); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, 32, 1, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ float scale; ++ float max_bias; ++ ++ memcpy(&scale, ((const int32_t *) op->op_params) + 0, sizeof(scale)); ++ memcpy(&max_bias, ((const int32_t *) op->op_params) + 1, sizeof(max_bias)); ++ ++ const uint32_t n_head = op->src[0]->ne[2]; ++ const int32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head)); ++ ++ const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); ++ const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); ++ ++ // softmax ++ ++ ggml_metal_kargs_soft_max args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.scale =*/ scale, ++ /*.max_bias =*/ max_bias, ++ /*.m0 =*/ m0, ++ /*.m1 =*/ m1, ++ /*.n_head_log2 =*/ n_head_log2, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_soft_max(lib, op); ++ ++ int nth = 32; // SIMD width ++ ++ if (ne00%4 == 0) { ++ while (nth < ne00/4 && nth*ne01*ne02*ne03 < 256) { ++ nth *= 2; ++ } ++ } else { ++ while (nth < ne00 && nth*ne01*ne02*ne03 < 256) { ++ nth *= 2; ++ } ++ } ++ ++ const size_t smem = pipeline.smem; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ if (op->src[1]) { ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ } else { ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 2); + } ++ if (op->src[2]) { ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[2]), 3); ++ } else { ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 3); ++ } ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_ssm_conv(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- if (ctx->debug_graph > 0) { +- if (n_fuse > 1) { +- GGML_LOG_DEBUG("%s: fuse %d ops\n", __func__, n_fuse); +- } +- } ++ ggml_metal_kargs_ssm_conv args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ }; + +- // update the mem ranges in the encoding context +- for (int i = 0; i < n_fuse; ++i) { +- if (!ggml_metal_op_concurrency_add(ctx, ctx->node(idx + i))) { +- ggml_metal_op_concurrency_reset(ctx); +- } ++ // Use batched kernel for prefill (ne1 > 1) to reduce threadgroup dispatch overhead ++ const bool use_batched = (ne1 > 1); ++ ++ if (use_batched) { ++ // Determine the smallest power of 2 that's >= ne1, but <= 256 ++ int BATCH_SIZE; ++ if (ne1 > 128) BATCH_SIZE = 256; ++ else if (ne1 > 64 ) BATCH_SIZE = 128; ++ else if (ne1 > 32 ) BATCH_SIZE = 64; ++ else if (ne1 > 16 ) BATCH_SIZE = 32; ++ else if (ne1 > 8 ) BATCH_SIZE = 16; ++ else if (ne1 > 4 ) BATCH_SIZE = 8; ++ else BATCH_SIZE = 2; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_ssm_conv_batched(lib, op, BATCH_SIZE); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3); ++ ++ // Dispatch: ne01 rows, ceil(ne1/BATCH_SIZE) token batches, ne02 sequences ++ // Each threadgroup has BATCH_SIZE threads, each handling one token ++ const int n_token_batches = (ne1 + BATCH_SIZE - 1) / BATCH_SIZE; ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, n_token_batches, ne02, BATCH_SIZE, 1, 1); ++ } else { ++ auto pipeline = ggml_metal_library_get_pipeline_ssm_conv(lib, op); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne1, ne02, 1, 1, 1); + } + +- return n_fuse; ++ return 1; + } + +-int ggml_metal_op_encode(ggml_metal_op_t ctx, int idx) { +- if (ctx->use_capture) { +- ggml_metal_encoder_debug_group_push(ctx->enc, ggml_op_desc(ctx->node(idx))); +- } ++int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); + +- int res = ggml_metal_op_encode_impl(ctx, idx); +- if (idx + res > ctx->n_nodes()) { +- GGML_ABORT("fusion error: nodes spanning multiple encoders have been fused. this indicates a bug in the fusion logic %s", +- "https://github.com/ggml-org/llama.cpp/pull/14849"); +- } ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +- if (ctx->use_capture) { +- ggml_metal_encoder_debug_group_pop(ctx->enc); ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne4, op->src[4], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb4, op->src[4], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne5, op->src[5], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb5, op->src[5], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne6, op->src[6], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb6, op->src[6], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ const ggml_tensor * src3 = op->src[3]; ++ const ggml_tensor * src4 = op->src[4]; ++ const ggml_tensor * src5 = op->src[5]; ++ const ggml_tensor * src6 = op->src[6]; ++ ++ GGML_ASSERT(src3); ++ GGML_ASSERT(src4); ++ GGML_ASSERT(src5); ++ GGML_ASSERT(src6); ++ ++ const int64_t d_state = ne00; ++ const int64_t d_inner = ne01; ++ const int64_t n_head = ne02; ++ const int64_t n_group = ne41; ++ const int64_t n_seq_tokens = ne12; ++ const int64_t n_seqs = ne13; ++ ++ ggml_metal_kargs_ssm_scan args = { ++ /*.d_state =*/ d_state, ++ /*.d_inner =*/ d_inner, ++ /*.n_head =*/ n_head, ++ /*.n_group =*/ n_group, ++ /*.n_seq_tokens =*/ n_seq_tokens, ++ /*.n_seqs =*/ n_seqs, ++ /*.s_off =*/ ggml_nelements(op->src[1]) * sizeof(float), ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.ns12 =*/ nb12/nb10, ++ /*.nb13 =*/ nb13, ++ /*.nb20 =*/ nb20, ++ /*.nb21 =*/ nb21, ++ /*.ns21 =*/ nb21/nb20, ++ /*.nb22 =*/ nb22, ++ /*.ne30 =*/ ne30, ++ /*.nb31 =*/ nb31, ++ /*.nb41 =*/ nb41, ++ /*.nb42 =*/ nb42, ++ /*.ns42 =*/ nb42/nb40, ++ /*.nb43 =*/ nb43, ++ /*.nb51 =*/ nb51, ++ /*.nb52 =*/ nb52, ++ /*.ns52 =*/ nb52/nb50, ++ /*.nb53 =*/ nb53, ++ /*.nb0 =*/ nb0, ++ }; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_ssm_scan(lib, op); ++ ++ GGML_ASSERT(d_state <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ ++ const size_t smem = pipeline.smem; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), 4); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), 5); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), 6); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), 7); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 8); ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, d_inner, n_head, n_seqs, d_state, 1, 1); ++ ++ return 1; ++} ++ ++int ggml_metal_op_rwkv(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ ++ const int64_t B = op->op == GGML_OP_RWKV_WKV6 ? op->src[5]->ne[1] : op->src[6]->ne[1]; ++ const int64_t T = op->src[0]->ne[2]; ++ const int64_t C = op->ne[0]; ++ const int64_t H = op->src[0]->ne[1]; ++ ++ auto pipeline = ggml_metal_library_get_pipeline_rwkv(lib, op); ++ ++ int ida = 0; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); ++ if (op->op == GGML_OP_RWKV_WKV7) { ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), ida++); + } ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); ++ ggml_metal_encoder_set_bytes (enc, (void *) &B, sizeof(B), ida++); ++ ggml_metal_encoder_set_bytes (enc, (void *) &T, sizeof(T), ida++); ++ ggml_metal_encoder_set_bytes (enc, (void *) &C, sizeof(C), ida++); ++ ggml_metal_encoder_set_bytes (enc, (void *) &H, sizeof(H), ida++); + +- return res; ++ ggml_metal_encoder_dispatch_threadgroups(enc, B * H, 1, 1, C/H, 1, 1); ++ ++ return 1; + } + +-int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + ++ + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const int32_t dim = ((const int32_t *) op->op_params)[0]; ++ auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op); + +- ggml_metal_kargs_concat args = { ++ int ida = 0; ++ ++ ggml_metal_kargs_gated_delta_net args = { + /*.ne00 =*/ ne00, + /*.ne01 =*/ ne01, + /*.ne02 =*/ ne02, +@@ -554,6 +7137,17 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) { + /*.nb11 =*/ nb11, + /*.nb12 =*/ nb12, + /*.nb13 =*/ nb13, ++ /*.ne20 =*/ ne20, ++ /*.ne21 =*/ ne21, ++ /*.ne22 =*/ ne22, ++ /*.ne23 =*/ ne23, ++ /*.nb20 =*/ nb20, ++ /*.nb21 =*/ nb21, ++ /*.nb22 =*/ nb22, ++ /*.nb23 =*/ nb23, ++ /*.ns02 =*/ (int32_t) (nb02/sizeof(float)), ++ /*.ns12 =*/ (int32_t) (nb12/sizeof(float)), ++ /*.ns22 =*/ (int32_t) (nb22/sizeof(float)), + /*.ne0 =*/ ne0, + /*.ne1 =*/ ne1, + /*.ne2 =*/ ne2, +@@ -561,37 +7155,27 @@ int ggml_metal_op_concat(ggml_metal_op_t ctx, int idx) { + /*.nb0 =*/ nb0, + /*.nb1 =*/ nb1, + /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.dim =*/ dim, +- }; +- +- auto pipeline = ggml_metal_library_get_pipeline_concat(lib, op->type); +- +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); +- +- int nth = std::min(256, ne0); +- +- // when rows are small, we can batch them together in a single threadgroup +- int nrptg = 1; +- if (nth < 256) { +- nrptg = std::min((256 + nth - 1) / nth, ne1); +- if (nrptg * nth > 256) { +- nrptg = 256 / nth; +- } +- } ++ /*.nb3 =*/ nb3, ++ }; + +- const int nw0 = (ne1 + nrptg - 1) / nrptg; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // k ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // v ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // gate ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); // beta ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); // state ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst + +- ggml_metal_encoder_dispatch_threadgroups(enc, nw0, ne2, ne3, nth, nrptg, 1); ++ const int nsg = pipeline.nsg; ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, op->src[2]->ne[0]/nsg, op->src[2]->ne[1], op->src[2]->ne[3], 32, nsg, 1); + + return 1; + } + +-int ggml_metal_op_repeat(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -599,12 +7183,12 @@ int ggml_metal_op_repeat(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- auto pipeline = ggml_metal_library_get_pipeline_repeat(lib, op->type); +- +- ggml_metal_kargs_repeat args = { ++ ggml_metal_kargs_solve_tri args = { + /*.ne00 =*/ ne00, + /*.ne01 =*/ ne01, + /*.ne02 =*/ ne02, +@@ -613,6 +7197,14 @@ int ggml_metal_op_repeat(ggml_metal_op_t ctx, int idx) { + /*.nb01 =*/ nb01, + /*.nb02 =*/ nb02, + /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, + /*.ne0 =*/ ne0, + /*.ne1 =*/ ne1, + /*.ne2 =*/ ne2, +@@ -623,19 +7215,24 @@ int ggml_metal_op_repeat(ggml_metal_op_t ctx, int idx) { + /*.nb3 =*/ nb3, + }; + ++ auto pipeline = ggml_metal_library_get_pipeline_solve_tri(lib, op); ++ + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne0); ++ const int nsg = pipeline.nsg; + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, (ne10 + nsg - 1)/nsg, ne02, ne03, 32, nsg, 1); + + return 1; + } + +-int ggml_metal_op_acc(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_set(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -648,12 +7245,9 @@ int ggml_metal_op_acc(ggml_metal_op_t ctx, int idx) { + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->type == GGML_TYPE_F32); +- +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[1])); ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); + + const size_t pnb1 = ((const int32_t *) op->op_params)[0]; + const size_t pnb2 = ((const int32_t *) op->op_params)[1]; +@@ -692,8 +7286,8 @@ int ggml_metal_op_acc(ggml_metal_op_t ctx, int idx) { + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); + + const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00); + +@@ -702,245 +7296,71 @@ int ggml_metal_op_acc(ggml_metal_op_t ctx, int idx) { + ggml_metal_op_concurrency_reset(ctx); + } + +- ggml_metal_kargs_bin args = { +- /*.ne00 =*/ ne10, +- /*.ne01 =*/ ne11, +- /*.ne02 =*/ ne12, +- /*.ne03 =*/ ne13, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ pnb1, +- /*.nb02 =*/ pnb2, +- /*.nb03 =*/ pnb3, +- /*.ne10 =*/ ne10, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.ne13 =*/ ne13, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ne0 =*/ ne10, +- /*.ne1 =*/ ne11, +- /*.ne2 =*/ ne12, +- /*.ne3 =*/ ne13, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ pnb1, +- /*.nb2 =*/ pnb2, +- /*.nb3 =*/ pnb3, +- /*.offs =*/ offs, +- /*.o1 =*/ { 0 }, +- }; +- +- auto pipeline = ggml_metal_library_get_pipeline_bin_one(lib, GGML_OP_ADD); +- +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); +- +- const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); +- +- int nth = 1; +- +- while (2*nth < args.ne0 && nth < nth_max) { +- nth *= 2; +- } +- +- ggml_metal_encoder_dispatch_threadgroups(enc, ne11, ne12, ne13, nth, 1, 1); +- +- return 1; +-} +- +-int ggml_metal_op_unary(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); +- +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; +- +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); +- +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); +- +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); +- +- ggml_metal_kargs_unary args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.slope =*/ 0.0, +- /*.scale =*/ 0.0, +- /*.bias =*/ 0.0, +- /*.val =*/ 0.0, +- /*.min =*/ 0.0, +- /*.max =*/ 0.0, +- }; +- +- if (op->op == GGML_OP_LEAKY_RELU) { +- args.slope = ggml_get_op_params_f32(op, 0); +- } +- +- if (op->op == GGML_OP_SCALE) { +- args.scale = ggml_get_op_params_f32(op, 0); +- args.bias = ggml_get_op_params_f32(op, 1); +- } +- +- if (op->op == GGML_OP_FILL) { +- args.val = ggml_get_op_params_f32(op, 0); +- } +- +- if (op->op == GGML_OP_CLAMP) { +- args.min = ggml_get_op_params_f32(op, 0); +- args.max = ggml_get_op_params_f32(op, 1); +- } +- +- if (op->op == GGML_OP_UNARY && ggml_get_unary_op(op) == GGML_UNARY_OP_XIELU) { +- args.slope = ggml_get_op_params_f32(op, 1); // alpha_n +- args.scale = ggml_get_op_params_f32(op, 2); // alpha_p +- args.bias = ggml_get_op_params_f32(op, 3); // beta +- args.val = ggml_get_op_params_f32(op, 4); // eps +- } +- +- auto pipeline = ggml_metal_library_get_pipeline_unary(lib, op); +- +- if (pipeline.c4) { +- args.ne00 = ne00/4; +- args.ne0 = ne0/4; +- } +- +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); +- +- if (pipeline.cnt) { +- const int n = pipeline.c4 ? ggml_nelements(op)/4 : ggml_nelements(op); +- +- ggml_metal_encoder_dispatch_threadgroups(enc, n, 1, 1, 1, 1, 1); +- } else { +- const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); +- const int nth = MIN(args.ne00, nth_max); +- const int nk0 = (args.ne00 + nth - 1)/nth; +- +- ggml_metal_encoder_dispatch_threadgroups(enc, nk0*ne01, ne02, ne03, nth, 1, 1); +- } +- +- return 1; +-} +- +-int ggml_metal_op_glu(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); +- +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; +- +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); +- +- if (op->src[1]) { +- GGML_ASSERT(ggml_are_same_shape(op->src[0], op->src[1])); +- } +- +- auto pipeline = ggml_metal_library_get_pipeline_glu(lib, op); +- +- const int32_t swp = ggml_get_op_params_i32(op, 1); +- const float alpha = ggml_get_op_params_f32(op, 2); +- const float limit = ggml_get_op_params_f32(op, 3); +- +- const int32_t i00 = swp ? ne0 : 0; +- const int32_t i10 = swp ? 0 : ne0; +- +- ggml_metal_kargs_glu args = { +- /*.ne00 =*/ ne00, +- /*.nb01 =*/ nb01, +- /*.ne10 =*/ op->src[1] ? ne10 : ne00, +- /*.nb11 =*/ op->src[1] ? nb11 : nb01, +- /*.ne0 =*/ ne0, +- /*.nb1 =*/ nb1, +- /*.i00 =*/ op->src[1] ? 0 : i00, +- /*.i10 =*/ op->src[1] ? 0 : i10, +- /*.alpha=*/ alpha, +- /*.limit=*/ limit +- }; +- +- const int64_t nrows = ggml_nrows(op->src[0]); ++ auto pipeline = ggml_metal_library_get_pipeline_cpy(lib, op->src[1]->type, op->type); + +- const int32_t nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00/2); ++ GGML_ASSERT(ne10 % ggml_blck_size(op->src[1]->type) == 0); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- if (op->src[1]) { +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- } else { +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 2); ++ int64_t nk0 = ne10; ++ if (ggml_is_quantized(op->src[1]->type)) { ++ nk0 = ne10/16; ++ } else if (ggml_is_quantized(op->type)) { ++ nk0 = ne10/ggml_blck_size(op->type); + } +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1); ++ int nth = std::min(nk0*ne11, 256); + +- return 1; +-} ++ // when rows are small, we can batch them together in a single threadgroup ++ int nrptg = 1; + +-int ggml_metal_op_sum(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ // TODO: relax this constraint in the future ++ if (ggml_blck_size(op->src[1]->type) == 1 && ggml_blck_size(op->type) == 1) { ++ if (nth > nk0) { ++ nrptg = (nth + nk0 - 1)/nk0; ++ nth = nk0; + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ if (nrptg*nth > 256) { ++ nrptg--; ++ } ++ } ++ } + +- const uint64_t n = (uint64_t) ggml_nelements(op->src[0]); ++ nth = std::min(nth, nk0); + +- ggml_metal_kargs_sum args = { +- /*.np =*/ n, ++ ggml_metal_kargs_cpy args = { ++ /*.nk0 =*/ nk0, ++ /*.ne00 =*/ ne10, ++ /*.ne01 =*/ ne11, ++ /*.ne02 =*/ ne12, ++ /*.ne03 =*/ ne13, ++ /*.nb00 =*/ nb10, ++ /*.nb01 =*/ nb11, ++ /*.nb02 =*/ nb12, ++ /*.nb03 =*/ nb13, ++ /*.ne0 =*/ ne10, ++ /*.ne1 =*/ ne11, ++ /*.ne2 =*/ ne12, ++ /*.ne3 =*/ ne13, ++ /*.nb0 =*/ ggml_element_size(op), ++ /*.nb1 =*/ pnb1, ++ /*.nb2 =*/ pnb2, ++ /*.nb3 =*/ pnb3, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_sum(lib, op); +- +- int nth = 32; // SIMD width +- +- while (nth < (int) n && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- nth *= 2; +- } +- +- nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); +- nth = std::min(nth, (int) n); ++ const int nw0 = nrptg == 1 ? (nk0 + nth - 1)/nth : 1; + +- const int nsg = (nth + 31) / 32; ++ bid_dst.offs += offs; + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, nsg * sizeof(float), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nw0*(ne11 + nrptg - 1)/nrptg, ne12, ne13, nth, nrptg, 1); + + return 1; + } + +-int ggml_metal_op_sum_rows(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_cpy(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -951,12 +7371,38 @@ int ggml_metal_op_sum_rows(ggml_metal_op_t ctx, int idx) { + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ auto pipeline = ggml_metal_library_get_pipeline_cpy(lib, op->src[0]->type, op->type); + +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ GGML_ASSERT(ne00 % ggml_blck_size(op->src[0]->type) == 0); + +- ggml_metal_kargs_sum_rows args = { ++ int64_t nk0 = ne00; ++ if (ggml_is_quantized(op->src[0]->type)) { ++ nk0 = ne00/16; ++ } else if (ggml_is_quantized(op->type)) { ++ nk0 = ne00/ggml_blck_size(op->type); ++ } ++ ++ int nth = std::min(nk0*ne01, 256); ++ ++ // when rows are small, we can batch them together in a single threadgroup ++ int nrptg = 1; ++ ++ // TODO: relax this constraint in the future ++ if (ggml_blck_size(op->src[0]->type) == 1 && ggml_blck_size(op->type) == 1) { ++ if (nth > nk0) { ++ nrptg = (nth + nk0 - 1)/nk0; ++ nth = nk0; ++ ++ if (nrptg*nth > 256) { ++ nrptg--; ++ } ++ } ++ } ++ ++ nth = std::min(nth, nk0); ++ ++ ggml_metal_kargs_cpy args = { ++ /*.nk0 =*/ nk0, + /*.ne00 =*/ ne00, + /*.ne01 =*/ ne01, + /*.ne02 =*/ ne02, +@@ -975,180 +7421,67 @@ int ggml_metal_op_sum_rows(ggml_metal_op_t ctx, int idx) { + /*.nb3 =*/ nb3, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_sum_rows(lib, op); +- +- if (pipeline.c4) { +- args.ne00 = ne00/4; +- args.ne0 = ne0/4; +- } +- +- int nth = 32; // SIMD width +- +- while (nth < args.ne00 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- nth *= 2; +- } +- +- nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); +- nth = std::min(nth, (int) args.ne00); +- +- const size_t smem = pipeline.smem; ++ const int nw0 = nrptg == 1 ? (nk0 + nth - 1)/nth : 1; + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nw0*(ne01 + nrptg - 1)/nrptg, ne02, ne03, nth, nrptg, 1); + + return 1; + } + +-int ggml_metal_op_cumsum(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_pool_1d(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); +- + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- auto pipeline_blk = ggml_metal_library_get_pipeline_cumsum_blk(lib, op); +- +- int nth = 1; +- while (nth < ne00 && 2*nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_blk)) { +- nth *= 2; +- } +- +- GGML_ASSERT(ne00 <= nth*nth); +- +- const int64_t net0 = (ne00 + nth - 1) / nth; +- const int64_t net1 = ne01; +- const int64_t net2 = ne02; +- const int64_t net3 = ne03; +- +- const uint64_t nbt0 = sizeof(float); +- const uint64_t nbt1 = net0*nbt0; +- const uint64_t nbt2 = net1*nbt1; +- const uint64_t nbt3 = net2*nbt2; +- +- const size_t smem = GGML_PAD(32*sizeof(float), 16); +- +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); +- +- ggml_metal_buffer_id bid_tmp = bid_dst; +- bid_tmp.offs += ggml_nbytes(op); +- +- { +- ggml_metal_kargs_cumsum_blk args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.net0 =*/ net0, +- /*.net1 =*/ net1, +- /*.net2 =*/ net2, +- /*.net3 =*/ net3, +- /*.nbt0 =*/ nbt0, +- /*.nbt1 =*/ nbt1, +- /*.nbt2 =*/ nbt2, +- /*.nbt3 =*/ nbt3, +- /*.outb =*/ ne00 > nth, +- }; +- +- ggml_metal_encoder_set_pipeline(enc, pipeline_blk); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_tmp, 2); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 3); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); +- +- ggml_metal_encoder_dispatch_threadgroups(enc, net0*ne01, ne02, ne03, nth, 1, 1); +- } +- +- if (ne00 > nth) { +- ggml_metal_op_concurrency_reset(ctx); +- +- { +- ggml_metal_kargs_cumsum_blk args = { +- /*.ne00 =*/ net0, +- /*.ne01 =*/ net1, +- /*.ne02 =*/ net2, +- /*.ne03 =*/ net3, +- /*.nb00 =*/ nbt0, +- /*.nb01 =*/ nbt1, +- /*.nb02 =*/ nbt2, +- /*.nb03 =*/ nbt3, +- /*.net0 =*/ net0, +- /*.net1 =*/ net1, +- /*.net2 =*/ net2, +- /*.net3 =*/ net3, +- /*.nbt0 =*/ nbt0, +- /*.nbt1 =*/ nbt1, +- /*.nbt2 =*/ nbt2, +- /*.nbt3 =*/ nbt3, +- /*.outb =*/ false, +- }; +- +- ggml_metal_encoder_set_pipeline(enc, pipeline_blk); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_tmp, 1); +- ggml_metal_encoder_set_buffer (enc, bid_tmp, 2); +- ggml_metal_encoder_set_buffer (enc, bid_tmp, 3); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ const int32_t * opts = op->op_params; ++ ggml_op_pool op_pool = (ggml_op_pool) opts[0]; + +- ggml_metal_encoder_dispatch_threadgroups(enc, net1, net2, net3, nth, 1, 1); +- } ++ const int32_t k0 = opts[1]; ++ const int32_t s0 = opts[2]; ++ const int32_t p0 = opts[3]; + +- ggml_metal_op_concurrency_reset(ctx); ++ const int64_t IW = op->src[0]->ne[0]; ++ const int64_t OW = op->ne[0]; + +- { +- auto pipeline_add = ggml_metal_library_get_pipeline_cumsum_add(lib, op); ++ const int64_t np = ggml_nelements(op); + +- ggml_metal_kargs_cumsum_add args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.net0 =*/ net0, +- /*.net1 =*/ net1, +- /*.net2 =*/ net2, +- /*.net3 =*/ net3, +- /*.nbt0 =*/ nbt0, +- /*.nbt1 =*/ nbt1, +- /*.nbt2 =*/ nbt2, +- /*.nbt3 =*/ nbt3, +- }; ++ ggml_metal_kargs_pool_1d args_pool_1d = { ++ /* .k0 = */ k0, ++ /* .s0 = */ s0, ++ /* .p0 = */ p0, ++ /* .IW = */ IW, ++ /* .OW = */ OW, ++ /* .np = */ np ++ }; + +- ggml_metal_encoder_set_pipeline(enc, pipeline_add); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_tmp, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ auto pipeline = ggml_metal_library_get_pipeline_pool_1d(lib, op, op_pool); + +- ggml_metal_encoder_dispatch_threadgroups(enc, net0*ne01, ne02, ne03, nth, 1, 1); +- } +- } ++ const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), (int) np); ++ const int ntg = (np + nth - 1) / nth; ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args_pool_1d, sizeof(args_pool_1d), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_get_rows(ggml_metal_op_t ctx, int idx) { ++ ++int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -1156,49 +7489,153 @@ int ggml_metal_op_get_rows(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- auto pipeline = ggml_metal_library_get_pipeline_get_rows(lib, op->src[0]->type); ++ const int32_t * opts = op->op_params; ++ ggml_op_pool op_pool = (ggml_op_pool) opts[0]; + +- ggml_metal_kargs_get_rows args = { +- /*.ne00t =*/ ggml_is_quantized(op->src[0]->type) ? ne00/16 : ne00, +- /*.ne00 =*/ ne00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne10 =*/ ne10, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, ++ const int32_t k0 = opts[1]; ++ const int32_t k1 = opts[2]; ++ const int32_t s0 = opts[3]; ++ const int32_t s1 = opts[4]; ++ const int32_t p0 = opts[5]; ++ const int32_t p1 = opts[6]; ++ ++ const int64_t IH = op->src[0]->ne[1]; ++ const int64_t IW = op->src[0]->ne[0]; ++ ++ const int64_t N = op->ne[3]; ++ const int64_t OC = op->ne[2]; ++ const int64_t OH = op->ne[1]; ++ const int64_t OW = op->ne[0]; ++ ++ const int64_t np = N * OC * OH * OW; ++ ++ ggml_metal_kargs_pool_2d args_pool_2d = { ++ /* .k0 = */ k0, ++ /* .k1 = */ k1, ++ /* .s0 = */ s0, ++ /* .s1 = */ s1, ++ /* .p0 = */ p0, ++ /* .p1 = */ p1, ++ /* .IH = */ IH, ++ /* .IW = */ IW, ++ /* .OH = */ OH, ++ /* .OW = */ OW, ++ /* .np = */ np + }; + +- const int nth = std::min(args.ne00t, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ auto pipeline = ggml_metal_library_get_pipeline_pool_2d(lib, op, op_pool); + +- const int nw0 = (args.ne00t + nth - 1)/nth; ++ const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), (int) np); ++ const int ntg = (np + nth - 1) / nth; + + ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_bytes (enc, &args_pool_2d, sizeof(args_pool_2d), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, nw0*ne10, ne11, ne12, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_set_rows(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + ++ const char * block_ceiling = getenv("GGML_METAL_EXPERIMENTAL_GLM_DECODE_BLOCK_BYTE_CEILING"); ++ if (block_ceiling != nullptr && atoi(block_ceiling) != 0 && ++ strcmp(ggml_metal_tensor_name(op->src[0]), "glm_decode_block_byte_ceiling_weights") == 0) { ++ const char * kernel = "kernel_glm_decode_block_phase_scan"; ++ auto pipeline = ggml_metal_library_get_pipeline(lib, kernel); ++ if (!pipeline.pipeline) { ++ pipeline = ggml_metal_library_compile_pipeline(lib, kernel, kernel, nullptr); ++ } ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 0); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 1); ++ ++ struct scan_args { ++ uint64_t offset; ++ uint64_t nbytes; ++ }; ++ constexpr uint64_t chunk_bytes = 64*1024; ++ constexpr uint64_t layer_bytes = 238569728; ++ constexpr uint64_t phase_bytes[4] = { ++ 17129472, ++ 35651584, ++ 49580032, ++ 136208640, ++ }; ++ constexpr uint64_t phase_prefix[4] = { ++ 0, ++ 17129472, ++ 52781056, ++ 102361088, ++ }; ++ ++ const int mode = atoi(block_ceiling); ++ if (mode == 2) { ++ scan_args args = { 0, 8*layer_bytes }; ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 2); ++ ggml_metal_encoder_dispatch_threadgroups( ++ enc, (args.nbytes + chunk_bytes - 1)/chunk_bytes, 1, 1, 32, 4, 1); ++ } else { ++ for (uint64_t layer = 0; layer < 8; ++layer) { ++ for (uint64_t phase = 0; phase < 4; ++phase) { ++ scan_args args = { ++ layer*layer_bytes + phase_prefix[phase], ++ phase_bytes[phase], ++ }; ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 2); ++ ggml_metal_encoder_dispatch_threadgroups( ++ enc, (args.nbytes + chunk_bytes - 1)/chunk_bytes, 1, 1, 32, 4, 1); ++ } ++ } ++ } ++ return 1; ++ } ++ ++ const bool shared_expert_weight = ++ ggml_metal_tensor_name_contains(op->src[0], "ffn_gate_shexp") || ++ ggml_metal_tensor_name_contains(op->src[0], "ffn_up_shexp") || ++ ggml_metal_tensor_name_contains(op->src[0], "ffn_down_shexp"); ++ const bool routed_expert_weight = ++ ggml_metal_tensor_name_contains(op->src[0], "ffn_gate_exps") || ++ ggml_metal_tensor_name_contains(op->src[0], "ffn_up_exps") || ++ ggml_metal_tensor_name_contains(op->src[0], "ffn_down_exps"); ++ const bool attention_projection_weight = ++ ggml_metal_tensor_name_contains(op->src[0], "attn_kv_a_mqa") || ++ ggml_metal_tensor_name_contains(op->src[0], "attn_q_a") || ++ ggml_metal_tensor_name_contains(op->src[0], "attn_q_b") || ++ ggml_metal_tensor_name_contains(op->src[0], "attn_k_b") || ++ ggml_metal_tensor_name_contains(op->src[0], "attn_v_b") || ++ ggml_metal_tensor_name_contains(op->src[0], "attn_output"); ++ const bool indexer_projection_weight = ++ ggml_metal_tensor_name_contains(op->src[0], "indexer.attn_q_b") || ++ ggml_metal_tensor_name_contains(op->src[0], "indexer.attn_k") || ++ ggml_metal_tensor_name_contains(op->src[0], "indexer.proj"); ++ const bool zero_projection = ++ (ggml_metal_glm_shared_expert_noop_enabled() && shared_expert_weight) || ++ (ggml_metal_glm_routed_expert_noop_enabled() && routed_expert_weight) || ++ (ggml_metal_glm_attention_projection_noop_enabled() && attention_projection_weight) || ++ (ggml_metal_glm_indexer_projection_noop_enabled() && indexer_projection_weight); ++ if (zero_projection) { ++ auto pipeline = ggml_metal_library_get_pipeline_zero_f32(lib); ++ const int nth = 256; ++ const int64_t n_tg = (ggml_nelements(op) + nth - 1)/nth; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, nth, 1, 1); ++ return 1; ++ } ++ ++ const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +@@ -1206,475 +7643,518 @@ int ggml_metal_op_set_rows(ggml_metal_op_t ctx, int idx) { + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- auto pipeline = ggml_metal_library_get_pipeline_set_rows(lib, op); ++ GGML_ASSERT(ne00 == ne10); + +- const int32_t nk0 = ne0/ggml_blck_size(op->type); ++ GGML_ASSERT(ne12 % ne02 == 0); ++ GGML_ASSERT(ne13 % ne03 == 0); + +- int nth = 32; // SIMD width ++ const int16_t r2 = ne12/ne02; ++ const int16_t r3 = ne13/ne03; + +- while (nth < nk0 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- nth *= 2; +- } ++ // find the break-even point where the matrix-matrix kernel becomes more efficient compared ++ // to the matrix-vector kernel ++ const int ne11_mm_min = 8; + +- int nrptg = 1; +- if (nth > nk0) { +- nrptg = (nth + nk0 - 1)/nk0; +- nth = nk0; ++ // first try to use small-batch mat-mv kernels ++ // these should be efficient for BS [2, ~8] ++ if (op->src[1]->type == GGML_TYPE_F32 && (ne00%128 == 0) && ++ ( ++ ( ++ ( ++ op->src[0]->type == GGML_TYPE_F32 || // TODO: helper function ++ op->src[0]->type == GGML_TYPE_F16 || ++ op->src[0]->type == GGML_TYPE_BF16 || ++ op->src[0]->type == GGML_TYPE_Q1_0 || ++ op->src[0]->type == GGML_TYPE_Q2_0 || ++ op->src[0]->type == GGML_TYPE_Q4_0 || ++ op->src[0]->type == GGML_TYPE_Q4_1 || ++ op->src[0]->type == GGML_TYPE_Q5_0 || ++ op->src[0]->type == GGML_TYPE_Q5_1 || ++ op->src[0]->type == GGML_TYPE_Q8_0 || ++ op->src[0]->type == GGML_TYPE_MXFP4 || ++ op->src[0]->type == GGML_TYPE_IQ4_NL || ++ false) && (ne11 >= 2 && ne11 <= 8) ++ ) || ++ ( ++ ( ++ op->src[0]->type == GGML_TYPE_Q4_K || ++ op->src[0]->type == GGML_TYPE_Q5_K || ++ op->src[0]->type == GGML_TYPE_Q6_K || ++ op->src[0]->type == GGML_TYPE_Q2_K || ++ op->src[0]->type == GGML_TYPE_Q3_K || ++ false) && (ne11 >= 4 && ne11 <= 8) ++ ) ++ ) ++ ) { ++ // TODO: determine the optimal parameters based on grid utilization ++ // I still don't know why we should not always use the maximum available threads: ++ // ++ // nsg = pipeline.maxTotalThreadsPerThreadgroup / 32 ++ // ++ // my current hypothesis is that the work grid is not evenly divisible for different nsg ++ // values and there can be some tail effects when nsg is high. need to confirm this ++ // ++ const int nsg = 2; // num simdgroups per threadgroup + +- if (nrptg*nth > ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- nrptg--; ++ // num threads along row per simdgroup ++ int16_t nxpsg = 0; ++ if (ne00 % 256 == 0 && ne11 < 3) { ++ nxpsg = 16; ++ } else if (ne00 % 128 == 0) { ++ nxpsg = 8; ++ } else { ++ nxpsg = 4; + } +- } +- +- nth = std::min(nth, nk0); +- +- ggml_metal_kargs_set_rows args = { +- /*.nk0 =*/ nk0, +- /*.ne01 =*/ ne01, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- }; +- +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); +- +- ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nrptg - 1)/nrptg, ne02, ne03, nth, nrptg, 1); + +- return 1; +-} +- +-int ggml_metal_op_diag(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ const int16_t nypsg = 32/nxpsg; // num threads along col per simdgroup (i.e. a simdgroup processes that many src0 rows at a time) ++ const int16_t r0ptg = nypsg*nsg; // num src0 rows per threadgroup ++ int16_t r1ptg = 4; // num src1 rows per threadgroup + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ // note: not sure how optimal are those across all different hardware. there might be something cleverer ++ switch (ne11) { ++ case 2: ++ r1ptg = 2; break; ++ case 3: ++ case 6: ++ r1ptg = 3; break; ++ case 4: ++ case 7: ++ case 8: ++ r1ptg = 4; break; ++ case 5: ++ r1ptg = 5; break; ++ default: ++ GGML_ABORT("unsupported ne11"); ++ }; + +- GGML_TENSOR_LOCALS(int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS(int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ auto pipeline = ggml_metal_library_get_pipeline_mul_mv_ext(lib, op, nsg, nxpsg, r1ptg); + +- ggml_metal_kargs_diag args = { +- /*.ne00 =*/ne00, +- /*.ne01 =*/ne01, +- /*.ne02 =*/ne02, +- /*.ne03 =*/ne03, +- /*.nb00 =*/nb00, +- /*.nb01 =*/nb01, +- /*.nb02 =*/nb02, +- /*.nb03 =*/nb03, +- /*.ne0 =*/ne0, +- /*.ne1 =*/ne1, +- /*.ne2 =*/ne2, +- /*.ne3 =*/ne3, +- /*.nb0 =*/nb0, +- /*.nb1 =*/nb1, +- /*.nb2 =*/nb2, +- /*.nb3 =*/nb3, +- }; ++ ggml_metal_kargs_mul_mv_ext args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.r2 =*/ r2, ++ /*.r3 =*/ r3, ++ }; + +- auto pipeline = ggml_metal_library_get_pipeline_diag(lib, op); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + r0ptg - 1)/r0ptg), ((ne11 + r1ptg - 1)/r1ptg), ne12*ne13, 32, nsg, 1); ++ } else if ( ++ !ggml_is_transposed(op->src[0]) && ++ !ggml_is_transposed(op->src[1]) && ++ // for now the matrix-matrix multiplication kernel only works on A14+/M1+ SoCs ++ // AMD GPU and older A-chips will reuse matrix-vector multiplication kernel ++ props_dev->has_simdgroup_mm && ne00 >= 64 && ne11 > ne11_mm_min) { ++ //GGML_LOG_INFO("matrix: ne00 = %6d, ne01 = %6d, ne02 = %6d, ne11 = %6d, ne12 = %6d\n", ne00, ne01, ne02, ne11, ne12); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, 32, 1, 1); ++ // some Metal matrix data types require aligned pointers ++ // ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf (Table 2.5) ++ //switch (op->src[0]->type) { ++ // case GGML_TYPE_F32: GGML_ASSERT(nb01 % 16 == 0); break; ++ // case GGML_TYPE_F16: GGML_ASSERT(nb01 % 8 == 0); break; ++ // case GGML_TYPE_BF16: GGML_ASSERT(nb01 % 8 == 0); break; ++ // default: break; ++ //} + +- return 1; +-} ++ auto pipeline = ggml_metal_library_get_pipeline_mul_mm(lib, op); + +-int ggml_metal_op_soft_max(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ ggml_metal_kargs_mul_mm args = { ++ /*.ne00 =*/ ne00, ++ /*.ne02 =*/ ne02, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne12 =*/ ne12, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.r2 =*/ r2, ++ /*.r3 =*/ r3, ++ }; + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ const size_t smem = pipeline.smem; + +- float scale; +- float max_bias; ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- memcpy(&scale, ((const int32_t *) op->op_params) + 0, sizeof(scale)); +- memcpy(&max_bias, ((const int32_t *) op->op_params) + 1, sizeof(max_bias)); ++ const int nr0 = pipeline.nr0; ++ const int nr1 = pipeline.nr1; ++ const int nsg = pipeline.nsg; + +- const uint32_t n_head = op->src[0]->ne[2]; +- const int32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head)); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ((ne11 + nr1 - 1) / nr1), ((ne01 + nr0 - 1) / nr0), ne12 * ne13, 32, nsg, 1); ++ } else { ++ auto pipeline = ggml_metal_library_get_pipeline_mul_mv(lib, op); + +- const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); +- const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); ++ const int nr0 = pipeline.nr0; ++ const int nr1 = pipeline.nr1; ++ const int nsg = pipeline.nsg; + +- // softmax ++ const size_t smem = pipeline.smem; + +- ggml_metal_kargs_soft_max args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.ne13 =*/ ne13, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.scale =*/ scale, +- /*.max_bias =*/ max_bias, +- /*.m0 =*/ m0, +- /*.m1 =*/ m1, +- /*.n_head_log2 =*/ n_head_log2, +- }; ++ ggml_metal_kargs_mul_mv args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nr0 =*/ nr0, ++ /*.r2 =*/ r2, ++ /*.r3 =*/ r3, ++ }; + +- auto pipeline = ggml_metal_library_get_pipeline_soft_max(lib, op); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- int nth = 32; // SIMD width ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- if (ne00%4 == 0) { +- while (nth < ne00/4 && nth*ne01*ne02*ne03 < 256) { +- nth *= 2; +- } +- } else { +- while (nth < ne00 && nth*ne01*ne02*ne03 < 256) { +- nth *= 2; ++ if (op->src[0]->type == GGML_TYPE_F32 || ++ op->src[0]->type == GGML_TYPE_F16 || ++ op->src[0]->type == GGML_TYPE_BF16 || ++ op->src[0]->type == GGML_TYPE_Q8_0) { ++ ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + nr0 - 1)/(nr0)), ((ne11 + nr1 - 1)/nr1), ne12*ne13, 32, nsg, 1); ++ } else { ++ ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + nr0*nsg - 1)/(nr0*nsg)), ((ne11 + nr1 - 1)/nr1), ne12*ne13, 32, nsg, 1); + } + } + +- const size_t smem = pipeline.smem; +- +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- if (op->src[1]) { +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- } else { +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 2); +- } +- if (op->src[2]) { +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[2]), 3); +- } else { +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 3); +- } +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 4); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); +- +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); +- + return 1; + } + +-int ggml_metal_op_ssm_conv(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); +- +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; +- +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); +- +- ggml_metal_kargs_ssm_conv args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.ne10 =*/ ne10, +- /*.ne11 =*/ ne11, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- }; ++size_t ggml_metal_op_mul_mat_id_extra_tpe(const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MUL_MAT_ID); + +- // Use batched kernel for prefill (ne1 > 1) to reduce threadgroup dispatch overhead +- const bool use_batched = (ne1 > 1); ++ const int64_t ne02 = op->src[0]->ne[2]; // n_expert + +- if (use_batched) { +- // Determine the smallest power of 2 that's >= ne1, but <= 256 +- int BATCH_SIZE; +- if (ne1 > 128) BATCH_SIZE = 256; +- else if (ne1 > 64 ) BATCH_SIZE = 128; +- else if (ne1 > 32 ) BATCH_SIZE = 64; +- else if (ne1 > 16 ) BATCH_SIZE = 32; +- else if (ne1 > 8 ) BATCH_SIZE = 16; +- else if (ne1 > 4 ) BATCH_SIZE = 8; +- else BATCH_SIZE = 2; ++ return ggml_type_size(GGML_TYPE_I32)*ne02; ++} + +- auto pipeline = ggml_metal_library_get_pipeline_ssm_conv_batched(lib, op, BATCH_SIZE); ++size_t ggml_metal_op_mul_mat_id_extra_ids(const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MUL_MAT_ID); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3); ++ const int64_t ne02 = op->src[0]->ne[2]; // n_expert ++ const int64_t ne21 = op->src[2]->ne[1]; // n_token + +- // Dispatch: ne01 rows, ceil(ne1/BATCH_SIZE) token batches, ne02 sequences +- // Each threadgroup has BATCH_SIZE threads, each handling one token +- const int n_token_batches = (ne1 + BATCH_SIZE - 1) / BATCH_SIZE; +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, n_token_batches, ne02, BATCH_SIZE, 1, 1); +- } else { +- auto pipeline = ggml_metal_library_get_pipeline_ssm_conv(lib, op); ++ return ggml_type_size(GGML_TYPE_I32)*ne02*ne21; ++} + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3); ++size_t ggml_metal_op_mul_mat_id_extra_src1_scratch(const ggml_tensor * op) { ++ assert(op->op == GGML_OP_MUL_MAT_ID); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne1, ne02, 1, 1, 1); ++ if (!ggml_metal_tensor_name_contains(op, "ffn_moe_down") || ++ (op->src[0]->type != GGML_TYPE_Q2_K && op->src[0]->type != GGML_TYPE_Q3_K) || ++ (op->src[1]->type != GGML_TYPE_F32 && op->src[1]->type != GGML_TYPE_F16) || ++ op->src[2]->ne[0] != 8) { ++ return 0; + } + +- return 1; ++ return ggml_nbytes(op->src[1]); + } + +-int ggml_metal_op_ssm_scan(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + ++ const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); + GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); +- GGML_TENSOR_LOCALS( int32_t, ne4, op->src[4], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb4, op->src[4], nb); +- GGML_TENSOR_LOCALS( int32_t, ne5, op->src[5], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb5, op->src[5], nb); +- GGML_TENSOR_LOCALS( int32_t, ne6, op->src[6], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb6, op->src[6], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const ggml_tensor * src3 = op->src[3]; +- const ggml_tensor * src4 = op->src[4]; +- const ggml_tensor * src5 = op->src[5]; +- const ggml_tensor * src6 = op->src[6]; ++ // src2 = ids ++ GGML_ASSERT(op->src[2]->type == GGML_TYPE_I32); + +- GGML_ASSERT(src3); +- GGML_ASSERT(src4); +- GGML_ASSERT(src5); +- GGML_ASSERT(src6); ++ GGML_ASSERT(!ggml_is_transposed(op->src[0])); ++ GGML_ASSERT(!ggml_is_transposed(op->src[1])); + +- const int64_t d_state = ne00; +- const int64_t d_inner = ne01; +- const int64_t n_head = ne02; +- const int64_t n_group = ne41; +- const int64_t n_seq_tokens = ne12; +- const int64_t n_seqs = ne13; ++ GGML_ASSERT(ne03 == 1); ++ GGML_ASSERT(ne13 == 1); + +- ggml_metal_kargs_ssm_scan args = { +- /*.d_state =*/ d_state, +- /*.d_inner =*/ d_inner, +- /*.n_head =*/ n_head, +- /*.n_group =*/ n_group, +- /*.n_seq_tokens =*/ n_seq_tokens, +- /*.n_seqs =*/ n_seqs, +- /*.s_off =*/ ggml_nelements(op->src[1]) * sizeof(float), +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.ns12 =*/ nb12/nb10, +- /*.nb13 =*/ nb13, +- /*.nb20 =*/ nb20, +- /*.nb21 =*/ nb21, +- /*.ns21 =*/ nb21/nb20, +- /*.nb22 =*/ nb22, +- /*.ne30 =*/ ne30, +- /*.nb31 =*/ nb31, +- /*.nb41 =*/ nb41, +- /*.nb42 =*/ nb42, +- /*.ns42 =*/ nb42/nb40, +- /*.nb43 =*/ nb43, +- /*.nb51 =*/ nb51, +- /*.nb52 =*/ nb52, +- /*.ns52 =*/ nb52/nb50, +- /*.nb53 =*/ nb53, +- /*.nb0 =*/ nb0, +- }; ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); ++ ggml_metal_buffer_id bid_src2 = ggml_metal_get_buffer_id(op->src[2]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); + +- auto pipeline = ggml_metal_library_get_pipeline_ssm_scan(lib, op); ++ const uint32_t r2 = 1; ++ const uint32_t r3 = 1; + +- GGML_ASSERT(d_state <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ // find the break-even point where the matrix-matrix kernel becomes more efficient compared ++ // to the matrix-vector kernel ++ // ne20 = n_used_experts ++ // ne21 = n_rows (batch size) ++ const int ne21_mm_id_min = ggml_metal_glm_dsa_mul_mm_id_min_tokens_requested(); + +- const size_t smem = pipeline.smem; ++ if (props_dev->has_simdgroup_mm && ne00 >= 64 && (ne21 >= ne21_mm_id_min)) { ++ // some Metal matrix data types require aligned pointers ++ // ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf (Table 2.5) ++ //switch (op->src[0]->type) { ++ // case GGML_TYPE_F32: GGML_ASSERT(nb01 % 16 == 0); break; ++ // case GGML_TYPE_F16: GGML_ASSERT(nb01 % 8 == 0); break; ++ // case GGML_TYPE_BF16: GGML_ASSERT(nb01 % 8 == 0); break; ++ // default: break; ++ //} + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), 4); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), 5); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), 6); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), 7); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 8); ++ // extra buffers for intermediate id mapping ++ ggml_metal_buffer_id bid_tpe = bid_dst; ++ bid_tpe.offs += ggml_nbytes(op); + +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_buffer_id bid_ids = bid_tpe; ++ bid_ids.offs += ggml_metal_op_mul_mat_id_extra_tpe(op); + +- ggml_metal_encoder_dispatch_threadgroups(enc, d_inner, n_head, n_seqs, d_state, 1, 1); ++ { ++ ggml_metal_kargs_mul_mm_id_map0 args = { ++ ne02, ++ ne10, ++ ne11, // n_expert_used (bcast) ++ nb11, ++ nb12, ++ ne21, // n_tokens ++ ne20, // n_expert_used ++ nb21, ++ }; + +- return 1; +-} ++ auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_map0(lib, ne02, ne20); + +-int ggml_metal_op_rwkv(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ const size_t smem = pipeline.smem; + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ GGML_ASSERT(ne02 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size); + +- const int64_t B = op->op == GGML_OP_RWKV_WKV6 ? op->src[5]->ne[1] : op->src[6]->ne[1]; +- const int64_t T = op->src[0]->ne[2]; +- const int64_t C = op->ne[0]; +- const int64_t H = op->src[0]->ne[1]; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_tpe, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_ids, 3); + +- auto pipeline = ggml_metal_library_get_pipeline_rwkv(lib, op); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- int ida = 0; ++ ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, ne02, 1, 1); ++ } + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); +- if (op->op == GGML_OP_RWKV_WKV7) { +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[6]), ida++); +- } +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); +- ggml_metal_encoder_set_bytes (enc, (void *) &B, sizeof(B), ida++); +- ggml_metal_encoder_set_bytes (enc, (void *) &T, sizeof(T), ida++); +- ggml_metal_encoder_set_bytes (enc, (void *) &C, sizeof(C), ida++); +- ggml_metal_encoder_set_bytes (enc, (void *) &H, sizeof(H), ida++); ++ // this barrier is always needed because the next kernel has to wait for the id maps to be computed ++ ggml_metal_op_concurrency_reset(ctx); + +- ggml_metal_encoder_dispatch_threadgroups(enc, B * H, 1, 1, C/H, 1, 1); ++ { ++ auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id(lib, op); + +- return 1; +-} ++ ggml_metal_kargs_mul_mm_id args = { ++ /*.ne00 =*/ ne00, ++ /*.ne02 =*/ ne02, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne11 =*/ ne11, // n_expert_used (bcast) ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne20 =*/ ne20, // n_expert_used ++ /*.ne21 =*/ ne21, // n_tokens ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.r2 =*/ r2, ++ /*.r3 =*/ r3, ++ }; + +-int ggml_metal_op_gated_delta_net(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_tpe, 3); ++ ggml_metal_encoder_set_buffer (enc, bid_ids, 4); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 5); + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ const size_t smem = pipeline.smem; + ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ const int grid_x = (ne21 + 31)/32; ++ const int grid_y = (ne01 + 63)/64; ++ const int grid_z = ne02; ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 128, 1, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=mul_mat_id kernel=mul_mm_id tensor=%s src0_type=%s src1_type=%s ids_type=%s dst_type=%s ne00=%d ne01=%d experts=%d used_experts=%d tokens=%d min_tokens=%d grid_x=%d grid_y=%d grid_z=%d threads_x=%d threads_y=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->type), ++ ne00, ++ ne01, ++ ne02, ++ ne20, ++ ne21, ++ ne21_mm_id_min, ++ grid_x, ++ grid_y, ++ grid_z, ++ 128, ++ 1); ++ } ++ } ++ } else { ++ auto pipeline = ggml_metal_library_get_pipeline_mul_mv_id(lib, op); + +- auto pipeline = ggml_metal_library_get_pipeline_gated_delta_net(lib, op); ++ const int nr0 = pipeline.nr0; ++ const int nr1 = pipeline.nr1; ++ const int nsg = pipeline.nsg; + +- int ida = 0; ++ const size_t smem = pipeline.smem; + +- ggml_metal_kargs_gated_delta_net args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne10 =*/ ne10, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.ne13 =*/ ne13, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ne20 =*/ ne20, +- /*.ne21 =*/ ne21, +- /*.ne22 =*/ ne22, +- /*.ne23 =*/ ne23, +- /*.nb20 =*/ nb20, +- /*.nb21 =*/ nb21, +- /*.nb22 =*/ nb22, +- /*.nb23 =*/ nb23, +- /*.ns02 =*/ (int32_t) (nb02/sizeof(float)), +- /*.ns12 =*/ (int32_t) (nb12/sizeof(float)), +- /*.ns22 =*/ (int32_t) (nb22/sizeof(float)), +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- }; ++ ggml_metal_kargs_mul_mv_id args = { ++ /*.nei0 =*/ ne20, ++ /*.nei1 =*/ ne21, ++ /*.nbi1 =*/ nb21, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nb1 =*/ nb1, ++ /*.nr0 =*/ nr0, ++ }; + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // k +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // v +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // gate +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); // beta +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[5]), ida++); // state +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst ++ if (ggml_is_quantized(op->src[0]->type)) { ++ GGML_ASSERT(ne00 >= nsg*nr0); ++ } + +- const int nsg = pipeline.nsg; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer(enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer(enc, bid_src1, 2); ++ ggml_metal_encoder_set_buffer(enc, bid_dst, 3); ++ ggml_metal_encoder_set_buffer(enc, bid_src2, 4); + +- ggml_metal_encoder_dispatch_threadgroups(enc, op->src[2]->ne[0]/nsg, op->src[2]->ne[1], op->src[2]->ne[3], 32, nsg, 1); ++ const int64_t _ne1 = 1; ++ const int64_t ne123 = ne20*ne21; ++ ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ++ int grid_x = 0; ++ int grid_y = 0; ++ const int grid_z = ne123; ++ if (op->src[0]->type == GGML_TYPE_F32 || ++ op->src[0]->type == GGML_TYPE_F16 || ++ op->src[0]->type == GGML_TYPE_BF16 || ++ op->src[0]->type == GGML_TYPE_Q8_0) { ++ grid_x = (ne01 + nr0 - 1)/(nr0); ++ grid_y = (_ne1 + nr1 - 1)/nr1; ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, nsg, 1); ++ } else { ++ grid_x = (ne01 + nr0*nsg - 1)/(nr0*nsg); ++ grid_y = (_ne1 + nr1 - 1)/nr1; ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, nsg, 1); ++ } ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ const bool weighted_slots = ggml_get_op_params_i32(op, 0) != 0; ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=mul_mat_id kernel=mul_mv_id tensor=%s src0_type=%s src1_type=%s ids_type=%s dst_type=%s ne00=%d ne01=%d experts=%d used_experts=%d tokens=%d min_tokens=%d nr0=%d nr1=%d nsg=%d grid_x=%d grid_y=%d grid_z=%d threads_x=%d threads_y=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->type), ++ ne00, ++ ne01, ++ ne02, ++ ne20, ++ ne21, ++ ne21_mm_id_min, ++ nr0, ++ nr1, ++ nsg, ++ grid_x, ++ grid_y, ++ grid_z, ++ 32, ++ nsg); ++ if (weighted_slots) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=mul_mv_id_weighted_slots kernel=%s tensor=%s src0_type=%s src1_type=%s ids_type=%s dst_type=%s ne00=%d ne01=%d experts=%d used_experts=%d tokens=%d grid_x=%d grid_y=%d grid_z=%d threads_x=%d threads_y=%d\n", ++ ggml_type_name(op->src[0]->type), ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->type), ++ ne00, ++ ne01, ++ ne02, ++ ne20, ++ ne21, ++ grid_x, ++ grid_y, ++ grid_z, ++ 32, ++ nsg); ++ } ++ } ++ } + + return 1; + } + +-int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_add_id(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -1684,1610 +8164,2148 @@ int ggml_metal_op_solve_tri(ggml_metal_op_t ctx, int idx) { + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- ggml_metal_kargs_solve_tri args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[2]->type == GGML_TYPE_I32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ ++ ggml_metal_kargs_add_id args = { ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, + /*.nb01 =*/ nb01, + /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne10 =*/ ne10, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.ne13 =*/ ne13, +- /*.nb10 =*/ nb10, + /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, ++ /*.nb21 =*/ nb21, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_solve_tri(lib, op); ++ auto pipeline = ggml_metal_library_get_pipeline_base(lib, GGML_OP_ADD_ID); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); +- +- const int nsg = pipeline.nsg; ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 4); + +- ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); ++ const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00); + +- ggml_metal_encoder_dispatch_threadgroups(enc, (ne10 + nsg - 1)/nsg, ne02, ne03, 32, nsg, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, 1, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_set(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++bool ggml_metal_op_flash_attn_ext_use_vec(const ggml_tensor * op) { ++ assert(op->op == GGML_OP_FLASH_ATTN_EXT); + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ const int64_t ne00 = op->src[0]->ne[0]; // head size ++ const int64_t ne01 = op->src[0]->ne[1]; // batch size + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ // use vec kernel if the batch size is small and if the head size is supported ++ return (ne01 < 20) && (ne00 % 32 == 0); ++} + +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++static int32_t ggml_metal_glm_dsa_compact_flash_nwg_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_COMPACT_FLASH_NWG"); ++ if (value == nullptr || value[0] == '\0') { ++ return 4; ++ } + +- const size_t pnb1 = ((const int32_t *) op->op_params)[0]; +- const size_t pnb2 = ((const int32_t *) op->op_params)[1]; +- const size_t pnb3 = ((const int32_t *) op->op_params)[2]; +- const size_t offs = ((const int32_t *) op->op_params)[3]; ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ case 8: ++ case 16: ++ case 32: ++ return requested; ++ default: ++ return 4; ++ } ++} + +- const bool inplace = (bool) ((const int32_t *) op->op_params)[4]; ++static int32_t ggml_metal_glm_dsa_selected_row_flash_nwg_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_SELECTED_ROW_FLASH_NWG"); ++ if (value == nullptr || value[0] == '\0') { ++ return 16; ++ } + +- if (!inplace) { +- // run a separate kernel to cpy src->dst +- // not sure how to avoid this +- // TODO: make a simpler cpy_bytes kernel ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ case 2: ++ case 4: ++ case 8: ++ case 16: ++ case 32: ++ return requested; ++ default: ++ return 16; ++ } ++} + +- //const id pipeline = ctx->pipelines[GGML_METAL_PIPELINE_TYPE_CPY_F32_F32].obj; +- auto pipeline = ggml_metal_library_get_pipeline_cpy(lib, op->src[0]->type, op->type); ++static int32_t ggml_metal_glm_dsa_selected_row_flash_heads_per_tg_requested() { ++ const char * value = getenv("LLAMA_GLM_DSA_SELECTED_ROW_FLASH_HEADS_PER_TG"); ++ if (value == nullptr || value[0] == '\0') { ++ return 1; ++ } + +- ggml_metal_kargs_cpy args = { +- /*.nk0 =*/ ne00, +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- }; ++ const int requested = atoi(value); ++ switch (requested) { ++ case 1: ++ return requested; ++ default: ++ return 1; ++ } ++} + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++static bool ggml_metal_glm_dsa_selected_row_flash_tiled_enabled() { ++ const char * value = getenv("LLAMA_GLM_DSA_EXPERIMENTAL_SELECTED_ROW_FLASH_TILED"); ++ return value == nullptr || value[0] == '\0' || atoi(value) != 0; ++} + +- const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00); ++static bool ggml_metal_glm_dsa_selected_row_flash_tiled_selected(const ggml_tensor * op) { ++ if (!ggml_metal_glm_dsa_selected_row_flash_tiled_enabled()) { ++ return false; ++ } + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); ++ const char * tensor = getenv("LLAMA_GLM_DSA_SELECTED_ROW_FLASH_TILED_TENSOR"); ++ return tensor == nullptr || tensor[0] == '\0' || strcmp(tensor, ggml_metal_tensor_name(op)) == 0; ++} + +- ggml_metal_op_concurrency_reset(ctx); ++static bool ggml_metal_glm_dsa_selected_row_flash_tiled_shape(const ggml_tensor * op) { ++ if (!ggml_metal_selected_row_flash_vec_shape_ok(op)) { ++ return false; + } + +- auto pipeline = ggml_metal_library_get_pipeline_cpy(lib, op->src[1]->type, op->type); ++ const ggml_tensor * q = op->src[0]; ++ const ggml_tensor * rows = op->src[1]; ++ const ggml_tensor * view = op->src[2]; ++ return q->ne[1] == 1 && q->ne[2] >= 8 && q->ne[2] % 8 == 0 && ++ rows->src[0]->ne[2] == 1 && view->ne[2] == 1; ++} + +- GGML_ASSERT(ne10 % ggml_blck_size(op->src[1]->type) == 0); ++static int32_t ggml_metal_glm_dsa_selected_row_flash_tiled_nwg(const ggml_tensor * op) { ++ const char * value = getenv("LLAMA_GLM_DSA_SELECTED_ROW_FLASH_NWG"); ++ if (value != nullptr && value[0] != '\0') { ++ return ggml_metal_glm_dsa_selected_row_flash_nwg_requested(); ++ } + +- int64_t nk0 = ne10; +- if (ggml_is_quantized(op->src[1]->type)) { +- nk0 = ne10/16; +- } else if (ggml_is_quantized(op->type)) { +- nk0 = ne10/ggml_blck_size(op->type); ++ const int64_t tiles = (op->src[1]->src[1]->ne[0] + 63)/64; ++ if (tiles >= 32) { ++ return 32; ++ } ++ if (tiles >= 16) { ++ return 16; ++ } ++ if (tiles >= 8) { ++ return 8; ++ } ++ if (tiles >= 4) { ++ return 4; ++ } ++ if (tiles >= 2) { ++ return 2; ++ } ++ return 1; ++} ++ ++static int32_t ggml_metal_glm_dsa_selected_row_flash_tiled_nwg_for_top_k(int64_t top_k) { ++ const char * value = getenv("LLAMA_GLM_DSA_SELECTED_ROW_FLASH_NWG"); ++ if (value != nullptr && value[0] != '\0') { ++ return ggml_metal_glm_dsa_selected_row_flash_nwg_requested(); ++ } ++ ++ const int64_t tiles = (top_k + 63)/64; ++ if (tiles >= 32) { ++ return 32; ++ } ++ if (tiles >= 16) { ++ return 16; ++ } ++ if (tiles >= 8) { ++ return 8; ++ } ++ if (tiles >= 4) { ++ return 4; ++ } ++ if (tiles >= 2) { ++ return 2; ++ } ++ return 1; ++} ++ ++static bool ggml_metal_glm_dsa_compact_flash_shape( ++ const ggml_tensor * op, ++ bool has_mask, ++ bool has_sinks, ++ bool has_bias, ++ bool has_scap, ++ bool has_kvpad) { ++ return !has_mask && ++ !has_sinks && ++ !has_bias && ++ !has_scap && ++ !has_kvpad && ++ op->src[1]->type == GGML_TYPE_F16 && ++ op->src[2]->type == GGML_TYPE_F16 && ++ op->src[0]->ne[0] == 576 && ++ op->src[2]->ne[0] == 512 && ++ op->src[0]->ne[1] == 1 && ++ op->src[1]->ne[1] == 2048; ++} ++ ++static int32_t ggml_metal_flash_attn_ext_vec_nwg( ++ const ggml_tensor * op, ++ bool has_mask, ++ bool has_sinks, ++ bool has_bias, ++ bool has_scap, ++ bool has_kvpad) { ++ if (ggml_metal_selected_row_flash_vec_shape_ok(op)) { ++ if (ggml_metal_glm_dsa_selected_row_flash_tiled_selected(op) && ++ ggml_metal_glm_dsa_selected_row_flash_tiled_shape(op)) { ++ return ggml_metal_glm_dsa_selected_row_flash_tiled_nwg(op); ++ } ++ return ggml_metal_glm_dsa_selected_row_flash_nwg_requested(); ++ } ++ ++ if (ggml_metal_glm_dsa_compact_flash_shape(op, has_mask, has_sinks, has_bias, has_scap, has_kvpad)) { ++ const char * nwg8_tensor = getenv("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_NWG8_TENSOR"); ++ if (nwg8_tensor != nullptr && nwg8_tensor[0] != '\0' && ++ strcmp(nwg8_tensor, ggml_metal_tensor_name(op)) == 0) { ++ return 8; ++ } ++ return ggml_metal_glm_dsa_compact_flash_nwg_requested(); ++ } ++ ++ return 32; ++} ++ ++size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) { ++ assert(op->op == GGML_OP_FLASH_ATTN_EXT); ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); ++ ++ size_t res = 0; ++ ++ const bool has_mask = op->src[3] != nullptr; ++ ++ // note: the non-vec kernel requires more extra memory, so always reserve for it ++ GGML_ASSERT(OP_FLASH_ATTN_EXT_NCPSG >= OP_FLASH_ATTN_EXT_VEC_NCPSG); ++ ++ //if (ggml_metal_op_flash_attn_ext_use_vec(op)) { ++ if (false) { ++ // note: always reserve the padding space to avoid graph reallocations ++ //const bool has_kvpad = ne11 % OP_FLASH_ATTN_EXT_VEC_NCPSG != 0; ++ const bool has_kvpad = true; ++ ++ if (has_kvpad) { ++ res += OP_FLASH_ATTN_EXT_VEC_NCPSG*( ++ nb11*ne12*ne13 + ++ nb21*ne22*ne23 + ++ (has_mask ? ggml_type_size(GGML_TYPE_F16)*ne31*ne32*ne33 : 0)); ++ } ++ } else { ++ //const bool has_kvpad = ne11 % OP_FLASH_ATTN_EXT_NCPSG != 0; ++ const bool has_kvpad = true; ++ ++ if (has_kvpad) { ++ res += OP_FLASH_ATTN_EXT_NCPSG*( ++ nb11*ne12*ne13 + ++ nb21*ne22*ne23 + ++ (has_mask ? ggml_type_size(GGML_TYPE_F16)*ne31*ne32*ne33 : 0)); ++ } + } + +- int nth = std::min(nk0*ne11, 256); ++ return res; ++} ++ ++size_t ggml_metal_op_flash_attn_ext_extra_blk(const ggml_tensor * op) { ++ assert(op->op == GGML_OP_FLASH_ATTN_EXT); ++ ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ //GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ //GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ //GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ //GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ //GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); + +- // when rows are small, we can batch them together in a single threadgroup +- int nrptg = 1; ++ size_t res = 0; + +- // TODO: relax this constraint in the future +- if (ggml_blck_size(op->src[1]->type) == 1 && ggml_blck_size(op->type) == 1) { +- if (nth > nk0) { +- nrptg = (nth + nk0 - 1)/nk0; +- nth = nk0; ++ const bool has_mask = op->src[3] != nullptr; + +- if (nrptg*nth > 256) { +- nrptg--; +- } +- } ++ if (!has_mask) { ++ return res; + } + +- nth = std::min(nth, nk0); +- +- ggml_metal_kargs_cpy args = { +- /*.nk0 =*/ nk0, +- /*.ne00 =*/ ne10, +- /*.ne01 =*/ ne11, +- /*.ne02 =*/ ne12, +- /*.ne03 =*/ ne13, +- /*.nb00 =*/ nb10, +- /*.nb01 =*/ nb11, +- /*.nb02 =*/ nb12, +- /*.nb03 =*/ nb13, +- /*.ne0 =*/ ne10, +- /*.ne1 =*/ ne11, +- /*.ne2 =*/ ne12, +- /*.ne3 =*/ ne13, +- /*.nb0 =*/ ggml_element_size(op), +- /*.nb1 =*/ pnb1, +- /*.nb2 =*/ pnb2, +- /*.nb3 =*/ pnb3, +- }; ++ const bool is_vec = ggml_metal_op_flash_attn_ext_use_vec(op); + +- const int nw0 = nrptg == 1 ? (nk0 + nth - 1)/nth : 1; ++ // this optimization is not useful for the vector kernels ++ // note: always reserve the blk buffer to avoid graph reallocations ++ //if (is_vec) { ++ // return res; ++ //} + +- bid_dst.offs += offs; ++ const int nqptg = is_vec ? OP_FLASH_ATTN_EXT_VEC_NQPSG : OP_FLASH_ATTN_EXT_NQPSG; ++ const int ncpsg = is_vec ? OP_FLASH_ATTN_EXT_VEC_NCPSG : OP_FLASH_ATTN_EXT_NCPSG; + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src1, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ const int64_t ne1 = (ne01 + nqptg - 1)/nqptg; ++ const int64_t ne0 = (ne30 + ncpsg - 1)/ncpsg; + +- ggml_metal_encoder_dispatch_threadgroups(enc, nw0*(ne11 + nrptg - 1)/nrptg, ne12, ne13, nth, nrptg, 1); ++ res += GGML_PAD(ggml_type_size(GGML_TYPE_I8)*ne0*ne1*ne32*ne33, 32); + +- return 1; ++ return res; + } + +-int ggml_metal_op_cpy(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); +- +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++size_t ggml_metal_op_flash_attn_ext_extra_tmp(const ggml_tensor * op) { ++ assert(op->op == GGML_OP_FLASH_ATTN_EXT); + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); +- +- auto pipeline = ggml_metal_library_get_pipeline_cpy(lib, op->src[0]->type, op->type); +- +- GGML_ASSERT(ne00 % ggml_blck_size(op->src[0]->type) == 0); ++ //GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ //GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ //GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); ++ //GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); + +- int64_t nk0 = ne00; +- if (ggml_is_quantized(op->src[0]->type)) { +- nk0 = ne00/16; +- } else if (ggml_is_quantized(op->type)) { +- nk0 = ne00/ggml_blck_size(op->type); +- } ++ size_t res = 0; + +- int nth = std::min(nk0*ne01, 256); ++ // note: always reserve the temp buffer to avoid graph reallocations ++ //if (ggml_metal_op_flash_attn_ext_use_vec(op)) { ++ if (true) { ++ float max_bias; ++ float logit_softcap; ++ ++ memcpy(&max_bias, ((const int32_t *) op->op_params) + 1, sizeof(max_bias)); ++ memcpy(&logit_softcap, ((const int32_t *) op->op_params) + 2, sizeof(logit_softcap)); ++ ++ const bool has_mask = op->src[3] != nullptr; ++ const bool has_sinks = op->src[4] != nullptr; ++ const bool has_bias = max_bias != 0.0f; ++ const bool has_scap = logit_softcap != 0.0f; ++ const bool has_kvpad = op->src[1]->ne[1] % OP_FLASH_ATTN_EXT_VEC_NCPSG != 0; ++ const int64_t nwg = ggml_metal_flash_attn_ext_vec_nwg( ++ op, has_mask, has_sinks, has_bias, has_scap, has_kvpad); ++ if (nwg == 1) { ++ return res; ++ } + +- // when rows are small, we can batch them together in a single threadgroup +- int nrptg = 1; ++ const int64_t ne01_max = std::min(ne01, 32); + +- // TODO: relax this constraint in the future +- if (ggml_blck_size(op->src[0]->type) == 1 && ggml_blck_size(op->type) == 1) { +- if (nth > nk0) { +- nrptg = (nth + nk0 - 1)/nk0; +- nth = nk0; ++ // temp buffer for writing the results from each workgroup ++ // - ne20: the size of the Value head ++ // - + 2: the S and M values for each intermediate result ++ res += ggml_type_size(GGML_TYPE_F32)*(ne01_max*ne02*ne03*nwg*(ne20 + 2)); + +- if (nrptg*nth > 256) { +- nrptg--; +- } ++ if (ggml_metal_glm_compact_split_exact_enabled() && ++ !has_mask && !has_sinks && !has_bias && !has_scap && !has_kvpad && ++ op->src[1]->type == GGML_TYPE_F16 && ++ op->src[2]->type == GGML_TYPE_F16 && ++ op->src[0]->ne[0] == 576 && ++ op->src[0]->ne[1] == 1 && ++ op->src[0]->ne[2] == 64 && ++ op->src[0]->ne[3] == 1 && ++ op->src[1]->ne[1] == 2048 && ++ op->src[1]->ne[2] == 1 && ++ op->src[2]->ne[0] == 512 && ++ op->src[2]->ne[2] == 1) { ++ constexpr size_t chunk_rows = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ const size_t heads = size_t(op->src[0]->ne[2])*size_t(op->src[0]->ne[3]); ++ const size_t rows = size_t(op->src[1]->ne[1]); ++ const size_t chunks = (rows + chunk_rows - 1)/chunk_rows; ++ const size_t score_values = heads*rows; ++ const size_t chunk_ms_values = heads*chunks; ++ const size_t chunk_v_values = heads*chunks*size_t(op->src[2]->ne[0]); ++ res += ggml_type_size(GGML_TYPE_F32)* ++ (score_values + chunk_ms_values + chunk_v_values); + } + } + +- nth = std::min(nth, nk0); +- +- ggml_metal_kargs_cpy args = { +- /*.nk0 =*/ nk0, +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- }; +- +- const int nw0 = nrptg == 1 ? (nk0 + nth - 1)/nth : 1; +- +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ return res; ++} + +- ggml_metal_encoder_dispatch_threadgroups(enc, nw0*(ne01 + nrptg - 1)/nrptg, ne02, ne03, nth, nrptg, 1); ++size_t ggml_metal_op_dsa_sparse_attn_extra_tmp(const ggml_tensor * op) { ++ assert(op->op == GGML_OP_DSA_SPARSE_ATTN); ++ ++ if (op->src[0] == nullptr || op->src[1] == nullptr || op->src[2] == nullptr || op->src[4] == nullptr || ++ op->type != GGML_TYPE_F32 || ++ op->src[0]->type != GGML_TYPE_F32 || ++ op->src[1]->type != GGML_TYPE_F16 || ++ op->src[2]->type != GGML_TYPE_F16 || ++ op->src[4]->type != GGML_TYPE_I32 || ++ op->src[0]->ne[0] != 576 || op->ne[0] != 512 || ++ op->src[4]->ne[0] <= 0 || op->src[4]->ne[0] > 4096) { ++ return 0; ++ } + +- return 1; ++ constexpr size_t max_nwg = 32; ++ const size_t nrows = size_t(op->ne[1])*size_t(op->ne[2])*size_t(op->ne[3]); ++ return sizeof(float)*nrows*max_nwg*size_t(op->ne[0] + 2); + } + +-int ggml_metal_op_pool_1d(ggml_metal_op_t ctx, int idx) { ++static int ggml_metal_op_selected_row_flash_vec(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); ++ if (!ggml_metal_selected_row_flash_vec_shape_ok(op)) { ++ return 0; ++ } ++ ++ ggml_tensor * q = op->src[0]; ++ ggml_tensor * rows = op->src[1]; ++ ggml_tensor * view = op->src[2]; ++ ggml_tensor * packed_kv = rows->src[0]; ++ ggml_tensor * top_k = rows->src[1]; + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ if (ggml_metal_glm_selected_row_flash_noop_enabled()) { ++ auto pipeline = ggml_metal_library_get_pipeline_zero_f32(lib); ++ const int nth = 256; ++ const int64_t n_tg = (ggml_nelements(op) + nth - 1)/nth; ++ ggml_metal_op_concurrency_reset(ctx); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, nth, 1, 1); ++ return 1; ++ } + +- const int32_t * opts = op->op_params; +- ggml_op_pool op_pool = (ggml_op_pool) opts[0]; ++ ggml_metal_kargs_selected_row_flash args = { ++ /*.ne00 =*/ (int32_t) q->ne[0], ++ /*.ne01 =*/ (int32_t) q->ne[1], ++ /*.ne02 =*/ (int32_t) q->ne[2], ++ /*.ne03 =*/ (int32_t) q->ne[3], ++ /*.nb00 =*/ q->nb[0], ++ /*.nb01 =*/ q->nb[1], ++ /*.nb02 =*/ q->nb[2], ++ /*.nb03 =*/ q->nb[3], ++ /*.ne10 =*/ (int32_t) packed_kv->ne[0], ++ /*.ne11 =*/ (int32_t) packed_kv->ne[1], ++ /*.ne12 =*/ (int32_t) packed_kv->ne[2], ++ /*.ne13 =*/ (int32_t) packed_kv->ne[3], ++ /*.nb10 =*/ packed_kv->nb[0], ++ /*.nb11 =*/ packed_kv->nb[1], ++ /*.nb12 =*/ packed_kv->nb[2], ++ /*.nb13 =*/ packed_kv->nb[3], ++ /*.ne20 =*/ (int32_t) view->ne[0], ++ /*.ne21 =*/ (int32_t) view->ne[1], ++ /*.ne22 =*/ (int32_t) view->ne[2], ++ /*.ne23 =*/ (int32_t) view->ne[3], ++ /*.nb20 =*/ view->nb[0], ++ /*.nb21 =*/ view->nb[1], ++ /*.nb22 =*/ view->nb[2], ++ /*.nb23 =*/ view->nb[3], ++ /*.ne30 =*/ 0, ++ /*.ne31 =*/ 0, ++ /*.ne32 =*/ 0, ++ /*.ne33 =*/ 0, ++ /*.nb30 =*/ 0, ++ /*.nb31 =*/ 0, ++ /*.nb32 =*/ 0, ++ /*.nb33 =*/ 0, ++ /*.ne40 =*/ (int32_t) top_k->ne[0], ++ /*.ne41 =*/ (int32_t) top_k->ne[1], ++ /*.ne42 =*/ (int32_t) top_k->ne[2], ++ /*.ne43 =*/ (int32_t) top_k->ne[3], ++ /*.nb40 =*/ top_k->nb[0], ++ /*.nb41 =*/ top_k->nb[1], ++ /*.nb42 =*/ top_k->nb[2], ++ /*.nb43 =*/ top_k->nb[3], ++ /*.ne0 =*/ (int32_t) op->ne[0], ++ /*.ne1 =*/ (int32_t) op->ne[1], ++ /*.ne2 =*/ (int32_t) op->ne[2], ++ /*.ne3 =*/ (int32_t) op->ne[3], ++ /*.nb0 =*/ op->nb[0], ++ /*.nb1 =*/ op->nb[1], ++ /*.nb2 =*/ op->nb[2], ++ /*.nb3 =*/ op->nb[3], ++ /*.scale =*/ ggml_get_op_params_f32(op, 0), ++ }; + +- const int32_t k0 = opts[1]; +- const int32_t s0 = opts[2]; +- const int32_t p0 = opts[3]; ++ const bool use_tiled = ++ ggml_metal_glm_dsa_selected_row_flash_tiled_selected(op) && ++ ggml_metal_glm_dsa_selected_row_flash_tiled_shape(op); ++ const int32_t requested_nsg = ++ ggml_metal_glm_dsa_selected_row_flash_heads_per_tg_requested(); ++ const bool use_pair = !use_tiled && requested_nsg == 2 && args.ne02 % 2 == 0; ++ const int32_t nsg = use_tiled ? 4 : (use_pair ? 2 : 1); ++ // Selected-row flash fuses the compact gather into attention. On the ++ // GLM-5.2 top_k=768 shape it wants the generic flash width, not the ++ // lower all-KV compact-flash default. ++ const int32_t nwg = ggml_metal_flash_attn_ext_vec_nwg(op, false, false, false, false, false); ++ ++ auto pipeline = use_tiled ? ++ ggml_metal_library_get_pipeline_selected_row_flash_tiled(lib, nwg) : ++ (use_pair ? ++ ggml_metal_library_get_pipeline_selected_row_flash_pair(lib, nwg) : ++ ggml_metal_library_get_pipeline_selected_row_flash_vec(lib, 1, nwg)); ++ const int32_t dispatch_nsg = use_tiled ? 4 : 1; ++ GGML_ASSERT(dispatch_nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + +- const int64_t IW = op->src[0]->ne[0]; +- const int64_t OW = op->ne[0]; ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ ggml_metal_buffer_id bid_pad = bid_dst; ++ bid_pad.offs += ggml_nbytes(op); + +- const int64_t np = ggml_nelements(op); ++ ggml_metal_buffer_id bid_blk = bid_pad; ++ bid_blk.offs += ggml_metal_op_flash_attn_ext_extra_pad(op); + +- ggml_metal_kargs_pool_1d args_pool_1d = { +- /* .k0 = */ k0, +- /* .s0 = */ s0, +- /* .p0 = */ p0, +- /* .IW = */ IW, +- /* .OW = */ OW, +- /* .np = */ np +- }; ++ ggml_metal_buffer_id bid_tmp = bid_blk; ++ bid_tmp.offs += ggml_metal_op_flash_attn_ext_extra_blk(op); + +- auto pipeline = ggml_metal_library_get_pipeline_pool_1d(lib, op, op_pool); ++ ggml_metal_buffer_id bid_v = ggml_metal_get_buffer_id(packed_kv); ++ bid_v.offs += view->view_offs; + +- const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), (int) np); +- const int ntg = (np + nth - 1) / nth; ++#define SELECTED_ROW_FLASH_VEC_SMEM(nsg_) (GGML_PAD(((GGML_PAD(args.ne00, 128) + 4*OP_FLASH_ATTN_EXT_VEC_NCPSG + 2*GGML_PAD(args.ne20, 128))*(nsg_))*(sizeof(float)/2), 16)) ++#define SELECTED_ROW_FLASH_TILED_SMEM (GGML_PAD((8*576 + 2*8*512 + 2*8*64 + 4*4*16*8)*sizeof(uint16_t), 16)) ++ const size_t smem = use_tiled ? SELECTED_ROW_FLASH_TILED_SMEM : SELECTED_ROW_FLASH_VEC_SMEM(nsg); ++#undef SELECTED_ROW_FLASH_TILED_SMEM ++#undef SELECTED_ROW_FLASH_VEC_SMEM ++ ++ const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size); + ++ ggml_metal_op_concurrency_reset(ctx); + ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args_pool_1d, sizeof(args_pool_1d), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(q), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(packed_kv), 2); ++ ggml_metal_encoder_set_buffer (enc, bid_v, 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(top_k), 4); ++ ++ const int grid_x = use_tiled ? (args.ne02 + 7)/8 : ++ (args.ne01 + OP_FLASH_ATTN_EXT_VEC_NQPSG - 1)/OP_FLASH_ATTN_EXT_VEC_NQPSG; ++ const int grid_y = use_tiled ? args.ne01 : (args.ne02 + nsg - 1)/nsg; ++ const int grid_z = args.ne03*nwg; ++ ++ if (nwg == 1) { ++ GGML_ASSERT(ggml_metal_op_flash_attn_ext_extra_tmp(op) == 0); ++ ggml_metal_encoder_set_buffer(enc, bid_dst, 5); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(top_k), 6); // unused mask ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, dispatch_nsg, 1); ++ } else { ++ GGML_ASSERT(ggml_metal_op_flash_attn_ext_extra_tmp(op) != 0); ++ ggml_metal_encoder_set_buffer(enc, bid_tmp, 5); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(top_k), 6); // unused mask ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, dispatch_nsg, 1); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1); ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ ggml_metal_kargs_flash_attn_ext_vec_reduce args0 = { ++ (int32_t) (args.ne1*args.ne2*args.ne3), ++ }; ++ ++ auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce(lib, op, args.ne20, nwg); ++ ggml_metal_encoder_set_pipeline(enc, pipeline0); ++ ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, args0.nrows, 1, 1, 32*nwg, 1, 1); ++ } ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=selected_row_flash kernel=%s tensor=%s q_type=%s k_type=%s top_k_type=%s dst_type=%s q_width=%lld v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld top_k=%lld nwg=%d smem=%zu grid_x=%d grid_y=%d grid_z=%d threads_x=32 threads_y=%d\n", ++ use_tiled ? "indirect_tiled" : (use_pair ? "gather_pair" : "gather_vec"), ++ ggml_metal_tensor_name(op), ++ ggml_type_name(q->type), ++ ggml_type_name(packed_kv->type), ++ ggml_type_name(top_k->type), ++ ggml_type_name(op->type), ++ (long long) args.ne00, ++ (long long) args.ne20, ++ (long long) args.ne01, ++ (long long) args.ne02, ++ (long long) args.ne03, ++ (long long) args.ne11, ++ (long long) args.ne40, ++ nwg, ++ smem, ++ grid_x, ++ grid_y, ++ grid_z, ++ dispatch_nsg); ++ } + + return 1; + } + +- +-int ggml_metal_op_pool_2d(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + ++ if (int selected_row_flash = ggml_metal_op_selected_row_flash_vec(ctx, idx)) { ++ return selected_row_flash; ++ } ++ + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + ++ const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ GGML_TENSOR_LOCALS( int32_t, nb, op, nb); + +- const int32_t * opts = op->op_params; +- ggml_op_pool op_pool = (ggml_op_pool) opts[0]; ++ GGML_ASSERT(ne00 % 4 == 0); + +- const int32_t k0 = opts[1]; +- const int32_t k1 = opts[2]; +- const int32_t s0 = opts[3]; +- const int32_t s1 = opts[4]; +- const int32_t p0 = opts[5]; +- const int32_t p1 = opts[6]; ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == op->src[2]->type); + +- const int64_t IH = op->src[0]->ne[1]; +- const int64_t IW = op->src[0]->ne[0]; ++ //GGML_ASSERT(ggml_are_same_shape (src1, src2)); ++ GGML_ASSERT(ne11 == ne21); ++ GGML_ASSERT(ne12 == ne22); + +- const int64_t N = op->ne[3]; +- const int64_t OC = op->ne[2]; +- const int64_t OH = op->ne[1]; +- const int64_t OW = op->ne[0]; ++ GGML_ASSERT(!op->src[3] || op->src[3]->type == GGML_TYPE_F16); ++ GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] >= op->src[0]->ne[1] && ++ "the Flash-Attention Metal kernel requires the mask to be at least n_queries big"); + +- const int64_t np = N * OC * OH * OW; ++ float scale; ++ float max_bias; ++ float logit_softcap; ++ ++ memcpy(&scale, ((const int32_t *) op->op_params) + 0, sizeof(scale)); ++ memcpy(&max_bias, ((const int32_t *) op->op_params) + 1, sizeof(max_bias)); ++ memcpy(&logit_softcap, ((const int32_t *) op->op_params) + 2, sizeof(logit_softcap)); ++ ++ if (logit_softcap != 0.0f) { ++ scale /= logit_softcap; ++ } ++ ++ const bool has_mask = op->src[3] != NULL; ++ const bool has_sinks = op->src[4] != NULL; ++ const bool has_bias = max_bias != 0.0f; ++ const bool has_scap = logit_softcap != 0.0f; + +- ggml_metal_kargs_pool_2d args_pool_2d = { +- /* .k0 = */ k0, +- /* .k1 = */ k1, +- /* .s0 = */ s0, +- /* .s1 = */ s1, +- /* .p0 = */ p0, +- /* .p1 = */ p1, +- /* .IH = */ IH, +- /* .IW = */ IW, +- /* .OH = */ OH, +- /* .OW = */ OW, +- /* .np = */ np +- }; ++ const uint32_t n_head = op->src[0]->ne[2]; ++ const int32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head)); + +- auto pipeline = ggml_metal_library_get_pipeline_pool_2d(lib, op, op_pool); ++ const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); ++ const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); + +- const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), (int) np); +- const int ntg = (np + nth - 1) / nth; ++ GGML_ASSERT(ne01 < 65536); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args_pool_2d, sizeof(args_pool_2d), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); ++ ggml_metal_buffer_id bid_src2 = ggml_metal_get_buffer_id(op->src[2]); ++ ggml_metal_buffer_id bid_src3 = has_mask ? ggml_metal_get_buffer_id(op->src[3]) : bid_src0; ++ ggml_metal_buffer_id bid_src4 = has_sinks ? ggml_metal_get_buffer_id(op->src[4]) : bid_src0; + +- ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); + +- return 1; +-} ++ ggml_metal_buffer_id bid_pad = bid_dst; ++ bid_pad.offs += ggml_nbytes(op); + +-int ggml_metal_op_mul_mat(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ ggml_metal_buffer_id bid_blk = bid_pad; ++ bid_blk.offs += ggml_metal_op_flash_attn_ext_extra_pad(op); + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ ggml_metal_buffer_id bid_tmp = bid_blk; ++ bid_tmp.offs += ggml_metal_op_flash_attn_ext_extra_blk(op); + +- const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ if (!ggml_metal_op_flash_attn_ext_use_vec(op)) { ++ // half8x8 kernel ++ const int nqptg = OP_FLASH_ATTN_EXT_NQPSG; // queries per threadgroup ++ const int ncpsg = OP_FLASH_ATTN_EXT_NCPSG; // cache values per simdgroup + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ GGML_ASSERT(nqptg <= 32); ++ GGML_ASSERT(nqptg % 8 == 0); ++ GGML_ASSERT(ncpsg % 32 == 0); + +- GGML_ASSERT(ne00 == ne10); ++ bool need_sync = false; + +- GGML_ASSERT(ne12 % ne02 == 0); +- GGML_ASSERT(ne13 % ne03 == 0); ++ const bool has_kvpad = ne11 % ncpsg != 0; + +- const int16_t r2 = ne12/ne02; +- const int16_t r3 = ne13/ne03; ++ if (has_kvpad) { ++ assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0); + +- // find the break-even point where the matrix-matrix kernel becomes more efficient compared +- // to the matrix-vector kernel +- const int ne11_mm_min = 8; ++ ggml_metal_kargs_flash_attn_ext_pad args0 = { ++ /*.ne11 =*/ne11, ++ /*.ne_12_2 =*/ne12, ++ /*.ne_12_3 =*/ne13, ++ /*.nb11 =*/nb11, ++ /*.nb12 =*/nb12, ++ /*.nb13 =*/nb13, ++ /*.nb21 =*/nb21, ++ /*.nb22 =*/nb22, ++ /*.nb23 =*/nb23, ++ /*.ne31 =*/ne31, ++ /*.ne32 =*/ne32, ++ /*.ne33 =*/ne33, ++ /*.nb31 =*/nb31, ++ /*.nb32 =*/nb32, ++ /*.nb33 =*/nb33, ++ }; + +- // first try to use small-batch mat-mv kernels +- // these should be efficient for BS [2, ~8] +- if (op->src[1]->type == GGML_TYPE_F32 && (ne00%128 == 0) && +- ( +- ( +- ( +- op->src[0]->type == GGML_TYPE_F32 || // TODO: helper function +- op->src[0]->type == GGML_TYPE_F16 || +- op->src[0]->type == GGML_TYPE_BF16 || +- op->src[0]->type == GGML_TYPE_Q1_0 || +- op->src[0]->type == GGML_TYPE_Q2_0 || +- op->src[0]->type == GGML_TYPE_Q4_0 || +- op->src[0]->type == GGML_TYPE_Q4_1 || +- op->src[0]->type == GGML_TYPE_Q5_0 || +- op->src[0]->type == GGML_TYPE_Q5_1 || +- op->src[0]->type == GGML_TYPE_Q8_0 || +- op->src[0]->type == GGML_TYPE_MXFP4 || +- op->src[0]->type == GGML_TYPE_IQ4_NL || +- false) && (ne11 >= 2 && ne11 <= 8) +- ) || +- ( +- ( +- op->src[0]->type == GGML_TYPE_Q4_K || +- op->src[0]->type == GGML_TYPE_Q5_K || +- op->src[0]->type == GGML_TYPE_Q6_K || +- op->src[0]->type == GGML_TYPE_Q2_K || +- op->src[0]->type == GGML_TYPE_Q3_K || +- false) && (ne11 >= 4 && ne11 <= 8) +- ) +- ) +- ) { +- // TODO: determine the optimal parameters based on grid utilization +- // I still don't know why we should not always use the maximum available threads: +- // +- // nsg = pipeline.maxTotalThreadsPerThreadgroup / 32 +- // +- // my current hypothesis is that the work grid is not evenly divisible for different nsg +- // values and there can be some tail effects when nsg is high. need to confirm this +- // +- const int nsg = 2; // num simdgroups per threadgroup ++ auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_pad(lib, op, has_mask, ncpsg); + +- // num threads along row per simdgroup +- int16_t nxpsg = 0; +- if (ne00 % 256 == 0 && ne11 < 3) { +- nxpsg = 16; +- } else if (ne00 % 128 == 0) { +- nxpsg = 8; +- } else { +- nxpsg = 4; +- } ++ ggml_metal_encoder_set_pipeline(enc, pipeline0); ++ ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_src3, 3); ++ ggml_metal_encoder_set_buffer (enc, bid_pad, 4); + +- const int16_t nypsg = 32/nxpsg; // num threads along col per simdgroup (i.e. a simdgroup processes that many src0 rows at a time) +- const int16_t r0ptg = nypsg*nsg; // num src0 rows per threadgroup +- int16_t r1ptg = 4; // num src1 rows per threadgroup ++ assert(ne12 == ne22); ++ assert(ne13 == ne23); + +- // note: not sure how optimal are those across all different hardware. there might be something cleverer +- switch (ne11) { +- case 2: +- r1ptg = 2; break; +- case 3: +- case 6: +- r1ptg = 3; break; +- case 4: +- case 7: +- case 8: +- r1ptg = 4; break; +- case 5: +- r1ptg = 5; break; +- default: +- GGML_ABORT("unsupported ne11"); +- }; ++ ggml_metal_encoder_dispatch_threadgroups(enc, ncpsg, std::max(ne12, ne32), std::max(ne13, ne33), 32, 1, 1); + +- auto pipeline = ggml_metal_library_get_pipeline_mul_mv_ext(lib, op, nsg, nxpsg, r1ptg); ++ need_sync = true; ++ } + +- ggml_metal_kargs_mul_mv_ext args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne10 =*/ ne10, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.r2 =*/ r2, +- /*.r3 =*/ r3, +- }; ++ if (has_mask) { ++ assert(ggml_metal_op_flash_attn_ext_extra_blk(op) != 0); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_kargs_flash_attn_ext_blk args0 = { ++ /*.ne01 =*/ ne01, ++ /*.ne30 =*/ ne30, ++ /*.ne31 =*/ ne31, ++ /*.ne32 =*/ ne32, ++ /*.ne33 =*/ ne33, ++ /*.nb31 =*/ nb31, ++ /*.nb32 =*/ nb32, ++ /*.nb33 =*/ nb33, ++ }; + +- ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + r0ptg - 1)/r0ptg), ((ne11 + r1ptg - 1)/r1ptg), ne12*ne13, 32, nsg, 1); +- } else if ( +- !ggml_is_transposed(op->src[0]) && +- !ggml_is_transposed(op->src[1]) && +- // for now the matrix-matrix multiplication kernel only works on A14+/M1+ SoCs +- // AMD GPU and older A-chips will reuse matrix-vector multiplication kernel +- props_dev->has_simdgroup_mm && ne00 >= 64 && ne11 > ne11_mm_min) { +- //GGML_LOG_INFO("matrix: ne00 = %6d, ne01 = %6d, ne02 = %6d, ne11 = %6d, ne12 = %6d\n", ne00, ne01, ne02, ne11, ne12); ++ auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_blk(lib, op, nqptg, ncpsg); + +- // some Metal matrix data types require aligned pointers +- // ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf (Table 2.5) +- //switch (op->src[0]->type) { +- // case GGML_TYPE_F32: GGML_ASSERT(nb01 % 16 == 0); break; +- // case GGML_TYPE_F16: GGML_ASSERT(nb01 % 8 == 0); break; +- // case GGML_TYPE_BF16: GGML_ASSERT(nb01 % 8 == 0); break; +- // default: break; +- //} ++ 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_blk, 2); + +- auto pipeline = ggml_metal_library_get_pipeline_mul_mm(lib, op); ++ const int32_t nblk1 = ((ne01 + nqptg - 1)/nqptg); ++ const int32_t nblk0 = ((ne30 + ncpsg - 1)/ncpsg); + +- ggml_metal_kargs_mul_mm args = { +- /*.ne00 =*/ ne00, +- /*.ne02 =*/ ne02, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne12 =*/ ne12, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.r2 =*/ r2, +- /*.r3 =*/ r3, +- }; ++ ggml_metal_encoder_dispatch_threadgroups(enc, nblk0, nblk1, ne32*ne33, 32, 1, 1); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ need_sync = true; ++ } + +- const size_t smem = pipeline.smem; ++ if (need_sync) { ++ ggml_metal_op_concurrency_reset(ctx); ++ } + +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ const int is_q = ggml_is_quantized(op->src[1]->type) ? 1 : 0; + +- const int nr0 = pipeline.nr0; +- const int nr1 = pipeline.nr1; +- const int nsg = pipeline.nsg; ++ // 2*(2*ncpsg) ++ // ncpsg soft_max values + ncpsg mask values ++ // ++ // 16*32*(nsg) ++ // the shared memory needed for the simdgroups to load the KV cache ++ // each thread loads (dequantizes) 16 head elements, there are 32 threads in th SG ++ // ++#define FATTN_SMEM(nsg) (GGML_PAD((nqptg*(ne00 + 2*GGML_PAD(ne20, 64) + 2*(2*ncpsg)) + is_q*(16*32*(nsg)))*(sizeof(float)/2), 16)) + +- ggml_metal_encoder_dispatch_threadgroups(enc, ((ne11 + nr1 - 1) / nr1), ((ne01 + nr0 - 1) / nr0), ne12 * ne13, 32, nsg, 1); +- } else { +- auto pipeline = ggml_metal_library_get_pipeline_mul_mv(lib, op); ++ //int64_t nsgmax = 4; ++ // ++ //if (is_q) { ++ // nsgmax = 2; ++ // while (true) { ++ // const size_t smem = FATTN_SMEM(nsgmax); ++ // if (smem > props_dev->max_theadgroup_memory_size) { ++ // break; ++ // } ++ // nsgmax *= 2; ++ // } ++ // nsgmax /= 2; ++ //} + +- const int nr0 = pipeline.nr0; +- const int nr1 = pipeline.nr1; +- const int nsg = pipeline.nsg; ++ // simdgroups per threadgroup (a.k.a. warps) ++ //nsg = ne01 <= nqptg ? MAX(4, MIN(nsgmax, MIN(ne11/ncpsg, (int64_t) pipeline.maxTotalThreadsPerThreadgroup/32))) : 4; ++ int32_t nsg = ne00 >= 512 ? 8 : 4; + +- const size_t smem = pipeline.smem; ++ const size_t smem = FATTN_SMEM(nsg); + +- ggml_metal_kargs_mul_mv args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne10 =*/ ne10, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.nr0 =*/ nr0, +- /*.r2 =*/ r2, +- /*.r3 =*/ r3, ++ ggml_metal_kargs_flash_attn_ext args = { ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne11 =*/ ne11, ++ /*.ne_12_2 =*/ ne12, ++ /*.ne_12_3 =*/ ne13, ++ /*.ns10 =*/ int32_t(nb11/nb10), ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ns20 =*/ int32_t(nb21/nb20), ++ /*.nb21 =*/ nb21, ++ /*.nb22 =*/ nb22, ++ /*.nb23 =*/ nb23, ++ /*.ne31 =*/ ne31, ++ /*.ne32 =*/ ne32, ++ /*.ne33 =*/ ne33, ++ /*.nb31 =*/ nb31, ++ /*.nb32 =*/ nb32, ++ /*.nb33 =*/ nb33, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.scale =*/ scale, ++ /*.max_bias =*/ max_bias, ++ /*.m0 =*/ m0, ++ /*.m1 =*/ m1, ++ /*.n_head_log2 =*/ n_head_log2, ++ /*.logit_softcap =*/ logit_softcap, + }; + ++ auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg); ++ + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 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, bid_pad, 6); ++ ggml_metal_encoder_set_buffer (enc, bid_blk, 7); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 8); + + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- if (op->src[0]->type == GGML_TYPE_F32 || +- op->src[0]->type == GGML_TYPE_F16 || +- op->src[0]->type == GGML_TYPE_BF16 || +- op->src[0]->type == GGML_TYPE_Q8_0) { +- ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + nr0 - 1)/(nr0)), ((ne11 + nr1 - 1)/nr1), ne12*ne13, 32, nsg, 1); +- } else { +- ggml_metal_encoder_dispatch_threadgroups(enc, ((ne01 + nr0*nsg - 1)/(nr0*nsg)), ((ne11 + nr1 - 1)/nr1), ne12*ne13, 32, nsg, 1); ++ const int grid_x = (ne01 + nqptg - 1)/nqptg; ++ const int grid_y = ne02; ++ const int grid_z = ne03; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=flash_attn_ext kernel=tile tensor=%s q_type=%s k_type=%s v_type=%s mask_type=%s dst_type=%s q_width=%lld v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=32 threads_y=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ has_mask ? ggml_type_name(op->src[3]->type) : "none", ++ ggml_type_name(op->type), ++ (long long) ne00, ++ (long long) ne20, ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ (long long) ne11, ++ grid_x, ++ grid_y, ++ grid_z, ++ nsg); + } +- } ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, nsg, 1); ++#undef FATTN_SMEM ++ } else { ++ // half4x4 kernel ++ const int nqptg = OP_FLASH_ATTN_EXT_VEC_NQPSG; // 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 + +- return 1; +-} ++ GGML_ASSERT(nqptg <= 32); ++ GGML_ASSERT(nqptg % 1 == 0); ++ GGML_ASSERT(ncpsg % 32 == 0); + +-size_t ggml_metal_op_mul_mat_id_extra_tpe(const ggml_tensor * op) { +- assert(op->op == GGML_OP_MUL_MAT_ID); ++ bool need_sync = false; + +- const int64_t ne02 = op->src[0]->ne[2]; // n_expert ++ const bool has_kvpad = ne11 % ncpsg != 0; + +- return ggml_type_size(GGML_TYPE_I32)*ne02; +-} ++ if (has_kvpad) { ++ assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0); + +-size_t ggml_metal_op_mul_mat_id_extra_ids(const ggml_tensor * op) { +- assert(op->op == GGML_OP_MUL_MAT_ID); ++ ggml_metal_kargs_flash_attn_ext_pad args0 = { ++ /*.ne11 =*/ne11, ++ /*.ne_12_2 =*/ne12, ++ /*.ne_12_3 =*/ne13, ++ /*.nb11 =*/nb11, ++ /*.nb12 =*/nb12, ++ /*.nb13 =*/nb13, ++ /*.nb21 =*/nb21, ++ /*.nb22 =*/nb22, ++ /*.nb23 =*/nb23, ++ /*.ne31 =*/ne31, ++ /*.ne32 =*/ne32, ++ /*.ne33 =*/ne33, ++ /*.nb31 =*/nb31, ++ /*.nb32 =*/nb32, ++ /*.nb33 =*/nb33, ++ }; + +- const int64_t ne02 = op->src[0]->ne[2]; // n_expert +- const int64_t ne21 = op->src[2]->ne[1]; // n_token ++ auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_pad(lib, op, has_mask, ncpsg); + +- return ggml_type_size(GGML_TYPE_I32)*ne02*ne21; +-} ++ ggml_metal_encoder_set_pipeline(enc, pipeline0); ++ ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_src3, 3); ++ ggml_metal_encoder_set_buffer (enc, bid_pad, 4); + +-int ggml_metal_op_mul_mat_id(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ assert(ne12 == ne22); ++ assert(ne13 == ne23); + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ ggml_metal_encoder_dispatch_threadgroups(enc, ncpsg, std::max(ne12, ne32), std::max(ne13, ne33), 32, 1, 1); + +- const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ need_sync = true; ++ } + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ if (need_sync) { ++ ggml_metal_op_concurrency_reset(ctx); ++ } + +- // src2 = ids +- GGML_ASSERT(op->src[2]->type == GGML_TYPE_I32); ++ // note: for simplicity assume the K is larger or equal than V ++ GGML_ASSERT(ne10 >= ne20); + +- GGML_ASSERT(!ggml_is_transposed(op->src[0])); +- GGML_ASSERT(!ggml_is_transposed(op->src[1])); ++ // ne00 + 2*ncpsg*(nsg) ++ // for each query, we load it as f16 in shared memory (ne00) ++ // and store the soft_max values and the mask ++ // ++ // ne20*(nsg) ++ // each simdgroup has a full f32 head vector in shared mem to accumulate results ++ // ++#define FATTN_SMEM(nsg) (GGML_PAD(((GGML_PAD(ne00, 128) + 4*ncpsg + 2*GGML_PAD(ne20, 128))*(nsg))*(sizeof(float)/2), 16)) + +- GGML_ASSERT(ne03 == 1); +- GGML_ASSERT(ne13 == 1); ++ int64_t nsg = 1; + +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); +- ggml_metal_buffer_id bid_src2 = ggml_metal_get_buffer_id(op->src[2]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ // workgroups ++ // each workgroup handles nsg*nkpsg cache values ++ int32_t nwg = 1; ++ if (false) { ++ // for small KV caches, we could launch a single workgroup and write the results directly to dst/ ++ // however, this does not lead to significant improvement, so disabled ++ nwg = 1; ++ nsg = 4; ++ } else { ++ nwg = 32; ++ nsg = 1; ++ while (2*nwg*nsg*ncpsg < ne11 && nsg < 4) { ++ nsg *= 2; ++ } ++ } ++ if (!has_mask && ++ !has_sinks && ++ !has_bias && ++ !has_scap && ++ !has_kvpad && ++ op->src[1]->type == GGML_TYPE_F16 && ++ op->src[2]->type == GGML_TYPE_F16 && ++ ne00 == 576 && ++ ne20 == 512 && ++ ne01 == 1 && ++ ne11 == 2048) { ++ nwg = ggml_metal_flash_attn_ext_vec_nwg( ++ op, has_mask, has_sinks, has_bias, has_scap, has_kvpad); ++ nsg = 1; ++ } ++ ++ ggml_metal_kargs_flash_attn_ext_vec args = { ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne11 =*/ ne11, ++ /*.ne_12_2 =*/ ne12, ++ /*.ne_12_3 =*/ ne13, ++ /*.ns10 =*/ int32_t(nb11/nb10), ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ns20 =*/ int32_t(nb21/nb20), ++ /*.nb21 =*/ nb21, ++ /*.nb22 =*/ nb22, ++ /*.nb23 =*/ nb23, ++ /*.ne31 =*/ ne31, ++ /*.ne32 =*/ ne32, ++ /*.ne33 =*/ ne33, ++ /*.nb31 =*/ nb31, ++ /*.nb32 =*/ nb32, ++ /*.nb33 =*/ nb33, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.scale =*/ scale, ++ /*.max_bias =*/ max_bias, ++ /*.m0 =*/ m0, ++ /*.m1 =*/ m1, ++ /*.n_head_log2 =*/ n_head_log2, ++ /*.logit_softcap =*/ logit_softcap, ++ }; + +- const uint32_t r2 = 1; +- const uint32_t r3 = 1; ++ const char * split_exact_tensor = ++ getenv("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_SPLIT_EXACT_TENSOR"); ++ const bool split_exact_tensor_selected = ++ split_exact_tensor == nullptr || split_exact_tensor[0] == '\0' || ++ strcmp(split_exact_tensor, ggml_metal_tensor_name(op)) == 0; ++ const bool glm_compact_split_exact = ++ ggml_metal_glm_compact_split_exact_enabled() && ++ split_exact_tensor_selected && ++ !has_mask && ++ !has_sinks && ++ !has_bias && ++ !has_scap && ++ !has_kvpad && ++ op->src[1]->type == GGML_TYPE_F16 && ++ op->src[2]->type == GGML_TYPE_F16 && ++ ne00 == 576 && ++ ne20 == 512 && ++ ne01 == 1 && ++ ne02 == 64 && ++ ne03 == 1 && ++ ne11 == 2048 && ++ ne12 == 1 && ++ ne22 == 1 && ++ nwg == 4; ++ if (glm_compact_split_exact) { ++ constexpr int32_t exact_nwg = 4; ++ constexpr size_t chunk_rows = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ const size_t nrows = size_t(ne1)*size_t(ne2)*size_t(ne3); ++ const size_t chunk_count = (size_t(ne11) + chunk_rows - 1)/chunk_rows; ++ const size_t score_bytes = sizeof(float)*size_t(ne02)*size_t(ne03)*size_t(ne11); ++ const size_t chunk_ms_bytes = sizeof(float)*size_t(ne02)*size_t(ne03)*chunk_count; ++ const size_t chunk_v_bytes = sizeof(float)*size_t(ne02)*size_t(ne03)*chunk_count*size_t(ne20); ++ const size_t partial_bytes = sizeof(float)*nrows*exact_nwg*size_t(ne20 + 2); ++ GGML_ASSERT(chunk_count == size_t(ne11)/OP_FLASH_ATTN_EXT_VEC_NCPSG); ++ GGML_ASSERT(ggml_metal_op_flash_attn_ext_extra_tmp(op) >= ++ score_bytes + chunk_ms_bytes + chunk_v_bytes + partial_bytes); ++ ++ ggml_metal_buffer_id bid_scores = bid_tmp; ++ ggml_metal_buffer_id bid_chunk_ms = bid_scores; ++ bid_chunk_ms.offs += score_bytes; ++ ggml_metal_buffer_id bid_chunk_v = bid_chunk_ms; ++ bid_chunk_v.offs += chunk_ms_bytes; ++ ggml_metal_buffer_id bid_partials = bid_chunk_v; ++ bid_partials.offs += chunk_v_bytes; ++ ggml_metal_buffer_id bid_reduce_partials = bid_partials; ++ ++ auto qk_pipeline = ggml_metal_library_get_pipeline_glm_compact_qk_scores(lib); ++ constexpr size_t qk_smem = GGML_PAD(576, 128)*sizeof(uint16_t); ++ GGML_ASSERT(32 <= ggml_metal_pipeline_max_theads_per_threadgroup(qk_pipeline)); ++ GGML_ASSERT(qk_smem <= props_dev->max_theadgroup_memory_size); ++ ++ ggml_metal_encoder_set_pipeline(enc, qk_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_scores, 3); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 4); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, qk_smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups( ++ enc, (ne11 + OP_FLASH_ATTN_EXT_VEC_NCPSG - 1)/OP_FLASH_ATTN_EXT_VEC_NCPSG, ++ ne02, ne03, 32, 1, 1); + +- // find the break-even point where the matrix-matrix kernel becomes more efficient compared +- // to the matrix-vector kernel +- // ne20 = n_used_experts +- // ne21 = n_rows (batch size) +- const int ne21_mm_id_min = 32; ++ ggml_metal_op_internal_phase_barrier(ctx); + +- if (props_dev->has_simdgroup_mm && ne00 >= 64 && (ne21 >= ne21_mm_id_min)) { +- // some Metal matrix data types require aligned pointers +- // ref: https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf (Table 2.5) +- //switch (op->src[0]->type) { +- // case GGML_TYPE_F32: GGML_ASSERT(nb01 % 16 == 0); break; +- // case GGML_TYPE_F16: GGML_ASSERT(nb01 % 8 == 0); break; +- // case GGML_TYPE_BF16: GGML_ASSERT(nb01 % 8 == 0); break; +- // default: break; +- //} ++ if (ggml_metal_glm_compact_dump_scores_diagnostic_enabled()) { ++ return 1; ++ } + +- // extra buffers for intermediate id mapping +- ggml_metal_buffer_id bid_tpe = bid_dst; +- bid_tpe.offs += ggml_nbytes(op); ++ const bool legacy_scores_v = ++ ggml_metal_glm_compact_legacy_scores_v_diagnostic_enabled(); ++ if (legacy_scores_v) { ++ bid_reduce_partials = bid_scores; ++ bid_reduce_partials.offs += score_bytes; ++ auto legacy_v_pipeline = ggml_metal_library_get_pipeline_glm_compact_scores_v(lib); ++ constexpr size_t legacy_v_smem = (OP_FLASH_ATTN_EXT_VEC_NCPSG + 512)*sizeof(float); ++ GGML_ASSERT(32 <= ggml_metal_pipeline_max_theads_per_threadgroup(legacy_v_pipeline)); ++ GGML_ASSERT(legacy_v_smem <= props_dev->max_theadgroup_memory_size); ++ ++ ggml_metal_encoder_set_pipeline(enc, legacy_v_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_scores, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_reduce_partials, 3); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, legacy_v_smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03*exact_nwg, 32, 1, 1); ++ } else { ++ auto prefix_pipeline = ggml_metal_library_get_pipeline_glm_compact_softmax_prefix(lib); ++ GGML_ASSERT(32 <= ggml_metal_pipeline_max_theads_per_threadgroup(prefix_pipeline)); ++ ++ ggml_metal_encoder_set_pipeline(enc, prefix_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_scores, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_chunk_ms, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_partials, 3); ++ ggml_metal_encoder_set_threadgroup_memory_size( ++ enc, 2*OP_FLASH_ATTN_EXT_VEC_NCPSG*sizeof(float), 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, exact_nwg, ne02, ne03, 32, 1, 1); ++ ++ ggml_metal_op_internal_phase_barrier(ctx); ++ } + +- ggml_metal_buffer_id bid_ids = bid_tpe; +- bid_ids.offs += ggml_metal_op_mul_mat_id_extra_tpe(op); ++ if (!legacy_scores_v && ggml_metal_glm_compact_sequential_v_diagnostic_enabled()) { ++ auto sequential_v_pipeline = ++ ggml_metal_library_get_pipeline_glm_compact_probs_v_sequential(lib); ++ GGML_ASSERT(32 <= ggml_metal_pipeline_max_theads_per_threadgroup(sequential_v_pipeline)); ++ ++ ggml_metal_encoder_set_pipeline(enc, sequential_v_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_scores, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_chunk_ms, 3); ++ ggml_metal_encoder_set_buffer (enc, bid_partials, 4); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, 512*sizeof(float), 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, exact_nwg, ne02, ne03, 32, 1, 1); ++ } else if (!legacy_scores_v) { ++ auto chunk_v_pipeline = ggml_metal_library_get_pipeline_glm_compact_chunk_v(lib); ++ GGML_ASSERT(32 <= ggml_metal_pipeline_max_theads_per_threadgroup(chunk_v_pipeline)); ++ ++ ggml_metal_encoder_set_pipeline(enc, chunk_v_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_scores, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_chunk_v, 3); ++ ggml_metal_encoder_dispatch_threadgroups(enc, chunk_count, ne02, ne03, 32, 1, 1); ++ ++ ggml_metal_op_internal_phase_barrier(ctx); ++ ++ auto fold_pipeline = ggml_metal_library_get_pipeline_glm_compact_chunk_fold(lib); ++ GGML_ASSERT(32 <= ggml_metal_pipeline_max_theads_per_threadgroup(fold_pipeline)); ++ ++ ggml_metal_encoder_set_pipeline(enc, fold_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_chunk_ms, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_chunk_v, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_partials, 3); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, 512*sizeof(float), 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, exact_nwg, ne02, ne03, 32, 1, 1); ++ } + +- { +- ggml_metal_kargs_mul_mm_id_map0 args = { +- ne02, +- ne10, +- ne11, // n_expert_used (bcast) +- nb11, +- nb12, +- ne21, // n_tokens +- ne20, // n_expert_used +- nb21, +- }; ++ ggml_metal_op_internal_phase_barrier(ctx); + +- auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id_map0(lib, ne02, ne20); ++ const int32_t nrows_i32 = ne1*ne2*ne3; ++ ggml_metal_kargs_flash_attn_ext_vec_reduce reduce_args = { ++ nrows_i32, ++ }; ++ auto reduce_pipeline = ++ ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce(lib, op, ne20, exact_nwg); ++ ggml_metal_encoder_set_pipeline(enc, reduce_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &reduce_args, sizeof(reduce_args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_reduce_partials, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nrows_i32, 1, 1, 32*exact_nwg, 1, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=flash_attn_ext kernel=glm_compact_chunk_exact tensor=%s q_type=%s k_type=%s v_type=%s mask_type=none dst_type=%s q_width=%lld v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld chunks=%zu score_bytes=%zu chunk_v_bytes=%zu reduction_strategy=scale_%.9g q_nb01=%llu q_nb02=%llu q_nb03=%llu k_nb11=%llu k_nb12=%llu k_nb13=%llu v_nb21=%llu v_nb22=%llu v_nb23=%llu out_ne1=%lld out_ne2=%lld out_ne3=%lld grid_x=%zu grid_y=%lld grid_z=%lld threads_x=32 threads_y=1 nwg=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->type), ++ (long long) ne00, ++ (long long) ne20, ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ (long long) ne11, ++ chunk_count, ++ score_bytes, ++ chunk_v_bytes, ++ (double) scale, ++ (unsigned long long) nb01, ++ (unsigned long long) nb02, ++ (unsigned long long) nb03, ++ (unsigned long long) nb11, ++ (unsigned long long) nb12, ++ (unsigned long long) nb13, ++ (unsigned long long) nb21, ++ (unsigned long long) nb22, ++ (unsigned long long) nb23, ++ (long long) ne1, ++ (long long) ne2, ++ (long long) ne3, ++ chunk_count, ++ (long long) ne02, ++ (long long) ne03, ++ exact_nwg); ++ } + +- const size_t smem = pipeline.smem; ++ return 1; ++ } + +- GGML_ASSERT(ne02 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ const bool glm_compact_multihead = ++ ggml_metal_glm_compact_multihead_flash_enabled(op) && ++ !has_mask && ++ !has_sinks && ++ !has_bias && ++ !has_scap && ++ !has_kvpad && ++ op->src[1]->type == GGML_TYPE_F16 && ++ op->src[2]->type == GGML_TYPE_F16 && ++ ne00 == 576 && ++ ne20 == 512 && ++ ne01 == 1 && ++ ne02 >= 8 && ++ ne02%8 == 0 && ++ ne11 >= 64 && ++ ne11%64 == 0 && ++ ne12 == 1 && ++ ne22 == 1; ++ if (glm_compact_multihead) { ++ auto glm_pipeline = ggml_metal_library_get_pipeline_glm_compact_multihead_flash(lib, nwg); ++ GGML_ASSERT(2*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(glm_pipeline)); ++ ++ constexpr size_t glm_smem = ++ GGML_PAD(576*sizeof(uint16_t) + (32 + 4 + 512)*sizeof(float), 16); ++ GGML_ASSERT(glm_smem <= props_dev->max_theadgroup_memory_size); ++ ++ ggml_metal_encoder_set_pipeline(enc, glm_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 3); ++ ggml_metal_encoder_set_buffer (enc, nwg == 1 ? bid_dst : bid_tmp, 4); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, glm_smem, 0); ++ ++ const int grid_x = ne01; ++ const int grid_y = ne02; ++ const int grid_z = ne03*nwg; ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, 2, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=flash_attn_ext kernel=glm_compact_multihead tensor=%s q_type=%s k_type=%s v_type=%s dst_type=%s q_width=%lld v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=32 threads_y=4 nwg=%d smem=%zu\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->type), ++ (long long) ne00, ++ (long long) ne20, ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ (long long) ne11, ++ grid_x, ++ grid_y, ++ grid_z, ++ nwg, ++ glm_smem); ++ } + +- GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size); ++ if (nwg > 1) { ++ GGML_ASSERT(ggml_metal_op_flash_attn_ext_extra_tmp(op) != 0); ++ ggml_metal_op_concurrency_reset(ctx); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src2, 1); +- ggml_metal_encoder_set_buffer (enc, bid_tpe, 2); +- ggml_metal_encoder_set_buffer (enc, bid_ids, 3); ++ const int32_t nrows = ne1*ne2*ne3; ++ ggml_metal_kargs_flash_attn_ext_vec_reduce reduce_args = { ++ nrows, ++ }; ++ auto reduce_pipeline = ++ ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce(lib, op, ne20, nwg); + +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_encoder_set_pipeline(enc, reduce_pipeline); ++ ggml_metal_encoder_set_bytes (enc, &reduce_args, sizeof(reduce_args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, 32*nwg, 1, 1); ++ } + +- ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, ne02, 1, 1); ++ return 1; + } + +- // this barrier is always needed because the next kernel has to wait for the id maps to be computed +- ggml_metal_op_concurrency_reset(ctx); ++ auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg); + +- { +- auto pipeline = ggml_metal_library_get_pipeline_mul_mm_id(lib, op); ++ GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + +- ggml_metal_kargs_mul_mm_id args = { +- /*.ne00 =*/ ne00, +- /*.ne02 =*/ ne02, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne11 =*/ ne11, // n_expert_used (bcast) +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ne20 =*/ ne20, // n_expert_used +- /*.ne21 =*/ ne21, // n_tokens +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.r2 =*/ r2, +- /*.r3 =*/ r3, +- }; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_src2, 3); ++ ggml_metal_encoder_set_buffer (enc, bid_src3, 4); ++ ggml_metal_encoder_set_buffer (enc, bid_src4, 5); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_src1, 2); +- ggml_metal_encoder_set_buffer (enc, bid_tpe, 3); +- ggml_metal_encoder_set_buffer (enc, bid_ids, 4); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 5); ++ const size_t smem = FATTN_SMEM(nsg); + +- const size_t smem = pipeline.smem; ++ //printf("smem: %zu, max: %zu, nsg = %d, nsgmax = %d\n", smem, props_dev->max_theadgroup_memory_size, (int) nsg, (int) nsgmax); ++ GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size); ++ ++ if (nwg == 1) { ++ assert(ggml_metal_op_flash_attn_ext_extra_tmp(op) == 0); ++ ++ // using 1 workgroup -> write the result directly into dst ++ ggml_metal_encoder_set_buffer(enc, bid_pad, 6); ++ ggml_metal_encoder_set_buffer(enc, bid_dst, 7); + + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- ggml_metal_encoder_dispatch_threadgroups(enc, (ne21 + 31)/32, (ne01 + 63)/64, ne02, 128, 1, 1); +- } +- } else { +- auto pipeline = ggml_metal_library_get_pipeline_mul_mv_id(lib, op); ++ const int grid_x = (ne01 + nqptg - 1)/nqptg; ++ const int grid_y = (ne02 + nhptg - 1)/nhptg; ++ const int grid_z = ne03*nwg; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=flash_attn_ext kernel=vec tensor=%s q_type=%s k_type=%s v_type=%s mask_type=%s dst_type=%s q_width=%lld v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=32 threads_y=%lld nwg=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ has_mask ? ggml_type_name(op->src[3]->type) : "none", ++ ggml_type_name(op->type), ++ (long long) ne00, ++ (long long) ne20, ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ (long long) ne11, ++ grid_x, ++ grid_y, ++ grid_z, ++ nsg, ++ nwg); ++ } ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, nsg, 1); ++ } else { ++ // sanity checks ++ assert(ggml_metal_op_flash_attn_ext_extra_tmp(op) != 0); + +- const int nr0 = pipeline.nr0; +- const int nr1 = pipeline.nr1; +- const int nsg = pipeline.nsg; ++ GGML_ASSERT(ne01*ne02*ne03 == ne1*ne2*ne3); ++ GGML_ASSERT((uint64_t)ne1*ne2*ne3 <= (1u << 31)); + +- const size_t smem = pipeline.smem; ++ // write the results from each workgroup into a temp buffer ++ ggml_metal_encoder_set_buffer(enc, bid_pad, 6); ++ ggml_metal_encoder_set_buffer(enc, bid_tmp, 7); + +- ggml_metal_kargs_mul_mv_id args = { +- /*.nei0 =*/ ne20, +- /*.nei1 =*/ ne21, +- /*.nbi1 =*/ nb21, +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.ne10 =*/ ne10, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.ne13 =*/ ne13, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.nb1 =*/ nb1, +- /*.nr0 =*/ nr0, +- }; ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ const int grid_x = (ne01 + nqptg - 1)/nqptg; ++ const int grid_y = (ne02 + nhptg - 1)/nhptg; ++ const int grid_z = ne03*nwg; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=flash_attn_ext kernel=vec tensor=%s q_type=%s k_type=%s v_type=%s mask_type=%s dst_type=%s q_width=%lld v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=32 threads_y=%lld nwg=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ has_mask ? ggml_type_name(op->src[3]->type) : "none", ++ ggml_type_name(op->type), ++ (long long) ne00, ++ (long long) ne20, ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ (long long) ne11, ++ grid_x, ++ grid_y, ++ grid_z, ++ nsg, ++ nwg); ++ } ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, nsg, 1); + +- if (ggml_is_quantized(op->src[0]->type)) { +- GGML_ASSERT(ne00 >= nsg*nr0); +- } ++ // sync the 2 kernels ++ ggml_metal_op_concurrency_reset(ctx); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer(enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer(enc, bid_src1, 2); +- ggml_metal_encoder_set_buffer(enc, bid_dst, 3); +- ggml_metal_encoder_set_buffer(enc, bid_src2, 4); ++ // reduce the results from the workgroups ++ { ++ const int32_t nrows = ne1*ne2*ne3; + +- const int64_t _ne1 = 1; +- const int64_t ne123 = ne20*ne21; ++ ggml_metal_kargs_flash_attn_ext_vec_reduce args0 = { ++ nrows, ++ }; + +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce(lib, op, ne20, nwg); + +- if (op->src[0]->type == GGML_TYPE_F32 || +- op->src[0]->type == GGML_TYPE_F16 || +- op->src[0]->type == GGML_TYPE_BF16 || +- op->src[0]->type == GGML_TYPE_Q8_0) { +- ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nr0 - 1)/(nr0), (_ne1 + nr1 - 1)/nr1, ne123, 32, nsg, 1); +- } else { +- ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nr0*nsg - 1)/(nr0*nsg), (_ne1 + nr1 - 1)/nr1, ne123, 32, nsg, 1); ++ ggml_metal_encoder_set_pipeline(enc, pipeline0); ++ ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, 32*nwg, 1, 1); ++ } + } ++#undef FATTN_SMEM + } + + return 1; + } + +-int ggml_metal_op_add_id(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); +- +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; +- +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- +- GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[2]->type == GGML_TYPE_I32); +- GGML_ASSERT(op->type == GGML_TYPE_F32); ++// Snake activation autofuse: mul -> sin -> sqr -> mul -> add ++static bool ggml_metal_op_can_fuse_snake(ggml_metal_op_t ctx, int idx) { ++ static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; + +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ if (ctx->node(idx)->op != GGML_OP_MUL || !ctx->can_fuse(idx, snake_ops, 5)) { ++ return false; ++ } + +- ggml_metal_kargs_add_id args = { +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb11 =*/ nb11, +- /*.nb21 =*/ nb21, +- }; ++ const ggml_tensor * mul0 = ctx->node(idx + 0); ++ const ggml_tensor * sin_node = ctx->node(idx + 1); ++ const ggml_tensor * sqr = ctx->node(idx + 2); ++ const ggml_tensor * mul1 = ctx->node(idx + 3); ++ const ggml_tensor * add = ctx->node(idx + 4); + +- auto pipeline = ggml_metal_library_get_pipeline_base(lib, GGML_OP_ADD_ID); ++ // x carries the full activation shape, a is the broadcast operand ++ const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; ++ const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 4); ++ // mul1 reads sqr and inv_b in either operand order ++ const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; + +- const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne00); ++ // closure check: the trailing add reads the same x as the leading mul ++ const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0]; + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, 1, nth, 1, 1); ++ // x is in the supported whitelist and every chain intermediate shares x's type. ++ // a and inv_b bind as device const float * in the kernel, so they stay F32. ++ const bool types_ok = ++ (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && ++ (a->type == GGML_TYPE_F32) && (inv_b->type == GGML_TYPE_F32) && ++ (mul0->type == x->type) && (sin_node->type == x->type) && ++ (sqr->type == x->type) && (mul1->type == x->type) && ++ (add->type == x->type); ++ // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D ++ const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; ++ const bool dim_ok = ++ (x->ne[2] == 1) && (x->ne[3] == 1) && ++ (add->ne[2] == 1) && (add->ne[3] == 1) && ++ (a->ne[2] == 1) && (a->ne[3] == 1) && ++ (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1); ++ // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous ++ const bool contig_ok = ++ ggml_is_contiguous(x) && ggml_is_contiguous(add) && ++ ggml_is_contiguous(a) && ggml_is_contiguous(inv_b); + +- return 1; ++ return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x; + } + +-bool ggml_metal_op_flash_attn_ext_use_vec(const ggml_tensor * op) { +- assert(op->op == GGML_OP_FLASH_ATTN_EXT); ++int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { ++ if (ctx->use_fusion && ggml_metal_op_can_fuse_snake(ctx, idx)) { ++ return ggml_metal_op_snake_fused(ctx, idx); ++ } + +- const int64_t ne00 = op->src[0]->ne[0]; // head size +- const int64_t ne01 = op->src[0]->ne[1]; // batch size ++ ggml_tensor * op = ctx->node(idx); + +- // use vec kernel if the batch size is small and if the head size is supported +- return (ne01 < 20) && (ne00 % 32 == 0); +-} ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +-size_t ggml_metal_op_flash_attn_ext_extra_pad(const ggml_tensor * op) { +- assert(op->op == GGML_OP_FLASH_ATTN_EXT); ++ const bool use_fusion = ctx->use_fusion; ++ ++ const int debug_fusion = ctx->debug_fusion; + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); + GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); +- +- size_t res = 0; +- +- const bool has_mask = op->src[3] != nullptr; +- +- // note: the non-vec kernel requires more extra memory, so always reserve for it +- GGML_ASSERT(OP_FLASH_ATTN_EXT_NCPSG >= OP_FLASH_ATTN_EXT_VEC_NCPSG); +- +- //if (ggml_metal_op_flash_attn_ext_use_vec(op)) { +- if (false) { +- // note: always reserve the padding space to avoid graph reallocations +- //const bool has_kvpad = ne11 % OP_FLASH_ATTN_EXT_VEC_NCPSG != 0; +- const bool has_kvpad = true; +- +- if (has_kvpad) { +- res += OP_FLASH_ATTN_EXT_VEC_NCPSG*( +- nb11*ne12*ne13 + +- nb21*ne22*ne23 + +- (has_mask ? ggml_type_size(GGML_TYPE_F16)*ne31*ne32*ne33 : 0)); +- } +- } else { +- //const bool has_kvpad = ne11 % OP_FLASH_ATTN_EXT_NCPSG != 0; +- const bool has_kvpad = true; +- +- if (has_kvpad) { +- res += OP_FLASH_ATTN_EXT_NCPSG*( +- nb11*ne12*ne13 + +- nb21*ne22*ne23 + +- (has_mask ? ggml_type_size(GGML_TYPE_F16)*ne31*ne32*ne33 : 0)); +- } +- } +- +- return res; +-} ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +-size_t ggml_metal_op_flash_attn_ext_extra_blk(const ggml_tensor * op) { +- assert(op->op == GGML_OP_FLASH_ATTN_EXT); ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- //GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- //GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- //GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- //GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); +- //GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[1])); + +- size_t res = 0; ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); + +- const bool has_mask = op->src[3] != nullptr; ++ ggml_metal_kargs_bin args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.offs =*/ 0, ++ /*.o1 =*/ { bid_src1.offs }, ++ }; + +- if (!has_mask) { +- return res; +- } ++ ggml_op fops[8]; + +- const bool is_vec = ggml_metal_op_flash_attn_ext_use_vec(op); ++ int n_fuse = 1; + +- // this optimization is not useful for the vector kernels +- // note: always reserve the blk buffer to avoid graph reallocations +- //if (is_vec) { +- // return res; +- //} ++ // c[0] = add(a, b[0]) ++ // c[1] = add(c[0], b[1]) ++ // c[2] = add(c[1], b[2]) ++ // ... ++ if (use_fusion) { ++ fops[0] = GGML_OP_ADD; ++ fops[1] = GGML_OP_ADD; ++ fops[2] = GGML_OP_ADD; ++ fops[3] = GGML_OP_ADD; ++ fops[4] = GGML_OP_ADD; ++ fops[5] = GGML_OP_ADD; ++ fops[6] = GGML_OP_ADD; ++ fops[7] = GGML_OP_ADD; + +- const int nqptg = is_vec ? OP_FLASH_ATTN_EXT_VEC_NQPSG : OP_FLASH_ATTN_EXT_NQPSG; +- const int ncpsg = is_vec ? OP_FLASH_ATTN_EXT_VEC_NCPSG : OP_FLASH_ATTN_EXT_NCPSG; ++ // note: in metal, we sometimes encode the graph in parallel so we have to avoid fusing ops ++ // across splits. idx_end indicates the last node in the current split ++ for (n_fuse = 0; n_fuse <= 6; ++n_fuse) { ++ if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) { ++ break; ++ } + +- const int64_t ne1 = (ne01 + nqptg - 1)/nqptg; +- const int64_t ne0 = (ne30 + ncpsg - 1)/ncpsg; ++ ggml_tensor * f0 = ctx->node(idx + n_fuse); ++ ggml_tensor * f1 = ctx->node(idx + n_fuse + 1); + +- res += GGML_PAD(ggml_type_size(GGML_TYPE_I8)*ne0*ne1*ne32*ne33, 32); ++ if (f0 != f1->src[0]) { ++ break; ++ } + +- return res; +-} ++ // b[0] === b[1] === ... ++ if (!ggml_are_same_layout(f0->src[1], f1->src[1])) { ++ break; ++ } + +-size_t ggml_metal_op_flash_attn_ext_extra_tmp(const ggml_tensor * op) { +- assert(op->op == GGML_OP_FLASH_ATTN_EXT); ++ // only fuse ops if src1 is in the same Metal buffer ++ ggml_metal_buffer_id bid_fuse = ggml_metal_get_buffer_id(f1->src[1]); ++ if (bid_fuse.metal != bid_src1.metal) { ++ break; ++ } + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- //GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- //GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- //GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); +- //GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); ++ //ctx->fuse_cnt[ops[n_fuse + 1]->op]++; + +- size_t res = 0; ++ args.o1[n_fuse + 1] = bid_fuse.offs; ++ } + +- // note: always reserve the temp buffer to avoid graph reallocations +- //if (ggml_metal_op_flash_attn_ext_use_vec(op)) { +- if (true) { +- const int64_t nwg = 32; +- const int64_t ne01_max = std::min(ne01, 32); ++ ++n_fuse; + +- // temp buffer for writing the results from each workgroup +- // - ne20: the size of the Value head +- // - + 2: the S and M values for each intermediate result +- res += ggml_type_size(GGML_TYPE_F32)*(ne01_max*ne02*ne03*nwg*(ne20 + 2)); ++ if (debug_fusion > 1 && n_fuse > 1) { ++ GGML_LOG_DEBUG("%s: fuse: ADD x %d\n", __func__, n_fuse); ++ } + } + +- return res; +-} ++ // the offsets of src1 and all fused buffers are relative to the start of the src1 buffer ++ bid_src1.offs = 0; + +-int ggml_metal_op_flash_attn_ext(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ struct ggml_metal_pipeline_with_params pipeline; + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ pipeline = ggml_metal_library_get_pipeline_bin(lib, op, n_fuse); + +- const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ if (n_fuse > 1) { ++ bid_dst = ggml_metal_get_buffer_id(ctx->node(idx + n_fuse - 1)); + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); +- GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS( int32_t, nb, op, nb); ++ for (int i = 1; i < n_fuse; ++i) { ++ if (!ggml_metal_op_concurrency_check(ctx, ctx->node(idx + i))) { ++ ggml_metal_op_concurrency_reset(ctx); + +- GGML_ASSERT(ne00 % 4 == 0); ++ break; ++ } ++ } ++ } + +- GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[1]->type == op->src[2]->type); ++ if (pipeline.c4) { ++ args.ne00 = ne00/4; ++ args.ne10 = ne10/4; ++ args.ne0 = ne0/4; ++ } + +- //GGML_ASSERT(ggml_are_same_shape (src1, src2)); +- GGML_ASSERT(ne11 == ne21); +- GGML_ASSERT(ne12 == ne22); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_src1, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 3); + +- GGML_ASSERT(!op->src[3] || op->src[3]->type == GGML_TYPE_F16); +- GGML_ASSERT(!op->src[3] || op->src[3]->ne[1] >= op->src[0]->ne[1] && +- "the Flash-Attention Metal kernel requires the mask to be at least n_queries big"); ++ int grid_x = 0; ++ int grid_y = 0; ++ int grid_z = 0; ++ int nth = 0; ++ if (pipeline.cnt) { ++ grid_x = args.ne0; ++ grid_y = ggml_nrows(op); ++ grid_z = 1; ++ nth = 1; ++ ggml_metal_encoder_dispatch_threadgroups(enc, args.ne0, ggml_nrows(op), 1, 1, 1, 1); ++ } else { ++ const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + +- float scale; +- float max_bias; +- float logit_softcap; ++ nth = 1; + +- memcpy(&scale, ((const int32_t *) op->op_params) + 0, sizeof(scale)); +- memcpy(&max_bias, ((const int32_t *) op->op_params) + 1, sizeof(max_bias)); +- memcpy(&logit_softcap, ((const int32_t *) op->op_params) + 2, sizeof(logit_softcap)); ++ while (2*nth < args.ne0 && nth < nth_max) { ++ nth *= 2; ++ } + +- if (logit_softcap != 0.0f) { +- scale /= logit_softcap; ++ grid_x = ne01; ++ grid_y = ne02; ++ grid_z = ne03; ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); + } + +- const bool has_mask = op->src[3] != NULL; +- const bool has_sinks = op->src[4] != NULL; +- const bool has_bias = max_bias != 0.0f; +- const bool has_scap = logit_softcap != 0.0f; +- +- const uint32_t n_head = op->src[0]->ne[2]; +- const int32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head)); ++ if (ggml_metal_glm_dsa_dispatch_log_enabled() && ++ ggml_metal_tensor_name_contains(op, "ffn_moe_down_weighted_input")) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=bin kernel=%s tensor=%s src0=%s src1=%s src0_type=%s src1_type=%s dst_type=%s ne0=%d rows=%lld n_fuse=%d c4=%d cnt=%d grid_x=%d grid_y=%d grid_z=%d threads_x=%d\n", ++ ggml_op_name(op->op), ++ ggml_metal_tensor_name(op), ++ ggml_metal_tensor_name(op->src[0]), ++ ggml_metal_tensor_name(op->src[1]), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->type), ++ ne0, ++ (long long) ggml_nrows(op), ++ n_fuse, ++ pipeline.c4 ? 1 : 0, ++ pipeline.cnt ? 1 : 0, ++ grid_x, ++ grid_y, ++ grid_z, ++ nth); ++ } + +- const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); +- const float m1 = powf(2.0f, -(max_bias / 2.0f) / n_head_log2); ++ return n_fuse; ++} + +- GGML_ASSERT(ne01 < 65536); ++int ggml_metal_op_l2_norm(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); + +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); +- ggml_metal_buffer_id bid_src2 = ggml_metal_get_buffer_id(op->src[2]); +- ggml_metal_buffer_id bid_src3 = has_mask ? ggml_metal_get_buffer_id(op->src[3]) : bid_src0; +- ggml_metal_buffer_id bid_src4 = has_sinks ? ggml_metal_get_buffer_id(op->src[4]) : bid_src0; ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ if (ggml_metal_glm_lightning_indexer_noop_enabled()) { ++ auto pipeline = ggml_metal_library_get_pipeline_zero_f32(lib); ++ const int nth = 256; ++ const int64_t n_tg = (ggml_nelements(op) + nth - 1)/nth; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, n_tg, 1, 1, nth, 1, 1); ++ return 1; ++ } + +- ggml_metal_buffer_id bid_pad = bid_dst; +- bid_pad.offs += ggml_nbytes(op); ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- ggml_metal_buffer_id bid_blk = bid_pad; +- bid_blk.offs += ggml_metal_op_flash_attn_ext_extra_pad(op); ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); + +- ggml_metal_buffer_id bid_tmp = bid_blk; +- bid_tmp.offs += ggml_metal_op_flash_attn_ext_extra_blk(op); ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); + +- if (!ggml_metal_op_flash_attn_ext_use_vec(op)) { +- // half8x8 kernel +- const int nqptg = OP_FLASH_ATTN_EXT_NQPSG; // queries per threadgroup +- const int ncpsg = OP_FLASH_ATTN_EXT_NCPSG; // cache values per simdgroup ++ float eps; ++ memcpy(&eps, op->op_params, sizeof(float)); + +- GGML_ASSERT(nqptg <= 32); +- GGML_ASSERT(nqptg % 8 == 0); +- GGML_ASSERT(ncpsg % 32 == 0); ++ ggml_metal_kargs_l2_norm args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.eps =*/ eps, ++ }; + +- bool need_sync = false; ++ auto pipeline = ggml_metal_library_get_pipeline_l2_norm(lib, op); + +- const bool has_kvpad = ne11 % ncpsg != 0; ++ if (pipeline.c4) { ++ args.ne00 = ne00/4; ++ args.ne0 = ne0/4; ++ } + +- if (has_kvpad) { +- assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0); ++ int nth = 32; // SIMD width + +- ggml_metal_kargs_flash_attn_ext_pad args0 = { +- /*.ne11 =*/ne11, +- /*.ne_12_2 =*/ne12, +- /*.ne_12_3 =*/ne13, +- /*.nb11 =*/nb11, +- /*.nb12 =*/nb12, +- /*.nb13 =*/nb13, +- /*.nb21 =*/nb21, +- /*.nb22 =*/nb22, +- /*.nb23 =*/nb23, +- /*.ne31 =*/ne31, +- /*.ne32 =*/ne32, +- /*.ne33 =*/ne33, +- /*.nb31 =*/nb31, +- /*.nb32 =*/nb32, +- /*.nb33 =*/nb33, +- }; ++ while (nth < ne00 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ nth *= 2; ++ } + +- auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_pad(lib, op, has_mask, ncpsg); ++ nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + +- ggml_metal_encoder_set_pipeline(enc, pipeline0); +- ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src1, 1); +- ggml_metal_encoder_set_buffer (enc, bid_src2, 2); +- ggml_metal_encoder_set_buffer (enc, bid_src3, 3); +- ggml_metal_encoder_set_buffer (enc, bid_pad, 4); ++ const size_t smem = pipeline.smem; + +- assert(ne12 == ne22); +- assert(ne13 == ne23); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ncpsg, std::max(ne12, ne32), std::max(ne13, ne33), 32, 1, 1); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- need_sync = true; +- } ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); + +- if (has_mask) { +- assert(ggml_metal_op_flash_attn_ext_extra_blk(op) != 0); ++ return 1; ++} + +- ggml_metal_kargs_flash_attn_ext_blk args0 = { +- /*.ne01 =*/ ne01, +- /*.ne30 =*/ ne30, +- /*.ne31 =*/ ne31, +- /*.ne32 =*/ ne32, +- /*.ne33 =*/ ne33, +- /*.nb31 =*/ nb31, +- /*.nb32 =*/ nb32, +- /*.nb33 =*/ nb33, +- }; ++int ggml_metal_op_group_norm(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); + +- auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_blk(lib, op, nqptg, ncpsg); ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +- 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_blk, 2); ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const int32_t nblk1 = ((ne01 + nqptg - 1)/nqptg); +- const int32_t nblk0 = ((ne30 + ncpsg - 1)/ncpsg); ++ const int32_t ngrp = ((const int32_t *) op->op_params)[0]; + +- ggml_metal_encoder_dispatch_threadgroups(enc, nblk0, nblk1, ne32*ne33, 32, 1, 1); ++ float eps; ++ memcpy(&eps, op->op_params + 1, sizeof(float)); + +- need_sync = true; +- } ++ ggml_metal_kargs_group_norm args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ngrp =*/ ngrp, ++ /*.eps =*/ eps, ++ }; + +- if (need_sync) { +- ggml_metal_op_concurrency_reset(ctx); +- } ++ auto pipeline = ggml_metal_library_get_pipeline_group_norm(lib, op); + +- const int is_q = ggml_is_quantized(op->src[1]->type) ? 1 : 0; ++ int nth = 32; // SIMD width ++ //while (nth < ne00/4 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ // nth *= 2; ++ //} + +- // 2*(2*ncpsg) +- // ncpsg soft_max values + ncpsg mask values +- // +- // 16*32*(nsg) +- // the shared memory needed for the simdgroups to load the KV cache +- // each thread loads (dequantizes) 16 head elements, there are 32 threads in th SG +- // +-#define FATTN_SMEM(nsg) (GGML_PAD((nqptg*(ne00 + 2*GGML_PAD(ne20, 64) + 2*(2*ncpsg)) + is_q*(16*32*(nsg)))*(sizeof(float)/2), 16)) ++ //nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ //nth = std::min(nth, ne00/4); + +- //int64_t nsgmax = 4; +- // +- //if (is_q) { +- // nsgmax = 2; +- // while (true) { +- // const size_t smem = FATTN_SMEM(nsgmax); +- // if (smem > props_dev->max_theadgroup_memory_size) { +- // break; +- // } +- // nsgmax *= 2; +- // } +- // nsgmax /= 2; +- //} ++ const size_t smem = pipeline.smem; + +- // simdgroups per threadgroup (a.k.a. warps) +- //nsg = ne01 <= nqptg ? MAX(4, MIN(nsgmax, MIN(ne11/ncpsg, (int64_t) pipeline.maxTotalThreadsPerThreadgroup/32))) : 4; +- int32_t nsg = ne00 >= 512 ? 8 : 4; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- const size_t smem = FATTN_SMEM(nsg); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- ggml_metal_kargs_flash_attn_ext args = { +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne11 =*/ ne11, +- /*.ne_12_2 =*/ ne12, +- /*.ne_12_3 =*/ ne13, +- /*.ns10 =*/ int32_t(nb11/nb10), +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ns20 =*/ int32_t(nb21/nb20), +- /*.nb21 =*/ nb21, +- /*.nb22 =*/ nb22, +- /*.nb23 =*/ nb23, +- /*.ne31 =*/ ne31, +- /*.ne32 =*/ ne32, +- /*.ne33 =*/ ne33, +- /*.nb31 =*/ nb31, +- /*.nb32 =*/ nb32, +- /*.nb33 =*/ nb33, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.scale =*/ scale, +- /*.max_bias =*/ max_bias, +- /*.m0 =*/ m0, +- /*.m1 =*/ m1, +- /*.n_head_log2 =*/ n_head_log2, +- /*.logit_softcap =*/ logit_softcap, +- }; ++ ggml_metal_encoder_dispatch_threadgroups(enc, ngrp, 1, 1, nth, 1, 1); + +- auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg); ++ return 1; ++} + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_src1, 2); +- ggml_metal_encoder_set_buffer (enc, bid_src2, 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, bid_pad, 6); +- ggml_metal_encoder_set_buffer (enc, bid_blk, 7); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 8); ++int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); + +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +- ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nqptg - 1)/nqptg, ne02, ne03, 32, nsg, 1); +-#undef FATTN_SMEM +- } else { +- // half4x4 kernel +- const int nqptg = OP_FLASH_ATTN_EXT_VEC_NQPSG; // 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 ++ const bool use_fusion = ctx->use_fusion; + +- GGML_ASSERT(nqptg <= 32); +- GGML_ASSERT(nqptg % 1 == 0); +- GGML_ASSERT(ncpsg % 32 == 0); ++ const int debug_fusion = ctx->debug_fusion; + +- bool need_sync = false; ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const bool has_kvpad = ne11 % ncpsg != 0; ++ float eps; ++ memcpy(&eps, op->op_params, sizeof(float)); + +- if (has_kvpad) { +- assert(ggml_metal_op_flash_attn_ext_extra_pad(op) != 0); ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); + +- ggml_metal_kargs_flash_attn_ext_pad args0 = { +- /*.ne11 =*/ne11, +- /*.ne_12_2 =*/ne12, +- /*.ne_12_3 =*/ne13, +- /*.nb11 =*/nb11, +- /*.nb12 =*/nb12, +- /*.nb13 =*/nb13, +- /*.nb21 =*/nb21, +- /*.nb22 =*/nb22, +- /*.nb23 =*/nb23, +- /*.ne31 =*/ne31, +- /*.ne32 =*/ne32, +- /*.ne33 =*/ne33, +- /*.nb31 =*/nb31, +- /*.nb32 =*/nb32, +- /*.nb33 =*/nb33, +- }; ++ ggml_metal_kargs_norm args = { ++ /*.ne00 =*/ ne00, ++ /*.ne00_t =*/ ne00 % 4 == 0 ? ne00/4 : ne00, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.eps =*/ eps, ++ /*.nef1 =*/ { ne01 }, ++ /*.nef2 =*/ { ne02 }, ++ /*.nef3 =*/ { ne03 }, ++ /*.nbf1 =*/ { nb01 }, ++ /*.nbf2 =*/ { nb02 }, ++ /*.nbf3 =*/ { nb03 }, ++ }; + +- auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_pad(lib, op, has_mask, ncpsg); ++ ggml_op fops[8]; + +- ggml_metal_encoder_set_pipeline(enc, pipeline0); +- ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src1, 1); +- ggml_metal_encoder_set_buffer (enc, bid_src2, 2); +- ggml_metal_encoder_set_buffer (enc, bid_src3, 3); +- ggml_metal_encoder_set_buffer (enc, bid_pad, 4); ++ int n_fuse = 1; + +- assert(ne12 == ne22); +- assert(ne13 == ne23); ++ ggml_metal_buffer_id bid_fuse[2] = { bid_src0, bid_src0 }; + +- ggml_metal_encoder_dispatch_threadgroups(enc, ncpsg, std::max(ne12, ne32), std::max(ne13, ne33), 32, 1, 1); ++ // d[0] = norm(a) ++ // d[1] = mul(d[0], b) ++ // d[2] = add(d[1], c) ++ if (use_fusion) { ++ fops[0] = op->op; ++ fops[1] = GGML_OP_MUL; ++ fops[2] = GGML_OP_ADD; + +- need_sync = true; +- } ++ for (n_fuse = 0; n_fuse <= 1; ++n_fuse) { ++ if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) { ++ break; ++ } + +- if (need_sync) { +- ggml_metal_op_concurrency_reset(ctx); +- } ++ ggml_tensor * f0 = ctx->node(idx + n_fuse); ++ ggml_tensor * f1 = ctx->node(idx + n_fuse + 1); + +- // note: for simplicity assume the K is larger or equal than V +- GGML_ASSERT(ne10 >= ne20); ++ if (f0 != f1->src[0]) { ++ break; ++ } + +- // ne00 + 2*ncpsg*(nsg) +- // for each query, we load it as f16 in shared memory (ne00) +- // and store the soft_max values and the mask +- // +- // ne20*(nsg) +- // each simdgroup has a full f32 head vector in shared mem to accumulate results +- // +-#define FATTN_SMEM(nsg) (GGML_PAD(((GGML_PAD(ne00, 128) + 4*ncpsg + 2*GGML_PAD(ne20, 128))*(nsg))*(sizeof(float)/2), 16)) ++ if (f1->src[1]->ne[0] != op->ne[0]) { ++ break; ++ } + +- int64_t nsg = 1; ++ if (!ggml_is_contiguous_rows(f1->src[1])) { ++ break; ++ } + +- // workgroups +- // each workgroup handles nsg*nkpsg cache values +- int32_t nwg = 1; +- if (false) { +- // for small KV caches, we could launch a single workgroup and write the results directly to dst/ +- // however, this does not lead to significant improvement, so disabled +- nwg = 1; +- nsg = 4; +- } else { +- nwg = 32; +- nsg = 1; +- while (2*nwg*nsg*ncpsg < ne11 && nsg < 4) { +- nsg *= 2; ++ if (f1->type != GGML_TYPE_F32) { ++ break; + } ++ ++ //ctx->fuse_cnt[f1->op]++; ++ ++ bid_fuse[n_fuse] = ggml_metal_get_buffer_id(f1->src[1]); ++ ++ args.nef1[n_fuse + 1] = f1->src[1]->ne[1]; ++ args.nef2[n_fuse + 1] = f1->src[1]->ne[2]; ++ args.nef3[n_fuse + 1] = f1->src[1]->ne[3]; ++ ++ args.nbf1[n_fuse + 1] = f1->src[1]->nb[1]; ++ args.nbf2[n_fuse + 1] = f1->src[1]->nb[2]; ++ args.nbf3[n_fuse + 1] = f1->src[1]->nb[3]; + } + +- ggml_metal_kargs_flash_attn_ext_vec args = { +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne11 =*/ ne11, +- /*.ne_12_2 =*/ ne12, +- /*.ne_12_3 =*/ ne13, +- /*.ns10 =*/ int32_t(nb11/nb10), +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ns20 =*/ int32_t(nb21/nb20), +- /*.nb21 =*/ nb21, +- /*.nb22 =*/ nb22, +- /*.nb23 =*/ nb23, +- /*.ne31 =*/ ne31, +- /*.ne32 =*/ ne32, +- /*.ne33 =*/ ne33, +- /*.nb31 =*/ nb31, +- /*.nb32 =*/ nb32, +- /*.nb33 =*/ nb33, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.scale =*/ scale, +- /*.max_bias =*/ max_bias, +- /*.m0 =*/ m0, +- /*.m1 =*/ m1, +- /*.n_head_log2 =*/ n_head_log2, +- /*.logit_softcap =*/ logit_softcap, +- }; ++ ++n_fuse; + +- auto pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec(lib, op, has_mask, has_sinks, has_bias, has_scap, has_kvpad, nsg, nwg); ++ if (debug_fusion > 1 && n_fuse > 1) { ++ if (n_fuse == 2) { ++ GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op)); ++ } ++ if (n_fuse == 3) { ++ GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op)); ++ } ++ } ++ } + +- GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ if (n_fuse > 1) { ++ bid_dst = ggml_metal_get_buffer_id(ctx->node(idx + n_fuse - 1)); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_src1, 2); +- ggml_metal_encoder_set_buffer (enc, bid_src2, 3); +- ggml_metal_encoder_set_buffer (enc, bid_src3, 4); +- ggml_metal_encoder_set_buffer (enc, bid_src4, 5); ++ for (int i = 1; i < n_fuse; ++i) { ++ if (!ggml_metal_op_concurrency_check(ctx, ctx->node(idx + i))) { ++ ggml_metal_op_concurrency_reset(ctx); + +- const size_t smem = FATTN_SMEM(nsg); ++ break; ++ } ++ } ++ } + +- //printf("smem: %zu, max: %zu, nsg = %d, nsgmax = %d\n", smem, props_dev->max_theadgroup_memory_size, (int) nsg, (int) nsgmax); +- GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size); ++ auto pipeline = ggml_metal_library_get_pipeline_norm(lib, op, n_fuse); + +- if (nwg == 1) { +- assert(ggml_metal_op_flash_attn_ext_extra_tmp(op) == 0); ++ int nth = 32; // SIMD width + +- // using 1 workgroup -> write the result directly into dst +- ggml_metal_encoder_set_buffer(enc, bid_pad, 6); +- ggml_metal_encoder_set_buffer(enc, bid_dst, 7); ++ while (nth < args.ne00_t && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ nth *= 2; ++ } + +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ nth = std::min(nth, args.ne00_t); + +- ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nqptg - 1)/nqptg, (ne02 + nhptg - 1)/nhptg, ne03*nwg, 32, nsg, 1); +- } else { +- // sanity checks +- assert(ggml_metal_op_flash_attn_ext_extra_tmp(op) != 0); ++ const size_t smem = pipeline.smem; + +- GGML_ASSERT(ne01*ne02*ne03 == ne1*ne2*ne3); +- GGML_ASSERT((uint64_t)ne1*ne2*ne3 <= (1u << 31)); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_fuse[0], 2); ++ ggml_metal_encoder_set_buffer (enc, bid_fuse[1], 3); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 4); + +- // write the results from each workgroup into a temp buffer +- ggml_metal_encoder_set_buffer(enc, bid_pad, 6); +- ggml_metal_encoder_set_buffer(enc, bid_tmp, 7); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); +- ggml_metal_encoder_dispatch_threadgroups(enc, (ne01 + nqptg - 1)/nqptg, (ne02 + nhptg - 1)/nhptg, ne03*nwg, 32, nsg, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); + +- // sync the 2 kernels +- ggml_metal_op_concurrency_reset(ctx); ++ return n_fuse; ++} + +- // reduce the results from the workgroups +- { +- const int32_t nrows = ne1*ne2*ne3; ++int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); + +- ggml_metal_kargs_flash_attn_ext_vec_reduce args0 = { +- nrows, +- }; ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +- auto pipeline0 = ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce(lib, op, ne20, nwg); ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- ggml_metal_encoder_set_pipeline(enc, pipeline0); +- ggml_metal_encoder_set_bytes (enc, &args0, sizeof(args0), 0); +- ggml_metal_encoder_set_buffer (enc, bid_tmp, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ // make sure we have one or more position id(ne10) per token(ne02) ++ GGML_ASSERT(ne10 % ne02 == 0); ++ GGML_ASSERT(ne10 >= ne02); + +- ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, 32*nwg, 1, 1); +- } +- } +-#undef FATTN_SMEM +- } ++ const int nth = std::min(1024, ne00); + +- return 1; +-} ++ const int n_past = ((const int32_t *) op->op_params)[0]; ++ const int n_dims = ((const int32_t *) op->op_params)[1]; ++ //const int mode = ((const int32_t *) op->op_params)[2]; ++ // skip 3, n_ctx, used in GLM RoPE, unimplemented in metal ++ const int n_ctx_orig = ((const int32_t *) op->op_params)[4]; + +-// Snake activation autofuse: mul -> sin -> sqr -> mul -> add +-static bool ggml_metal_op_can_fuse_snake(ggml_metal_op_t ctx, int idx) { +- static constexpr ggml_op snake_ops[5] = { GGML_OP_MUL, GGML_OP_SIN, GGML_OP_SQR, GGML_OP_MUL, GGML_OP_ADD }; ++ float freq_base; ++ float freq_scale; ++ float ext_factor; ++ float attn_factor; ++ float beta_fast; ++ float beta_slow; + +- if (ctx->node(idx)->op != GGML_OP_MUL || !ctx->can_fuse(idx, snake_ops, 5)) { +- return false; +- } ++ memcpy(&freq_base, (const int32_t *) op->op_params + 5, sizeof(float)); ++ memcpy(&freq_scale, (const int32_t *) op->op_params + 6, sizeof(float)); ++ memcpy(&ext_factor, (const int32_t *) op->op_params + 7, sizeof(float)); ++ memcpy(&attn_factor, (const int32_t *) op->op_params + 8, sizeof(float)); ++ memcpy(&beta_fast, (const int32_t *) op->op_params + 9, sizeof(float)); ++ memcpy(&beta_slow, (const int32_t *) op->op_params + 10, sizeof(float)); + +- const ggml_tensor * mul0 = ctx->node(idx + 0); +- const ggml_tensor * sin_node = ctx->node(idx + 1); +- const ggml_tensor * sqr = ctx->node(idx + 2); +- const ggml_tensor * mul1 = ctx->node(idx + 3); +- const ggml_tensor * add = ctx->node(idx + 4); ++ // mrope ++ const int sect_0 = ((const int32_t *) op->op_params)[11]; ++ const int sect_1 = ((const int32_t *) op->op_params)[12]; ++ const int sect_2 = ((const int32_t *) op->op_params)[13]; ++ const int sect_3 = ((const int32_t *) op->op_params)[14]; + +- // x carries the full activation shape, a is the broadcast operand +- const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; +- const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; ++ ggml_metal_kargs_rope args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.n_past =*/ n_past, ++ /*.n_dims =*/ n_dims, ++ /*.n_ctx_orig =*/ n_ctx_orig, ++ /*.freq_base =*/ freq_base, ++ /*.freq_scale =*/ freq_scale, ++ /*.ext_factor =*/ ext_factor, ++ /*.attn_factor =*/ attn_factor, ++ /*.beta_fast =*/ beta_fast, ++ /*.beta_slow =*/ beta_slow, ++ /* sect_0 =*/ sect_0, ++ /* sect_1 =*/ sect_1, ++ /* sect_2 =*/ sect_2, ++ /* sect_3 =*/ sect_3, ++ /* src2 =*/ op->src[2] != nullptr, ++ }; + +- // mul1 reads sqr and inv_b in either operand order +- const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; ++ auto pipeline = ggml_metal_library_get_pipeline_rope(lib, op); + +- // closure check: the trailing add reads the same x as the leading mul +- const ggml_tensor * x_in_add = (add->src[0] == mul1) ? add->src[1] : add->src[0]; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ if (op->src[2]) { ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); ++ } else { ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 3); ++ } ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 4); + +- // x is in the supported whitelist and every chain intermediate shares x's type. +- // a and inv_b bind as device const float * in the kernel, so they stay F32. +- const bool types_ok = +- (x->type == GGML_TYPE_F32 || x->type == GGML_TYPE_F16 || x->type == GGML_TYPE_BF16) && +- (a->type == GGML_TYPE_F32) && (inv_b->type == GGML_TYPE_F32) && +- (mul0->type == x->type) && (sin_node->type == x->type) && +- (sqr->type == x->type) && (mul1->type == x->type) && +- (add->type == x->type); +- // a / inv_b collapse to [1, C, 1, 1], x and add stay 2D +- const bool shape_ok = ggml_are_same_shape(a, inv_b) && a->ne[0] == 1 && a->ne[1] == x->ne[1]; +- const bool dim_ok = +- (x->ne[2] == 1) && (x->ne[3] == 1) && +- (add->ne[2] == 1) && (add->ne[3] == 1) && +- (a->ne[2] == 1) && (a->ne[3] == 1) && +- (inv_b->ne[2] == 1) && (inv_b->ne[3] == 1); +- // kernel reads x[idx] and a[c] / inv_b[c] linearly, so every operand is contiguous +- const bool contig_ok = +- ggml_is_contiguous(x) && ggml_is_contiguous(add) && +- ggml_is_contiguous(a) && ggml_is_contiguous(inv_b); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); + +- return types_ok && shape_ok && dim_ok && contig_ok && x_in_add == x; ++ return 1; + } + +-int ggml_metal_op_bin(ggml_metal_op_t ctx, int idx) { +- if (ctx->use_fusion && ggml_metal_op_can_fuse_snake(ctx, idx)) { +- return ggml_metal_op_snake_fused(ctx, idx); +- } +- ++int ggml_metal_op_im2col(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- const bool use_fusion = ctx->use_fusion; +- +- const int debug_fusion = ctx->debug_fusion; +- + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); +- +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[1])); +- +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_src1 = ggml_metal_get_buffer_id(op->src[1]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); +- +- ggml_metal_kargs_bin args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne10 =*/ ne10, +- /*.ne11 =*/ ne11, +- /*.ne12 =*/ ne12, +- /*.ne13 =*/ ne13, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.offs =*/ 0, +- /*.o1 =*/ { bid_src1.offs }, +- }; ++ const int32_t s0 = ((const int32_t *)(op->op_params))[0]; ++ const int32_t s1 = ((const int32_t *)(op->op_params))[1]; ++ const int32_t p0 = ((const int32_t *)(op->op_params))[2]; ++ const int32_t p1 = ((const int32_t *)(op->op_params))[3]; ++ const int32_t d0 = ((const int32_t *)(op->op_params))[4]; ++ const int32_t d1 = ((const int32_t *)(op->op_params))[5]; + +- ggml_op fops[8]; ++ const bool is_2D = ((const int32_t *)(op->op_params))[6] == 1; + +- int n_fuse = 1; ++ const int32_t N = op->src[1]->ne[is_2D ? 3 : 2]; ++ const int32_t IC = op->src[1]->ne[is_2D ? 2 : 1]; ++ const int32_t IH = is_2D ? op->src[1]->ne[1] : 1; ++ const int32_t IW = op->src[1]->ne[0]; + +- // c[0] = add(a, b[0]) +- // c[1] = add(c[0], b[1]) +- // c[2] = add(c[1], b[2]) +- // ... +- if (use_fusion) { +- fops[0] = GGML_OP_ADD; +- fops[1] = GGML_OP_ADD; +- fops[2] = GGML_OP_ADD; +- fops[3] = GGML_OP_ADD; +- fops[4] = GGML_OP_ADD; +- fops[5] = GGML_OP_ADD; +- fops[6] = GGML_OP_ADD; +- fops[7] = GGML_OP_ADD; ++ const int32_t KH = is_2D ? op->src[0]->ne[1] : 1; ++ const int32_t KW = op->src[0]->ne[0]; + +- // note: in metal, we sometimes encode the graph in parallel so we have to avoid fusing ops +- // across splits. idx_end indicates the last node in the current split +- for (n_fuse = 0; n_fuse <= 6; ++n_fuse) { +- if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) { +- break; +- } ++ const int32_t OH = is_2D ? op->ne[2] : 1; ++ const int32_t OW = op->ne[1]; + +- ggml_tensor * f0 = ctx->node(idx + n_fuse); +- ggml_tensor * f1 = ctx->node(idx + n_fuse + 1); ++ const int32_t CHW = IC * KH * KW; + +- if (f0 != f1->src[0]) { +- break; +- } ++ const uint64_t ofs0 = op->src[1]->nb[is_2D ? 3 : 2] / 4; ++ const uint64_t ofs1 = op->src[1]->nb[is_2D ? 2 : 1] / 4; + +- // b[0] === b[1] === ... +- if (!ggml_are_same_layout(f0->src[1], f1->src[1])) { +- break; +- } ++ ggml_metal_kargs_im2col args = { ++ /*.ofs0 =*/ ofs0, ++ /*.ofs1 =*/ ofs1, ++ /*.IW =*/ IW, ++ /*.IH =*/ IH, ++ /*.CHW =*/ CHW, ++ /*.s0 =*/ s0, ++ /*.s1 =*/ s1, ++ /*.p0 =*/ p0, ++ /*.p1 =*/ p1, ++ /*.d0 =*/ d0, ++ /*.d1 =*/ d1, ++ /*.N =*/ N, ++ /*.KH =*/ KH, ++ /*.KW =*/ KW, ++ /*.KHW =*/ KH * KW, ++ }; + +- // only fuse ops if src1 is in the same Metal buffer +- ggml_metal_buffer_id bid_fuse = ggml_metal_get_buffer_id(f1->src[1]); +- if (bid_fuse.metal != bid_src1.metal) { +- break; +- } ++ auto pipeline = ggml_metal_library_get_pipeline_im2col(lib, op); + +- //ctx->fuse_cnt[ops[n_fuse + 1]->op]++; ++ if (KH*KW <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ const uint64_t ntptg0 = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)/(KH*KW), N); + +- args.o1[n_fuse + 1] = bid_fuse.offs; +- } ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- ++n_fuse; ++ ggml_metal_encoder_dispatch_threadgroups(enc, IC, OH, OW, ntptg0, KH, KW); ++ } else { ++ const uint64_t n_threads = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), N); ++ const int64_t quotient = N / n_threads + (N % n_threads > 0 ? 1 : 0); + +- if (debug_fusion > 1 && n_fuse > 1) { +- GGML_LOG_DEBUG("%s: fuse: ADD x %d\n", __func__, n_fuse); +- } ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ++ ggml_metal_encoder_dispatch_threadgroups(enc, quotient * CHW, OH, OW, n_threads, 1, 1); + } + +- // the offsets of src1 and all fused buffers are relative to the start of the src1 buffer +- bid_src1.offs = 0; ++ return 1; ++} + +- struct ggml_metal_pipeline_with_params pipeline; ++int ggml_metal_op_conv_2d(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); + +- pipeline = ggml_metal_library_get_pipeline_bin(lib, op, n_fuse); ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +- if (n_fuse > 1) { +- bid_dst = ggml_metal_get_buffer_id(ctx->node(idx + n_fuse - 1)); ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- for (int i = 1; i < n_fuse; ++i) { +- if (!ggml_metal_op_concurrency_check(ctx, ctx->node(idx + i))) { +- ggml_metal_op_concurrency_reset(ctx); ++ GGML_ASSERT(ggml_is_contiguous(op->src[0])); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); + +- break; +- } +- } +- } ++ const int32_t s0 = ((const int32_t *) op->op_params)[0]; ++ const int32_t s1 = ((const int32_t *) op->op_params)[1]; ++ const int32_t p0 = ((const int32_t *) op->op_params)[2]; ++ const int32_t p1 = ((const int32_t *) op->op_params)[3]; ++ const int32_t d0 = ((const int32_t *) op->op_params)[4]; ++ const int32_t d1 = ((const int32_t *) op->op_params)[5]; + +- if (pipeline.c4) { +- args.ne00 = ne00/4; +- args.ne10 = ne10/4; +- args.ne0 = ne0/4; +- } ++ ggml_metal_kargs_conv_2d args = { ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.IW =*/ ne10, ++ /*.IH =*/ ne11, ++ /*.KW =*/ ne00, ++ /*.KH =*/ ne01, ++ /*.IC =*/ ne02, ++ /*.OC =*/ ne03, ++ /*.OW =*/ ne0, ++ /*.OH =*/ ne1, ++ /*.N =*/ ne3, ++ /*.s0 =*/ s0, ++ /*.s1 =*/ s1, ++ /*.p0 =*/ p0, ++ /*.p1 =*/ p1, ++ /*.d0 =*/ d0, ++ /*.d1 =*/ d1, ++ }; + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_src1, 2); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 3); ++ auto pipeline = ggml_metal_library_get_pipeline_conv_2d(lib, op); + +- if (pipeline.cnt) { +- ggml_metal_encoder_dispatch_threadgroups(enc, args.ne0, ggml_nrows(op), 1, 1, 1, 1); +- } else { +- const int nth_max = MIN(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ int nth = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); ++ nth = std::min(nth, 256); ++ nth = std::max(nth, 1); + +- int nth = 1; ++ const uint64_t n_out = ggml_nelements(op); + +- while (2*nth < args.ne0 && nth < nth_max) { +- nth *= 2; +- } ++ uint64_t tg = (n_out + nth - 1)/nth; ++ tg = std::max(tg, 1); ++ tg = std::min(tg, (uint64_t) std::numeric_limits::max()); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); +- } ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- return n_fuse; ++ ggml_metal_encoder_dispatch_threadgroups(enc, tg, 1, 1, nth, 1, 1); ++ ++ return 1; + } + +-int ggml_metal_op_l2_norm(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_conv_2d_dw(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -3295,352 +10313,285 @@ int ggml_metal_op_l2_norm(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); +- +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->type == GGML_TYPE_F32); ++ GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); + +- float eps; +- memcpy(&eps, op->op_params, sizeof(float)); ++ const int32_t s0 = ((const int32_t *) op->op_params)[0]; ++ const int32_t s1 = ((const int32_t *) op->op_params)[1]; ++ const int32_t p0 = ((const int32_t *) op->op_params)[2]; ++ const int32_t p1 = ((const int32_t *) op->op_params)[3]; ++ const int32_t d0 = ((const int32_t *) op->op_params)[4]; ++ const int32_t d1 = ((const int32_t *) op->op_params)[5]; + +- ggml_metal_kargs_l2_norm args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.eps =*/ eps, ++ ggml_metal_kargs_conv_2d_dw args = { ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb03, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.IW =*/ ne10, ++ /*.IH =*/ ne11, ++ /*.KW =*/ ne00, ++ /*.KH =*/ ne01, ++ /*.C =*/ ne12, ++ /*.OW =*/ ne0, ++ /*.OH =*/ ne1, ++ /*.N =*/ ne13, ++ /*.s0 =*/ s0, ++ /*.s1 =*/ s1, ++ /*.p0 =*/ p0, ++ /*.p1 =*/ p1, ++ /*.d0 =*/ d0, ++ /*.d1 =*/ d1, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_l2_norm(lib, op); +- +- if (pipeline.c4) { +- args.ne00 = ne00/4; +- args.ne0 = ne0/4; +- } ++ const bool use_tiled = (nb12 < nb10); + +- int nth = 32; // SIMD width ++ auto pipeline = ggml_metal_library_get_pipeline_conv_2d_dw(lib, op, use_tiled); + +- while (nth < ne00 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- nth *= 2; +- } ++ int nth = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); ++ nth = std::min(nth, 256); ++ nth = std::max(nth, 1); + +- nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ const int32_t OW = ne0; ++ const int32_t OH = ne1; ++ const int32_t C = ne12; ++ const int32_t N = ne13; + +- const size_t smem = pipeline.smem; ++ const int tg_x = use_tiled ? (C + nth - 1) / nth : (OW + nth - 1) / nth; ++ const int tg_y = OH; ++ const int tg_z = use_tiled ? OW * N : C * N; + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, tg_x, tg_y, tg_z, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_group_norm(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_conv_3d(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ // 1. Extract standard dimensions and byte strides + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const int32_t ngrp = ((const int32_t *) op->op_params)[0]; +- +- float eps; +- memcpy(&eps, op->op_params + 1, sizeof(float)); ++ // 2. Extract hyperparams from op_params ++ const int32_t s0 = ((const int32_t *)(op->op_params))[0]; ++ const int32_t s1 = ((const int32_t *)(op->op_params))[1]; ++ const int32_t s2 = ((const int32_t *)(op->op_params))[2]; ++ const int32_t p0 = ((const int32_t *)(op->op_params))[3]; ++ const int32_t p1 = ((const int32_t *)(op->op_params))[4]; ++ const int32_t p2 = ((const int32_t *)(op->op_params))[5]; ++ const int32_t d0 = ((const int32_t *)(op->op_params))[6]; ++ const int32_t d1 = ((const int32_t *)(op->op_params))[7]; ++ const int32_t d2 = ((const int32_t *)(op->op_params))[8]; ++ const int32_t IC = ((const int32_t *)(op->op_params))[9]; ++ const int32_t N = ((const int32_t *)(op->op_params))[10]; ++ const int32_t OC = ((const int32_t *)(op->op_params))[11]; + +- ggml_metal_kargs_group_norm args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.ngrp =*/ ngrp, +- /*.eps =*/ eps, ++ // 3. Build the parameter struct using the macro-generated variables ++ ggml_metal_kargs_conv_3d args = { ++ /*.IW =*/ (int32_t)op->src[1]->ne[0], ++ /*.IH =*/ (int32_t)op->src[1]->ne[1], ++ /*.ID =*/ (int32_t)op->src[1]->ne[2], ++ /*.OW =*/ (int32_t)op->ne[0], ++ /*.OH =*/ (int32_t)op->ne[1], ++ /*.OD =*/ (int32_t)op->ne[2], ++ /*.KW =*/ (int32_t)op->src[0]->ne[0], ++ /*.KH =*/ (int32_t)op->src[0]->ne[1], ++ /*.KD =*/ (int32_t)op->src[0]->ne[2], ++ s0, s1, s2, ++ p0, p1, p2, ++ d0, d1, d2, ++ IC, N, OC, ++ nb00, nb01, nb02, nb03, // Weight strides ++ nb10, nb11, nb12, nb13, // Input strides ++ nb0, nb1, nb2, nb3 // Output strides + }; + +- auto pipeline = ggml_metal_library_get_pipeline_group_norm(lib, op); ++ // 4. Fetch the JIT pipeline ++ auto pipeline = ggml_metal_library_get_pipeline_conv_3d(lib, op); + +- int nth = 32; // SIMD width +- //while (nth < ne00/4 && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- // nth *= 2; +- //} ++ // 5. Grid mapping ++ int nth0 = 32; // Standard SIMD width for Apple Silicon ++ int nth1 = 1; ++ int nth2 = 1; + +- //nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); +- //nth = std::min(nth, ne00/4); ++ int64_t spatial_volume = args.OW * args.OH * args.OD; + +- const size_t smem = pipeline.smem; ++ int ntg0 = (spatial_volume + nth0 - 1) / nth0; ++ int ntg1 = args.OC; ++ int ntg2 = args.N; + ++ // 6. Bind and Dispatch via the ggml C wrapper + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ngrp, 1, 1, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ntg0, ntg1, ntg2, nth0, nth1, nth2); + + return 1; + } + +-int ggml_metal_op_norm(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- const bool use_fusion = ctx->use_fusion; +- +- const int debug_fusion = ctx->debug_fusion; +- + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- float eps; +- memcpy(&eps, op->op_params, sizeof(float)); +- +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); +- +- ggml_metal_kargs_norm args = { +- /*.ne00 =*/ ne00, +- /*.ne00_t =*/ ne00 % 4 == 0 ? ne00/4 : ne00, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.eps =*/ eps, +- /*.nef1 =*/ { ne01 }, +- /*.nef2 =*/ { ne02 }, +- /*.nef3 =*/ { ne03 }, +- /*.nbf1 =*/ { nb01 }, +- /*.nbf2 =*/ { nb02 }, +- /*.nbf3 =*/ { nb03 }, +- }; +- +- ggml_op fops[8]; +- +- int n_fuse = 1; +- +- ggml_metal_buffer_id bid_fuse[2] = { bid_src0, bid_src0 }; +- +- // d[0] = norm(a) +- // d[1] = mul(d[0], b) +- // d[2] = add(d[1], c) +- if (use_fusion) { +- fops[0] = op->op; +- fops[1] = GGML_OP_MUL; +- fops[2] = GGML_OP_ADD; +- +- for (n_fuse = 0; n_fuse <= 1; ++n_fuse) { +- if (!ctx->can_fuse(idx + n_fuse, fops + n_fuse, 2)) { +- break; +- } +- +- ggml_tensor * f0 = ctx->node(idx + n_fuse); +- ggml_tensor * f1 = ctx->node(idx + n_fuse + 1); +- +- if (f0 != f1->src[0]) { +- break; +- } +- +- if (f1->src[1]->ne[0] != op->ne[0]) { +- break; +- } +- +- if (!ggml_is_contiguous_rows(f1->src[1])) { +- break; +- } ++ const int32_t s0 = ((const int32_t *)(op->op_params))[0]; + +- if (f1->type != GGML_TYPE_F32) { +- break; +- } ++ const int32_t IC = op->src[1]->ne[1]; ++ const int32_t IL = op->src[1]->ne[0]; + +- //ctx->fuse_cnt[f1->op]++; ++ const int32_t K = op->src[0]->ne[0]; + +- bid_fuse[n_fuse] = ggml_metal_get_buffer_id(f1->src[1]); ++ const int32_t OL = op->ne[0]; ++ const int32_t OC = op->ne[1]; + +- args.nef1[n_fuse + 1] = f1->src[1]->ne[1]; +- args.nef2[n_fuse + 1] = f1->src[1]->ne[2]; +- args.nef3[n_fuse + 1] = f1->src[1]->ne[3]; ++ ggml_metal_kargs_conv_transpose_1d args = { ++ /*.IC =*/ IC, ++ /*.IL =*/ IL, ++ /*.K =*/ K, ++ /*.s0 =*/ s0, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ }; + +- args.nbf1[n_fuse + 1] = f1->src[1]->nb[1]; +- args.nbf2[n_fuse + 1] = f1->src[1]->nb[2]; +- args.nbf3[n_fuse + 1] = f1->src[1]->nb[3]; +- } ++ auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_1d(lib, op); + +- ++n_fuse; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- if (debug_fusion > 1 && n_fuse > 1) { +- if (n_fuse == 2) { +- GGML_LOG_DEBUG("%s: fuse: %s + MUL\n", __func__, ggml_op_name(op->op)); +- } +- if (n_fuse == 3) { +- GGML_LOG_DEBUG("%s: fuse: %s + MUL + ADD\n", __func__, ggml_op_name(op->op)); +- } +- } +- } ++ ggml_metal_encoder_dispatch_threadgroups(enc, OL, OC, 1, 1, 1, 1); + +- if (n_fuse > 1) { +- bid_dst = ggml_metal_get_buffer_id(ctx->node(idx + n_fuse - 1)); ++ return 1; ++} + +- for (int i = 1; i < n_fuse; ++i) { +- if (!ggml_metal_op_concurrency_check(ctx, ctx->node(idx + i))) { +- ggml_metal_op_concurrency_reset(ctx); ++int ggml_metal_op_col2im_1d(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); + +- break; +- } +- } +- } ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +- auto pipeline = ggml_metal_library_get_pipeline_norm(lib, op, n_fuse); ++ const int32_t s0 = ((const int32_t *)(op->op_params))[0]; ++ const int32_t OC = ((const int32_t *)(op->op_params))[1]; ++ const int32_t p0 = ((const int32_t *)(op->op_params))[2]; + +- int nth = 32; // SIMD width ++ const int32_t K_OC = (int32_t) op->src[0]->ne[0]; ++ const int32_t T_in = (int32_t) op->src[0]->ne[1]; ++ const int32_t K = K_OC / OC; ++ const int32_t T_out = (int32_t) op->ne[0]; + +- while (nth < args.ne00_t && nth < ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- nth *= 2; +- } ++ ggml_metal_kargs_col2im_1d args = { ++ /*.T_in =*/ T_in, ++ /*.T_out =*/ T_out, ++ /*.OC =*/ OC, ++ /*.K =*/ K, ++ /*.K_OC =*/ K_OC, ++ /*.s0 =*/ s0, ++ /*.p0 =*/ p0, ++ }; + +- nth = std::min(nth, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); +- nth = std::min(nth, args.ne00_t); ++ auto pipeline = ggml_metal_library_get_pipeline_col2im_1d(lib, op); + +- const size_t smem = pipeline.smem; ++ const int total = T_out * OC; ++ const int nth = 256; ++ const int ntg = (total + nth - 1) / nth; + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_fuse[0], 2); +- ggml_metal_encoder_set_buffer (enc, bid_fuse[1], 3); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 4); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1); + +- return n_fuse; ++ return 1; + } + +-int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); +- ++// Dispatch the fused snake kernel from the matched mul -> sin -> sqr -> mul -> add chain. ++// idx points at the leading mul. The caller has validated the chain. ++int ggml_metal_op_snake_fused(ggml_metal_op_t ctx, int idx) { + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); +- +- // make sure we have one or more position id(ne10) per token(ne02) +- GGML_ASSERT(ne10 % ne02 == 0); +- GGML_ASSERT(ne10 >= ne02); +- +- const int nth = std::min(1024, ne00); +- +- const int n_past = ((const int32_t *) op->op_params)[0]; +- const int n_dims = ((const int32_t *) op->op_params)[1]; +- //const int mode = ((const int32_t *) op->op_params)[2]; +- // skip 3, n_ctx, used in GLM RoPE, unimplemented in metal +- const int n_ctx_orig = ((const int32_t *) op->op_params)[4]; +- +- float freq_base; +- float freq_scale; +- float ext_factor; +- float attn_factor; +- float beta_fast; +- float beta_slow; +- +- memcpy(&freq_base, (const int32_t *) op->op_params + 5, sizeof(float)); +- memcpy(&freq_scale, (const int32_t *) op->op_params + 6, sizeof(float)); +- memcpy(&ext_factor, (const int32_t *) op->op_params + 7, sizeof(float)); +- memcpy(&attn_factor, (const int32_t *) op->op_params + 8, sizeof(float)); +- memcpy(&beta_fast, (const int32_t *) op->op_params + 9, sizeof(float)); +- memcpy(&beta_slow, (const int32_t *) op->op_params + 10, sizeof(float)); +- +- // mrope +- const int sect_0 = ((const int32_t *) op->op_params)[11]; +- const int sect_1 = ((const int32_t *) op->op_params)[12]; +- const int sect_2 = ((const int32_t *) op->op_params)[13]; +- const int sect_3 = ((const int32_t *) op->op_params)[14]; +- +- ggml_metal_kargs_rope args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.n_past =*/ n_past, +- /*.n_dims =*/ n_dims, +- /*.n_ctx_orig =*/ n_ctx_orig, +- /*.freq_base =*/ freq_base, +- /*.freq_scale =*/ freq_scale, +- /*.ext_factor =*/ ext_factor, +- /*.attn_factor =*/ attn_factor, +- /*.beta_fast =*/ beta_fast, +- /*.beta_slow =*/ beta_slow, +- /* sect_0 =*/ sect_0, +- /* sect_1 =*/ sect_1, +- /* sect_2 =*/ sect_2, +- /* sect_3 =*/ sect_3, +- /* src2 =*/ op->src[2] != nullptr, +- }; ++ const ggml_tensor * mul0 = ctx->node(idx + 0); ++ const ggml_tensor * sqr = ctx->node(idx + 2); ++ const ggml_tensor * mul1 = ctx->node(idx + 3); ++ ggml_tensor * add = ctx->node(idx + 4); + +- auto pipeline = ggml_metal_library_get_pipeline_rope(lib, op); ++ const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; ++ const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; ++ const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; ++ ++ const int T = (int) x->ne[0]; ++ const int C = (int) x->ne[1]; ++ const int total = T * C; ++ ++ // the encode loop pre-checked the leading mul only, check the rest of the chain ++ for (int i = 1; i < 5; ++i) { ++ if (!ggml_metal_op_concurrency_check(ctx, ctx->node(idx + i))) { ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ break; ++ } ++ } ++ ++ auto pipeline = ggml_metal_library_get_pipeline_snake(lib, x->type); ++ ++ ggml_metal_kargs_snake args = { ++ /*.T =*/ T, ++ /*.C =*/ C, ++ }; + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- if (op->src[2]) { +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); +- } else { +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 3); +- } +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 4); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(x), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(a), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(inv_b), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(add), 4); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne01, ne02, ne03, nth, 1, 1); ++ const int nth = 256; ++ const int ntg = (total + nth - 1) / nth; ++ ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1); + +- return 1; ++ return 5; + } + +-int ggml_metal_op_im2col(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -3648,79 +10599,55 @@ int ggml_metal_op_im2col(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + + const int32_t s0 = ((const int32_t *)(op->op_params))[0]; +- const int32_t s1 = ((const int32_t *)(op->op_params))[1]; +- const int32_t p0 = ((const int32_t *)(op->op_params))[2]; +- const int32_t p1 = ((const int32_t *)(op->op_params))[3]; +- const int32_t d0 = ((const int32_t *)(op->op_params))[4]; +- const int32_t d1 = ((const int32_t *)(op->op_params))[5]; +- +- const bool is_2D = ((const int32_t *)(op->op_params))[6] == 1; +- +- const int32_t N = op->src[1]->ne[is_2D ? 3 : 2]; +- const int32_t IC = op->src[1]->ne[is_2D ? 2 : 1]; +- const int32_t IH = is_2D ? op->src[1]->ne[1] : 1; +- const int32_t IW = op->src[1]->ne[0]; +- +- const int32_t KH = is_2D ? op->src[0]->ne[1] : 1; +- const int32_t KW = op->src[0]->ne[0]; + +- const int32_t OH = is_2D ? op->ne[2] : 1; +- const int32_t OW = op->ne[1]; ++ const int32_t IC = op->src[1]->ne[2]; ++ const int32_t IH = op->src[1]->ne[1]; ++ const int32_t IW = op->src[1]->ne[0]; + +- const int32_t CHW = IC * KH * KW; ++ const int32_t KH = op->src[0]->ne[1]; ++ const int32_t KW = op->src[0]->ne[0]; + +- const uint64_t ofs0 = op->src[1]->nb[is_2D ? 3 : 2] / 4; +- const uint64_t ofs1 = op->src[1]->nb[is_2D ? 2 : 1] / 4; ++ const int32_t OW = op->ne[0]; ++ const int32_t OH = op->ne[1]; ++ const int32_t OC = op->ne[2]; + +- ggml_metal_kargs_im2col args = { +- /*.ofs0 =*/ ofs0, +- /*.ofs1 =*/ ofs1, +- /*.IW =*/ IW, +- /*.IH =*/ IH, +- /*.CHW =*/ CHW, +- /*.s0 =*/ s0, +- /*.s1 =*/ s1, +- /*.p0 =*/ p0, +- /*.p1 =*/ p1, +- /*.d0 =*/ d0, +- /*.d1 =*/ d1, +- /*.N =*/ N, +- /*.KH =*/ KH, +- /*.KW =*/ KW, +- /*.KHW =*/ KH * KW, ++ ggml_metal_kargs_conv_transpose_2d args = { ++ /*.IC =*/ IC, ++ /*.IH =*/ IH, ++ /*.IW =*/ IW, ++ /*.KH =*/ KH, ++ /*.KW =*/ KW, ++ /*.OC =*/ OC, ++ /*.s0 =*/ s0, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_im2col(lib, op); +- +- if (KH*KW <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- const uint64_t ntptg0 = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)/(KH*KW), N); +- +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_2d(lib, op); + +- ggml_metal_encoder_dispatch_threadgroups(enc, IC, OH, OW, ntptg0, KH, KW); +- } else { +- const uint64_t n_threads = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), N); +- const int64_t quotient = N / n_threads + (N % n_threads > 0 ? 1 : 0); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ // Metal requires buffer size to be multiple of 16 bytes ++ const size_t smem = GGML_PAD(KW * KH * sizeof(float), 16); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- ggml_metal_encoder_dispatch_threadgroups(enc, quotient * CHW, OH, OW, n_threads, 1, 1); +- } ++ ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC, KW, KH, 1); + + return 1; + } + +-int ggml_metal_op_conv_2d(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_upscale(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -3728,77 +10655,63 @@ int ggml_metal_op_conv_2d(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- GGML_ASSERT(ggml_is_contiguous(op->src[0])); +- GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); ++ float sf0 = (float)ne0/op->src[0]->ne[0]; ++ float sf1 = (float)ne1/op->src[0]->ne[1]; ++ float sf2 = (float)ne2/op->src[0]->ne[2]; ++ float sf3 = (float)ne3/op->src[0]->ne[3]; + +- const int32_t s0 = ((const int32_t *) op->op_params)[0]; +- const int32_t s1 = ((const int32_t *) op->op_params)[1]; +- const int32_t p0 = ((const int32_t *) op->op_params)[2]; +- const int32_t p1 = ((const int32_t *) op->op_params)[3]; +- const int32_t d0 = ((const int32_t *) op->op_params)[4]; +- const int32_t d1 = ((const int32_t *) op->op_params)[5]; ++ const int32_t mode_flags = ggml_get_op_params_i32(op, 0); + +- ggml_metal_kargs_conv_2d args = { +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.IW =*/ ne10, +- /*.IH =*/ ne11, +- /*.KW =*/ ne00, +- /*.KH =*/ ne01, +- /*.IC =*/ ne02, +- /*.OC =*/ ne03, +- /*.OW =*/ ne0, +- /*.OH =*/ ne1, +- /*.N =*/ ne3, +- /*.s0 =*/ s0, +- /*.s1 =*/ s1, +- /*.p0 =*/ p0, +- /*.p1 =*/ p1, +- /*.d0 =*/ d0, +- /*.d1 =*/ d1, +- }; ++ float poffs = 0.5f; + +- auto pipeline = ggml_metal_library_get_pipeline_conv_2d(lib, op); ++ if (mode_flags & GGML_SCALE_FLAG_ALIGN_CORNERS) { ++ poffs = 0.0f; ++ sf0 = ne0 > 1 && ne00 > 1 ? (float)(ne0 - 1) / (ne00 - 1) : sf0; ++ sf1 = ne1 > 1 && ne01 > 1 ? (float)(ne1 - 1) / (ne01 - 1) : sf1; ++ } + +- int nth = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); +- nth = std::min(nth, 256); +- nth = std::max(nth, 1); ++ ggml_metal_kargs_upscale args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.sf0 =*/ sf0, ++ /*.sf1 =*/ sf1, ++ /*.sf2 =*/ sf2, ++ /*.sf3 =*/ sf3, ++ /*.poffs =*/ poffs, ++ }; + +- const uint64_t n_out = ggml_nelements(op); ++ auto pipeline = ggml_metal_library_get_pipeline_upscale(lib, op); + +- uint64_t tg = (n_out + nth - 1)/nth; +- tg = std::max(tg, 1); +- tg = std::min(tg, (uint64_t) std::numeric_limits::max()); ++ const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne0); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, tg, 1, 1, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_conv_2d_dw(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_roll(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -3806,150 +10719,103 @@ int ggml_metal_op_conv_2d_dw(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32); +- GGML_ASSERT(op->type == GGML_TYPE_F32); +- GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32); +- +- const int32_t s0 = ((const int32_t *) op->op_params)[0]; +- const int32_t s1 = ((const int32_t *) op->op_params)[1]; +- const int32_t p0 = ((const int32_t *) op->op_params)[2]; +- const int32_t p1 = ((const int32_t *) op->op_params)[3]; +- const int32_t d0 = ((const int32_t *) op->op_params)[4]; +- const int32_t d1 = ((const int32_t *) op->op_params)[5]; ++ const int32_t s0 = ggml_get_op_params_i32(op, 0); ++ const int32_t s1 = ggml_get_op_params_i32(op, 1); ++ const int32_t s2 = ggml_get_op_params_i32(op, 2); ++ const int32_t s3 = ggml_get_op_params_i32(op, 3); + +- ggml_metal_kargs_conv_2d_dw args = { ++ ggml_metal_kargs_roll args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, + /*.nb00 =*/ nb00, + /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb03, +- /*.nb10 =*/ nb10, +- /*.nb11 =*/ nb11, +- /*.nb12 =*/ nb12, +- /*.nb13 =*/ nb13, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, + /*.nb0 =*/ nb0, + /*.nb1 =*/ nb1, + /*.nb2 =*/ nb2, + /*.nb3 =*/ nb3, +- /*.IW =*/ ne10, +- /*.IH =*/ ne11, +- /*.KW =*/ ne00, +- /*.KH =*/ ne01, +- /*.C =*/ ne12, +- /*.OW =*/ ne0, +- /*.OH =*/ ne1, +- /*.N =*/ ne13, + /*.s0 =*/ s0, + /*.s1 =*/ s1, +- /*.p0 =*/ p0, +- /*.p1 =*/ p1, +- /*.d0 =*/ d0, +- /*.d1 =*/ d1, ++ /*.s2 =*/ s2, ++ /*.s3 =*/ s3 + }; + +- const bool use_tiled = (nb12 < nb10); +- +- auto pipeline = ggml_metal_library_get_pipeline_conv_2d_dw(lib, op, use_tiled); +- +- int nth = ggml_metal_pipeline_max_theads_per_threadgroup(pipeline); +- nth = std::min(nth, 256); +- nth = std::max(nth, 1); +- +- const int32_t OW = ne0; +- const int32_t OH = ne1; +- const int32_t C = ne12; +- const int32_t N = ne13; ++ auto pipeline = ggml_metal_library_get_pipeline_roll(lib, op); + +- const int tg_x = use_tiled ? (C + nth - 1) / nth : (OW + nth - 1) / nth; +- const int tg_y = OH; +- const int tg_z = use_tiled ? OW * N : C * N; ++ const int nth = std::min(1024, ne0); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, tg_x, tg_y, tg_z, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_conv_3d(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- // 1. Extract standard dimensions and byte strides ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- // 2. Extract hyperparams from op_params +- const int32_t s0 = ((const int32_t *)(op->op_params))[0]; +- const int32_t s1 = ((const int32_t *)(op->op_params))[1]; +- const int32_t s2 = ((const int32_t *)(op->op_params))[2]; +- const int32_t p0 = ((const int32_t *)(op->op_params))[3]; +- const int32_t p1 = ((const int32_t *)(op->op_params))[4]; +- const int32_t p2 = ((const int32_t *)(op->op_params))[5]; +- const int32_t d0 = ((const int32_t *)(op->op_params))[6]; +- const int32_t d1 = ((const int32_t *)(op->op_params))[7]; +- const int32_t d2 = ((const int32_t *)(op->op_params))[8]; +- const int32_t IC = ((const int32_t *)(op->op_params))[9]; +- const int32_t N = ((const int32_t *)(op->op_params))[10]; +- const int32_t OC = ((const int32_t *)(op->op_params))[11]; +- +- // 3. Build the parameter struct using the macro-generated variables +- ggml_metal_kargs_conv_3d args = { +- /*.IW =*/ (int32_t)op->src[1]->ne[0], +- /*.IH =*/ (int32_t)op->src[1]->ne[1], +- /*.ID =*/ (int32_t)op->src[1]->ne[2], +- /*.OW =*/ (int32_t)op->ne[0], +- /*.OH =*/ (int32_t)op->ne[1], +- /*.OD =*/ (int32_t)op->ne[2], +- /*.KW =*/ (int32_t)op->src[0]->ne[0], +- /*.KH =*/ (int32_t)op->src[0]->ne[1], +- /*.KD =*/ (int32_t)op->src[0]->ne[2], +- s0, s1, s2, +- p0, p1, p2, +- d0, d1, d2, +- IC, N, OC, +- nb00, nb01, nb02, nb03, // Weight strides +- nb10, nb11, nb12, nb13, // Input strides +- nb0, nb1, nb2, nb3 // Output strides ++ ggml_metal_kargs_pad args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3 + }; + +- // 4. Fetch the JIT pipeline +- auto pipeline = ggml_metal_library_get_pipeline_conv_3d(lib, op); +- +- // 5. Grid mapping +- int nth0 = 32; // Standard SIMD width for Apple Silicon +- int nth1 = 1; +- int nth2 = 1; ++ auto pipeline = ggml_metal_library_get_pipeline_pad(lib, op); + +- int64_t spatial_volume = args.OW * args.OH * args.OD; ++ if (pipeline.c4) { ++ args.ne00 = ne00/4; ++ args.ne0 = ne0/4; ++ } + +- int ntg0 = (spatial_volume + nth0 - 1) / nth0; +- int ntg1 = args.OC; +- int ntg2 = args.N; ++ const int nth_max = MIN(64, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ const int nth = MIN(args.ne0, nth_max); ++ const int nk0 = (args.ne0 + 1024 - 1)/1024; // note: 1024 is hardcoded in the kernel! + +- // 6. Bind and Dispatch via the ggml C wrapper + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ntg0, ntg1, ntg2, nth0, nth1, nth2); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nk0*ne1, ne2, ne3, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_pad_reflect_1d(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -3957,134 +10823,113 @@ int ggml_metal_op_conv_transpose_1d(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const int32_t s0 = ((const int32_t *)(op->op_params))[0]; +- +- const int32_t IC = op->src[1]->ne[1]; +- const int32_t IL = op->src[1]->ne[0]; +- +- const int32_t K = op->src[0]->ne[0]; +- +- const int32_t OL = op->ne[0]; +- const int32_t OC = op->ne[1]; +- +- ggml_metal_kargs_conv_transpose_1d args = { +- /*.IC =*/ IC, +- /*.IL =*/ IL, +- /*.K =*/ K, +- /*.s0 =*/ s0, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, ++ ggml_metal_kargs_pad_reflect_1d args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ /*.p0 =*/ ((const int32_t *)(op->op_params))[0], ++ /*.p1 =*/ ((const int32_t *)(op->op_params))[1] + }; + +- auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_1d(lib, op); ++ auto pipeline = ggml_metal_library_get_pipeline_pad_reflect_1d(lib, op); ++ ++ const int nth = std::min(1024, ne0); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, OL, OC, 1, 1, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_col2im_1d(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_arange(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- const int32_t s0 = ((const int32_t *)(op->op_params))[0]; +- const int32_t OC = ((const int32_t *)(op->op_params))[1]; +- const int32_t p0 = ((const int32_t *)(op->op_params))[2]; ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const int32_t K_OC = (int32_t) op->src[0]->ne[0]; +- const int32_t T_in = (int32_t) op->src[0]->ne[1]; +- const int32_t K = K_OC / OC; +- const int32_t T_out = (int32_t) op->ne[0]; ++ float start; ++ float step; + +- ggml_metal_kargs_col2im_1d args = { +- /*.T_in =*/ T_in, +- /*.T_out =*/ T_out, +- /*.OC =*/ OC, +- /*.K =*/ K, +- /*.K_OC =*/ K_OC, +- /*.s0 =*/ s0, +- /*.p0 =*/ p0, ++ memcpy(&start, ((const int32_t *) op->op_params) + 0, sizeof(float)); ++ memcpy(&step, ((const int32_t *) op->op_params) + 2, sizeof(float)); ++ ++ ggml_metal_kargs_arange args = { ++ /*.ne0 =*/ ne0, ++ /*.start =*/ start, ++ /*.step =*/ step + }; + +- auto pipeline = ggml_metal_library_get_pipeline_col2im_1d(lib, op); ++ const int nth = std::min(1024, ne0); + +- const int total = T_out * OC; +- const int nth = 256; +- const int ntg = (total + nth - 1) / nth; ++ auto pipeline = ggml_metal_library_get_pipeline_arange(lib, op); + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 1); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, nth, 1, 1); + + return 1; + } + +-// Dispatch the fused snake kernel from the matched mul -> sin -> sqr -> mul -> add chain. +-// idx points at the leading mul. The caller has validated the chain. +-int ggml_metal_op_snake_fused(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx) { ++ ggml_tensor * op = ctx->node(idx); ++ + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- const ggml_tensor * mul0 = ctx->node(idx + 0); +- const ggml_tensor * sqr = ctx->node(idx + 2); +- const ggml_tensor * mul1 = ctx->node(idx + 3); +- ggml_tensor * add = ctx->node(idx + 4); +- +- const ggml_tensor * x = ggml_are_same_shape(mul0, mul0->src[0]) ? mul0->src[0] : mul0->src[1]; +- const ggml_tensor * a = (x == mul0->src[0]) ? mul0->src[1] : mul0->src[0]; +- const ggml_tensor * inv_b = (mul1->src[0] == sqr) ? mul1->src[1] : mul1->src[0]; +- +- const int T = (int) x->ne[0]; +- const int C = (int) x->ne[1]; +- const int total = T * C; ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- // the encode loop pre-checked the leading mul only, check the rest of the chain +- for (int i = 1; i < 5; ++i) { +- if (!ggml_metal_op_concurrency_check(ctx, ctx->node(idx + i))) { +- ggml_metal_op_concurrency_reset(ctx); ++ const int dim = op->op_params[0]; ++ const int max_period = op->op_params[1]; + +- break; +- } +- } ++ ggml_metal_kargs_timestep_embedding args = { ++ /*.nb1 =*/ nb1, ++ /*.dim =*/ dim, ++ /*.max_period =*/ max_period, ++ }; + +- auto pipeline = ggml_metal_library_get_pipeline_snake(lib, x->type); ++ auto pipeline = ggml_metal_library_get_pipeline_timestep_embedding(lib, op); + +- ggml_metal_kargs_snake args = { +- /*.T =*/ T, +- /*.C =*/ C, +- }; ++ const int nth = std::max(1, std::min(1024, dim/2)); + + ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(x), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(a), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(inv_b), 3); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(add), 4); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- const int nth = 256; +- const int ntg = (total + nth - 1) / nth; +- ggml_metal_encoder_dispatch_threadgroups(enc, ntg, 1, 1, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, ne00, 1, 1, nth, 1, 1); + +- return 5; ++ return 1; + } + +-int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -4092,81 +10937,75 @@ int ggml_metal_op_conv_transpose_2d(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const int32_t s0 = ((const int32_t *)(op->op_params))[0]; +- +- const int32_t IC = op->src[1]->ne[2]; +- const int32_t IH = op->src[1]->ne[1]; +- const int32_t IW = op->src[1]->ne[0]; ++ ggml_metal_kargs_argmax args = { ++ /*.ne00 = */ ne00, ++ /*.nb01 = */ nb01, ++ }; + +- const int32_t KH = op->src[0]->ne[1]; +- const int32_t KW = op->src[0]->ne[0]; ++ auto pipeline = ggml_metal_library_get_pipeline_argmax(lib, op); + +- const int32_t OW = op->ne[0]; +- const int32_t OH = op->ne[1]; +- const int32_t OC = op->ne[2]; ++ const int64_t nrows = ggml_nrows(op->src[0]); + +- ggml_metal_kargs_conv_transpose_2d args = { +- /*.IC =*/ IC, +- /*.IH =*/ IH, +- /*.IW =*/ IW, +- /*.KH =*/ KH, +- /*.KW =*/ KW, +- /*.OC =*/ OC, +- /*.s0 =*/ s0, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- }; ++ int nth = 32; // SIMD width ++ while (nth < ne00 && nth*ne01*ne02*ne03 < 256) { ++ nth *= 2; ++ } + +- auto pipeline = ggml_metal_library_get_pipeline_conv_transpose_2d(lib, op); ++ const size_t smem = pipeline.smem; + + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); + +- // Metal requires buffer size to be multiple of 16 bytes +- const size_t smem = GGML_PAD(KW * KH * sizeof(float), 16); + ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- ggml_metal_encoder_dispatch_threadgroups(enc, OW, OH, OC, KW, KH, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_upscale(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_argsort(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + ++ GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); ++ + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- float sf0 = (float)ne0/op->src[0]->ne[0]; +- float sf1 = (float)ne1/op->src[0]->ne[1]; +- float sf2 = (float)ne2/op->src[0]->ne[2]; +- float sf3 = (float)ne3/op->src[0]->ne[3]; ++ auto pipeline = ggml_metal_library_get_pipeline_argsort(lib, op); + +- const int32_t mode_flags = ggml_get_op_params_i32(op, 0); ++ // bitonic sort requires the number of elements to be power of 2 ++ int nth = 1; ++ while (nth < ne00 && 2*nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { ++ nth *= 2; ++ } + +- float poffs = 0.5f; ++ const int npr = (ne00 + nth - 1)/nth; + +- if (mode_flags & GGML_SCALE_FLAG_ALIGN_CORNERS) { +- poffs = 0.0f; +- sf0 = ne0 > 1 && ne00 > 1 ? (float)(ne0 - 1) / (ne00 - 1) : sf0; +- sf1 = ne1 > 1 && ne01 > 1 ? (float)(ne1 - 1) / (ne01 - 1) : sf1; ++ // Metal kernels require the buffer size to be multiple of 16 bytes ++ // https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/1443142-setthreadgroupmemorylength ++ const size_t smem = GGML_PAD(nth*sizeof(int32_t), 16); ++ ++ ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ ++ ggml_metal_buffer_id bid_tmp = bid_dst; ++ bid_tmp.offs += ggml_nbytes(op); ++ ++ if ((int) ceil(std::log(npr) / std::log(2)) % 2 == 1) { ++ std::swap(bid_dst, bid_tmp); + } + +- ggml_metal_kargs_upscale args = { ++ ggml_metal_kargs_argsort args = { + /*.ne00 =*/ ne00, + /*.ne01 =*/ ne01, + /*.ne02 =*/ ne02, +@@ -4179,85 +11018,64 @@ int ggml_metal_op_upscale(ggml_metal_op_t ctx, int idx) { + /*.ne1 =*/ ne1, + /*.ne2 =*/ ne2, + /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.sf0 =*/ sf0, +- /*.sf1 =*/ sf1, +- /*.sf2 =*/ sf2, +- /*.sf3 =*/ sf3, +- /*.poffs =*/ poffs, ++ /*.top_k =*/ nth, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_upscale(lib, op); +- +- const int nth = std::min(ggml_metal_pipeline_max_theads_per_threadgroup(pipeline), ne0); +- + ggml_metal_encoder_set_pipeline(enc, pipeline); + ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); + +- return 1; +-} ++ ggml_metal_encoder_dispatch_threadgroups(enc, npr*ne01, ne02, ne03, nth, 1, 1); + +-int ggml_metal_op_roll(ggml_metal_op_t ctx, int idx) { +- ggml_tensor * op = ctx->node(idx); ++ auto pipeline_merge = ggml_metal_library_get_pipeline_argsort_merge(lib, op); + +- ggml_metal_library_t lib = ctx->lib; +- ggml_metal_encoder_t enc = ctx->enc; ++ int len = nth; + +- GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); +- GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); +- GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); ++ while (len < ne00) { ++ ggml_metal_op_concurrency_reset(ctx); + +- const int32_t s0 = ggml_get_op_params_i32(op, 0); +- const int32_t s1 = ggml_get_op_params_i32(op, 1); +- const int32_t s2 = ggml_get_op_params_i32(op, 2); +- const int32_t s3 = ggml_get_op_params_i32(op, 3); ++ ggml_metal_kargs_argsort_merge args_merge = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.top_k =*/ ne00, ++ /*.len =*/ len, ++ }; + +- ggml_metal_kargs_roll args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.s0 =*/ s0, +- /*.s1 =*/ s1, +- /*.s2 =*/ s2, +- /*.s3 =*/ s3 +- }; ++ // merges per row ++ const int nm = (ne00 + 2*len - 1) / (2*len); + +- auto pipeline = ggml_metal_library_get_pipeline_roll(lib, op); ++ const int nth = std::min(512, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_merge)); + +- const int nth = std::min(1024, ne0); ++ ggml_metal_encoder_set_pipeline(enc, pipeline_merge); ++ ggml_metal_encoder_set_bytes (enc, &args_merge, sizeof(args_merge), 0); ++ ggml_metal_encoder_set_buffer (enc, bid_src0, 1); ++ ggml_metal_encoder_set_buffer (enc, bid_dst, 2); ++ ggml_metal_encoder_set_buffer (enc, bid_tmp, 3); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_dispatch_threadgroups(enc, nm*ne01, ne02, ne03, nth, 1, 1); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); ++ std::swap(bid_dst, bid_tmp); ++ ++ len <<= 1; ++ } + + return 1; + } + +-int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_lightning_indexer(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -4265,50 +11083,146 @@ int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- ggml_metal_kargs_pad args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3 +- }; +- +- auto pipeline = ggml_metal_library_get_pipeline_pad(lib, op); ++ const bool parallel_requested = ggml_metal_lightning_indexer_parallel_requested(); ++ bool parallel = parallel_requested && ne01 <= 1024; ++ const size_t staged_q_bytes = size_t(ne00)*size_t(ne01)*sizeof(float); ++ const bool staged_q_requested = ggml_metal_lightning_indexer_staged_q_requested(); ++ const bool staged_q = staged_q_requested && !parallel && ++ op->src[0]->type == GGML_TYPE_F32 && ++ (op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32) && ++ staged_q_bytes <= 32*1024; ++ const int requested_parallel_threads = ggml_metal_lightning_indexer_parallel_threads_requested(); ++ auto pipeline = ggml_metal_library_get_pipeline_lightning_indexer(lib, op, parallel, staged_q); ++ int nth = std::min(parallel ? requested_parallel_threads : 64, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ if (parallel && ne01 > nth) { ++ parallel = false; ++ pipeline = ggml_metal_library_get_pipeline_lightning_indexer(lib, op, parallel, false); ++ nth = std::min(64, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ } + +- if (pipeline.c4) { +- args.ne00 = ne00/4; +- args.ne0 = ne0/4; ++ static bool logged = false; ++ if (!logged) { ++ logged = true; ++ const char * parallel_env = getenv("LLAMA_GLM_DSA_PARALLEL_LIGHTNING_INDEXER"); ++ GGML_LOG_INFO( ++ "%s: env=%s requested=%d staged_q_requested=%d staged_q_selected=%d staged_q_bytes=%zu q_ne=(%lld,%lld,%lld,%lld) k_type=%s selected=%d nth=%d\n", ++ __func__, ++ parallel_env ? parallel_env : "", ++ parallel_requested, ++ staged_q_requested, ++ staged_q, ++ staged_q_bytes, ++ (long long) ne00, ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ ggml_type_name(op->src[1]->type), ++ parallel, ++ nth); + } + +- const int nth_max = MIN(64, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); +- const int nth = MIN(args.ne0, nth_max); +- const int nk0 = (args.ne0 + 1024 - 1)/1024; // note: 1024 is hardcoded in the kernel! ++ ggml_metal_kargs_lightning_indexer args = { ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb13 =*/ nb13, ++ /*.ne20 =*/ ne20, ++ /*.ne21 =*/ ne21, ++ /*.ne22 =*/ ne22, ++ /*.ne23 =*/ ne23, ++ /*.nb20 =*/ nb20, ++ /*.nb21 =*/ nb21, ++ /*.nb22 =*/ nb22, ++ /*.nb23 =*/ nb23, ++ /*.ne30 =*/ ne30, ++ /*.ne31 =*/ ne31, ++ /*.ne32 =*/ ne32, ++ /*.ne33 =*/ ne33, ++ /*.nb30 =*/ nb30, ++ /*.nb31 =*/ nb31, ++ /*.nb32 =*/ nb32, ++ /*.nb33 =*/ nb33, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.ne2 =*/ ne2, ++ /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, ++ }; ++ ++ int ida = 0; + + ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // k ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // weights ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // mask ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst ++ if (staged_q) { ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, staged_q_bytes, 0); ++ } + +- ggml_metal_encoder_dispatch_threadgroups(enc, nk0*ne1, ne2, ne3, nth, 1, 1); ++ const int grid_x = parallel ? ne0 : (ne0 + nth - 1)/nth; ++ const int grid_y = ne1; ++ const int grid_z = ne3; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=lightning_indexer kernel=%s tensor=%s parallel=%d staged_q=%d staged_q_bytes=%zu q_type=%s k_type=%s dst_type=%s q_ne0=%lld q_ne1=%lld q_ne2=%lld q_ne3=%lld k_ne0=%lld k_ne1=%lld k_ne2=%lld k_ne3=%lld dst_ne0=%lld dst_ne1=%lld dst_ne2=%lld dst_ne3=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=%d\n", ++ staged_q ? "staged_q" : parallel ? "parallel" : "serial", ++ ggml_metal_tensor_name(op), ++ parallel ? 1 : 0, ++ staged_q ? 1 : 0, ++ staged_q_bytes, ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->type), ++ (long long) ne00, ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ (long long) ne10, ++ (long long) ne11, ++ (long long) ne12, ++ (long long) ne13, ++ (long long) ne0, ++ (long long) ne1, ++ (long long) ne2, ++ (long long) ne3, ++ grid_x, ++ grid_y, ++ grid_z, ++ nth); ++ } ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_pad_reflect_1d(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_dsa_sparse_mask(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -4316,79 +11230,318 @@ int ggml_metal_op_pad_reflect_1d(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- ggml_metal_kargs_pad_reflect_1d args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.nb0 =*/ nb0, +- /*.nb1 =*/ nb1, +- /*.nb2 =*/ nb2, +- /*.nb3 =*/ nb3, +- /*.p0 =*/ ((const int32_t *)(op->op_params))[0], +- /*.p1 =*/ ((const int32_t *)(op->op_params))[1] ++ ggml_metal_kargs_dsa_sparse_mask args = { ++ /*.n_kv =*/ ne01, ++ /*.n_batch =*/ ne02, ++ /*.n_stream =*/ ne03, ++ /*.n_top_k =*/ ne10, ++ /*.n_top_stream =*/ ne12, ++ /*.elem_size =*/ (int32_t) ggml_type_size(op->type), ++ /*._pad1 =*/ 0, ++ /*._pad2 =*/ 0, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.nb03 =*/ nb03, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.nb0 =*/ nb0, ++ /*.nb1 =*/ nb1, ++ /*.nb2 =*/ nb2, ++ /*.nb3 =*/ nb3, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_pad_reflect_1d(lib, op); ++ auto pipeline_fill = ggml_metal_library_get_pipeline_dsa_sparse_mask_fill(lib); ++ auto pipeline_set = ggml_metal_library_get_pipeline_dsa_sparse_mask_set(lib); + +- const int nth = std::min(1024, ne0); ++ ggml_metal_encoder_set_pipeline(enc, pipeline_fill); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 1); + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ const int nth_fill = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_fill)); ++ const int fill_grid_x = (ne01 + nth_fill - 1)/nth_fill; ++ const int fill_grid_y = ne02; ++ const int fill_grid_z = ne03; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=dsa_sparse_mask kernel=fill tensor=%s src_type=%s top_k_type=%s dst_type=%s kv=%lld batch=%lld stream=%lld top_k=%lld top_stream=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->type), ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ (long long) ne10, ++ (long long) ne12, ++ fill_grid_x, ++ fill_grid_y, ++ fill_grid_z, ++ nth_fill); ++ } ++ ggml_metal_encoder_dispatch_threadgroups(enc, fill_grid_x, fill_grid_y, fill_grid_z, nth_fill, 1, 1); ++ ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ ggml_metal_encoder_set_pipeline(enc, pipeline_set); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); + ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne1, ne2, ne3, nth, 1, 1); ++ const int nth_set = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_set)); ++ const int set_grid_x = (ne10 + nth_set - 1)/nth_set; ++ const int set_grid_y = ne11; ++ const int set_grid_z = ne03; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=dsa_sparse_mask kernel=set tensor=%s src_type=%s top_k_type=%s dst_type=%s kv=%lld batch=%lld stream=%lld top_k=%lld top_stream=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->type), ++ (long long) ne01, ++ (long long) ne02, ++ (long long) ne03, ++ (long long) ne10, ++ (long long) ne12, ++ set_grid_x, ++ set_grid_y, ++ set_grid_z, ++ nth_set); ++ } ++ ggml_metal_encoder_dispatch_threadgroups(enc, set_grid_x, set_grid_y, set_grid_z, nth_set, 1, 1); + + return 1; + } + +-int ggml_metal_op_arange(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_dsa_sparse_attn(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne3, op->src[3], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne4, op->src[4], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb4, op->src[4], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- float start; +- float step; ++ ggml_metal_kargs_dsa_sparse_attn args = { ++ /*.ne00 =*/ ne00, /*.ne01 =*/ ne01, /*.ne02 =*/ ne02, /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, /*.nb01 =*/ nb01, /*.nb02 =*/ nb02, /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, /*.ne11 =*/ ne11, /*.ne12 =*/ ne12, /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, /*.nb11 =*/ nb11, /*.nb12 =*/ nb12, /*.nb13 =*/ nb13, ++ /*.ne20 =*/ ne20, /*.ne21 =*/ ne21, /*.ne22 =*/ ne22, /*.ne23 =*/ ne23, ++ /*.nb20 =*/ nb20, /*.nb21 =*/ nb21, /*.nb22 =*/ nb22, /*.nb23 =*/ nb23, ++ /*.ne30 =*/ ne30, /*.ne31 =*/ ne31, /*.ne32 =*/ ne32, /*.ne33 =*/ ne33, ++ /*.nb30 =*/ nb30, /*.nb31 =*/ nb31, /*.nb32 =*/ nb32, /*.nb33 =*/ nb33, ++ /*.ne40 =*/ ne40, /*.ne41 =*/ ne41, /*.ne42 =*/ ne42, /*.ne43 =*/ ne43, ++ /*.nb40 =*/ nb40, /*.nb41 =*/ nb41, /*.nb42 =*/ nb42, /*.nb43 =*/ nb43, ++ /*.ne0 =*/ ne0, /*.ne1 =*/ ne1, /*.ne2 =*/ ne2, /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, /*.nb3 =*/ nb3, ++ /*.scale =*/ ggml_get_op_params_f32(op, 0), ++ }; + +- memcpy(&start, ((const int32_t *) op->op_params) + 0, sizeof(float)); +- memcpy(&step, ((const int32_t *) op->op_params) + 2, sizeof(float)); ++ const bool use_selected_row = ++ ggml_metal_glm_dsa_selected_row_flash_enabled() && ++ op->type == GGML_TYPE_F32 && ++ op->src[0]->type == GGML_TYPE_F32 && ++ op->src[1]->type == GGML_TYPE_F16 && ++ op->src[2]->type == GGML_TYPE_F16 && ++ op->src[4]->type == GGML_TYPE_I32 && ++ ne00 == 576 && ++ ne20 == 512 && ++ ne01 > 0 && ++ ne02 > 0 && ++ ne03 > 0 && ++ ne12 == 1 && ++ ne22 == 1 && ++ ne40 > 0 && ++ ne40 <= 4096 && ++ ne41 == ne01 && ++ ne42 > 0 && ++ ne03 % ne42 == 0; ++ if (use_selected_row) { ++ const bool use_tiled = ++ ggml_metal_glm_dsa_selected_row_flash_tiled_selected(op) && ++ ne02 >= 8 && ne02 % 8 == 0; ++ const int32_t nwg = use_tiled ++ ? ggml_metal_glm_dsa_selected_row_flash_tiled_nwg_for_top_k(ne40) ++ : 1; ++ const int32_t nsg = use_tiled ? 4 : 1; ++ auto pipeline = use_tiled ++ ? ggml_metal_library_get_pipeline_selected_row_flash_tiled(lib, nwg) ++ : ggml_metal_library_get_pipeline_selected_row_flash_vec(lib, nsg, nwg); ++ GGML_ASSERT(nsg*32 <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); + +- ggml_metal_kargs_arange args = { +- /*.ne0 =*/ ne0, +- /*.start =*/ start, +- /*.step =*/ step +- }; ++#define DSA_SELECTED_ROW_VEC_SMEM (GGML_PAD(((GGML_PAD(args.ne00, 128) + 4*OP_FLASH_ATTN_EXT_VEC_NCPSG + 2*GGML_PAD(args.ne20, 128))*nsg)*(sizeof(float)/2), 16)) ++#define DSA_SELECTED_ROW_TILED_SMEM (GGML_PAD((8*576 + 2*8*512 + 2*8*64 + 4*4*16*8)*sizeof(uint16_t), 16)) ++ const size_t smem = use_tiled ? DSA_SELECTED_ROW_TILED_SMEM : DSA_SELECTED_ROW_VEC_SMEM; ++#undef DSA_SELECTED_ROW_TILED_SMEM ++#undef DSA_SELECTED_ROW_VEC_SMEM + +- const int nth = std::min(1024, ne0); ++ const ggml_metal_device_props * props_dev = ggml_metal_device_get_props(ctx->dev); ++ GGML_ASSERT(smem <= props_dev->max_theadgroup_memory_size); + +- auto pipeline = ggml_metal_library_get_pipeline_arange(lib, op); ++ ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); ++ ggml_metal_buffer_id bid_tmp = bid_dst; ++ bid_tmp.offs += ggml_nbytes(op); ++ ++ ggml_metal_op_concurrency_reset(ctx); ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); // q ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); // k ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), 3); // v ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), 4); // top_k ++ ggml_metal_encoder_set_buffer (enc, nwg == 1 ? bid_dst : bid_tmp, 5); // dst ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), 6); // kq_mask_rows ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ const int grid_x = use_tiled ? (ne02 + 7)/8 : ne01; ++ const int grid_y = use_tiled ? ne01 : ne02; ++ const int grid_z = ne03*nwg; ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, 32, nsg, 1); ++ ++ if (nwg > 1) { ++ ggml_metal_op_concurrency_reset(ctx); ++ ++ ggml_metal_kargs_flash_attn_ext_vec_reduce reduce_args = { ++ (int32_t) (ne1*ne2*ne3), ++ }; ++ auto reduce_pipeline = ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce(lib, op, ne20, nwg); ++ GGML_ASSERT(32*nwg <= ggml_metal_pipeline_max_theads_per_threadgroup(reduce_pipeline)); ++ ggml_metal_encoder_set_pipeline(enc, reduce_pipeline); ++ ggml_metal_encoder_set_bytes(enc, &reduce_args, sizeof(reduce_args), 0); ++ ggml_metal_encoder_set_buffer(enc, bid_tmp, 1); ++ ggml_metal_encoder_set_buffer(enc, bid_dst, 2); ++ ggml_metal_encoder_dispatch_threadgroups(enc, reduce_args.nrows, 1, 1, 32*nwg, 1, 1); ++ } ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=dsa_sparse_attn kernel=%s tensor=%s q_type=%s k_type=%s v_type=%s mask_type=%s top_k_type=%s dst_type=%s q_width=%lld v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld top_k=%lld top_stream=%lld nwg=%d q_nb=%llu,%llu,%llu,%llu k_nb=%llu,%llu,%llu,%llu v_nb=%llu,%llu,%llu,%llu top_k_nb=%llu,%llu,%llu,%llu grid_x=%d grid_y=%d grid_z=%d threads_x=32 threads_y=%d\n", ++ use_tiled ? "selected_row_tiled" : "selected_row", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->src[3]->type), ++ ggml_type_name(op->src[4]->type), ++ ggml_type_name(op->type), ++ (long long) ne00, ++ (long long) ne20, ++ (long long) ne1, ++ (long long) ne2, ++ (long long) ne3, ++ (long long) ne11, ++ (long long) ne40, ++ (long long) ne42, ++ nwg, ++ (unsigned long long) nb00, ++ (unsigned long long) nb01, ++ (unsigned long long) nb02, ++ (unsigned long long) nb03, ++ (unsigned long long) nb10, ++ (unsigned long long) nb11, ++ (unsigned long long) nb12, ++ (unsigned long long) nb13, ++ (unsigned long long) nb20, ++ (unsigned long long) nb21, ++ (unsigned long long) nb22, ++ (unsigned long long) nb23, ++ (unsigned long long) nb40, ++ (unsigned long long) nb41, ++ (unsigned long long) nb42, ++ (unsigned long long) nb43, ++ grid_x, ++ grid_y, ++ grid_z, ++ nsg); ++ } ++ ++ return 1; ++ } ++ ++ const int requested_head_group = ggml_metal_glm_dsa_sparse_attn_decode_group_heads_requested(); ++ const bool use_decode_grouped = ++ requested_head_group > 1 && ++ ne1 == 1 && ++ ne3 == 1 && ++ ne40 <= 1024 && ++ ne2 % requested_head_group == 0; ++ const bool use_cached_topk = ++ !use_decode_grouped && ++ ggml_metal_glm_dsa_sparse_attn_cache_topk_enabled() && ++ ne40 <= 1024; ++ auto pipeline = use_decode_grouped ++ ? ggml_metal_library_get_pipeline_dsa_sparse_attn_decode_grouped(lib, op) ++ : use_cached_topk ++ ? ggml_metal_library_get_pipeline_dsa_sparse_attn_cached_topk(lib, op) ++ : ggml_metal_library_get_pipeline_dsa_sparse_attn(lib, op); ++ ++ int ida = 0; + + ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 1); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // k ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // v ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[3]), ida++); // kq_mask_rows ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[4]), ida++); // top_k ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst + +- ggml_metal_encoder_dispatch_threadgroups(enc, 1, 1, 1, nth, 1, 1); ++ const int nth_requested = ggml_metal_glm_dsa_sparse_attn_threads_for_shape( ++ ggml_metal_glm_dsa_sparse_attn_threads_requested(), ne1, ne40); ++ const int head_group = use_decode_grouped ? requested_head_group : 1; ++ const int max_threads_per_group = std::max(1, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)/head_group); ++ const int nth = std::min(nth_requested, max_threads_per_group); ++ const int grid_x = ne1; ++ const int grid_y = use_decode_grouped ? (ne2 + head_group - 1)/head_group : ne2; ++ const int grid_z = ne3; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=dsa_sparse_attn kernel=%s tensor=%s q_type=%s k_type=%s v_type=%s mask_type=%s top_k_type=%s dst_type=%s q_width=%lld v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld top_k=%lld top_stream=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=%d threads_y=%d\n", ++ use_decode_grouped ? "decode_grouped" : use_cached_topk ? "cached_topk" : "default", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->src[3]->type), ++ ggml_type_name(op->src[4]->type), ++ ggml_type_name(op->type), ++ (long long) ne00, ++ (long long) ne20, ++ (long long) ne1, ++ (long long) ne2, ++ (long long) ne3, ++ (long long) ne11, ++ (long long) ne40, ++ (long long) ne42, ++ grid_x, ++ grid_y, ++ grid_z, ++ nth, ++ head_group); ++ } ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, nth, head_group, 1); + + return 1; + } + +-int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_dsa_top1_attn(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -4396,33 +11549,65 @@ int ggml_metal_op_timestep_embedding(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- const int dim = op->op_params[0]; +- const int max_period = op->op_params[1]; +- +- ggml_metal_kargs_timestep_embedding args = { +- /*.nb1 =*/ nb1, +- /*.dim =*/ dim, +- /*.max_period =*/ max_period, ++ ggml_metal_kargs_dsa_top1_attn args = { ++ /*.ne00 =*/ ne00, /*.ne01 =*/ ne01, /*.ne02 =*/ ne02, /*.ne03 =*/ ne03, ++ /*.nb00 =*/ nb00, /*.nb01 =*/ nb01, /*.nb02 =*/ nb02, /*.nb03 =*/ nb03, ++ /*.ne10 =*/ ne10, /*.ne11 =*/ ne11, /*.ne12 =*/ ne12, /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, /*.nb11 =*/ nb11, /*.nb12 =*/ nb12, /*.nb13 =*/ nb13, ++ /*.ne20 =*/ ne20, /*.ne21 =*/ ne21, /*.ne22 =*/ ne22, /*.ne23 =*/ ne23, ++ /*.nb20 =*/ nb20, /*.nb21 =*/ nb21, /*.nb22 =*/ nb22, /*.nb23 =*/ nb23, ++ /*.ne0 =*/ ne0, /*.ne1 =*/ ne1, /*.ne2 =*/ ne2, /*.ne3 =*/ ne3, ++ /*.nb0 =*/ nb0, /*.nb1 =*/ nb1, /*.nb2 =*/ nb2, /*.nb3 =*/ nb3, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_timestep_embedding(lib, op); ++ auto pipeline = ggml_metal_library_get_pipeline_dsa_top1_attn(lib, op); + +- const int nth = std::max(1, std::min(1024, dim/2)); ++ int ida = 0; + + ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), ida++); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), ida++); // q ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), ida++); // v ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[2]), ida++); // top_k ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), ida++); // dst + +- ggml_metal_encoder_dispatch_threadgroups(enc, ne00, 1, 1, nth, 1, 1); ++ const int nth = std::min(64, std::max(1, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline))); ++ const int grid_x = (ne0 + nth - 1)/nth; ++ const int grid_y = ne1*ne2; ++ const int grid_z = ne3; ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=dsa_top1_attn tensor=%s q_type=%s v_type=%s top_k_type=%s dst_type=%s v_width=%lld batch=%lld heads=%lld stream=%lld kv=%lld value_heads=%lld top_stream=%lld grid_x=%d grid_y=%d grid_z=%d threads_x=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_type_name(op->src[0]->type), ++ ggml_type_name(op->src[1]->type), ++ ggml_type_name(op->src[2]->type), ++ ggml_type_name(op->type), ++ (long long) ne10, ++ (long long) ne1, ++ (long long) ne2, ++ (long long) ne3, ++ (long long) ne11, ++ (long long) ne12, ++ (long long) ne22, ++ grid_x, ++ grid_y, ++ grid_z, ++ nth); ++ } ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, grid_z, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_moe_route_weights(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; +@@ -4430,144 +11615,234 @@ int ggml_metal_op_argmax(ggml_metal_op_t ctx, int idx) { + + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); +- GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- ggml_metal_kargs_argmax args = { +- /*.ne00 = */ ne00, +- /*.nb01 = */ nb01, ++ ggml_metal_kargs_moe_route_weights args = { ++ /*.n_expert =*/ ne01, ++ /*.n_tokens =*/ ne11, ++ /*.n_expert_used =*/ ne10, ++ /*.norm =*/ ggml_get_op_params_i32(op, 2), ++ /*.clamp_min =*/ ggml_get_op_params_f32(op, 0), ++ /*.scale =*/ ggml_get_op_params_f32(op, 1), ++ /*._pad0 =*/ ggml_metal_glm_dsa_moe_route_weights_slot0_enabled() ? 1 : 0, ++ /*._pad1 =*/ 0, ++ /*.probs_nb1 =*/ nb01, ++ /*.probs_nb2 =*/ nb02, ++ /*.ids_nb0 =*/ nb10, ++ /*.ids_nb1 =*/ nb11, ++ /*.dst_nb1 =*/ nb1, ++ /*.dst_nb2 =*/ nb2, + }; + +- auto pipeline = ggml_metal_library_get_pipeline_argmax(lib, op); +- +- const int64_t nrows = ggml_nrows(op->src[0]); +- +- int nth = 32; // SIMD width +- while (nth < ne00 && nth*ne01*ne02*ne03 < 256) { +- nth *= 2; +- } +- +- const size_t smem = pipeline.smem; ++ auto pipeline = ggml_metal_library_get_pipeline_moe_route_weights(lib); + + ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); +- ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 2); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ++ const int nth = 1; ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=moe_route_weights tensor=%s probs=%s ids=%s experts=%d tokens=%d used_experts=%d norm=%d scale=%g force_slot0=%d grid_x=%d grid_y=1 grid_z=1 threads_x=%d\n", ++ ggml_metal_tensor_name(op), ++ ggml_metal_tensor_name(op->src[0]), ++ ggml_metal_tensor_name(op->src[1]), ++ args.n_expert, ++ args.n_tokens, ++ args.n_expert_used, ++ args.norm, ++ (double) args.scale, ++ args._pad0, ++ args.n_tokens, ++ nth); ++ } + +- ggml_metal_encoder_dispatch_threadgroups(enc, nrows, 1, 1, nth, 1, 1); ++ ggml_metal_encoder_dispatch_threadgroups(enc, args.n_tokens, 1, 1, nth, 1, 1); + + return 1; + } + +-int ggml_metal_op_argsort(ggml_metal_op_t ctx, int idx) { ++int ggml_metal_op_moe_weighted_sum(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + + ggml_metal_library_t lib = ctx->lib; + ggml_metal_encoder_t enc = ctx->enc; + +- GGML_ASSERT(ggml_is_contiguous_rows(op->src[0])); +- + GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); + GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); + GGML_TENSOR_LOCALS( int32_t, ne, op, ne); + GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- auto pipeline = ggml_metal_library_get_pipeline_argsort(lib, op); +- +- // bitonic sort requires the number of elements to be power of 2 +- int nth = 1; +- while (nth < ne00 && 2*nth <= ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)) { +- nth *= 2; +- } +- +- const int npr = (ne00 + nth - 1)/nth; +- +- // Metal kernels require the buffer size to be multiple of 16 bytes +- // https://developer.apple.com/documentation/metal/mtlcomputecommandencoder/1443142-setthreadgroupmemorylength +- const size_t smem = GGML_PAD(nth*sizeof(int32_t), 16); +- +- ggml_metal_buffer_id bid_src0 = ggml_metal_get_buffer_id(op->src[0]); +- ggml_metal_buffer_id bid_dst = ggml_metal_get_buffer_id(op); +- +- ggml_metal_buffer_id bid_tmp = bid_dst; +- bid_tmp.offs += ggml_nbytes(op); +- +- if ((int) ceil(std::log(npr) / std::log(2)) % 2 == 1) { +- std::swap(bid_dst, bid_tmp); +- } +- +- ggml_metal_kargs_argsort args = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.top_k =*/ nth, ++ ggml_metal_kargs_moe_weighted_sum args = { ++ /*.n_embd =*/ ne00, ++ /*.n_tokens =*/ ne02, ++ /*.n_expert_used =*/ ne01, ++ /*.already_weighted =*/ ggml_get_op_params_i32(op, 0) != 0 ? 1 : 0, ++ /*.experts_nb0 =*/ nb00, ++ /*.experts_nb1 =*/ nb01, ++ /*.experts_nb2 =*/ nb02, ++ /*.weights_nb1 =*/ nb11, ++ /*.weights_nb2 =*/ nb12, ++ /*.dst_nb0 =*/ nb0, ++ /*.dst_nb1 =*/ nb1, + }; + +- ggml_metal_encoder_set_pipeline(enc, pipeline); +- ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); +- +- ggml_metal_encoder_set_threadgroup_memory_size(enc, smem, 0); +- +- ggml_metal_encoder_dispatch_threadgroups(enc, npr*ne01, ne02, ne03, nth, 1, 1); ++ const bool use_x4 = ++ ne00 % 4 == 0 && ++ nb00 == sizeof(float) && ++ nb0 == sizeof(float) && ++ nb01 % (4*sizeof(float)) == 0 && ++ nb02 % (4*sizeof(float)) == 0 && ++ nb1 % (4*sizeof(float)) == 0; ++ auto pipeline = use_x4 ? ++ ggml_metal_library_get_pipeline_moe_weighted_sum_x4(lib) : ++ ggml_metal_library_get_pipeline_moe_weighted_sum(lib); + +- auto pipeline_merge = ggml_metal_library_get_pipeline_argsort_merge(lib, op); +- +- int len = nth; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer (enc, ggml_metal_get_buffer_id(op), 3); ++ ++ const int nth = std::min(256, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline)); ++ const int ncols = use_x4 ? ne00/4 : ne00; ++ const int grid_x = (ncols + nth - 1)/nth; ++ const int grid_y = ne02; ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=moe_weighted_sum kernel=%s tensor=%s experts=%s weights=%s embd=%lld tokens=%lld used_experts=%lld grid_x=%d grid_y=%d grid_z=1 threads_x=%d\n", ++ args.already_weighted ? "already_weighted" : (use_x4 ? "f32x4" : "f32"), ++ ggml_metal_tensor_name(op), ++ ggml_metal_tensor_name(op->src[0]), ++ ggml_metal_tensor_name(op->src[1]), ++ (long long) ne00, ++ (long long) ne02, ++ (long long) ne01, ++ grid_x, ++ grid_y, ++ nth); ++ } + +- while (len < ne00) { +- ggml_metal_op_concurrency_reset(ctx); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, 1, nth, 1, 1); + +- ggml_metal_kargs_argsort_merge args_merge = { +- /*.ne00 =*/ ne00, +- /*.ne01 =*/ ne01, +- /*.ne02 =*/ ne02, +- /*.ne03 =*/ ne03, +- /*.nb00 =*/ nb00, +- /*.nb01 =*/ nb01, +- /*.nb02 =*/ nb02, +- /*.nb03 =*/ nb03, +- /*.ne0 =*/ ne0, +- /*.ne1 =*/ ne1, +- /*.ne2 =*/ ne2, +- /*.ne3 =*/ ne3, +- /*.top_k =*/ ne00, +- /*.len =*/ len, +- }; ++ return 1; ++} + +- // merges per row +- const int nm = (ne00 + 2*len - 1) / (2*len); ++static int ggml_metal_encode_moe_mul_mat_id( ++ ggml_metal_op_t ctx, ++ ggml_tensor * op, ++ const ggml_metal_glm_moe_private_bindings * private_bindings) { ++ const bool private_path = private_bindings != nullptr; + +- const int nth = std::min(512, ggml_metal_pipeline_max_theads_per_threadgroup(pipeline_merge)); ++ ggml_metal_library_t lib = ctx->lib; ++ ggml_metal_encoder_t enc = ctx->enc; + +- ggml_metal_encoder_set_pipeline(enc, pipeline_merge); +- ggml_metal_encoder_set_bytes (enc, &args_merge, sizeof(args_merge), 0); +- ggml_metal_encoder_set_buffer (enc, bid_src0, 1); +- ggml_metal_encoder_set_buffer (enc, bid_dst, 2); +- ggml_metal_encoder_set_buffer (enc, bid_tmp, 3); ++ GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb0, op->src[0], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb1, op->src[1], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne2, op->src[2], ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb2, op->src[2], nb); ++ GGML_TENSOR_LOCALS(uint64_t, nb3, op->src[3], nb); ++ GGML_TENSOR_LOCALS( int32_t, ne, op, ne); ++ GGML_TENSOR_LOCALS(uint64_t, nb, op, nb); + +- ggml_metal_encoder_dispatch_threadgroups(enc, nm*ne01, ne02, ne03, nth, 1, 1); ++ const bool use_glm52_slot_parallel = ++ op->src[0]->type == GGML_TYPE_Q3_K && ++ ne00 == 2048 && ne01 == 6144 && ne02 == 256 && ++ op->src[1]->type == GGML_TYPE_F32 && ++ ne20 == 8 && ne21 == 1; ++ auto pipeline = use_glm52_slot_parallel ? ++ ggml_metal_library_get_pipeline_mul_mv_id_q3_weighted_reduce_slots_sg_r8_nb8_w0(lib, op) : ++ ggml_metal_library_get_pipeline_mul_mv_id_weighted_reduce(lib, op); ++ const int nr0 = pipeline.nr0; ++ const int nsg = pipeline.nsg; + +- std::swap(bid_dst, bid_tmp); ++ ggml_metal_kargs_mul_mv_id args = { ++ /*.nei0 =*/ ne20, ++ /*.nei1 =*/ ne21, ++ /*.nbi1 =*/ nb21, ++ /*.ne00 =*/ ne00, ++ /*.ne01 =*/ ne01, ++ /*.ne02 =*/ ne02, ++ /*.nb00 =*/ nb00, ++ /*.nb01 =*/ nb01, ++ /*.nb02 =*/ nb02, ++ /*.ne10 =*/ ne10, ++ /*.ne11 =*/ ne11, ++ /*.ne12 =*/ ne12, ++ /*.ne13 =*/ ne13, ++ /*.nb10 =*/ nb10, ++ /*.nb11 =*/ nb11, ++ /*.nb12 =*/ nb12, ++ /*.ne0 =*/ ne0, ++ /*.ne1 =*/ ne1, ++ /*.nb1 =*/ nb1, ++ /*.nr0 =*/ nr0, ++ }; ++ ggml_metal_kargs_mul_mv_id_weighted_reduce_extra extra = { ++ /*.weights_nb1 =*/ nb31, ++ /*.weights_nb2 =*/ nb32, ++ /*.dst_nb0 =*/ nb0, ++ /*.dst_nb1 =*/ nb1, ++ /*.already_weighted =*/ 0, ++ /*._pad0 =*/ 0, ++ }; + +- len <<= 1; ++ ggml_metal_encoder_set_pipeline(enc, pipeline); ++ ggml_metal_encoder_set_bytes(enc, &args, sizeof(args), 0); ++ ggml_metal_encoder_set_bytes(enc, &extra, sizeof(extra), 6); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op->src[0]), 1); ++ ggml_metal_encoder_set_buffer(enc, private_path ? private_bindings->activation : ggml_metal_get_buffer_id(op->src[1]), 2); ++ ggml_metal_encoder_set_buffer(enc, ggml_metal_get_buffer_id(op), 3); ++ ggml_metal_encoder_set_buffer(enc, private_path ? private_bindings->ids : ggml_metal_get_buffer_id(op->src[2]), 4); ++ ggml_metal_encoder_set_buffer(enc, private_path ? private_bindings->weights : ggml_metal_get_buffer_id(op->src[3]), 5); ++ ++ const int grid_x = use_glm52_slot_parallel ? ++ (ne01 + nr0 - 1)/nr0 : ++ (ne01 + nr0*nsg - 1)/(nr0*nsg); ++ const int grid_y = ne21; ++ ggml_metal_encoder_set_threadgroup_memory_size(enc, pipeline.smem, 0); ++ ggml_metal_encoder_dispatch_threadgroups(enc, grid_x, grid_y, 1, 32, nsg, 1); ++ ++ if (ggml_metal_glm_dsa_dispatch_log_enabled()) { ++ GGML_LOG_INFO( ++ "skippy: glm_dsa_metal_dispatch op=moe_mul_mat_id tensor=%s experts=%s input=%s ids=%s weights=%s kernel=%s private_scratch=%d src0_type=%s n_ff=%d n_embd=%d experts_total=%d used_experts=%d tokens=%d nr0=%d nsg=%d grid_x=%d grid_y=%d grid_z=1 threads_x=32 threads_y=%d fused_nodes=1\n", ++ ggml_metal_tensor_name(op), ++ ggml_metal_tensor_name(op->src[0]), ++ ggml_metal_tensor_name(op->src[1]), ++ ggml_metal_tensor_name(op->src[2]), ++ ggml_metal_tensor_name(op->src[3]), ++ use_glm52_slot_parallel ? "q3_r8_nb8_w0" : "sequential_slots", ++ private_path ? 1 : 0, ++ ggml_type_name(op->src[0]->type), ++ ne00, ++ ne01, ++ ne02, ++ ne20, ++ ne21, ++ nr0, ++ nsg, ++ grid_x, ++ grid_y, ++ nsg); + } + + return 1; + } + ++int ggml_metal_op_moe_mul_mat_id(ggml_metal_op_t ctx, int idx) { ++ return ggml_metal_encode_moe_mul_mat_id(ctx, ctx->node(idx), nullptr); ++} ++ + int ggml_metal_op_top_k(ggml_metal_op_t ctx, int idx) { + ggml_tensor * op = ctx->node(idx); + +diff --git a/ggml/src/ggml-metal/ggml-metal-ops.h b/ggml/src/ggml-metal/ggml-metal-ops.h +index 89a6ad82..5d22c181 100644 +--- a/ggml/src/ggml-metal/ggml-metal-ops.h ++++ b/ggml/src/ggml-metal/ggml-metal-ops.h +@@ -11,6 +11,8 @@ typedef struct ggml_metal_op * ggml_metal_op_t; + ggml_metal_op_t ggml_metal_op_init( + ggml_metal_device_t dev, + ggml_metal_cmd_buf_t cmd_buf, ++ struct ggml_metal_buffer_id fusion_scratch, ++ size_t fusion_scratch_size, + struct ggml_cgraph * gf, + int idx_start, + int idx_end, +@@ -36,12 +38,16 @@ size_t ggml_metal_op_mul_mat_id_extra_tpe(const struct ggml_tensor * op); + // id map [n_tokens, n_expert] + size_t ggml_metal_op_mul_mat_id_extra_ids(const struct ggml_tensor * op); + ++// routed activation preserved across shared-expert work ++size_t ggml_metal_op_mul_mat_id_extra_src1_scratch(const struct ggml_tensor * op); ++ + // return true if we should use the FA vector kernel for this op + bool ggml_metal_op_flash_attn_ext_use_vec(const struct ggml_tensor * op); + + 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_dsa_sparse_attn_extra_tmp(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); +@@ -59,6 +65,13 @@ int ggml_metal_op_ssm_conv (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_ssm_scan (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_rwkv (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_gated_delta_net (ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_lightning_indexer (ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_dsa_sparse_mask (ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_dsa_sparse_attn (ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_dsa_top1_attn (ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_moe_route_weights (ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_moe_weighted_sum (ggml_metal_op_t ctx, int idx); ++int ggml_metal_op_moe_mul_mat_id (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_solve_tri (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_set (ggml_metal_op_t ctx, int idx); + int ggml_metal_op_cpy (ggml_metal_op_t ctx, int idx); +diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp +index a1003b3a..3807bf76 100644 +--- a/ggml/src/ggml-metal/ggml-metal.cpp ++++ b/ggml/src/ggml-metal/ggml-metal.cpp +@@ -219,6 +219,7 @@ static size_t ggml_backend_metal_buffer_type_get_alloc_size(ggml_backend_buffer_ + { + res += ggml_metal_op_mul_mat_id_extra_tpe(tensor); + res += ggml_metal_op_mul_mat_id_extra_ids(tensor); ++ res += ggml_metal_op_mul_mat_id_extra_src1_scratch(tensor); + } break; + case GGML_OP_FLASH_ATTN_EXT: + { +@@ -226,6 +227,10 @@ 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); + } break; ++ case GGML_OP_DSA_SPARSE_ATTN: ++ { ++ res += ggml_metal_op_dsa_sparse_attn_extra_tmp(tensor); ++ } break; + case GGML_OP_CUMSUM: + case GGML_OP_ARGSORT: + { +diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal +index 969fddfa..b7b7e065 100644 +--- a/ggml/src/ggml-metal/ggml-metal.metal ++++ b/ggml/src/ggml-metal/ggml-metal.metal +@@ -1562,6 +1562,60 @@ typedef decltype(kernel_swiglu) kernel_swiglu_t; + template [[host_name("kernel_swiglu_f32")]] kernel kernel_swiglu_t kernel_swiglu; + template [[host_name("kernel_swiglu_f16")]] kernel kernel_swiglu_t kernel_swiglu; + ++kernel void kernel_swiglu_weighted_f32( ++ constant ggml_metal_kargs_glu_weighted & args, ++ device const char * src0, ++ device const char * src1, ++ device const char * weights, ++ device char * dst, ++ uint tgpig[[threadgroup_position_in_grid]], ++ uint tpitg[[thread_position_in_threadgroup]], ++ uint ntg[[threads_per_threadgroup]]) { ++ const int32_t slot = int32_t(tgpig % uint(args.ne1)); ++ const int32_t token = int32_t(tgpig / uint(args.ne1)); ++ ++ device const float * src0_row = (device const float *) ((device const char *) src0 + tgpig*args.nb01) + args.i00; ++ device const float * src1_row = (device const float *) ((device const char *) src1 + tgpig*args.nb11) + args.i10; ++ device float * dst_row = (device float *) ((device char *) dst + tgpig*args.dst_nb1); ++ ++ const float route_weight = ((device const float *) (weights + ++ uint64_t(slot)*args.weights_nb1 + uint64_t(token)*args.weights_nb2))[0]; ++ ++ for (int i0 = tpitg; i0 < args.ne0; i0 += ntg) { ++ const float x0 = src0_row[i0]; ++ const float x1 = src1_row[i0]; ++ const float silu = x0 / (1.0f + exp(-x0)); ++ dst_row[i0] = silu*x1*route_weight; ++ } ++} ++ ++kernel void kernel_swiglu_weighted_f32_4( ++ constant ggml_metal_kargs_glu_weighted & args, ++ device const char * src0, ++ device const char * src1, ++ device const char * weights, ++ device char * dst, ++ uint tgpig[[threadgroup_position_in_grid]], ++ uint tpitg[[thread_position_in_threadgroup]], ++ uint ntg[[threads_per_threadgroup]]) { ++ const int32_t slot = int32_t(tgpig % uint(args.ne1)); ++ const int32_t token = int32_t(tgpig / uint(args.ne1)); ++ ++ device const float4 * src0_row = (device const float4 *) ((device const char *) src0 + tgpig*args.nb01) + args.i00/4; ++ device const float4 * src1_row = (device const float4 *) ((device const char *) src1 + tgpig*args.nb11) + args.i10/4; ++ device float4 * dst_row = (device float4 *) ((device char *) dst + tgpig*args.dst_nb1); ++ ++ const float route_weight = ((device const float *) (weights + ++ uint64_t(slot)*args.weights_nb1 + uint64_t(token)*args.weights_nb2))[0]; ++ ++ for (int i0 = tpitg; i0 < args.ne0/4; i0 += ntg) { ++ const float4 x0 = src0_row[i0]; ++ const float4 x1 = src1_row[i0]; ++ const float4 silu = x0 / (1.0f + exp(-x0)); ++ dst_row[i0] = silu*x1*route_weight; ++ } ++} ++ + template + kernel void kernel_swiglu_oai( + constant ggml_metal_kargs_glu & args, +@@ -2873,6 +2927,347 @@ template [[host_name("kernel_gated_delta_net_f32_2")]] kernel kernel_gated_delta + template [[host_name("kernel_gated_delta_net_f32_4")]] kernel kernel_gated_delta_net_t kernel_gated_delta_net_impl; + #endif + ++static inline float lightning_indexer_mask( ++ constant ggml_metal_kargs_lightning_indexer & args, ++ device const char * mask, ++ int i_kv, ++ int i_batch, ++ int i_stream) { ++ const int i_mask_stream = i_stream % args.ne33; ++ device const half * value = (device const half *) ( ++ mask + i_kv*args.nb30 + i_batch*args.nb31 + i_mask_stream*args.nb33); ++ return float(*value); ++} ++ ++template ++kernel void kernel_lightning_indexer_impl( ++ constant ggml_metal_kargs_lightning_indexer & args, ++ device const char * q, ++ device const char * k, ++ device const char * weights, ++ device const char * mask, ++ device char * dst, ++ uint3 gid[[thread_position_in_grid]]) { ++ const int i_kv = gid.x; ++ const int i_batch = gid.y; ++ const int i_stream = gid.z; ++ ++ if (i_kv >= args.ne0 || i_batch >= args.ne1 || i_stream >= args.ne3) { ++ return; ++ } ++ ++ float score = 0.0f; ++ ++ for (int i_head = 0; i_head < args.ne01; ++i_head) { ++ float qk = 0.0f; ++ ++ int i_embd = 0; ++ for (; i_embd + 7 < args.ne00; i_embd += 8) { ++ device const float * q_ptr = (device const float *) (q + i_embd*args.nb00 + i_head*args.nb01 + i_batch*args.nb02 + i_stream*args.nb03); ++ device const K * k_ptr = (device const K *) (k + i_embd*args.nb10 + i_kv*args.nb12 + i_stream*args.nb13); ++ ++ qk += q_ptr[0] * float(k_ptr[0]); ++ qk += q_ptr[1] * float(k_ptr[1]); ++ qk += q_ptr[2] * float(k_ptr[2]); ++ qk += q_ptr[3] * float(k_ptr[3]); ++ qk += q_ptr[4] * float(k_ptr[4]); ++ qk += q_ptr[5] * float(k_ptr[5]); ++ qk += q_ptr[6] * float(k_ptr[6]); ++ qk += q_ptr[7] * float(k_ptr[7]); ++ } ++ for (; i_embd < args.ne00; ++i_embd) { ++ device const float * q_ptr = (device const float *) (q + i_embd*args.nb00 + i_head*args.nb01 + i_batch*args.nb02 + i_stream*args.nb03); ++ device const K * k_ptr = (device const K *) (k + i_embd*args.nb10 + i_kv*args.nb12 + i_stream*args.nb13); ++ ++ qk += *q_ptr * float(*k_ptr); ++ } ++ ++ device const float * weight_ptr = (device const float *) (weights + i_head*args.nb20 + i_batch*args.nb21 + i_stream*args.nb23); ++ score += max(qk, 0.0f) * *weight_ptr; ++ } ++ ++ device float * dst_ptr = (device float *) (dst + i_kv*args.nb0 + i_batch*args.nb1 + i_stream*args.nb3); ++ *dst_ptr = score + lightning_indexer_mask(args, mask, i_kv, i_batch, i_stream); ++} ++ ++typedef decltype(kernel_lightning_indexer_impl) kernel_lightning_indexer_t; ++ ++template ++kernel void kernel_lightning_indexer_staged_q_impl( ++ constant ggml_metal_kargs_lightning_indexer & args, ++ device const char * q, ++ device const char * k, ++ device const char * weights, ++ device const char * mask, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ uint tiitg [[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ threadgroup float * q_stage = (threadgroup float *) shmem; ++ const int tid = tiitg; ++ const int n_threads = tptg.x; ++ const int i_batch = tgpig.y; ++ const int i_stream = tgpig.z; ++ ++ const int q_values = args.ne00*args.ne01; ++ for (int flat = tid; flat < q_values; flat += n_threads) { ++ const int i_head = flat/args.ne00; ++ const int i_embd = flat - i_head*args.ne00; ++ device const float * q_ptr = (device const float *) ++ (q + i_embd*args.nb00 + i_head*args.nb01 + i_batch*args.nb02 + i_stream*args.nb03); ++ q_stage[flat] = *q_ptr; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ const int i_kv = tgpig.x*n_threads + tid; ++ if (i_kv >= args.ne0 || i_batch >= args.ne1 || i_stream >= args.ne3) { ++ return; ++ } ++ ++ float score = 0.0f; ++ for (int i_head = 0; i_head < args.ne01; ++i_head) { ++ float qk = 0.0f; ++ threadgroup const float * q_ptr = q_stage + i_head*args.ne00; ++ ++ int i_embd = 0; ++ for (; i_embd + 7 < args.ne00; i_embd += 8) { ++ device const K * k_ptr = (device const K *) ++ (k + i_embd*args.nb10 + i_kv*args.nb12 + i_stream*args.nb13); ++ qk += q_ptr[i_embd + 0] * float(k_ptr[0]); ++ qk += q_ptr[i_embd + 1] * float(k_ptr[1]); ++ qk += q_ptr[i_embd + 2] * float(k_ptr[2]); ++ qk += q_ptr[i_embd + 3] * float(k_ptr[3]); ++ qk += q_ptr[i_embd + 4] * float(k_ptr[4]); ++ qk += q_ptr[i_embd + 5] * float(k_ptr[5]); ++ qk += q_ptr[i_embd + 6] * float(k_ptr[6]); ++ qk += q_ptr[i_embd + 7] * float(k_ptr[7]); ++ } ++ for (; i_embd < args.ne00; ++i_embd) { ++ device const K * k_ptr = (device const K *) ++ (k + i_embd*args.nb10 + i_kv*args.nb12 + i_stream*args.nb13); ++ qk += q_ptr[i_embd] * float(*k_ptr); ++ } ++ ++ device const float * weight_ptr = (device const float *) ++ (weights + i_head*args.nb20 + i_batch*args.nb21 + i_stream*args.nb23); ++ score += max(qk, 0.0f) * *weight_ptr; ++ } ++ ++ device float * dst_ptr = (device float *) ++ (dst + i_kv*args.nb0 + i_batch*args.nb1 + i_stream*args.nb3); ++ *dst_ptr = score + lightning_indexer_mask(args, mask, i_kv, i_batch, i_stream); ++} ++ ++typedef decltype(kernel_lightning_indexer_staged_q_impl) kernel_lightning_indexer_staged_q_t; ++ ++template ++kernel void kernel_lightning_indexer_quant( ++ constant ggml_metal_kargs_lightning_indexer & args, ++ device const char * q, ++ device const char * k, ++ device const char * weights, ++ device const char * mask, ++ device char * dst, ++ uint3 gid[[thread_position_in_grid]]) { ++ const int i_kv = gid.x; ++ const int i_batch = gid.y; ++ const int i_stream = gid.z; ++ ++ if (i_kv >= args.ne0 || i_batch >= args.ne1 || i_stream >= args.ne3) { ++ return; ++ } ++ ++ float score = 0.0f; ++ ++ device const block_q * k_row = (device const block_q *) (k + i_kv*args.nb12 + i_stream*args.nb13); ++ ++ for (int i_head = 0; i_head < args.ne01; ++i_head) { ++ float qk = 0.0f; ++ ++ for (int i_embd = 0; i_embd < args.ne00; i_embd += 16) { ++ device const float * q_ptr = (device const float *) (q + i_embd*args.nb00 + i_head*args.nb01 + i_batch*args.nb02 + i_stream*args.nb03); ++ ++ float4x4 k_reg; ++ const int i_block = i_embd / (16*nl); ++ const short il = (i_embd / 16) % nl; ++ dequantize_func(k_row + i_block, il, k_reg); ++ ++ const int n_chunk = min(16, args.ne00 - i_embd); ++ for (int i = 0; i < n_chunk; ++i) { ++ qk += q_ptr[i] * k_reg[i/4][i%4]; ++ } ++ } ++ ++ device const float * weight_ptr = (device const float *) (weights + i_head*args.nb20 + i_batch*args.nb21 + i_stream*args.nb23); ++ score += max(qk, 0.0f) * *weight_ptr; ++ } ++ ++ device float * dst_ptr = (device float *) (dst + i_kv*args.nb0 + i_batch*args.nb1 + i_stream*args.nb3); ++ *dst_ptr = score + lightning_indexer_mask(args, mask, i_kv, i_batch, i_stream); ++} ++ ++typedef decltype(kernel_lightning_indexer_quant) kernel_lightning_indexer_quant_t; ++ ++#define LIGHTNING_INDEXER_PARALLEL_MAX_THREADS 1024 ++ ++template ++kernel void kernel_lightning_indexer_parallel_impl( ++ constant ggml_metal_kargs_lightning_indexer & args, ++ device const char * q, ++ device const char * k, ++ device const char * weights, ++ device const char * mask, ++ device char * dst, ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ uint tiitg [[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ threadgroup float shared_qk[LIGHTNING_INDEXER_PARALLEL_MAX_THREADS]; ++ threadgroup float shared_scores[LIGHTNING_INDEXER_PARALLEL_MAX_THREADS]; ++ ++ const int i_kv = tgpig.x; ++ const int i_batch = tgpig.y; ++ const int i_stream = tgpig.z; ++ const int tid = tiitg; ++ const int n_threads = tptg.x; ++ ++ const int lanes_per_head = max(1, n_threads / args.ne01); ++ const int i_head = tid / lanes_per_head; ++ const int i_lane = tid - i_head * lanes_per_head; ++ const bool active = i_kv < args.ne0 && i_batch < args.ne1 && i_stream < args.ne3 && i_head < args.ne01; ++ ++ float qk = 0.0f; ++ if (active) { ++ for (int i_embd = i_lane; i_embd < args.ne00; i_embd += lanes_per_head) { ++ device const float * q_ptr = (device const float *) (q + i_embd*args.nb00 + i_head*args.nb01 + i_batch*args.nb02 + i_stream*args.nb03); ++ device const K * k_ptr = (device const K *) (k + i_embd*args.nb10 + i_kv*args.nb12 + i_stream*args.nb13); ++ ++ qk += *q_ptr * float(*k_ptr); ++ } ++ } ++ ++ shared_qk[tid] = qk; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (active && i_lane == 0) { ++ float head_qk = 0.0f; ++ const int head_start = i_head * lanes_per_head; ++ for (int lane = 0; lane < lanes_per_head; ++lane) { ++ head_qk += shared_qk[head_start + lane]; ++ } ++ ++ device const float * weight_ptr = (device const float *) (weights + i_head*args.nb20 + i_batch*args.nb21 + i_stream*args.nb23); ++ shared_scores[i_head] = max(head_qk, 0.0f) * *weight_ptr; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (tid == 0 && i_kv < args.ne0 && i_batch < args.ne1 && i_stream < args.ne3) { ++ float score = 0.0f; ++ for (int i = 0; i < args.ne01; ++i) { ++ score += shared_scores[i]; ++ } ++ ++ device float * dst_ptr = (device float *) (dst + i_kv*args.nb0 + i_batch*args.nb1 + i_stream*args.nb3); ++ *dst_ptr = score + lightning_indexer_mask(args, mask, i_kv, i_batch, i_stream); ++ } ++} ++ ++template ++kernel void kernel_lightning_indexer_parallel_quant( ++ constant ggml_metal_kargs_lightning_indexer & args, ++ device const char * q, ++ device const char * k, ++ device const char * weights, ++ device const char * mask, ++ device char * dst, ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ uint tiitg [[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ threadgroup float shared_qk[LIGHTNING_INDEXER_PARALLEL_MAX_THREADS]; ++ threadgroup float shared_scores[LIGHTNING_INDEXER_PARALLEL_MAX_THREADS]; ++ ++ const int i_kv = tgpig.x; ++ const int i_batch = tgpig.y; ++ const int i_stream = tgpig.z; ++ const int tid = tiitg; ++ const int n_threads = tptg.x; ++ ++ const int lanes_per_head = max(1, n_threads / args.ne01); ++ const int i_head = tid / lanes_per_head; ++ const int i_lane = tid - i_head * lanes_per_head; ++ const bool active = i_kv < args.ne0 && i_batch < args.ne1 && i_stream < args.ne3 && i_head < args.ne01; ++ ++ float qk = 0.0f; ++ if (active) { ++ device const block_q * k_row = (device const block_q *) (k + i_kv*args.nb12 + i_stream*args.nb13); ++ ++ for (int i_embd = i_lane * 16; i_embd < args.ne00; i_embd += lanes_per_head * 16) { ++ device const float * q_ptr = (device const float *) (q + i_embd*args.nb00 + i_head*args.nb01 + i_batch*args.nb02 + i_stream*args.nb03); ++ ++ float4x4 k_reg; ++ const int i_block = i_embd / (16*nl); ++ const short il = (i_embd / 16) % nl; ++ dequantize_func(k_row + i_block, il, k_reg); ++ ++ const int n_chunk = min(16, args.ne00 - i_embd); ++ for (int i = 0; i < n_chunk; ++i) { ++ qk += q_ptr[i] * k_reg[i/4][i%4]; ++ } ++ } ++ } ++ ++ shared_qk[tid] = qk; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (active && i_lane == 0) { ++ float head_qk = 0.0f; ++ const int head_start = i_head * lanes_per_head; ++ for (int lane = 0; lane < lanes_per_head; ++lane) { ++ head_qk += shared_qk[head_start + lane]; ++ } ++ ++ device const float * weight_ptr = (device const float *) (weights + i_head*args.nb20 + i_batch*args.nb21 + i_stream*args.nb23); ++ shared_scores[i_head] = max(head_qk, 0.0f) * *weight_ptr; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (tid == 0 && i_kv < args.ne0 && i_batch < args.ne1 && i_stream < args.ne3) { ++ float score = 0.0f; ++ for (int i = 0; i < args.ne01; ++i) { ++ score += shared_scores[i]; ++ } ++ ++ device float * dst_ptr = (device float *) (dst + i_kv*args.nb0 + i_batch*args.nb1 + i_stream*args.nb3); ++ *dst_ptr = score + lightning_indexer_mask(args, mask, i_kv, i_batch, i_stream); ++ } ++} ++ ++typedef decltype(kernel_lightning_indexer_parallel_impl) kernel_lightning_indexer_parallel_t; ++typedef decltype(kernel_lightning_indexer_parallel_quant) kernel_lightning_indexer_parallel_quant_t; ++ ++template [[host_name("kernel_lightning_indexer_q4_0")]] kernel kernel_lightning_indexer_quant_t kernel_lightning_indexer_quant; ++template [[host_name("kernel_lightning_indexer_q8_0")]] kernel kernel_lightning_indexer_quant_t kernel_lightning_indexer_quant; ++template [[host_name("kernel_lightning_indexer_q2_K")]] kernel kernel_lightning_indexer_quant_t kernel_lightning_indexer_quant; ++template [[host_name("kernel_lightning_indexer_q3_K")]] kernel kernel_lightning_indexer_quant_t kernel_lightning_indexer_quant; ++template [[host_name("kernel_lightning_indexer_q4_K")]] kernel kernel_lightning_indexer_quant_t kernel_lightning_indexer_quant; ++template [[host_name("kernel_lightning_indexer_q5_K")]] kernel kernel_lightning_indexer_quant_t kernel_lightning_indexer_quant; ++template [[host_name("kernel_lightning_indexer_q6_K")]] kernel kernel_lightning_indexer_quant_t kernel_lightning_indexer_quant; ++ ++template [[host_name("kernel_lightning_indexer_parallel_q4_0")]] kernel kernel_lightning_indexer_parallel_quant_t kernel_lightning_indexer_parallel_quant; ++template [[host_name("kernel_lightning_indexer_parallel_q8_0")]] kernel kernel_lightning_indexer_parallel_quant_t kernel_lightning_indexer_parallel_quant; ++template [[host_name("kernel_lightning_indexer_parallel_q2_K")]] kernel kernel_lightning_indexer_parallel_quant_t kernel_lightning_indexer_parallel_quant; ++template [[host_name("kernel_lightning_indexer_parallel_q3_K")]] kernel kernel_lightning_indexer_parallel_quant_t kernel_lightning_indexer_parallel_quant; ++template [[host_name("kernel_lightning_indexer_parallel_q4_K")]] kernel kernel_lightning_indexer_parallel_quant_t kernel_lightning_indexer_parallel_quant; ++template [[host_name("kernel_lightning_indexer_parallel_q5_K")]] kernel kernel_lightning_indexer_parallel_quant_t kernel_lightning_indexer_parallel_quant; ++template [[host_name("kernel_lightning_indexer_parallel_q6_K")]] kernel kernel_lightning_indexer_parallel_quant_t kernel_lightning_indexer_parallel_quant; ++ ++template [[host_name("kernel_lightning_indexer_f32")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer_impl; ++template [[host_name("kernel_lightning_indexer_f16")]] kernel kernel_lightning_indexer_t kernel_lightning_indexer_impl; ++template [[host_name("kernel_lightning_indexer_staged_q_f32")]] kernel kernel_lightning_indexer_staged_q_t kernel_lightning_indexer_staged_q_impl; ++template [[host_name("kernel_lightning_indexer_staged_q_f16")]] kernel kernel_lightning_indexer_staged_q_t kernel_lightning_indexer_staged_q_impl; ++template [[host_name("kernel_lightning_indexer_parallel_f32")]] kernel kernel_lightning_indexer_parallel_t kernel_lightning_indexer_parallel_impl; ++template [[host_name("kernel_lightning_indexer_parallel_f16")]] kernel kernel_lightning_indexer_parallel_t kernel_lightning_indexer_parallel_impl; ++ + constant short FC_solve_tri_nsg [[function_constant(FC_SOLVE_TRI + 0)]]; + constant short FC_solve_tri_n [[function_constant(FC_SOLVE_TRI + 1)]]; + constant short FC_solve_tri_k [[function_constant(FC_SOLVE_TRI + 2)]]; +@@ -3683,6 +4078,73 @@ void kernel_mul_mv_q1_0_f32_impl( + } + } + ++template ++void kernel_mul_mv_tq2_0_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ const int first_row = (r0*NSG + sgitg)*nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ const uint64_t offset1 = r1*args.nb11 + i12*args.nb12 + i13*args.nb13; ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ device const block_tq2_0 * ax[nr0]; ++ for (short row = 0; row < nr0; ++row) { ++ const uint64_t offset0 = (first_row + row)*args.nb01 + ++ (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ ax[row] = (device const block_tq2_0 *) (src0 + offset0); ++ } ++ ++ float sumf[nr0] = {0.f}; ++ for (int ib = 0; ib < nb; ++ib) { ++ device const float * yb = y + ib*QK_K; ++ const float4 y0 = {yb[tiisg], yb[32 + tiisg], yb[64 + tiisg], yb[96 + tiisg]}; ++ const float4 y1 = {yb[128 + tiisg], yb[160 + tiisg], yb[192 + tiisg], yb[224 + tiisg]}; ++ const float sumy0 = y0[0] + y0[1] + y0[2] + y0[3]; ++ const float sumy1 = y1[0] + y1[1] + y1[2] + y1[3]; ++ ++ for (short row = 0; row < nr0; ++row) { ++ device const block_tq2_0 * xb = ax[row] + ib; ++ const uint q0 = xb->qs[tiisg]; ++ const uint q1 = xb->qs[32 + tiisg]; ++ const float4 v0 = { ++ float( q0 & 3), float((q0 >> 2) & 3), ++ float((q0 >> 4) & 3), float((q0 >> 6) & 3), ++ }; ++ const float4 v1 = { ++ float( q1 & 3), float((q1 >> 2) & 3), ++ float((q1 >> 4) & 3), float((q1 >> 6) & 3), ++ }; ++ const float dot_q = dot(y0, v0) + dot(y1, v1) - sumy0 - sumy1; ++ sumf[row] += float(xb->d)*dot_q; ++ } ++ } ++ ++ device float * dst_f32 = (device float *) dst + ++ (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ for (short row = 0; row < nr0; ++row) { ++ const float total = simd_sum(sumf[row]); ++ if (tiisg == 0 && first_row + row < args.ne01) { ++ dst_f32[first_row + row] = total; ++ } ++ } ++ ++ (void) shmem; ++} ++ + [[host_name("kernel_mul_mv_q1_0_f32")]] + kernel void kernel_mul_mv_q1_0_f32( + constant ggml_metal_kargs_mul_mv & args, +@@ -3775,6 +4237,19 @@ kernel void kernel_mul_mv_q2_0_f32( + kernel_mul_mv_q2_0_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); + } + ++[[host_name("kernel_mul_mv_tq2_0_f32")]] ++kernel void kernel_mul_mv_tq2_0_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_tq2_0_f32_impl( ++ args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++} ++ + kernel void kernel_mul_mv_q4_0_f32( + constant ggml_metal_kargs_mul_mv & args, + device const char * src0, +@@ -3823,7 +4298,7 @@ kernel void kernel_mul_mv_q5_1_f32( + mul_vec_q_n_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + } + +-template ++template + void kernel_mul_mv_q8_0_f32_impl( + args_t args, + device const char * src0, +@@ -3881,9 +4356,19 @@ void kernel_mul_mv_q8_0_f32_impl( + for (short row = 0; row < NR0; row++) { + device const int8_t * qs = ax[row][ib].qs + il*NQ; + +- float sumq = 0.f; +- FOR_UNROLL (short i = 0; i < NQ; ++i) { +- sumq += qs[i] * yl[i]; ++ float sumq; ++ if (vector_dot) { ++ sumq = dot( ++ float4(qs[0], qs[1], qs[2], qs[3]), ++ float4(yl[0], yl[1], yl[2], yl[3])); ++ sumq += dot( ++ float4(qs[4], qs[5], qs[6], qs[7]), ++ float4(yl[4], yl[5], yl[6], yl[7])); ++ } else { ++ sumq = 0.f; ++ FOR_UNROLL (short i = 0; i < NQ; ++i) { ++ sumq += qs[i] * yl[i]; ++ } + } + + sumf[row] += sumq*ax[row][ib].d; +@@ -3910,37 +4395,180 @@ kernel void kernel_mul_mv_q8_0_f32( + kernel_mul_mv_q8_0_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); + } + +-// mat-vec kernel processing in chunks of float4 +-// chpb - chunks per quantization block +-template +-void kernel_mul_mv_ext_q4_f32_impl( +- constant ggml_metal_kargs_mul_mv_ext & args, ++[[host_name("kernel_mul_mv_q8_0_f32_vector_dot")]] ++kernel void kernel_mul_mv_q8_0_f32_vector_dot( ++ constant ggml_metal_kargs_mul_mv & args, + device const char * src0, + device const char * src1, + device char * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { +- const short NSG = FC_mul_mv_nsg; +- const short nxpsg = FC_mul_mv_nxpsg; +- +- const short chpt = 4; // chunks per thread ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_q8_0_f32_impl( ++ args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} + +- //const short nxpsg = (32); +- const short nypsg = (32/nxpsg); ++[[host_name("kernel_mul_mv_q8_0_f32_r1")]] ++kernel void kernel_mul_mv_q8_0_f32_r1( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_q8_0_f32_impl<1, constant ggml_metal_kargs_mul_mv &>( ++ args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} + +- const short tx = tiisg%nxpsg; +- const short ty = tiisg/nxpsg; ++[[host_name("kernel_mul_mv_q8_0_f32_r4")]] ++kernel void kernel_mul_mv_q8_0_f32_r4( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_q8_0_f32_impl<4, constant ggml_metal_kargs_mul_mv &>( ++ args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} + +- const int i01 = tgpig.x*(nypsg*NSG) + nypsg*sgitg + ty; +- const int i11 = tgpig.y*r1ptg; +- const int i1m = tgpig.z; ++[[host_name("kernel_mul_mv_q8_0_f32_r8")]] ++kernel void kernel_mul_mv_q8_0_f32_r8( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_q8_0_f32_impl<8, constant ggml_metal_kargs_mul_mv &>( ++ args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} + +- const int i12 = i1m%FC_mul_mv_ne12; +- const int i13 = i1m/FC_mul_mv_ne12; ++template ++void kernel_mul_mv_q8_0_f32_row_parallel_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; + +- const uint64_t offset0 = i01*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = i11*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ constexpr short NW = N_SIMDWIDTH; ++ constexpr short NQ = 8; ++ ++ const int nb = args.ne00/QK8_0; ++ const int r0 = (tgpig.x*NSG + sgitg)*NR0; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ const uint64_t offset1 = r1*args.nb11 + i12*args.nb12 + i13*args.nb13; ++ ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ device const block_q8_0 * ax[NR0]; ++ FOR_UNROLL (short row = 0; row < NR0; ++row) { ++ const int src_row = min(r0 + row, args.ne01 - 1); ++ const uint64_t offset0 = src_row*args.nb01 + ++ (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0); ++ } ++ ++ float sumf[NR0] = { 0.f }; ++ ++ const short ix = tiisg/(NW/NQ); ++ const short il = tiisg%(NW/NQ); ++ const int ib0 = ix; ++ ++ float yl[NQ]; ++ device const float * yb = y + ib0*QK8_0 + il*NQ; ++ ++ for (int ib = ib0; ib < nb; ib += NQ) { ++ FOR_UNROLL (short i = 0; i < NQ; ++i) { ++ yl[i] = yb[i]; ++ } ++ ++ FOR_UNROLL (short row = 0; row < NR0; ++row) { ++ device const int8_t * qs = ax[row][ib].qs + il*NQ; ++ ++ float sumq = 0.f; ++ FOR_UNROLL (short i = 0; i < NQ; ++i) { ++ sumq += qs[i]*yl[i]; ++ } ++ ++ sumf[row] += sumq*ax[row][ib].d; ++ } ++ ++ yb += NQ*QK8_0; ++ } ++ ++ device float * dst_f32 = (device float *) dst + ++ (uint64_t) im*args.ne0*args.ne1 + (uint64_t) r1*args.ne0; ++ ++ FOR_UNROLL (short row = 0; row < NR0; ++row) { ++ const float total = simd_sum(sumf[row]); ++ if (tiisg == 0 && r0 + row < args.ne01) { ++ dst_f32[r0 + row] = total; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_q8_0_f32_row_parallel")]] ++kernel void kernel_mul_mv_q8_0_f32_row_parallel( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ (void) shmem; ++ kernel_mul_mv_q8_0_f32_row_parallel_impl<4, constant ggml_metal_kargs_mul_mv &>( ++ args, src0, src1, dst, tgpig, tiisg, sgitg); ++} ++ ++// mat-vec kernel processing in chunks of float4 ++// chpb - chunks per quantization block ++template ++void kernel_mul_mv_ext_q4_f32_impl( ++ constant ggml_metal_kargs_mul_mv_ext & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const short NSG = FC_mul_mv_nsg; ++ const short nxpsg = FC_mul_mv_nxpsg; ++ ++ const short chpt = 4; // chunks per thread ++ ++ //const short nxpsg = (32); ++ const short nypsg = (32/nxpsg); ++ ++ const short tx = tiisg%nxpsg; ++ const short ty = tiisg/nxpsg; ++ ++ const int i01 = tgpig.x*(nypsg*NSG) + nypsg*sgitg + ty; ++ const int i11 = tgpig.y*r1ptg; ++ const int i1m = tgpig.z; ++ ++ const int i12 = i1m%FC_mul_mv_ne12; ++ const int i13 = i1m/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = i01*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = i11*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; + + device const q_t * xq = (i01 < args.ne01) ? (device const q_t *) (src0 + offset0) + tx/chpb : (device const q_t *) src0; + +@@ -4864,6 +5492,148 @@ template [[host_name("kernel_rope_multi_f16")]] kernel kernel_rope_multi_t kerne + template [[host_name("kernel_rope_vision_f32")]] kernel kernel_rope_vision_t kernel_rope_vision; + template [[host_name("kernel_rope_vision_f16")]] kernel kernel_rope_vision_t kernel_rope_vision; + ++[[host_name("kernel_glm_absorbed_q_q8_q4")]] ++kernel void kernel_glm_absorbed_q_q8_q4( ++ constant ggml_metal_kargs_glm_absorbed_q & fused_args, ++ constant ggml_metal_kargs_rope & rope_args, ++ device const char * q_b_bytes, ++ device const float * q_input, ++ device const char * wk_b_bytes, ++ device const int32_t * pos, ++ device const float * freq_factors, ++ device float * q_packed, ++ threadgroup float * q_shared [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort3 ntg [[threads_per_threadgroup]], ++ ushort tpitg [[thread_index_in_threadgroup]], ++ ushort sgitg [[simdgroup_index_in_threadgroup]], ++ ushort tiisg [[thread_index_in_simdgroup]]) { ++ const int head = tgpig.x; ++ if (head >= fused_args.n_head) { ++ return; ++ } ++ ++ // Reproduce the ordinary Q8 matvec's four-SIMD-group/two-row reduction, ++ // but retain all 256 per-head outputs in threadgroup memory. ++ threadgroup float * q_partials = q_shared + fused_args.q_head_dim; ++ const int q_group = sgitg/4; ++ const int q_local_sg = sgitg%4; ++ const int q_ix = tiisg/4; ++ const int q_il = tiisg%4; ++ for (int row_base = 0; row_base < fused_args.q_head_dim; row_base += 16) { ++ const int row0 = row_base + 2*q_group; ++ float sumf[2] = { 0.0f, 0.0f }; ++ const int block0 = q_local_sg*8 + q_ix; ++ for (int block = block0; block < fused_args.q_rank/QK8_0; block += 32) { ++ float values[8]; ++ FOR_UNROLL (short i = 0; i < 8; ++i) { ++ values[i] = q_input[block*QK8_0 + q_il*8 + i]; ++ } ++ FOR_UNROLL (short row = 0; row < 2; ++row) { ++ const int global_row = head*fused_args.q_head_dim + row0 + row; ++ device const block_q8_0 * weights = ++ (device const block_q8_0 *) (q_b_bytes + uint64_t(global_row)*fused_args.q_b_nb1); ++ device const int8_t * quants = weights[block].qs + q_il*8; ++ float block_sum = 0.0f; ++ FOR_UNROLL (short i = 0; i < 8; ++i) { ++ block_sum += float(quants[i])*values[i]; ++ } ++ sumf[row] += block_sum*float(weights[block].d); ++ } ++ } ++ FOR_UNROLL (short row = 0; row < 2; ++row) { ++ sumf[row] = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ q_partials[q_group*8 + row*4 + q_local_sg] = sumf[row]; ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (q_local_sg == 0) { ++ FOR_UNROLL (short row = 0; row < 2; ++row) { ++ const float partial = tiisg < 4 ? q_partials[q_group*8 + row*4 + tiisg] : 0.0f; ++ const float total = simd_sum(partial); ++ if (tiisg == 0) { ++ q_shared[row0 + row] = total; ++ } ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ // Reuse the ordinary Q4 matvec's four-row SIMD decomposition, consuming ++ // the on-chip no-position query and publishing directly into Qcur. ++ const int q4_ix = tiisg/2; ++ const int q4_il = (tiisg%2)*8; ++ for (int row_base = 0; row_base < fused_args.q_abs_dim; row_base += 128) { ++ const int row0 = row_base + 4*sgitg; ++ float sumf[4] = { 0.0f, 0.0f, 0.0f, 0.0f }; ++ if (q4_ix < fused_args.q_nope_dim/QK4_0) { ++ threadgroup const float * values = q_shared + q4_ix*QK4_0 + q4_il; ++ float yl[16]; ++ float sumy[2] = { 0.0f, 0.0f }; ++ FOR_UNROLL (short i = 0; i < 8; i += 2) { ++ sumy[0] += values[i + 0] + values[i + 1]; ++ yl[i + 0] = values[i + 0]; ++ yl[i + 1] = values[i + 1]/256.0f; ++ sumy[1] += values[i + 16] + values[i + 17]; ++ yl[i + 8] = values[i + 16]/16.0f; ++ yl[i + 9] = values[i + 17]/4096.0f; ++ } ++ FOR_UNROLL (short row = 0; row < 4; ++row) { ++ device const block_q4_0 * weights = ++ (device const block_q4_0 *) (wk_b_bytes + uint64_t(head)*fused_args.wk_b_nb2 + ++ uint64_t(row0 + row)*fused_args.wk_b_nb1); ++ sumf[row] = block_q_n_dot_y(weights + q4_ix, sumy[0] + sumy[1], yl, q4_il); ++ } ++ } ++ FOR_UNROLL (short row = 0; row < 4; ++row) { ++ const float total = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ q_packed[head*(fused_args.q_abs_dim + fused_args.rope_dim) + row0 + row] = total; ++ } ++ } ++ } ++ ++ if (tpitg < fused_args.rope_dim/2) { ++ const int i0 = 2*tpitg; ++ const int ic = tpitg; ++ float corr_dims[2]; ++ rope_yarn_corr_dims( ++ rope_args.n_dims, ++ rope_args.n_ctx_orig, ++ rope_args.freq_base, ++ rope_args.beta_fast, ++ rope_args.beta_slow, ++ corr_dims); ++ const float theta_base = float(pos[0]); ++ const float theta = theta_base*pow(rope_args.freq_base, -float(i0)/rope_args.n_dims); ++ const float freq_factor = rope_args.src2 ? freq_factors[ic] : 1.0f; ++ float cos_theta; ++ float sin_theta; ++ rope_yarn( ++ theta/freq_factor, ++ rope_args.freq_scale, ++ corr_dims, ++ i0, ++ rope_args.ext_factor, ++ rope_args.attn_factor, ++ &cos_theta, ++ &sin_theta); ++ ++ const bool is_neox = (fused_args.rope_mode & 2) != 0; ++ const int x0_index = is_neox ? ic : 2*ic; ++ const int x1_index = is_neox ? ic + fused_args.rope_dim/2 : 2*ic + 1; ++ const int q_rope_offset = fused_args.q_nope_dim; ++ const float x0 = q_shared[q_rope_offset + x0_index]; ++ const float x1 = q_shared[q_rope_offset + x1_index]; ++ const float y0 = x0*cos_theta - x1*sin_theta; ++ const float y1 = x0*sin_theta + x1*cos_theta; ++ const int packed_rope_offset = head*(fused_args.q_abs_dim + fused_args.rope_dim) + fused_args.q_abs_dim; ++ q_packed[packed_rope_offset + x0_index] = y0; ++ q_packed[packed_rope_offset + x1_index] = y1; ++ } ++} ++ + typedef void (im2col_t)( + constant ggml_metal_kargs_im2col & args, + device const float * x, +@@ -6016,136 +6786,624 @@ kernel void kernel_argsort_f32_i32( + template [[host_name("kernel_argsort_f32_i32_asc")]] kernel argsort_t kernel_argsort_f32_i32; + template [[host_name("kernel_argsort_f32_i32_desc")]] kernel argsort_t kernel_argsort_f32_i32; + +-typedef void (argsort_merge_t)( +- constant ggml_metal_kargs_argsort_merge & args, +- device const char * src0, +- device const int32_t * tmp, +- device int32_t * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort3 tpitg[[thread_position_in_threadgroup]], +- ushort3 ntg[[threads_per_threadgroup]]); +- +-template +-kernel void kernel_argsort_merge_f32_i32( +- constant ggml_metal_kargs_argsort_merge & args, +- device const char * src0, +- device const int32_t * tmp, +- device int32_t * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort3 tpitg[[thread_position_in_threadgroup]], +- ushort3 ntg[[threads_per_threadgroup]]) { +- +- const int im = tgpig[0] / args.ne01; +- const int i01 = tgpig[0] % args.ne01; +- const int i02 = tgpig[1]; +- const int i03 = tgpig[2]; ++kernel void kernel_topk_moe_route_f32_i32( ++ constant ggml_metal_kargs_topk_moe_route & args, ++ device const char * logits, ++ device const char * bias, ++ device char * ids, ++ device char * weights, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort3 tpitg[[thread_position_in_threadgroup]]) { ++ constexpr int32_t MAX_EXPERTS = 256; ++ constexpr int32_t MAX_TOP_K = 16; ++ threadgroup float scores[MAX_EXPERTS]; ++ threadgroup float probs[MAX_EXPERTS]; ++ threadgroup int32_t expert_ids[MAX_EXPERTS]; ++ threadgroup float work_scores[MAX_EXPERTS]; ++ threadgroup float work_probs[MAX_EXPERTS]; ++ threadgroup int32_t work_ids[MAX_EXPERTS]; ++ threadgroup float top_probs[MAX_TOP_K + 1]; ++ threadgroup int32_t top_ids[MAX_TOP_K]; ++ ++ const int32_t token = tgpig.x; ++ const int32_t tid = tpitg.x; ++ if (token >= args.n_tokens || args.top_k <= 0 || args.top_k > MAX_TOP_K || args.n_expert > MAX_EXPERTS) { ++ return; ++ } + +- const int start = im * (2 * args.len); ++ float prob = 0.0f; ++ float score = -INFINITY; ++ if (tid < args.n_expert) { ++ const float logit = ((device const float *) (logits + tid*args.logits_nb0 + token*args.logits_nb1))[0]; ++ prob = 1.0f/(1.0f + exp(-logit)); ++ score = prob + (args.has_bias ? ((device const float *) (bias + tid*args.bias_nb0))[0] : 0.0f); ++ } ++ scores[tid] = score; ++ probs[tid] = prob; ++ expert_ids[tid] = tid < args.n_expert ? tid : -1; + +- const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start))); +- const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len))); ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- const int total = len0 + len1; ++ for (int32_t top = 0; top < args.top_k; ++top) { ++ work_scores[tid] = scores[tid]; ++ work_probs[tid] = probs[tid]; ++ work_ids[tid] = expert_ids[tid]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- device const int32_t * tmp0 = tmp + start +- + i01*args.ne0 +- + i02*args.ne0*args.ne01 +- + i03*args.ne0*args.ne01*args.ne02; ++ for (int32_t stride = MAX_EXPERTS/2; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ const float lhs_score = work_scores[tid]; ++ const float rhs_score = work_scores[tid + stride]; ++ const int32_t lhs_id = work_ids[tid]; ++ const int32_t rhs_id = work_ids[tid + stride]; ++ const bool rhs_better = ++ rhs_score > lhs_score || ++ (rhs_score == lhs_score && rhs_id >= 0 && (lhs_id < 0 || rhs_id < lhs_id)); ++ ++ if (rhs_better) { ++ work_scores[tid] = rhs_score; ++ work_probs[tid] = work_probs[tid + stride]; ++ work_ids[tid] = rhs_id; ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +- device const int32_t * tmp1 = tmp0 + args.len; ++ if (tid == 0) { ++ top_probs[top] = work_probs[0]; ++ top_ids[top] = work_ids[0]; ++ if (work_ids[0] >= 0 && work_ids[0] < args.n_expert) { ++ scores[work_ids[0]] = -INFINITY; ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +- dst += start +- + i01*args.top_k +- + i02*args.top_k*args.ne01 +- + i03*args.top_k*args.ne01*args.ne02; ++ if (tid == 0) { ++ const int32_t active_top_k = args._pad1 > 0 ? min(args.top_k, args._pad1) : args.top_k; ++ float sum = 0.0f; ++ for (int32_t i = 0; i < active_top_k; ++i) { ++ sum += top_probs[i]; ++ } ++ if (args.norm) { ++ sum = max(sum, args.clamp_min); ++ } + +- device const float * src0_row = (device const float *)(src0 +- + args.nb01*i01 +- + args.nb02*i02 +- + args.nb03*i03); ++ for (int32_t i = 0; i < args.top_k; ++i) { ++ float weight = i < active_top_k ? top_probs[i] : 0.0f; ++ if (args.norm) { ++ weight /= sum; ++ } ++ top_probs[i] = weight * args.scale; ++ } + +- if (total == 0) { +- return; ++ if (args._pad0 != 0) { ++ for (int32_t i = 0; i < args.top_k; ++i) { ++ for (int32_t j = i + 1; j < args.top_k; ++j) { ++ const int32_t lhs = top_ids[i]; ++ const int32_t rhs = top_ids[j]; ++ const bool rhs_first = rhs >= 0 && (lhs < 0 || rhs < lhs); ++ if (rhs_first) { ++ const int32_t tmp_id = top_ids[i]; ++ const float tmp_prob = top_probs[i]; ++ top_ids[i] = top_ids[j]; ++ top_probs[i] = top_probs[j]; ++ top_ids[j] = tmp_id; ++ top_probs[j] = tmp_prob; ++ } ++ } ++ } ++ } + } ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- const int chunk = (total + ntg.x - 1) / ntg.x; +- +- const int k0 = tpitg.x * chunk; +- const int k1 = MIN(MIN(k0 + chunk, total), args.top_k); +- +- if (k0 >= args.top_k) { +- return; ++ if (tid < args.top_k) { ++ ((device int32_t *) (ids + tid*args.ids_nb0 + token*args.ids_nb1))[0] = top_ids[tid]; ++ ((device float *) (weights + tid*args.weights_nb1 + token*args.weights_nb2))[0] = top_probs[tid]; + } ++} + +- if (k0 >= total) { ++kernel void kernel_topk_moe_route_f32_i32_sg_reduce( ++ constant ggml_metal_kargs_topk_moe_route & args, ++ device const char * logits, ++ device const char * bias, ++ device char * ids, ++ device char * weights, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort3 tpitg[[thread_position_in_threadgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]], ++ ushort tiisg[[thread_index_in_simdgroup]]) { ++ constexpr int32_t MAX_EXPERTS = 256; ++ constexpr int32_t MAX_TOP_K = 16; ++ constexpr int32_t N_SIMDGROUPS = 8; ++ threadgroup float candidate_scores[N_SIMDGROUPS*MAX_TOP_K]; ++ threadgroup float candidate_probs[N_SIMDGROUPS*MAX_TOP_K]; ++ threadgroup int32_t candidate_ids[N_SIMDGROUPS*MAX_TOP_K]; ++ threadgroup float top_probs[MAX_TOP_K + 1]; ++ threadgroup int32_t top_ids[MAX_TOP_K]; ++ ++ const int32_t token = tgpig.x; ++ const int32_t tid = tpitg.x; ++ const int32_t lane = tiisg; ++ const int32_t simdgroup = sgitg; ++ ++ if (token >= args.n_tokens || args.top_k <= 0 || args.top_k > MAX_TOP_K || args.n_expert > MAX_EXPERTS) { + return; + } + +- int low = k0 > len1 ? k0 - len1 : 0; +- int high = MIN(k0, len0); +- +- // binary-search partition (i, j) such that i + j = k +- while (low < high) { +- const int mid = (low + high) >> 1; +- +- const int32_t idx0 = tmp0[mid]; +- const int32_t idx1 = tmp1[k0 - mid - 1]; ++ float lane_prob = 0.0f; ++ float lane_score = -INFINITY; ++ const int32_t expert_id = tid; ++ if (tid < args.n_expert) { ++ const float logit = ((device const float *) (logits + tid*args.logits_nb0 + token*args.logits_nb1))[0]; ++ lane_prob = 1.0f/(1.0f + exp(-logit)); ++ lane_score = lane_prob + (args.has_bias ? ((device const float *) (bias + tid*args.bias_nb0))[0] : 0.0f); ++ } + +- const float val0 = src0_row[idx0]; +- const float val1 = src0_row[idx1]; ++ for (int32_t top = 0; top < args.top_k; ++top) { ++ const float max_score = simd_max(lane_score); ++ const int32_t lane_id = lane_score == max_score ? expert_id : MAX_EXPERTS; ++ const int32_t max_id = simd_min(lane_id); ++ const float selected_prob = expert_id == max_id ? lane_prob : 0.0f; ++ const float max_prob = simd_sum(selected_prob); + +- bool take_left; +- if (order == GGML_SORT_ORDER_ASC) { +- take_left = (val0 <= val1); +- } else { +- take_left = (val0 >= val1); ++ if (lane == 0) { ++ const int32_t candidate = simdgroup*args.top_k + top; ++ candidate_scores[candidate] = max_score; ++ candidate_probs[candidate] = max_prob; ++ candidate_ids[candidate] = max_id == MAX_EXPERTS ? -1 : max_id; + } + +- if (take_left) { +- low = mid + 1; +- } else { +- high = mid; ++ if (expert_id == max_id) { ++ lane_score = -INFINITY; + } + } + +- int i = low; +- int j = k0 - i; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- // keep the merge fronts into registers +- int32_t idx0 = 0; +- float val0 = 0.0f; +- if (i < len0) { +- idx0 = tmp0[i]; +- val0 = src0_row[idx0]; +- } ++ if (simdgroup == 0) { ++ const int32_t n_candidates = N_SIMDGROUPS*args.top_k; ++ float local_scores[4]; ++ float local_probs[4]; ++ int32_t local_ids[4]; ++ ++ for (int32_t i = 0; i < 4; ++i) { ++ const int32_t candidate = lane + i*32; ++ if (candidate < n_candidates) { ++ local_scores[i] = candidate_scores[candidate]; ++ local_probs[i] = candidate_probs[candidate]; ++ local_ids[i] = candidate_ids[candidate]; ++ } else { ++ local_scores[i] = -INFINITY; ++ local_probs[i] = 0.0f; ++ local_ids[i] = -1; ++ } ++ } + +- int32_t idx1 = 0; +- float val1 = 0.0f; +- if (j < len1) { +- idx1 = tmp1[j]; +- val1 = src0_row[idx1]; +- } ++ for (int32_t top = 0; top < args.top_k; ++top) { ++ float best_score = local_scores[0]; ++ float best_prob = local_probs[0]; ++ int32_t best_id = local_ids[0]; ++ ++ for (int32_t i = 1; i < 4; ++i) { ++ const bool better = ++ local_scores[i] > best_score || ++ (local_scores[i] == best_score && local_ids[i] >= 0 && (best_id < 0 || local_ids[i] < best_id)); ++ if (better) { ++ best_score = local_scores[i]; ++ best_prob = local_probs[i]; ++ best_id = local_ids[i]; ++ } ++ } + +- for (int k = k0; k < k1; ++k) { +- int32_t out_idx; ++ const float max_score = simd_max(best_score); ++ const int32_t winner_id = best_score == max_score ? best_id : MAX_EXPERTS; ++ const int32_t max_id = simd_min(winner_id); ++ const float selected_prob = best_id == max_id ? best_prob : 0.0f; ++ const float max_prob = simd_sum(selected_prob); + +- if (i >= len0) { +- while (k < k1) { +- dst[k++] = tmp1[j++]; +- } +- break; +- } else if (j >= len1) { +- while (k < k1) { +- dst[k++] = tmp0[i++]; ++ if (lane == 0) { ++ top_probs[top] = max_prob; ++ top_ids[top] = max_id == MAX_EXPERTS ? -1 : max_id; + } +- break; +- } else { +- bool take_left; + +- if (order == GGML_SORT_ORDER_ASC) { +- take_left = (val0 <= val1); +- } else { +- take_left = (val0 >= val1); ++ for (int32_t i = 0; i < 4; ++i) { ++ if (local_ids[i] == max_id) { ++ local_scores[i] = -INFINITY; ++ } ++ } ++ } ++ ++ float sum = 0.0f; ++ if (lane == 0) { ++ const int32_t active_top_k = args._pad1 > 0 ? min(args.top_k, args._pad1) : args.top_k; ++ if (args.norm) { ++ for (int32_t i = 0; i < active_top_k; ++i) { ++ sum += top_probs[i]; ++ } ++ sum = max(sum, args.clamp_min); ++ } ++ ++ for (int32_t i = 0; i < args.top_k; ++i) { ++ float weight = i < active_top_k ? top_probs[i] : 0.0f; ++ if (args.norm) { ++ weight /= sum; ++ } ++ top_probs[i] = weight * args.scale; ++ } ++ ++ if (args._pad0 != 0) { ++ for (int32_t i = 0; i < args.top_k; ++i) { ++ for (int32_t j = i + 1; j < args.top_k; ++j) { ++ const int32_t lhs = top_ids[i]; ++ const int32_t rhs = top_ids[j]; ++ const bool rhs_first = rhs >= 0 && (lhs < 0 || rhs < lhs); ++ if (rhs_first) { ++ const int32_t tmp_id = top_ids[i]; ++ const float tmp_prob = top_probs[i]; ++ top_ids[i] = top_ids[j]; ++ top_probs[i] = top_probs[j]; ++ top_ids[j] = tmp_id; ++ top_probs[j] = tmp_prob; ++ } ++ } ++ } ++ } ++ ++ for (int32_t i = 0; i < args.top_k; ++i) { ++ ((device int32_t *) (ids + i*args.ids_nb0 + token*args.ids_nb1))[0] = top_ids[i]; ++ ((device float *) (weights + i*args.weights_nb1 + token*args.weights_nb2))[0] = top_probs[i]; ++ } ++ } ++ } ++} ++ ++kernel void kernel_topk_moe_route_glm_256_8_sg32( ++ constant ggml_metal_kargs_topk_moe_route & args, ++ device const char * logits, ++ device const char * bias, ++ device char * ids, ++ device char * weights, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort3 tpitg[[thread_position_in_threadgroup]], ++ ushort tiisg[[thread_index_in_simdgroup]]) { ++ constexpr int32_t N_EXPERTS = 256; ++ constexpr int32_t TOP_K = 8; ++ constexpr int32_t LANES = 32; ++ constexpr int32_t EXPERTS_PER_LANE = N_EXPERTS/LANES; ++ ++ const int32_t token = tgpig.x; ++ const int32_t lane = tiisg; ++ ++ if (token >= args.n_tokens || args.n_expert != N_EXPERTS || args.top_k != TOP_K || tpitg.x >= LANES) { ++ return; ++ } ++ ++ float local_scores[EXPERTS_PER_LANE]; ++ float local_probs[EXPERTS_PER_LANE]; ++ int32_t local_ids[EXPERTS_PER_LANE]; ++ ++ for (int32_t i = 0; i < EXPERTS_PER_LANE; ++i) { ++ const int32_t expert = lane + i*LANES; ++ const float logit = ((device const float *) (logits + expert*args.logits_nb0 + token*args.logits_nb1))[0]; ++ const float prob = 1.0f/(1.0f + exp(-logit)); ++ local_probs[i] = prob; ++ local_scores[i] = prob + (args.has_bias ? ((device const float *) (bias + expert*args.bias_nb0))[0] : 0.0f); ++ local_ids[i] = expert; ++ } ++ ++ float top_probs[TOP_K]; ++ int32_t top_ids[TOP_K]; ++ ++ for (int32_t top = 0; top < TOP_K; ++top) { ++ float best_score = local_scores[0]; ++ float best_prob = local_probs[0]; ++ int32_t best_id = local_ids[0]; ++ ++ for (int32_t i = 1; i < EXPERTS_PER_LANE; ++i) { ++ const bool better = ++ local_scores[i] > best_score || ++ (local_scores[i] == best_score && local_ids[i] >= 0 && (best_id < 0 || local_ids[i] < best_id)); ++ if (better) { ++ best_score = local_scores[i]; ++ best_prob = local_probs[i]; ++ best_id = local_ids[i]; ++ } ++ } ++ ++ const float max_score = simd_max(best_score); ++ const int32_t winner_id = best_score == max_score ? best_id : N_EXPERTS; ++ const int32_t max_id = simd_min(winner_id); ++ const float selected_prob = best_id == max_id ? best_prob : 0.0f; ++ const float max_prob = simd_sum(selected_prob); ++ ++ if (lane == 0) { ++ top_probs[top] = max_prob; ++ top_ids[top] = max_id == N_EXPERTS ? -1 : max_id; ++ } ++ ++ for (int32_t i = 0; i < EXPERTS_PER_LANE; ++i) { ++ if (local_ids[i] == max_id) { ++ local_scores[i] = -INFINITY; ++ } ++ } ++ } ++ ++ if (lane == 0) { ++ const int32_t active_top_k = args._pad1 > 0 ? min(TOP_K, args._pad1) : TOP_K; ++ float sum = 0.0f; ++ if (args.norm) { ++ for (int32_t i = 0; i < active_top_k; ++i) { ++ sum += top_probs[i]; ++ } ++ sum = max(sum, args.clamp_min); ++ } ++ ++ for (int32_t i = 0; i < TOP_K; ++i) { ++ float weight = i < active_top_k ? top_probs[i] : 0.0f; ++ if (args.norm) { ++ weight /= sum; ++ } ++ top_probs[i] = weight * args.scale; ++ } ++ ++ if (args._pad0 != 0) { ++ for (int32_t i = 0; i < TOP_K; ++i) { ++ for (int32_t j = i + 1; j < TOP_K; ++j) { ++ const int32_t lhs = top_ids[i]; ++ const int32_t rhs = top_ids[j]; ++ const bool rhs_first = rhs >= 0 && (lhs < 0 || rhs < lhs); ++ if (rhs_first) { ++ const int32_t tmp_id = top_ids[i]; ++ const float tmp_prob = top_probs[i]; ++ top_ids[i] = top_ids[j]; ++ top_probs[i] = top_probs[j]; ++ top_ids[j] = tmp_id; ++ top_probs[j] = tmp_prob; ++ } ++ } ++ } ++ } ++ ++ for (int32_t i = 0; i < TOP_K; ++i) { ++ ((device int32_t *) (ids + i*args.ids_nb0 + token*args.ids_nb1))[0] = top_ids[i]; ++ ((device float *) (weights + i*args.weights_nb1 + token*args.weights_nb2))[0] = top_probs[i]; ++ } ++ } ++} ++ ++kernel void kernel_moe_route_weights_f32_i32( ++ constant ggml_metal_kargs_moe_route_weights & args, ++ device const char * probs, ++ device const char * ids, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort3 tpitg[[thread_position_in_threadgroup]]) { ++ if (tpitg.x != 0) { ++ return; ++ } ++ ++ constexpr int32_t MAX_TOP_K = 16; ++ float weights[MAX_TOP_K]; ++ ++ const int32_t token = int32_t(tgpig.x); ++ if (token >= args.n_tokens || args.n_expert_used <= 0 || args.n_expert_used > MAX_TOP_K) { ++ return; ++ } ++ ++ float sum = 0.0f; ++ for (int32_t i = 0; i < args.n_expert_used; ++i) { ++ const int32_t expert = ((device const int32_t *) (ids + i*args.ids_nb0 + token*args.ids_nb1))[0]; ++ float weight = 0.0f; ++ if (expert >= 0 && expert < args.n_expert) { ++ weight = ((device const float *) (probs + expert*args.probs_nb1 + token*args.probs_nb2))[0]; ++ } ++ weights[i] = weight; ++ sum += weight; ++ } ++ ++ if (args.norm) { ++ sum = max(sum, args.clamp_min); ++ } ++ ++ for (int32_t i = 0; i < args.n_expert_used; ++i) { ++ float weight = weights[i]; ++ if (args.norm) { ++ weight /= sum; ++ } ++ weight *= args.scale; ++ if (args._pad0 != 0) { ++ weight = i == 0 ? 1.0f : 0.0f; ++ } ++ ((device float *) (dst + i*args.dst_nb1 + token*args.dst_nb2))[0] = weight; ++ } ++} ++ ++kernel void kernel_moe_weighted_sum_f32( ++ constant ggml_metal_kargs_moe_weighted_sum & args, ++ device const char * experts, ++ device const char * weights, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiitg[[thread_index_in_threadgroup]], ++ ushort3 tptg[[threads_per_threadgroup]]) { ++ const int32_t i_embd = int32_t(tgpig.x*tptg.x + tiitg); ++ const int32_t token = int32_t(tgpig.y); ++ ++ if (i_embd >= args.n_embd || token >= args.n_tokens) { ++ return; ++ } ++ ++ float acc = 0.0f; ++ for (int32_t expert = 0; expert < args.n_expert_used; ++expert) { ++ const float value = ((device const float *) (experts + ++ i_embd*args.experts_nb0 + expert*args.experts_nb1 + token*args.experts_nb2))[0]; ++ const float weight = args.already_weighted ? 1.0f : ((device const float *) (weights + ++ expert*args.weights_nb1 + token*args.weights_nb2))[0]; ++ acc += value*weight; ++ } ++ ++ ((device float *) (dst + i_embd*args.dst_nb0 + token*args.dst_nb1))[0] = acc; ++} ++ ++kernel void kernel_moe_weighted_sum_f32x4( ++ constant ggml_metal_kargs_moe_weighted_sum & args, ++ device const char * experts, ++ device const char * weights, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiitg[[thread_index_in_threadgroup]], ++ ushort3 tptg[[threads_per_threadgroup]]) { ++ const int32_t vec = int32_t(tgpig.x*tptg.x + tiitg); ++ const int32_t i_embd = 4*vec; ++ const int32_t token = int32_t(tgpig.y); ++ ++ if (i_embd >= args.n_embd || token >= args.n_tokens) { ++ return; ++ } ++ ++ float4 acc = float4(0.0f); ++ for (int32_t expert = 0; expert < args.n_expert_used; ++expert) { ++ const float4 value = ((device const float4 *) (experts + ++ i_embd*args.experts_nb0 + expert*args.experts_nb1 + token*args.experts_nb2))[0]; ++ const float weight = args.already_weighted ? 1.0f : ((device const float *) (weights + ++ expert*args.weights_nb1 + token*args.weights_nb2))[0]; ++ acc += value*weight; ++ } ++ ++ ((device float4 *) (dst + i_embd*args.dst_nb0 + token*args.dst_nb1))[0] = acc; ++} ++ ++typedef void (argsort_merge_t)( ++ constant ggml_metal_kargs_argsort_merge & args, ++ device const char * src0, ++ device const int32_t * tmp, ++ device int32_t * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort3 tpitg[[thread_position_in_threadgroup]], ++ ushort3 ntg[[threads_per_threadgroup]]); ++ ++template ++kernel void kernel_argsort_merge_f32_i32( ++ constant ggml_metal_kargs_argsort_merge & args, ++ device const char * src0, ++ device const int32_t * tmp, ++ device int32_t * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort3 tpitg[[thread_position_in_threadgroup]], ++ ushort3 ntg[[threads_per_threadgroup]]) { ++ ++ const int im = tgpig[0] / args.ne01; ++ const int i01 = tgpig[0] % args.ne01; ++ const int i02 = tgpig[1]; ++ const int i03 = tgpig[2]; ++ ++ const int start = im * (2 * args.len); ++ ++ const int len0 = MIN(args.len, MAX(0, args.ne0 - (int)(start))); ++ const int len1 = MIN(args.len, MAX(0, args.ne0 - (int)(start + args.len))); ++ ++ const int total = len0 + len1; ++ ++ device const int32_t * tmp0 = tmp + start ++ + i01*args.ne0 ++ + i02*args.ne0*args.ne01 ++ + i03*args.ne0*args.ne01*args.ne02; ++ ++ device const int32_t * tmp1 = tmp0 + args.len; ++ ++ dst += start ++ + i01*args.top_k ++ + i02*args.top_k*args.ne01 ++ + i03*args.top_k*args.ne01*args.ne02; ++ ++ device const float * src0_row = (device const float *)(src0 ++ + args.nb01*i01 ++ + args.nb02*i02 ++ + args.nb03*i03); ++ ++ if (total == 0) { ++ return; ++ } ++ ++ const int chunk = (total + ntg.x - 1) / ntg.x; ++ ++ const int k0 = tpitg.x * chunk; ++ const int k1 = MIN(MIN(k0 + chunk, total), args.top_k); ++ ++ if (k0 >= args.top_k) { ++ return; ++ } ++ ++ if (k0 >= total) { ++ return; ++ } ++ ++ int low = k0 > len1 ? k0 - len1 : 0; ++ int high = MIN(k0, len0); ++ ++ // binary-search partition (i, j) such that i + j = k ++ while (low < high) { ++ const int mid = (low + high) >> 1; ++ ++ const int32_t idx0 = tmp0[mid]; ++ const int32_t idx1 = tmp1[k0 - mid - 1]; ++ ++ const float val0 = src0_row[idx0]; ++ const float val1 = src0_row[idx1]; ++ ++ bool take_left; ++ if (order == GGML_SORT_ORDER_ASC) { ++ take_left = (val0 <= val1); ++ } else { ++ take_left = (val0 >= val1); ++ } ++ ++ if (take_left) { ++ low = mid + 1; ++ } else { ++ high = mid; ++ } ++ } ++ ++ int i = low; ++ int j = k0 - i; ++ ++ // keep the merge fronts into registers ++ int32_t idx0 = 0; ++ float val0 = 0.0f; ++ if (i < len0) { ++ idx0 = tmp0[i]; ++ val0 = src0_row[idx0]; ++ } ++ ++ int32_t idx1 = 0; ++ float val1 = 0.0f; ++ if (j < len1) { ++ idx1 = tmp1[j]; ++ val1 = src0_row[idx1]; ++ } ++ ++ for (int k = k0; k < k1; ++k) { ++ int32_t out_idx; ++ ++ if (i >= len0) { ++ while (k < k1) { ++ dst[k++] = tmp1[j++]; ++ } ++ break; ++ } else if (j >= len1) { ++ while (k < k1) { ++ dst[k++] = tmp0[i++]; ++ } ++ break; ++ } else { ++ bool take_left; ++ ++ if (order == GGML_SORT_ORDER_ASC) { ++ take_left = (val0 <= val1); ++ } else { ++ take_left = (val0 >= val1); + } + + if (take_left) { +@@ -7797,14 +9055,20 @@ kernel void kernel_flash_attn_ext_vec_reduce( + const uint64_t rid = tgpig; + + const short iwg = tiisg; ++ const bool active = iwg < NWG; + + device const float * ss = (device const float *) htmp + (uint64_t)args.nrows*DV*NWG; + +- float S = ss[rid*(2*NWG) + 2*iwg + 0]; +- float M = ss[rid*(2*NWG) + 2*iwg + 1]; ++ float S = 0.0f; ++ float M = -FLT_MAX/2; ++ ++ if (active) { ++ S = ss[rid*(2*NWG) + 2*iwg + 0]; ++ M = ss[rid*(2*NWG) + 2*iwg + 1]; ++ } + + const float m = simd_max(M); +- const float ms = exp(M - m); ++ const float ms = active ? exp(M - m) : 0.0f; + + S = simd_sum(S*ms); + S = S == 0.0f ? 0.0f : 1.0f/S; +@@ -7815,7 +9079,11 @@ kernel void kernel_flash_attn_ext_vec_reduce( + device float4 * dst4 = (device float4 *) dst + rid*DV4; + + for (short i = sgitg; i < DV4; i += NWG) { +- const float4 v = simd_sum(htmp4[i*NWG + iwg]*ms); ++ float4 partial = 0.0f; ++ if (active) { ++ partial = htmp4[i*NWG + iwg]*ms; ++ } ++ const float4 v = simd_sum(partial); + + if (iwg == 0) { + dst4[i] = v*S; +@@ -8026,7 +9294,7 @@ template [[host_name("kernel_concat_i16")]] kernel kernel_concat_t kernel_conca + template [[host_name("kernel_concat_i32")]] kernel kernel_concat_t kernel_concat; + template [[host_name("kernel_concat_i64")]] kernel kernel_concat_t kernel_concat; + +-template ++template + void kernel_mul_mv_q2_K_f32_impl( + args_t args, + device const char * src0, +@@ -8038,7 +9306,7 @@ void kernel_mul_mv_q2_K_f32_impl( + ushort sgitg) { + const short NSG = FC_mul_mv_nsg; + +- const int nb = args.ne00/QK_K; ++ const int nb = fixed_nb > 0 ? fixed_nb : args.ne00/QK_K; + + const int r0 = tgpig.x; + const int r1 = tgpig.y; +@@ -8131,299 +9399,615 @@ kernel void kernel_mul_mv_q2_K_f32( + kernel_mul_mv_q2_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); + } + +-template +-void kernel_mul_mv_q3_K_f32_impl( +- args_t args, +- device const char * src0, ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, + device const char * src1, + device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { + const short NSG = FC_mul_mv_nsg; +- ++ const int nr0 = 8; + const int nb = args.ne00/QK_K; + + const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; +- ++ const int token = tgpig.y; + const int first_row = (r0 * NSG + sgitg) * nr0; + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; +- +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; +- +- device const block_q3_K * x = (device const block_q3_K *) (src0 + offset0); +- device const float * yy = (device const float *) (src1 + offset1); +- ++ float sumf[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + float yl[32]; + +- //const uint16_t kmask1 = 0x3030; +- //const uint16_t kmask2 = 0x0f0f; +- +- const short tid = tiisg/4; +- const short ix = tiisg%4; +- const short ip = tid/4; // 0 or 1 +- const short il = 2*((tid%4)/2); // 0 or 2 +- const short ir = tid%2; +- const short l0 = 8*ir; +- +- // One would think that the Metal compiler would figure out that ip and il can only have +- // 4 possible states, and optimize accordingly. Well, no. It needs help, and we do it +- // with these two tales. +- // +- // Possible masks for the high bit +- const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200}, // ip = 0, il = 0 +- {0x0004, 0x0400, 0x0008, 0x0800}, // ip = 0, il = 2 +- {0x0010, 0x1000, 0x0020, 0x2000}, // ip = 1, il = 0 +- {0x0040, 0x4000, 0x0080, 0x8000}}; // ip = 1, il = 2 +- +- // Possible masks for the low 2 bits +- const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}}; +- +- const ushort4 hm = mm[2*ip + il/2]; +- +- const short shift = 2*il; +- +- const float v1 = il == 0 ? 4.f : 64.f; +- const float v2 = 4.f * v1; +- +- const uint16_t s_shift1 = 4*ip; +- const uint16_t s_shift2 = s_shift1 + il; +- +- const short q_offset = 32*ip + l0; +- const short y_offset = 128*ip + 32*il + l0; ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ const short iq = it/4; ++ const short ir = it%4; ++ const short is = (8*ir)/16; + +- device const float * y1 = yy + ix*QK_K + y_offset; +- +- uint32_t scales32, aux32; +- thread uint16_t * scales16 = (thread uint16_t *)&scales32; +- thread const int8_t * scales = (thread const int8_t *)&scales32; ++ if (extra._pad0 != 0) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[0]; ++ if (expert_id < 0 || expert_id >= args.ne02) { ++ return; ++ } + +- float sumf1[nr0] = {0.f}; +- float sumf2[nr0] = {0.f}; ++ device const char * src0_cur = src0s + uint64_t(expert_id)*args.nb02; ++ device const char * src1_cur = src1 + uint64_t(token)*args.nb12; ++ ggml_metal_kargs_mul_mv args0 = { ++ /*.ne00 =*/ args.ne00, ++ /*.ne01 =*/ args.ne01, ++ /*.ne02 =*/ 1, ++ /*.nb00 =*/ args.nb00, ++ /*.nb01 =*/ args.nb01, ++ /*.nb02 =*/ args.nb02, ++ /*.nb03 =*/ args.nb02, ++ /*.ne10 =*/ args.ne10, ++ /*.ne11 =*/ 1, ++ /*.ne12 =*/ 1, ++ /*.nb10 =*/ args.nb10, ++ /*.nb11 =*/ args.nb11, ++ /*.nb12 =*/ args.nb12, ++ /*.nb13 =*/ args.nb12, ++ /*.ne0 =*/ args.ne01, ++ /*.ne1 =*/ 1, ++ /*.nr0 =*/ args.nr0, ++ /*.r2 =*/ 1, ++ /*.r3 =*/ 1, ++ }; ++ kernel_mul_mv_q2_K_f32_impl<8, ggml_metal_kargs_mul_mv, 8>( ++ args0, ++ src0_cur, ++ src1_cur, ++ dst, ++ nullptr, ++ uint3(tgpig.x, 0, 0), ++ tiisg, ++ sgitg); ++ return; ++ } + +- for (int i = ix; i < nb; i += 4) { +- for (short l = 0; l < 8; ++l) { +- yl[l+ 0] = y1[l+ 0]; +- yl[l+ 8] = y1[l+16]; +- yl[l+16] = y1[l+32]; +- yl[l+24] = y1[l+48]; ++ for (int slot = 0; slot < args.nei0; ++slot) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id < 0 || expert_id >= args.ne02) { ++ continue; + } + +- device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset); +- device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0); +- device const uint16_t * a = (device const uint16_t *)(x[i].scales); +- device const half * dh = &x[i].d; ++ const float route_weight = extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0]; ++ if (route_weight == 0.0f) { ++ continue; ++ } + +- for (short row = 0; row < nr0; ++row) { +- const float d_all = (float)dh[0]; ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot)*args.nb11 + uint64_t(token)*args.nb12; + +- scales16[0] = a[4]; +- scales16[1] = a[5]; +- aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030; +- scales16[0] = a[il+0]; +- scales16[1] = a[il+1]; +- scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32; ++ device const block_q2_K * x = (device const block_q2_K *) (src0s + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir; + +- float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0; +- for (short l = 0; l < 8; l += 2) { +- const int32_t qs = q[l/2]; +- s1 += yl[l+0] * (qs & qm[il/2][0]); +- s2 += yl[l+1] * (qs & qm[il/2][1]); +- s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]); +- s4 += yl[l+16] * (qs & qm[il/2][2]); +- s5 += yl[l+17] * (qs & qm[il/2][3]); +- s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]); ++ for (int ib = ix; ib < nb; ib += 4) { ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ for (short i = 0; i < 8; ++i) { ++ yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0]; ++ yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8]; ++ yl[i+16] = y4[i+64]; sumy[2] += yl[i+16]; ++ yl[i+24] = y4[i+96]; sumy[3] += yl[i+24]; + } +- float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); +- float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); +- sumf1[row] += d1 * (scales[0] - 32); +- sumf2[row] += d2 * (scales[2] - 32); + +- s1 = s2 = s3 = s4 = s5 = s6 = 0; +- for (short l = 0; l < 8; l += 2) { +- const int32_t qs = q[l/2+8]; +- s1 += yl[l+8] * (qs & qm[il/2][0]); +- s2 += yl[l+9] * (qs & qm[il/2][1]); +- s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]); +- s4 += yl[l+24] * (qs & qm[il/2][2]); +- s5 += yl[l+25] * (qs & qm[il/2][3]); +- s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]); ++ device const uint8_t * sc = (device const uint8_t *)x[ib].scales + 8*iq + is; ++ device const uint16_t * qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; ++ device const half * dh = &x[ib].d; ++ ++ for (short row = 0; row < nr0; row++) { ++ float4 acc1 = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2 = {0.f, 0.f, 0.f, 0.f}; ++ for (int i = 0; i < 8; i += 2) { ++ acc1[0] += yl[i+ 0] * (qs[i/2] & 0x0003); ++ acc2[0] += yl[i+ 1] * (qs[i/2] & 0x0300); ++ acc1[1] += yl[i+ 8] * (qs[i/2] & 0x000c); ++ acc2[1] += yl[i+ 9] * (qs[i/2] & 0x0c00); ++ acc1[2] += yl[i+16] * (qs[i/2] & 0x0030); ++ acc2[2] += yl[i+17] * (qs[i/2] & 0x3000); ++ acc1[3] += yl[i+24] * (qs[i/2] & 0x00c0); ++ acc2[3] += yl[i+25] * (qs[i/2] & 0xc000); ++ } ++ float dall = dh[0]; ++ float dmin = dh[1] * 1.f/16.f; ++ const float dot = ++ dall * ((acc1[0] + 1.f/256.f * acc2[0]) * (sc[0] & 0xF) * 1.f/ 1.f + ++ (acc1[1] + 1.f/256.f * acc2[1]) * (sc[2] & 0xF) * 1.f/ 4.f + ++ (acc1[2] + 1.f/256.f * acc2[2]) * (sc[4] & 0xF) * 1.f/16.f + ++ (acc1[3] + 1.f/256.f * acc2[3]) * (sc[6] & 0xF) * 1.f/64.f) - ++ dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + ++ sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); ++ sumf[row] += route_weight * dot; ++ ++ qs += args.nb01/2; ++ sc += args.nb01; ++ dh += args.nb01/2; + } +- d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); +- d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); +- sumf1[row] += d1 * (scales[1] - 32); +- sumf2[row] += d2 * (scales[3] - 32); + +- q += args.nb01/2; +- h += args.nb01/2; +- a += args.nb01/2; +- dh += args.nb01/2; ++ y4 += 4 * QK_K; + } +- +- y1 += 4 * QK_K; +- } +- +- for (int row = 0; row < nr0; ++row) { +- const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift); +- sumf1[row] = simd_sum(sumf); + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; +- +- if (tiisg == 0) { +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- dst_f32[first_row + row] = sumf1[row]; ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ ((device float *) (dst + uint64_t(first_row + row)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum_all; + } + } + } + +-[[host_name("kernel_mul_mv_q3_K_f32")]] +-kernel void kernel_mul_mv_q3_K_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, ++template ++void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, + device const char * src1, + device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { +- +- kernel_mul_mv_q3_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); +-} +- +-template +-void kernel_mul_mv_q4_K_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; +- +- constexpr uint16_t kmask1 = 0x3f3f; +- constexpr uint16_t kmask2 = 0x0f0f; +- constexpr uint16_t kmask3 = 0xc0c0; +- +- const short ix = tiisg/8; // 0...3 +- const short it = tiisg%8; // 0...7 +- const short iq = it/4; // 0 or 1 +- const short ir = it%4; // 0...3 +- ++ constexpr int n_slots = 8; + const int nb = args.ne00/QK_K; + + const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; +- +- const int first_row = (r0 * NSG + sgitg) * nr0; +- +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; +- +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; +- +- device const block_q4_K * x = (device const block_q4_K *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); +- +- float yl[16]; +- float yh[16]; +- +- float sumf[nr0]={0.f}; +- +- device const float * y4 = y + ix * QK_K + 64 * iq + 8 * ir; +- +- uint16_t sc16[4]; +- thread const uint8_t * sc8 = (thread const uint8_t *)sc16; +- +- for (int ib = ix; ib < nb; ib += 4) { +- float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ const int token = tgpig.y; ++ const int slot = sgitg; ++ const int first_row = r0 * nr0; ++ ++ if (scan_only) { ++ float scan_sum[nr0] = {0.0f}; ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id >= 0 && expert_id < args.ne02) { ++ const float route_weight = extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0]; ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ ++ if (route_weight != 0.0f) { ++ constexpr int words_per_block = sizeof(block_q2_K)/sizeof(uint32_t); ++ for (int ib = ix; ib < nb; ib += 4) { ++ for (short row = 0; row < nr0; ++row) { ++ device const uint32_t * words = (device const uint32_t *) ++ (src0s + offset0 + uint64_t(row)*args.nb01 + uint64_t(ib)*sizeof(block_q2_K)); ++ uint32_t checksum = 0; ++ for (int word = it; word < words_per_block; word += 8) { ++ checksum ^= words[word]; ++ } ++ scan_sum[row] += route_weight * float(checksum & 0xffffu); ++ } ++ } ++ } ++ } ++ } + +- for (short i = 0; i < 8; ++i) { +- yl[i+0] = y4[i+ 0]; sumy[0] += yl[i+0]; +- yl[i+8] = y4[i+ 32]; sumy[1] += yl[i+8]; +- yh[i+0] = y4[i+128]; sumy[2] += yh[i+0]; +- yh[i+8] = y4[i+160]; sumy[3] += yh[i+8]; ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0; ++row) { ++ partials[slot*nr0 + row] = first_row + row < args.ne0 ? simd_sum(scan_sum[row]) : 0.0f; ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ if (sgitg == 0 && tiisg < nr0 && first_row + tiisg < args.ne0) { ++ float sum = 0.0f; ++ for (int s = 0; s < n_slots; ++s) { ++ sum += partials[s*nr0 + tiisg]; ++ } ++ ((device float *) (dst + uint64_t(first_row + tiisg)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum; + } ++ return; ++ } + +- device const uint16_t * sc = (device const uint16_t *)x[ib].scales + iq; +- device const uint16_t * q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; +- device const half * dh = &x[ib].d; ++ float sumf[nr0]; ++ for (int row = 0; row < nr0; ++row) { ++ sumf[row] = 0.0f; ++ } ++ float yl[32]; + +- for (short row = 0; row < nr0; row++) { +- sc16[0] = sc[0] & kmask1; +- sc16[1] = sc[2] & kmask1; +- sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); +- sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ const short iq = it/4; ++ const short ir = it%4; ++ const short is = (8*ir)/16; ++ ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id >= 0 && expert_id < args.ne02) { ++ const float route_weight = extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0]; ++ ++ if (route_weight != 0.0f) { ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot)*args.nb11 + uint64_t(token)*args.nb12; ++ ++ device const block_q2_K * x = (device const block_q2_K *) (src0s + offset0); ++ device const act_t * y = (device const act_t *) (src1 + offset1); ++ device const act_t * y4 = y + ix * QK_K + 128 * iq + 8 * ir; ++ ++ for (int ib = ix; ib < nb; ib += 4) { ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ for (short i = 0; i < 8; ++i) { ++ yl[i+ 0] = float(y4[i+ 0]); sumy[0] += yl[i+ 0]; ++ yl[i+ 8] = float(y4[i+32]); sumy[1] += yl[i+ 8]; ++ yl[i+16] = float(y4[i+64]); sumy[2] += yl[i+16]; ++ yl[i+24] = float(y4[i+96]); sumy[3] += yl[i+24]; ++ } + +- device const uint16_t * q2 = q1 + 32; ++ device const uint8_t * sc = (device const uint8_t *)x[ib].scales + 8*iq + is; ++ device const uint16_t * qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; ++ device const half * dh = &x[ib].d; ++ ++ for (short row = 0; row < nr0; row++) { ++ float4 acc1 = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2 = {0.f, 0.f, 0.f, 0.f}; ++ for (int i = 0; i < 8; i += 2) { ++ acc1[0] += yl[i+ 0] * (qs[i/2] & 0x0003); ++ acc1[1] += yl[i+ 8] * (qs[i/2] & 0x000c); ++ acc1[2] += yl[i+16] * (qs[i/2] & 0x0030); ++ acc1[3] += yl[i+24] * (qs[i/2] & 0x00c0); ++ if (shift_high_bits) { ++ const uint16_t qs_hi = qs[i/2] >> 8; ++ acc2[0] += yl[i+ 1] * (qs_hi & 0x0003); ++ acc2[1] += yl[i+ 9] * (qs_hi & 0x000c); ++ acc2[2] += yl[i+17] * (qs_hi & 0x0030); ++ acc2[3] += yl[i+25] * (qs_hi & 0x00c0); ++ } else { ++ acc2[0] += yl[i+ 1] * (qs[i/2] & 0x0300); ++ acc2[1] += yl[i+ 9] * (qs[i/2] & 0x0c00); ++ acc2[2] += yl[i+17] * (qs[i/2] & 0x3000); ++ acc2[3] += yl[i+25] * (qs[i/2] & 0xc000); ++ } ++ } ++ const float dall = dh[0]; ++ const float dmin = dh[1] * 1.f/16.f; ++ float dot_val; ++ if (use_vec_scale) { ++ const float4 qscale = float4( ++ float(sc[0] & 0x0F), ++ float(sc[2] & 0x0F) * 0.25f, ++ float(sc[4] & 0x0F) * 0.0625f, ++ float(sc[6] & 0x0F) * 0.015625f); ++ const float4 minscale = float4( ++ float(sc[0] & 0xF0), ++ float(sc[2] & 0xF0), ++ float(sc[4] & 0xF0), ++ float(sc[6] & 0xF0)); ++ dot_val = dall * dot(float4(acc1 + (shift_high_bits ? 1.f : 1.f/256.f) * acc2), qscale) - ++ dmin * dot(sumy, minscale); ++ } else { ++ const float acc2_scale = shift_high_bits ? 1.f : 1.f/256.f; ++ dot_val = ++ dall * ((acc1[0] + acc2_scale * acc2[0]) * (sc[0] & 0xF) * 1.f/ 1.f + ++ (acc1[1] + acc2_scale * acc2[1]) * (sc[2] & 0xF) * 1.f/ 4.f + ++ (acc1[2] + acc2_scale * acc2[2]) * (sc[4] & 0xF) * 1.f/16.f + ++ (acc1[3] + acc2_scale * acc2[3]) * (sc[6] & 0xF) * 1.f/64.f) - ++ dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + ++ sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); ++ } ++ sumf[row] += route_weight * dot_val; + +- float4 acc1 = {0.f, 0.f, 0.f, 0.f}; +- float4 acc2 = {0.f, 0.f, 0.f, 0.f}; ++ qs += args.nb01/2; ++ sc += args.nb01; ++ dh += args.nb01/2; ++ } + +- FOR_UNROLL (short i = 0; i < 4; ++i) { +- acc1[0] += yl[2*i + 0] * (q1[i] & 0x000F); +- acc1[1] += yl[2*i + 1] * (q1[i] & 0x0F00); +- acc1[2] += yl[2*i + 8] * (q1[i] & 0x00F0); +- acc1[3] += yl[2*i + 9] * (q1[i] & 0xF000); +- acc2[0] += yh[2*i + 0] * (q2[i] & 0x000F); +- acc2[1] += yh[2*i + 1] * (q2[i] & 0x0F00); +- acc2[2] += yh[2*i + 8] * (q2[i] & 0x00F0); +- acc2[3] += yh[2*i + 9] * (q2[i] & 0xF000); ++ y4 += 4 * QK_K; ++ } + } ++ } ++ } + +- sumf[row] += dh[0] * ((acc1[0] + 1.f/256.f * acc1[1]) * sc8[0] + +- (acc1[2] + 1.f/256.f * acc1[3]) * sc8[1] * 1.f/16.f + +- (acc2[0] + 1.f/256.f * acc2[1]) * sc8[4] + +- (acc2[2] + 1.f/256.f * acc2[3]) * sc8[5] * 1.f/16.f) - +- dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]); ++ float reduced[nr0]; ++ for (int row = 0; row < nr0; ++row) { ++ reduced[row] = simd_sum(sumf[row]); ++ } + +- q1 += args.nb01/2; +- sc += args.nb01/2; +- dh += args.nb01/2; ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0; ++row) { ++ partials[slot*nr0 + row] = first_row + row < args.ne0 ? reduced[row] : 0.0f; + } +- +- y4 += 4 * QK_K; + } + +- device float * dst_f32 = (device float *) dst + (int64_t)im*args.ne0*args.ne1 + (int64_t)r1*args.ne0; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all; ++ if (sgitg == 0 && tiisg < nr0 && first_row + tiisg < args.ne0) { ++ float sum = 0.0f; ++ for (int s = 0; s < n_slots; ++s) { ++ sum += partials[s*nr0 + tiisg]; + } ++ ((device float *) (dst + uint64_t(first_row + tiisg)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum; + } + } + +-[[host_name("kernel_mul_mv_q4_K_f32")]] +-kernel void kernel_mul_mv_q4_K_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, + device const char * src1, + device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<8>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} + +- kernel_mul_mv_q4_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_f16( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<8, half>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); + } + +-template +-void kernel_mul_mv_q5_K_f32_impl( +- args_t args, +- device const char * src0, ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_shifted( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, + device const char * src1, + device char * dst, +- threadgroup char * shmem, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<8, float, false, true>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_f16_shifted( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<8, half, false, true>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_vecscale( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<8, float, true>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r8_nb8_scan( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<8, float, false, false, true>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r4_nb8( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<4>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r4_nb8_vecscale( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<4, float, true>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_r16_nb8( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_weighted_reduce_slots_sg_impl<16>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_glm_moe_swiglu_q2_K_down_weighted( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * gate, ++ device const char * up, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ constant ggml_metal_kargs_glm_moe_swiglu_q3_down & swiglu_args, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ constexpr int nr0 = 8; ++ constexpr int n_slots = 8; ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int token = tgpig.y; ++ const int slot = sgitg; ++ const int first_row = r0 * nr0; ++ ++ float sumf[nr0]; ++ for (int row = 0; row < nr0; ++row) { ++ sumf[row] = 0.0f; ++ } ++ ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ const short iq = it/4; ++ const short ir = it%4; ++ const short is = (8*ir)/16; ++ ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id >= 0 && expert_id < args.ne02) { ++ const float route_weight = extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0]; ++ ++ if (route_weight != 0.0f) { ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ device const block_q2_K * x = (device const block_q2_K *) (src0s + offset0); ++ device const float * gate_base = (device const float *) (gate + ++ uint64_t(slot)*swiglu_args.gate_nb1 + uint64_t(token)*swiglu_args.gate_nb2); ++ device const float * up_base = (device const float *) (up + ++ uint64_t(slot)*swiglu_args.up_nb1 + uint64_t(token)*swiglu_args.up_nb2); ++ ++ float yl[32]; ++ for (int ib = ix; ib < nb; ib += 4) { ++ const int y0 = ib*QK_K + 128*iq + 8*ir; ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ for (short i = 0; i < 8; ++i) { ++ const float g0 = gate_base[y0 + i + 0]; ++ const float g1 = gate_base[y0 + i + 32]; ++ const float g2 = gate_base[y0 + i + 64]; ++ const float g3 = gate_base[y0 + i + 96]; ++ yl[i+ 0] = (g0 / (1.0f + exp(-g0))) * up_base[y0 + i + 0]; sumy[0] += yl[i+ 0]; ++ yl[i+ 8] = (g1 / (1.0f + exp(-g1))) * up_base[y0 + i + 32]; sumy[1] += yl[i+ 8]; ++ yl[i+16] = (g2 / (1.0f + exp(-g2))) * up_base[y0 + i + 64]; sumy[2] += yl[i+16]; ++ yl[i+24] = (g3 / (1.0f + exp(-g3))) * up_base[y0 + i + 96]; sumy[3] += yl[i+24]; ++ } ++ ++ device const uint8_t * sc = (device const uint8_t *)x[ib].scales + 8*iq + is; ++ device const uint16_t * qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; ++ device const half * dh = &x[ib].d; ++ ++ for (short row = 0; row < nr0; row++) { ++ float4 acc1 = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2 = {0.f, 0.f, 0.f, 0.f}; ++ for (int i = 0; i < 8; i += 2) { ++ acc1[0] += yl[i+ 0] * (qs[i/2] & 0x0003); ++ acc2[0] += yl[i+ 1] * (qs[i/2] & 0x0300); ++ acc1[1] += yl[i+ 8] * (qs[i/2] & 0x000c); ++ acc2[1] += yl[i+ 9] * (qs[i/2] & 0x0c00); ++ acc1[2] += yl[i+16] * (qs[i/2] & 0x0030); ++ acc2[2] += yl[i+17] * (qs[i/2] & 0x3000); ++ acc1[3] += yl[i+24] * (qs[i/2] & 0x00c0); ++ acc2[3] += yl[i+25] * (qs[i/2] & 0xc000); ++ } ++ ++ const float dall = dh[0]; ++ const float dmin = dh[1] * 1.f/16.f; ++ const float4 qscale = float4( ++ float(sc[0] & 0x0F), ++ float(sc[2] & 0x0F) * 0.25f, ++ float(sc[4] & 0x0F) * 0.0625f, ++ float(sc[6] & 0x0F) * 0.015625f); ++ const float4 minscale = float4( ++ float(sc[0] & 0xF0), ++ float(sc[2] & 0xF0), ++ float(sc[4] & 0xF0), ++ float(sc[6] & 0xF0)); ++ const float dot_val = dall * dot(float4(acc1 + (1.f/256.f) * acc2), qscale) - ++ dmin * dot(sumy, minscale); ++ sumf[row] += route_weight * dot_val; ++ ++ qs += args.nb01/2; ++ sc += args.nb01; ++ dh += args.nb01/2; ++ } ++ } ++ } ++ } ++ } ++ ++ float reduced[nr0]; ++ for (int row = 0; row < nr0; ++row) { ++ reduced[row] = simd_sum(sumf[row]); ++ } ++ ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0; ++row) { ++ partials[slot*nr0 + row] = first_row + row < args.ne0 ? reduced[row] : 0.0f; ++ } ++ } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (sgitg == 0 && tiisg < nr0 && first_row + tiisg < args.ne0) { ++ float sum = 0.0f; ++ for (int s = 0; s < n_slots; ++s) { ++ sum += partials[s*nr0 + tiisg]; ++ } ++ ((device float *) (dst + uint64_t(first_row + tiisg)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum; ++ } ++} ++ ++template ++void kernel_mul_mv_q2_K_gate_up_swiglu_impl( ++ ggml_metal_kargs_mul_mv args, ++ device const char * src0_up, ++ device const char * src0_gate, ++ device const char * src1, ++ device char * dst, ++ float route_weight, + uint3 tgpig, + ushort tiisg, + ushort sgitg) { +@@ -8443,1377 +10027,6793 @@ void kernel_mul_mv_q5_K_f32_impl( + const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; + const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; + +- device const block_q5_K * x = (device const block_q5_K *) (src0 + offset0); +- device const float * yy = (device const float *) (src1 + offset1); +- +- float sumf[nr0]={0.f}; +- +- float yl[16], yh[16]; +- +- constexpr uint16_t kmask1 = 0x3f3f; +- constexpr uint16_t kmask2 = 0x0f0f; +- constexpr uint16_t kmask3 = 0xc0c0; +- +- const short tid = tiisg/4; +- const short ix = tiisg%4; +- const short iq = tid/4; +- const short ir = tid%4; +- +- const short l0 = 8*ir; +- const short q_offset = 32*iq + l0; +- const short y_offset = 64*iq + l0; +- +- const uint8_t hm1 = 1u << (2*iq); +- const uint8_t hm2 = hm1 << 1; +- const uint8_t hm3 = hm1 << 4; +- const uint8_t hm4 = hm2 << 4; ++ device const block_q2_K * x_up = (device const block_q2_K *) (src0_up + offset0); ++ device const block_q2_K * x_gate = (device const block_q2_K *) (src0_gate + offset0); ++ device const float * y = (device const float *) (src1 + offset1); + +- uint16_t sc16[4]; +- thread const uint8_t * sc8 = (thread const uint8_t *)sc16; ++ float yl[32]; ++ float sum_up[nr0] = {0.f}; ++ float sum_gate[nr0] = {0.f}; + +- device const float * y1 = yy + ix*QK_K + y_offset; ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ const short iq = it/4; ++ const short ir = it%4; ++ const short is = (8*ir)/16; + +- for (int i = ix; i < nb; i += 4) { +- device const uint8_t * q1 = x[i].qs + q_offset; +- device const uint8_t * qh = x[i].qh + l0; +- device const half * dh = &x[i].d; +- device const uint16_t * a = (device const uint16_t *)x[i].scales + iq; ++ device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir; + +- device const float * y2 = y1 + 128; ++ for (int ib = ix; ib < nb; ib += 4) { + float4 sumy = {0.f, 0.f, 0.f, 0.f}; +- for (short l = 0; l < 8; ++l) { +- yl[l+0] = y1[l+ 0]; sumy[0] += yl[l+0]; +- yl[l+8] = y1[l+32]; sumy[1] += yl[l+8]; +- yh[l+0] = y2[l+ 0]; sumy[2] += yh[l+0]; +- yh[l+8] = y2[l+32]; sumy[3] += yh[l+8]; ++ for (short i = 0; i < 8; ++i) { ++ yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0]; ++ yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8]; ++ yl[i+16] = y4[i+64]; sumy[2] += yl[i+16]; ++ yl[i+24] = y4[i+96]; sumy[3] += yl[i+24]; + } + +- for (short row = 0; row < nr0; ++row) { +- device const uint8_t * q2 = q1 + 64; +- +- sc16[0] = a[0] & kmask1; +- sc16[1] = a[2] & kmask1; +- sc16[2] = ((a[4] >> 0) & kmask2) | ((a[0] & kmask3) >> 2); +- sc16[3] = ((a[4] >> 4) & kmask2) | ((a[2] & kmask3) >> 2); ++ device const uint8_t * sc_up = (device const uint8_t *)x_up[ib].scales + 8*iq + is; ++ device const uint16_t * qs_up = (device const uint16_t *)x_up[ib].qs + 16*iq + 4*ir; ++ device const half * dh_up = &x_up[ib].d; ++ device const uint8_t * sc_gate = (device const uint8_t *)x_gate[ib].scales + 8*iq + is; ++ device const uint16_t * qs_gate = (device const uint16_t *)x_gate[ib].qs + 16*iq + 4*ir; ++ device const half * dh_gate = &x_gate[ib].d; + +- float4 acc1 = {0.f}; +- float4 acc2 = {0.f}; +- FOR_UNROLL (short l = 0; l < 8; ++l) { +- uint8_t h = qh[l]; +- acc1[0] += yl[l+0] * (q1[l] & 0x0F); +- acc1[1] += yl[l+8] * (q1[l] & 0xF0); +- acc1[2] += yh[l+0] * (q2[l] & 0x0F); +- acc1[3] += yh[l+8] * (q2[l] & 0xF0); +- acc2[0] += h & hm1 ? yl[l+0] : 0.f; +- acc2[1] += h & hm2 ? yl[l+8] : 0.f; +- acc2[2] += h & hm3 ? yh[l+0] : 0.f; +- acc2[3] += h & hm4 ? yh[l+8] : 0.f; ++ for (short row = 0; row < nr0; row++) { ++ float4 acc1_up = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2_up = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc1_gate = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2_gate = {0.f, 0.f, 0.f, 0.f}; ++ for (int i = 0; i < 8; i += 2) { ++ acc1_up[0] += yl[i+ 0] * (qs_up[i/2] & 0x0003); ++ acc2_up[0] += yl[i+ 1] * (qs_up[i/2] & 0x0300); ++ acc1_up[1] += yl[i+ 8] * (qs_up[i/2] & 0x000c); ++ acc2_up[1] += yl[i+ 9] * (qs_up[i/2] & 0x0c00); ++ acc1_up[2] += yl[i+16] * (qs_up[i/2] & 0x0030); ++ acc2_up[2] += yl[i+17] * (qs_up[i/2] & 0x3000); ++ acc1_up[3] += yl[i+24] * (qs_up[i/2] & 0x00c0); ++ acc2_up[3] += yl[i+25] * (qs_up[i/2] & 0xc000); ++ ++ acc1_gate[0] += yl[i+ 0] * (qs_gate[i/2] & 0x0003); ++ acc2_gate[0] += yl[i+ 1] * (qs_gate[i/2] & 0x0300); ++ acc1_gate[1] += yl[i+ 8] * (qs_gate[i/2] & 0x000c); ++ acc2_gate[1] += yl[i+ 9] * (qs_gate[i/2] & 0x0c00); ++ acc1_gate[2] += yl[i+16] * (qs_gate[i/2] & 0x0030); ++ acc2_gate[2] += yl[i+17] * (qs_gate[i/2] & 0x3000); ++ acc1_gate[3] += yl[i+24] * (qs_gate[i/2] & 0x00c0); ++ acc2_gate[3] += yl[i+25] * (qs_gate[i/2] & 0xc000); ++ } ++ const float dall_up = dh_up[0]; ++ const float dmin_up = dh_up[1] * 1.f/16.f; ++ const float dall_gate = dh_gate[0]; ++ const float dmin_gate = dh_gate[1] * 1.f/16.f; ++ ++ if (use_vec_scale) { ++ const float4 qscale_up = float4( ++ float(sc_up[0] & 0x0F), ++ float(sc_up[2] & 0x0F) * 0.25f, ++ float(sc_up[4] & 0x0F) * 0.0625f, ++ float(sc_up[6] & 0x0F) * 0.015625f); ++ const float4 minscale_up = float4( ++ float(sc_up[0] & 0xF0), ++ float(sc_up[2] & 0xF0), ++ float(sc_up[4] & 0xF0), ++ float(sc_up[6] & 0xF0)); ++ const float4 qscale_gate = float4( ++ float(sc_gate[0] & 0x0F), ++ float(sc_gate[2] & 0x0F) * 0.25f, ++ float(sc_gate[4] & 0x0F) * 0.0625f, ++ float(sc_gate[6] & 0x0F) * 0.015625f); ++ const float4 minscale_gate = float4( ++ float(sc_gate[0] & 0xF0), ++ float(sc_gate[2] & 0xF0), ++ float(sc_gate[4] & 0xF0), ++ float(sc_gate[6] & 0xF0)); ++ sum_up[row] += dall_up * dot(float4(acc1_up + (1.f/256.f) * acc2_up), qscale_up) - ++ dmin_up * dot(sumy, minscale_up); ++ sum_gate[row] += dall_gate * dot(float4(acc1_gate + (1.f/256.f) * acc2_gate), qscale_gate) - ++ dmin_gate * dot(sumy, minscale_gate); ++ } else { ++ sum_up[row] += dall_up * ((acc1_up[0] + 1.f/256.f * acc2_up[0]) * (sc_up[0] & 0xF) * 1.f/ 1.f + ++ (acc1_up[1] + 1.f/256.f * acc2_up[1]) * (sc_up[2] & 0xF) * 1.f/ 4.f + ++ (acc1_up[2] + 1.f/256.f * acc2_up[2]) * (sc_up[4] & 0xF) * 1.f/16.f + ++ (acc1_up[3] + 1.f/256.f * acc2_up[3]) * (sc_up[6] & 0xF) * 1.f/64.f) - ++ dmin_up * (sumy[0] * (sc_up[0] & 0xF0) + sumy[1] * (sc_up[2] & 0xF0) + sumy[2] * (sc_up[4] & 0xF0) + sumy[3] * (sc_up[6] & 0xF0)); ++ ++ sum_gate[row] += dall_gate * ((acc1_gate[0] + 1.f/256.f * acc2_gate[0]) * (sc_gate[0] & 0xF) * 1.f/ 1.f + ++ (acc1_gate[1] + 1.f/256.f * acc2_gate[1]) * (sc_gate[2] & 0xF) * 1.f/ 4.f + ++ (acc1_gate[2] + 1.f/256.f * acc2_gate[2]) * (sc_gate[4] & 0xF) * 1.f/16.f + ++ (acc1_gate[3] + 1.f/256.f * acc2_gate[3]) * (sc_gate[6] & 0xF) * 1.f/64.f) - ++ dmin_gate * (sumy[0] * (sc_gate[0] & 0xF0) + sumy[1] * (sc_gate[2] & 0xF0) + sumy[2] * (sc_gate[4] & 0xF0) + sumy[3] * (sc_gate[6] & 0xF0)); + } + +- sumf[row] += dh[0] * (sc8[0] * (acc1[0] + 16.f*acc2[0]) + +- sc8[1] * (acc1[1]/16.f + 16.f*acc2[1]) + +- sc8[4] * (acc1[2] + 16.f*acc2[2]) + +- sc8[5] * (acc1[3]/16.f + 16.f*acc2[3])) - +- dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]); +- +- q1 += args.nb01; +- qh += args.nb01; +- dh += args.nb01/2; +- a += args.nb01/2; ++ qs_up += args.nb01/2; ++ sc_up += args.nb01; ++ dh_up += args.nb01/2; ++ qs_gate += args.nb01/2; ++ sc_gate += args.nb01; ++ dh_gate += args.nb01/2; + } + +- y1 += 4 * QK_K; ++ y4 += 4 * QK_K; + } + + device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; + + for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- const float tot = simd_sum(sumf[row]); ++ const float up_value = simd_sum(sum_up[row]); ++ const float gate_value = simd_sum(sum_gate[row]); + if (tiisg == 0) { +- dst_f32[first_row + row] = tot; ++ const float silu = gate_value / (1.0f + exp(-gate_value)); ++ dst_f32[first_row + row] = up_value * silu * route_weight; + } + } + } + +-[[host_name("kernel_mul_mv_q5_K_f32")]] +-kernel void kernel_mul_mv_q5_K_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, + device const char * src1, + device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup char * shmem [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiitg[[thread_index_in_threadgroup]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const int iid1 = tgpig.z/args.nei0; ++ const int idx = tgpig.z%args.nei0; + +- kernel_mul_mv_q5_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); +-} ++ tgpig.z = 0; + +-template +-void kernel_mul_mv_q6_K_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ const int32_t i02 = ((device const int32_t *) (ids + iid1*args.nbi1))[idx]; ++ const float route_weight = args.weighted != 0 ? ++ ((device const float *) (weights + uint64_t(idx)*args.weights_nb1 + uint64_t(iid1)*args.weights_nb2))[0] : ++ 1.0f; + +- constexpr uint8_t kmask1 = 0x03; +- constexpr uint8_t kmask2 = 0x0C; +- constexpr uint8_t kmask3 = 0x30; +- constexpr uint8_t kmask4 = 0xC0; ++ const int64_t i11 = idx % args.ne11; ++ const int64_t i12 = iid1; + +- const int nb = args.ne00/QK_K; ++ const int64_t i1 = idx; ++ const int64_t i2 = i12; + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; +- +- const int first_row = (r0 * NSG + sgitg) * nr0; +- +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; +- +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; +- +- device const block_q6_K * x = (device const block_q6_K *) (src0 + offset0); +- device const float * yy = (device const float *) (src1 + offset1); ++ device const char * src0_up_cur = src0_up_all + i02*args.nb02; ++ device const char * src0_gate_cur = src0_gate_all + i02*args.nb02; ++ device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12; ++ device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float); + +- float sumf[nr0] = { 0.f }; ++ ggml_metal_kargs_mul_mv args0 = { ++ /*.ne00 =*/ args.ne00, ++ /*.ne01 =*/ args.ne01, ++ /*.ne02 =*/ 1, ++ /*.nb00 =*/ args.nb00, ++ /*.nb01 =*/ args.nb01, ++ /*.nb02 =*/ args.nb02, ++ /*.nb03 =*/ args.nb02, ++ /*.ne10 =*/ args.ne10, ++ /*.ne11 =*/ 1, ++ /*.ne12 =*/ 1, ++ /*.nb10 =*/ args.nb10, ++ /*.nb11 =*/ args.nb11, ++ /*.nb12 =*/ args.nb12, ++ /*.nb13 =*/ args.nb12, ++ /*.ne0 =*/ args.ne0, ++ /*.ne1 =*/ 1, ++ /*.nr0 =*/ args.nr0, ++ /*.r2 =*/ 1, ++ /*.r3 =*/ 1, ++ }; + +- float yl[16]; ++ switch (args.nr0) { ++ case 1: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl<1>( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ case 2: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl<2>( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ case 4: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl<4>( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ case 8: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl<8>( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ default: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ } + +- const short tid = tiisg/2; +- const short ix = tiisg%2; +- const short ip = tid/8; // 0 or 1 +- const short il = tid%8; +- const short l0 = 4*il; +- const short is = 8*ip + l0/16; ++ (void) shmem; ++ (void) tiitg; ++} + +- const short y_offset = 128*ip + l0; +- const short q_offset_l = 64*ip + l0; +- const short q_offset_h = 32*ip + l0; ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_vecscale( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiitg[[thread_index_in_threadgroup]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const int iid1 = tgpig.z/args.nei0; ++ const int idx = tgpig.z%args.nei0; + +- for (int i = ix; i < nb; i += 2) { +- device const uint8_t * q1 = x[i].ql + q_offset_l; +- device const uint8_t * q2 = q1 + 32; +- device const uint8_t * qh = x[i].qh + q_offset_h; +- device const int8_t * sc = x[i].scales + is; +- device const half * dh = &x[i].d; ++ tgpig.z = 0; + +- device const float * y = yy + i * QK_K + y_offset; ++ const int32_t i02 = ((device const int32_t *) (ids + iid1*args.nbi1))[idx]; ++ const float route_weight = args.weighted != 0 ? ++ ((device const float *) (weights + uint64_t(idx)*args.weights_nb1 + uint64_t(iid1)*args.weights_nb2))[0] : ++ 1.0f; + +- for (short l = 0; l < 4; ++l) { +- yl[4*l + 0] = y[l + 0]; +- yl[4*l + 1] = y[l + 32]; +- yl[4*l + 2] = y[l + 64]; +- yl[4*l + 3] = y[l + 96]; +- } ++ const int64_t i11 = idx % args.ne11; ++ const int64_t i12 = iid1; + +- for (short row = 0; row < nr0; ++row) { +- float4 sums = {0.f, 0.f, 0.f, 0.f}; ++ const int64_t i1 = idx; ++ const int64_t i2 = i12; + +- FOR_UNROLL (short l = 0; l < 4; ++l) { +- sums[0] += yl[4*l + 0] * ((int8_t)((q1[l] & 0xF) | ((qh[l] & kmask1) << 4)) - 32); +- sums[1] += yl[4*l + 1] * ((int8_t)((q2[l] & 0xF) | ((qh[l] & kmask2) << 2)) - 32); +- sums[2] += yl[4*l + 2] * ((int8_t)((q1[l] >> 4) | ((qh[l] & kmask3) << 0)) - 32); +- sums[3] += yl[4*l + 3] * ((int8_t)((q2[l] >> 4) | ((qh[l] & kmask4) >> 2)) - 32); +- } ++ device const char * src0_up_cur = src0_up_all + i02*args.nb02; ++ device const char * src0_gate_cur = src0_gate_all + i02*args.nb02; ++ device const char * src1_cur = src1 + i11*args.nb11 + i12*args.nb12; ++ device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float); + +- sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + sums[2] * sc[4] + sums[3] * sc[6]); ++ ggml_metal_kargs_mul_mv args0 = { ++ /*.ne00 =*/ args.ne00, ++ /*.ne01 =*/ args.ne01, ++ /*.ne02 =*/ 1, ++ /*.nb00 =*/ args.nb00, ++ /*.nb01 =*/ args.nb01, ++ /*.nb02 =*/ args.nb02, ++ /*.nb03 =*/ args.nb02, ++ /*.ne10 =*/ args.ne10, ++ /*.ne11 =*/ 1, ++ /*.ne12 =*/ 1, ++ /*.nb10 =*/ args.nb10, ++ /*.nb11 =*/ args.nb11, ++ /*.nb12 =*/ args.nb12, ++ /*.nb13 =*/ args.nb12, ++ /*.ne0 =*/ args.ne0, ++ /*.ne1 =*/ 1, ++ /*.nr0 =*/ args.nr0, ++ /*.r2 =*/ 1, ++ /*.r3 =*/ 1, ++ }; + +- q1 += args.nb01; +- q2 += args.nb01; +- qh += args.nb01; +- sc += args.nb01; +- dh += args.nb01/2; +- } ++ switch (args.nr0) { ++ case 1: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl<1, true>( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ case 2: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl<2, true>( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ case 4: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl<4, true>( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ case 8: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl<8, true>( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; ++ default: ++ kernel_mul_mv_q2_K_gate_up_swiglu_impl( ++ args0, src0_up_cur, src0_gate_cur, src1_cur, dst_cur, route_weight, tgpig, tiisg, sgitg); ++ break; + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; +- +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all; +- } +- } ++ (void) shmem; ++ (void) tiitg; + } + +-[[host_name("kernel_mul_mv_q6_K_f32")]] +-kernel void kernel_mul_mv_q6_K_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, ++template ++void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, + device const char * src1, + device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { +- +- kernel_mul_mv_q6_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); +-} +- +-// ======================= "True" 2-bit +- +-template +-void kernel_mul_mv_iq2_xxs_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; +- + const int nb = args.ne00/QK_K; + + const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; +- +- const int first_row = (r0 * NSG + sgitg) * nr0; +- +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; +- +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; +- +- device const block_iq2_xxs * x = (device const block_iq2_xxs *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); +- +- float yl[32]; +- float sumf[nr0]={0.f}; ++ const int token = tgpig.z / ((args.nei0 + slots_per_tg - 1) / slots_per_tg); ++ const int slot_group = tgpig.z % ((args.nei0 + slots_per_tg - 1) / slots_per_tg); ++ const int slot = slot_group * slots_per_tg + int(sgitg / 2); ++ const bool is_gate = (sgitg & 1) != 0; ++ const int first_row = r0 * nr0; + +- const int nb32 = nb * (QK_K / 32); +- +- threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem); +- threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256); +- { +- int nval = 4; +- int pos = (32*sgitg + tiisg)*nval; +- for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2xxs_grid[pos + i]; +- nval = 2; +- pos = (32*sgitg + tiisg)*nval; +- for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i]; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- } +- +- const int ix = tiisg; +- +- device const float * y4 = y + 32 * ix; +- +- for (int ib32 = ix; ib32 < nb32; ib32 += 32) { +- for (short i = 0; i < 32; ++i) { +- yl[i] = y4[i]; +- } ++ float sumf[nr0] = {0.f}; + +- const int ibl = ib32 / (QK_K / 32); +- const int ib = ib32 % (QK_K / 32); ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ const bool active_slot = args._pad0 <= 0 || slot < args._pad0; ++ if (active_slot && expert_id >= 0 && expert_id < args.ne02) { ++ device const char * src0_all = is_gate ? src0_gate_all : src0_up_all; ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot % args.ne11)*args.nb11 + uint64_t(token)*args.nb12; ++ ++ device const block_q2_K * x = (device const block_q2_K *) (src0_all + offset0); ++ device const block_q2_K * x_expert = (device const block_q2_K *) (src0_all + uint64_t(expert_id)*args.nb02); ++ device const float * y = (device const float *) (src1 + offset1); ++ threadgroup float * y_shared = partials + 8*nr0; ++ ++ y_t yl[32]; ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ const short iq = it/4; ++ const short ir = it%4; ++ const short is = (8*ir)/16; ++ ++ device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir; ++ ++ for (int ib = ix; ib < nb; ib += 4) { ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ if (share_y) { ++ if (sgitg == 0) { ++ for (short i = 0; i < 8; ++i) { ++ y_shared[tiisg*32 + i + 0] = y4[i+ 0]; ++ y_shared[tiisg*32 + i + 8] = y4[i+32]; ++ y_shared[tiisg*32 + i + 16] = y4[i+64]; ++ y_shared[tiisg*32 + i + 24] = y4[i+96]; ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ for (short i = 0; i < 8; ++i) { ++ yl[i+ 0] = y_t(y_shared[tiisg*32 + i + 0]); sumy[0] += float(yl[i+ 0]); ++ yl[i+ 8] = y_t(y_shared[tiisg*32 + i + 8]); sumy[1] += float(yl[i+ 8]); ++ yl[i+16] = y_t(y_shared[tiisg*32 + i + 16]); sumy[2] += float(yl[i+16]); ++ yl[i+24] = y_t(y_shared[tiisg*32 + i + 24]); sumy[3] += float(yl[i+24]); ++ } ++ } else { ++ for (short i = 0; i < 8; ++i) { ++ yl[i+ 0] = y_t(y4[i+ 0]); sumy[0] += float(yl[i+ 0]); ++ yl[i+ 8] = y_t(y4[i+32]); sumy[1] += float(yl[i+ 8]); ++ yl[i+16] = y_t(y4[i+64]); sumy[2] += float(yl[i+16]); ++ yl[i+24] = y_t(y4[i+96]); sumy[3] += float(yl[i+24]); ++ } ++ } + +- device const block_iq2_xxs * xr = x + ibl; +- device const uint16_t * q2 = xr->qs + 4 * ib; +- device const half * dh = &xr->d; ++ int qy[32]; ++ int4 sumq = {0, 0, 0, 0}; ++ float4 d8 = {0.f, 0.f, 0.f, 0.f}; ++ if (use_q8_act) { ++ for (short group = 0; group < 4; ++group) { ++ const short base = 8 * group; ++ float amax = 0.f; ++ for (short i = 0; i < 8; ++i) { ++ amax = MAX(amax, fabs(float(yl[base + i]))); ++ } ++ const float d = amax == 0.f ? 0.f : amax / 127.f; ++ const float id = d == 0.f ? 0.f : 1.f / d; ++ d8[group] = d; ++ for (short i = 0; i < 8; ++i) { ++ const int q = int(round(clamp(float(yl[base + i]) * id, -127.f, 127.f))); ++ qy[base + i] = q; ++ sumq[group] += q; ++ } ++ } ++ } + +- for (short row = 0; row < nr0; row++) { +- const float db = dh[0]; +- device const uint8_t * aux8 = (device const uint8_t *)q2; +- const uint32_t aux32 = q2[2] | (q2[3] << 16); +- const float d = db * (0.5f + (aux32 >> 28)); ++ for (short row = 0; row < nr0; row++) { ++ device const block_q2_K * xb = use_row_tile ? ++ (x_expert + uint64_t(first_row/nr0)*uint64_t(nb*nr0) + uint64_t(ib)*uint64_t(nr0) + uint64_t(row)) : ++ (x + ib + uint64_t(row)*uint64_t(args.nb01/sizeof(block_q2_K))); ++ device const uint8_t * sc = (device const uint8_t *)xb->scales + 8*iq + is; ++ device const uint16_t * qs = (device const uint16_t *)xb->qs + 16 * iq + 4 * ir; ++ device const half * dh = &xb->d; ++ ++ float4 acc1 = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2 = {0.f, 0.f, 0.f, 0.f}; ++ for (int i = 0; i < 8; i += 2) { ++ acc1[0] += float(yl[i+ 0]) * (qs[i/2] & 0x0003); ++ acc2[0] += float(yl[i+ 1]) * (qs[i/2] & 0x0300); ++ acc1[1] += float(yl[i+ 8]) * (qs[i/2] & 0x000c); ++ acc2[1] += float(yl[i+ 9]) * (qs[i/2] & 0x0c00); ++ acc1[2] += float(yl[i+16]) * (qs[i/2] & 0x0030); ++ acc2[2] += float(yl[i+17]) * (qs[i/2] & 0x3000); ++ acc1[3] += float(yl[i+24]) * (qs[i/2] & 0x00c0); ++ acc2[3] += float(yl[i+25]) * (qs[i/2] & 0xc000); ++ } ++ const float dall = dh[0]; ++ const float dmin = dh[1] * 1.f/16.f; ++ float dot_val; ++ if (use_q8_act) { ++ int4 accq = {0, 0, 0, 0}; ++ for (int i = 0; i < 8; i += 2) { ++ const uint16_t q = qs[i/2]; ++ accq[0] += qy[i+ 0] * int( q & 0x0003) + qy[i+ 1] * int((q >> 8) & 0x0003); ++ accq[1] += qy[i+ 8] * int((q >> 2) & 0x0003) + qy[i+ 9] * int((q >> 10) & 0x0003); ++ accq[2] += qy[i+16] * int((q >> 4) & 0x0003) + qy[i+17] * int((q >> 12) & 0x0003); ++ accq[3] += qy[i+24] * int((q >> 6) & 0x0003) + qy[i+25] * int((q >> 14) & 0x0003); ++ } ++ const float4 qscale = float4( ++ float(sc[0] & 0x0F), ++ float(sc[2] & 0x0F), ++ float(sc[4] & 0x0F), ++ float(sc[6] & 0x0F)); ++ const float4 minscale = float4( ++ float(sc[0] & 0xF0), ++ float(sc[2] & 0xF0), ++ float(sc[4] & 0xF0), ++ float(sc[6] & 0xF0)); ++ const float4 accf = float4(float(accq[0]), float(accq[1]), float(accq[2]), float(accq[3])); ++ const float4 sumf_q = float4(float(sumq[0]), float(sumq[1]), float(sumq[2]), float(sumq[3])); ++ dot_val = dall * dot(accf, qscale*d8) - ++ dmin * dot(sumf_q, minscale*d8); ++ } else if (use_vec_scale) { ++ const float4 qscale = float4( ++ float(sc[0] & 0x0F), ++ float(sc[2] & 0x0F) * 0.25f, ++ float(sc[4] & 0x0F) * 0.0625f, ++ float(sc[6] & 0x0F) * 0.015625f); ++ const float4 minscale = float4( ++ float(sc[0] & 0xF0), ++ float(sc[2] & 0xF0), ++ float(sc[4] & 0xF0), ++ float(sc[6] & 0xF0)); ++ dot_val = dall * dot(float4(acc1 + (1.f/256.f) * acc2), qscale) - ++ dmin * dot(sumy, minscale); ++ } else { ++ dot_val = dall * ((acc1[0] + 1.f/256.f * acc2[0]) * (sc[0] & 0xF) * 1.f/ 1.f + ++ (acc1[1] + 1.f/256.f * acc2[1]) * (sc[2] & 0xF) * 1.f/ 4.f + ++ (acc1[2] + 1.f/256.f * acc2[2]) * (sc[4] & 0xF) * 1.f/16.f + ++ (acc1[3] + 1.f/256.f * acc2[3]) * (sc[6] & 0xF) * 1.f/64.f) - ++ dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + ++ sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0)); ++ } ++ sumf[row] += dot_val; ++ } + +- float sum = 0; +- for (short l = 0; l < 4; ++l) { +- const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + aux8[l]); +- const uint8_t signs = ssigns[(aux32 >> 7*l) & 127]; +- for (short j = 0; j < 8; ++j) { +- sum += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); ++ y4 += 4 * QK_K; ++ if (share_y) { ++ threadgroup_barrier(mem_flags::mem_threadgroup); + } + } +- sumf[row] += d * sum; +- +- dh += args.nb01/2; +- q2 += args.nb01/2; + } +- +- y4 += 32 * 32; + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; +- +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); ++ for (int row = 0; row < nr0; ++row) { ++ const float row_sum = simd_sum(sumf[row]); + if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all * 0.25f; ++ partials[sgitg*nr0 + row] = first_row + row < args.ne0 ? row_sum : 0.0f; + } + } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if ((sgitg & 1) == 0 && tiisg < nr0 && slot < args.nei0 && token < args.nei1 && first_row + tiisg < args.ne0) { ++ const float up_value = partials[sgitg*nr0 + tiisg]; ++ const float gate_value = partials[(sgitg + 1)*nr0 + tiisg]; ++ const float route_weight = args.weighted != 0 ? ++ ((device const float *) (weights + uint64_t(slot)*args.weights_nb1 + uint64_t(token)*args.weights_nb2))[0] : ++ 1.0f; ++ const float silu = gate_value / (1.0f + exp(-gate_value)); ++ ((device dst_t *) dst)[uint64_t(slot)*args.ne0 + uint64_t(token)*args.ne1*args.ne0 + first_row + tiisg] = ++ dst_t(up_value * silu * route_weight); ++ } + } + +-[[host_name("kernel_mul_mv_iq2_xxs_f32")]] +-kernel void kernel_mul_mv_iq2_xxs_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, + device const char * src1, + device char * dst, +- threadgroup char * shmem [[threadgroup(0)]], ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], + uint3 tgpig[[threadgroup_position_in_grid]], + ushort tiisg[[thread_index_in_simdgroup]], + ushort sgitg[[simdgroup_index_in_threadgroup]]) { +- kernel_mul_mv_iq2_xxs_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, float>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); + } + +-template +-void kernel_mul_mv_iq2_xs_f32_impl( +- args_t args, +- device const char * src0, ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_share_y( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, + device const char * src1, + device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, float, 4, true>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- const int nb = args.ne00/QK_K; ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_vecscale( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, float, 4, false, true>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_q8_act( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, float, 4, false, false, true>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- const int first_row = (r0 * NSG + sgitg) * nr0; ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_half_y( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, float, 4, false, false, false, half>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_rowtile( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, float, 4, false, false, false, float, true>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_r12( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<12, float>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- device const block_iq2_xs * x = (device const block_iq2_xs *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_r16( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<16, float>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- float yl[32]; +- float sumf[nr0]={0.f}; ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot2( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, float, 2>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- const int nb32 = nb * (QK_K / 32); ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_f16( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, half>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem); +- threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 512); +- { +- int nval = 8; +- int pos = (32*sgitg + tiisg)*nval; +- for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2xs_grid[pos + i]; +- nval = 2; +- pos = (32*sgitg + tiisg)*nval; +- for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i]; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- } ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot8_split( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_impl<8, float, 8>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, partials, tgpig, tiisg, sgitg); ++} + +- const int ix = tiisg; ++static inline float4 q2_repacked_unpack4(uint8_t packed) { ++ return float4( ++ float( packed & 0x3), ++ float((packed >> 2) & 0x3), ++ float((packed >> 4) & 0x3), ++ float((packed >> 6) & 0x3)); ++} + +- device const float * y4 = y + 32 * ix; ++static inline float q2_repacked_dot8(device const uint8_t * packed, thread const float * values) { ++ const float4 y0 = float4(values[0], values[1], values[2], values[3]); ++ const float4 y1 = float4(values[4], values[5], values[6], values[7]); ++ return dot(y0, q2_repacked_unpack4(packed[0])) + ++ dot(y1, q2_repacked_unpack4(packed[1])); ++} + +- for (int ib32 = ix; ib32 < nb32; ib32 += 32) { +- for (short i = 0; i < 32; ++i) { +- yl[i] = y4[i]; ++template< ++ int slots_per_tg, ++ int nr0 = 8, ++ bool scan_only = false, ++ bool use_prequant_q8 = false, ++ bool use_inblock_repack = false> ++void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int token = tgpig.z / ((args.nei0 + slots_per_tg - 1) / slots_per_tg); ++ const int slot_group = tgpig.z % ((args.nei0 + slots_per_tg - 1) / slots_per_tg); ++ const int slot = slot_group * slots_per_tg + sgitg; ++ const int first_row = r0 * nr0; ++ ++ float sum_up[nr0] = {0.f}; ++ float sum_gate[nr0] = {0.f}; ++ ++ if (scan_only) { ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id >= 0 && expert_id < args.ne02) { ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ constexpr int words_per_block = sizeof(block_q2_K)/sizeof(uint32_t); ++ for (int ib = ix; ib < nb; ib += 4) { ++ for (short row = 0; row < nr0; ++row) { ++ device const uint32_t * up_words = (device const uint32_t *) ++ (src0_up_all + offset0 + uint64_t(row)*args.nb01 + uint64_t(ib)*sizeof(block_q2_K)); ++ device const uint32_t * gate_words = (device const uint32_t *) ++ (src0_gate_all + offset0 + uint64_t(row)*args.nb01 + uint64_t(ib)*sizeof(block_q2_K)); ++ uint32_t up_checksum = 0; ++ uint32_t gate_checksum = 0; ++ for (int word = it; word < words_per_block; word += 8) { ++ up_checksum ^= up_words[word]; ++ gate_checksum ^= gate_words[word]; ++ } ++ sum_up[row] += float(up_checksum & 0xffffu); ++ sum_gate[row] += float(gate_checksum & 0xffffu); ++ } ++ } ++ } + } + +- const int ibl = ib32 / (QK_K / 32); +- const int ib = ib32 % (QK_K / 32); ++ const bool valid_output = slot < args.nei0 && token < args.nei1; ++ for (int row = 0; row < nr0; ++row) { ++ const float up_value = simd_sum(sum_up[row]); ++ const float gate_value = simd_sum(sum_gate[row]); ++ if (tiisg == 0 && valid_output && first_row + row < args.ne0) { ++ ((device float *) dst)[uint64_t(slot)*args.ne0 + uint64_t(token)*args.ne1*args.ne0 + first_row + row] = ++ up_value + gate_value; ++ } ++ } ++ (void) src1; ++ (void) weights; ++ (void) src1_q8; ++ (void) partials; ++ return; ++ } + +- device const block_iq2_xs * xr = x + ibl; +- device const uint16_t * q2 = xr->qs + 4 * ib; +- device const uint8_t * sc = xr->scales + ib; +- device const half * dh = &xr->d; ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ const bool active_slot = args._pad0 <= 0 || slot < args._pad0; ++ if (active_slot && expert_id >= 0 && expert_id < args.ne02) { ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot % args.ne11)*args.nb11 + uint64_t(token)*args.nb12; ++ ++ device const block_q2_K * x_up = (device const block_q2_K *) (src0_up_all + offset0); ++ device const block_q2_K * x_gate = (device const block_q2_K *) (src0_gate_all + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ device const block_q8_0 * y_q8 = (device const block_q8_0 *) ++ (src1_q8 + uint64_t(token)*uint64_t(args.ne10/QK8_0)*sizeof(block_q8_0)); ++ ++ float yl[32]; ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ const short iq = it/4; ++ const short ir = it%4; ++ const short is = (8*ir)/16; ++ ++ device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir; ++ ++ for (int ib = ix; ib < nb; ib += 4) { ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ int qy[32]; ++ int4 sumq = {0, 0, 0, 0}; ++ float4 d8 = {0.f, 0.f, 0.f, 0.f}; ++ if (use_prequant_q8) { ++ for (short group = 0; group < 4; ++group) { ++ device const block_q8_0 * yb = y_q8 + ib*8 + 4*iq + group; ++ device const int8_t * yqs = yb->qs + 8*ir; ++ d8[group] = float(yb->d); ++ for (short i = 0; i < 8; ++i) { ++ const int q = int(yqs[i]); ++ qy[8*group + i] = q; ++ sumq[group] += q; ++ } ++ } ++ } else { ++ for (short i = 0; i < 8; ++i) { ++ yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0]; ++ yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8]; ++ yl[i+16] = y4[i+64]; sumy[2] += yl[i+16]; ++ yl[i+24] = y4[i+96]; sumy[3] += yl[i+24]; ++ } ++ } + +- for (short row = 0; row < nr0; row++) { +- const float db = dh[0]; +- const uint8_t ls1 = sc[0] & 0xf; +- const uint8_t ls2 = sc[0] >> 4; +- const float d1 = db * (0.5f + ls1); +- const float d2 = db * (0.5f + ls2); ++ device const uint8_t * sc_up = (device const uint8_t *)x_up[ib].scales + 8*iq + is; ++ device const uint16_t * qs_up = (device const uint16_t *)x_up[ib].qs + 16 * iq + 4 * ir; ++ device const half * dh_up = &x_up[ib].d; ++ device const uint8_t * sc_gate = (device const uint8_t *)x_gate[ib].scales + 8*iq + is; ++ device const uint16_t * qs_gate = (device const uint16_t *)x_gate[ib].qs + 16 * iq + 4 * ir; ++ device const half * dh_gate = &x_gate[ib].d; ++ ++ for (short row = 0; row < nr0; row++) { ++ if (use_inblock_repack) { ++ device const block_q2_K * xb_up = (device const block_q2_K *)dh_up; ++ device const block_q2_K * xb_gate = (device const block_q2_K *)dh_gate; ++ float up_dot = 0.f; ++ float up_min = 0.f; ++ float gate_dot = 0.f; ++ float gate_min = 0.f; ++ for (short group = 0; group < 4; ++group) { ++ const int subblock = 8*iq + ir/2 + 2*group; ++ const int q_offset = 4*subblock + 2*(ir & 1); ++ const uint8_t up_sc = xb_up->scales[subblock]; ++ const uint8_t gate_sc = xb_gate->scales[subblock]; ++ up_dot += q2_repacked_dot8(xb_up->qs + q_offset, yl + 8*group) * float(up_sc & 0x0F); ++ gate_dot += q2_repacked_dot8(xb_gate->qs + q_offset, yl + 8*group) * float(gate_sc & 0x0F); ++ up_min += sumy[group] * float(up_sc >> 4); ++ gate_min += sumy[group] * float(gate_sc >> 4); ++ } ++ sum_up[row] += float(xb_up->d) * up_dot - float(xb_up->dmin) * up_min; ++ sum_gate[row] += float(xb_gate->d) * gate_dot - float(xb_gate->dmin) * gate_min; ++ } else if (use_prequant_q8) { ++ int4 accq_up = {0, 0, 0, 0}; ++ int4 accq_gate = {0, 0, 0, 0}; ++ for (int i = 0; i < 8; i += 2) { ++ const uint16_t qu = qs_up[i/2]; ++ const uint16_t qg = qs_gate[i/2]; ++ accq_up[0] += qy[i+ 0] * int( qu & 0x0003) + qy[i+ 1] * int((qu >> 8) & 0x0003); ++ accq_up[1] += qy[i+ 8] * int((qu >> 2) & 0x0003) + qy[i+ 9] * int((qu >> 10) & 0x0003); ++ accq_up[2] += qy[i+16] * int((qu >> 4) & 0x0003) + qy[i+17] * int((qu >> 12) & 0x0003); ++ accq_up[3] += qy[i+24] * int((qu >> 6) & 0x0003) + qy[i+25] * int((qu >> 14) & 0x0003); ++ accq_gate[0] += qy[i+ 0] * int( qg & 0x0003) + qy[i+ 1] * int((qg >> 8) & 0x0003); ++ accq_gate[1] += qy[i+ 8] * int((qg >> 2) & 0x0003) + qy[i+ 9] * int((qg >> 10) & 0x0003); ++ accq_gate[2] += qy[i+16] * int((qg >> 4) & 0x0003) + qy[i+17] * int((qg >> 12) & 0x0003); ++ accq_gate[3] += qy[i+24] * int((qg >> 6) & 0x0003) + qy[i+25] * int((qg >> 14) & 0x0003); ++ } ++ const float4 up_scale = float4(float(sc_up[0] & 0x0F), float(sc_up[2] & 0x0F), ++ float(sc_up[4] & 0x0F), float(sc_up[6] & 0x0F)); ++ const float4 up_min = float4(float(sc_up[0] & 0xF0), float(sc_up[2] & 0xF0), ++ float(sc_up[4] & 0xF0), float(sc_up[6] & 0xF0)); ++ const float4 gate_scale = float4(float(sc_gate[0] & 0x0F), float(sc_gate[2] & 0x0F), ++ float(sc_gate[4] & 0x0F), float(sc_gate[6] & 0x0F)); ++ const float4 gate_min = float4(float(sc_gate[0] & 0xF0), float(sc_gate[2] & 0xF0), ++ float(sc_gate[4] & 0xF0), float(sc_gate[6] & 0xF0)); ++ const float4 acc_up = float4(accq_up) * d8; ++ const float4 acc_gate = float4(accq_gate) * d8; ++ const float4 sums = float4(sumq) * d8; ++ sum_up[row] += float(dh_up[0]) * dot(acc_up, up_scale) - ++ float(dh_up[1]) * (1.f/16.f) * dot(sums, up_min); ++ sum_gate[row] += float(dh_gate[0]) * dot(acc_gate, gate_scale) - ++ float(dh_gate[1]) * (1.f/16.f) * dot(sums, gate_min); ++ } else { ++ float4 acc1_up = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2_up = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc1_gate = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2_gate = {0.f, 0.f, 0.f, 0.f}; ++ for (int i = 0; i < 8; i += 2) { ++ acc1_up[0] += yl[i+ 0] * (qs_up[i/2] & 0x0003); ++ acc2_up[0] += yl[i+ 1] * (qs_up[i/2] & 0x0300); ++ acc1_up[1] += yl[i+ 8] * (qs_up[i/2] & 0x000c); ++ acc2_up[1] += yl[i+ 9] * (qs_up[i/2] & 0x0c00); ++ acc1_up[2] += yl[i+16] * (qs_up[i/2] & 0x0030); ++ acc2_up[2] += yl[i+17] * (qs_up[i/2] & 0x3000); ++ acc1_up[3] += yl[i+24] * (qs_up[i/2] & 0x00c0); ++ acc2_up[3] += yl[i+25] * (qs_up[i/2] & 0xc000); ++ ++ acc1_gate[0] += yl[i+ 0] * (qs_gate[i/2] & 0x0003); ++ acc2_gate[0] += yl[i+ 1] * (qs_gate[i/2] & 0x0300); ++ acc1_gate[1] += yl[i+ 8] * (qs_gate[i/2] & 0x000c); ++ acc2_gate[1] += yl[i+ 9] * (qs_gate[i/2] & 0x0c00); ++ acc1_gate[2] += yl[i+16] * (qs_gate[i/2] & 0x0030); ++ acc2_gate[2] += yl[i+17] * (qs_gate[i/2] & 0x3000); ++ acc1_gate[3] += yl[i+24] * (qs_gate[i/2] & 0x00c0); ++ acc2_gate[3] += yl[i+25] * (qs_gate[i/2] & 0xc000); ++ } + +- float sum1 = 0, sum2 = 0; +- for (short l = 0; l < 2; ++l) { +- const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + (q2[l] & 511)); +- const uint8_t signs = ssigns[(q2[l] >> 9)]; +- for (short j = 0; j < 8; ++j) { +- sum1 += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); ++ float dall = dh_up[0]; ++ float dmin = dh_up[1] * 1.f/16.f; ++ sum_up[row] += dall * ((acc1_up[0] + 1.f/256.f * acc2_up[0]) * (sc_up[0] & 0xF) * 1.f/ 1.f + ++ (acc1_up[1] + 1.f/256.f * acc2_up[1]) * (sc_up[2] & 0xF) * 1.f/ 4.f + ++ (acc1_up[2] + 1.f/256.f * acc2_up[2]) * (sc_up[4] & 0xF) * 1.f/16.f + ++ (acc1_up[3] + 1.f/256.f * acc2_up[3]) * (sc_up[6] & 0xF) * 1.f/64.f) - ++ dmin * (sumy[0] * (sc_up[0] & 0xF0) + sumy[1] * (sc_up[2] & 0xF0) + ++ sumy[2] * (sc_up[4] & 0xF0) + sumy[3] * (sc_up[6] & 0xF0)); ++ ++ dall = dh_gate[0]; ++ dmin = dh_gate[1] * 1.f/16.f; ++ sum_gate[row] += dall * ((acc1_gate[0] + 1.f/256.f * acc2_gate[0]) * (sc_gate[0] & 0xF) * 1.f/ 1.f + ++ (acc1_gate[1] + 1.f/256.f * acc2_gate[1]) * (sc_gate[2] & 0xF) * 1.f/ 4.f + ++ (acc1_gate[2] + 1.f/256.f * acc2_gate[2]) * (sc_gate[4] & 0xF) * 1.f/16.f + ++ (acc1_gate[3] + 1.f/256.f * acc2_gate[3]) * (sc_gate[6] & 0xF) * 1.f/64.f) - ++ dmin * (sumy[0] * (sc_gate[0] & 0xF0) + sumy[1] * (sc_gate[2] & 0xF0) + ++ sumy[2] * (sc_gate[4] & 0xF0) + sumy[3] * (sc_gate[6] & 0xF0)); ++ } ++ ++ qs_up += args.nb01/2; ++ sc_up += args.nb01; ++ dh_up += args.nb01/2; ++ qs_gate += args.nb01/2; ++ sc_gate += args.nb01; ++ dh_gate += args.nb01/2; + } +- } +- for (short l = 2; l < 4; ++l) { +- const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + (q2[l] & 511)); +- const uint8_t signs = ssigns[(q2[l] >> 9)]; +- for (short j = 0; j < 8; ++j) { +- sum2 += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); ++ ++ if (!use_prequant_q8) { ++ y4 += 4 * QK_K; + } + } +- sumf[row] += d1 * sum1 + d2 * sum2; ++ } ++ } ++ ++ const bool valid_output = slot < args.nei0 && token < args.nei1; ++ const float route_weight = valid_output && args.weighted != 0 ? ++ ((device const float *) (weights + uint64_t(slot)*args.weights_nb1 + uint64_t(token)*args.weights_nb2))[0] : ++ 1.0f; ++ for (int row = 0; row < nr0; ++row) { ++ const float up_value = simd_sum(sum_up[row]); ++ const float gate_value = simd_sum(sum_gate[row]); ++ if (tiisg == 0 && valid_output && first_row + row < args.ne0) { ++ const float silu = gate_value / (1.0f + exp(-gate_value)); ++ ((device float *) dst)[uint64_t(slot)*args.ne0 + uint64_t(token)*args.ne1*args.ne0 + first_row + row] = ++ up_value * silu * route_weight; ++ } ++ } ++ ++ (void) partials; ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot4_dual( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<4>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot1_dual( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<1>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot1_dual_prequant_q8( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<1, 8, false, true>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot1_dual_inblock_q2( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<1, 8, false, false, true>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot1_dual_scan( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<1, 8, true>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot4_dual_r12( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<4, 12>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot4_dual_r16( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<4, 16>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot2_dual( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<2>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_slot8( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ device const char * src1_q8, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<8>( ++ args, src0_up_all, src0_gate_all, src1, dst, ids, weights, src1_q8, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_glm_moe_weights_q2_gate_up_swiglu_pair_sg_slot1( ++ constant ggml_metal_kargs_moe_route_weights & route_args, ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * probs, ++ device const char * ids, ++ device char * weights, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ constexpr int32_t TOP_K = 8; ++ ++ if (route_args.n_tokens != 1 || route_args.n_expert_used != TOP_K || ++ args.nei0 != TOP_K || args.nei1 != 1 || args.weighted != 0) { ++ return; ++ } ++ ++ // The route weights are independent of the output-row tile. One lane in ++ // the first tile computes them while every tile evaluates its Q2 rows. ++ if (tgpig.x == 0 && tgpig.z == 0 && sgitg == 0 && tiisg == 0) { ++ float selected[TOP_K]; ++ float sum = 0.0f; ++ for (int32_t slot = 0; slot < TOP_K; ++slot) { ++ const int32_t expert = ((device const int32_t *) ++ (ids + uint64_t(slot)*route_args.ids_nb0))[0]; ++ float value = 0.0f; ++ if (expert >= 0 && expert < route_args.n_expert) { ++ value = ((device const float *) ++ (probs + uint64_t(expert)*route_args.probs_nb1))[0]; ++ } ++ selected[slot] = value; ++ sum += value; ++ } ++ ++ if (route_args.norm) { ++ sum = max(sum, route_args.clamp_min); ++ } ++ for (int32_t slot = 0; slot < TOP_K; ++slot) { ++ float value = selected[slot]; ++ if (route_args.norm) { ++ value /= sum; ++ } ++ value *= route_args.scale; ++ if (route_args._pad0 != 0) { ++ value = slot == 0 ? 1.0f : 0.0f; ++ } ++ ((device float *) (weights + uint64_t(slot)*route_args.dst_nb1))[0] = value; ++ } ++ } ++ ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<1>( ++ args, ++ src0_up_all, ++ src0_gate_all, ++ src1, ++ dst, ++ ids, ++ weights, ++ src1, ++ partials, ++ tgpig, ++ tiisg, ++ sgitg); ++} ++ ++kernel void kernel_glm_moe_route_q2_gate_up_swiglu_pair_sg_slot8( ++ constant ggml_metal_kargs_topk_moe_route & route_args, ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & args, ++ device const char * logits, ++ device const char * bias, ++ device char * ids, ++ device char * weights, ++ device const char * src0_up_all, ++ device const char * src0_gate_all, ++ device const char * src1, ++ device char * dst, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ constexpr int32_t MAX_EXPERTS = 256; ++ constexpr int32_t TOP_K = 8; ++ constexpr int32_t N_SIMDGROUPS = 8; ++ constexpr int nr0 = 8; ++ ++ threadgroup float candidate_scores[N_SIMDGROUPS*TOP_K]; ++ threadgroup float candidate_probs[N_SIMDGROUPS*TOP_K]; ++ threadgroup int32_t candidate_ids[N_SIMDGROUPS*TOP_K]; ++ threadgroup float top_probs[TOP_K + 1]; ++ threadgroup int32_t top_ids[TOP_K]; ++ ++ const int32_t r0 = int32_t(tgpig.x); ++ const int32_t token = int32_t(tgpig.z); ++ const int32_t lane = int32_t(tiisg); ++ const int32_t simdgroup = int32_t(sgitg); ++ const int32_t tid = simdgroup*32 + lane; ++ ++ if (route_args.n_expert != MAX_EXPERTS || ++ route_args.top_k != TOP_K || ++ args.nei0 != TOP_K || ++ token >= route_args.n_tokens || ++ token >= args.nei1) { ++ return; ++ } ++ ++ float lane_prob = 0.0f; ++ float lane_score = -INFINITY; ++ const int32_t expert_id = tid; ++ if (tid < route_args.n_expert) { ++ const float logit = ((device const float *) (logits + uint64_t(tid)*route_args.logits_nb0 + uint64_t(token)*route_args.logits_nb1))[0]; ++ lane_prob = 1.0f/(1.0f + exp(-logit)); ++ lane_score = lane_prob + (route_args.has_bias ? ((device const float *) (bias + uint64_t(tid)*route_args.bias_nb0))[0] : 0.0f); ++ } ++ ++ for (int32_t top = 0; top < TOP_K; ++top) { ++ const float max_score = simd_max(lane_score); ++ const int32_t lane_id = lane_score == max_score ? expert_id : MAX_EXPERTS; ++ const int32_t max_id = simd_min(lane_id); ++ const float selected_prob = expert_id == max_id ? lane_prob : 0.0f; ++ const float max_prob = simd_sum(selected_prob); ++ ++ if (lane == 0) { ++ const int32_t candidate = simdgroup*TOP_K + top; ++ candidate_scores[candidate] = max_score; ++ candidate_probs[candidate] = max_prob; ++ candidate_ids[candidate] = max_id == MAX_EXPERTS ? -1 : max_id; ++ } ++ ++ if (expert_id == max_id) { ++ lane_score = -INFINITY; ++ } ++ } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (simdgroup == 0) { ++ constexpr int32_t N_CANDIDATES = N_SIMDGROUPS*TOP_K; ++ float local_scores[2]; ++ float local_probs[2]; ++ int32_t local_ids[2]; ++ ++ for (int32_t i = 0; i < 2; ++i) { ++ const int32_t candidate = lane + i*32; ++ local_scores[i] = candidate < N_CANDIDATES ? candidate_scores[candidate] : -INFINITY; ++ local_probs[i] = candidate < N_CANDIDATES ? candidate_probs[candidate] : 0.0f; ++ local_ids[i] = candidate < N_CANDIDATES ? candidate_ids[candidate] : -1; ++ } ++ ++ for (int32_t top = 0; top < TOP_K; ++top) { ++ const bool second_better = ++ local_scores[1] > local_scores[0] || ++ (local_scores[1] == local_scores[0] && local_ids[1] >= 0 && ++ (local_ids[0] < 0 || local_ids[1] < local_ids[0])); ++ float best_score = second_better ? local_scores[1] : local_scores[0]; ++ float best_prob = second_better ? local_probs[1] : local_probs[0]; ++ int32_t best_id = second_better ? local_ids[1] : local_ids[0]; ++ ++ const float max_score = simd_max(best_score); ++ const int32_t winner_id = best_score == max_score ? best_id : MAX_EXPERTS; ++ const int32_t max_id = simd_min(winner_id); ++ const float selected_prob = best_id == max_id ? best_prob : 0.0f; ++ const float max_prob = simd_sum(selected_prob); ++ ++ if (lane == 0) { ++ top_probs[top] = max_prob; ++ top_ids[top] = max_id == MAX_EXPERTS ? -1 : max_id; ++ } ++ ++ for (int32_t i = 0; i < 2; ++i) { ++ if (local_ids[i] == max_id) { ++ local_scores[i] = -INFINITY; ++ } ++ } ++ } ++ ++ if (lane == 0) { ++ const int32_t active_top_k = route_args._pad1 > 0 ? min(TOP_K, route_args._pad1) : TOP_K; ++ float sum = 0.0f; ++ if (route_args.norm) { ++ for (int32_t i = 0; i < active_top_k; ++i) { ++ sum += top_probs[i]; ++ } ++ sum = max(sum, route_args.clamp_min); ++ } ++ ++ for (int32_t i = 0; i < TOP_K; ++i) { ++ float weight = i < active_top_k ? top_probs[i] : 0.0f; ++ if (route_args.norm) { ++ weight /= sum; ++ } ++ weight *= route_args.scale; ++ top_probs[i] = weight; ++ } ++ ++ if (route_args._pad0 != 0) { ++ for (int32_t i = 0; i < TOP_K; ++i) { ++ for (int32_t j = i + 1; j < TOP_K; ++j) { ++ const int32_t lhs = top_ids[i]; ++ const int32_t rhs = top_ids[j]; ++ const bool rhs_first = rhs >= 0 && (lhs < 0 || rhs < lhs); ++ if (rhs_first) { ++ const int32_t tmp_id = top_ids[i]; ++ const float tmp_prob = top_probs[i]; ++ top_ids[i] = top_ids[j]; ++ top_probs[i] = top_probs[j]; ++ top_ids[j] = tmp_id; ++ top_probs[j] = tmp_prob; ++ } ++ } ++ } ++ } ++ ++ for (int32_t i = 0; i < TOP_K; ++i) { ++ if (r0 == 0) { ++ ((device int32_t *) (ids + uint64_t(i)*route_args.ids_nb0 + uint64_t(token)*route_args.ids_nb1))[0] = top_ids[i]; ++ ((device float *) (weights + uint64_t(i)*route_args.weights_nb1 + uint64_t(token)*route_args.weights_nb2))[0] = top_probs[i]; ++ } ++ } ++ } ++ } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ const int nb = args.ne00/QK_K; ++ const int slot = sgitg; ++ const int first_row = r0 * nr0; ++ ++ float sum_up[nr0] = {0.f}; ++ float sum_gate[nr0] = {0.f}; ++ ++ const int32_t active_top_k = route_args._pad1 > 0 ? min(TOP_K, route_args._pad1) : TOP_K; ++ const int32_t selected_expert_id = slot < active_top_k ? top_ids[slot] : -1; ++ if (slot < active_top_k && selected_expert_id >= 0 && selected_expert_id < args.ne02) { ++ const uint64_t offset0 = uint64_t(selected_expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot % args.ne11)*args.nb11 + uint64_t(token)*args.nb12; ++ ++ device const block_q2_K * x_up = (device const block_q2_K *) (src0_up_all + offset0); ++ device const block_q2_K * x_gate = (device const block_q2_K *) (src0_gate_all + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ const short iq = it/4; ++ const short ir = it%4; ++ const short is = (8*ir)/16; ++ ++ device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir; ++ ++ for (int ib = ix; ib < nb; ib += 4) { ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ for (short i = 0; i < 8; ++i) { ++ yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0]; ++ yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8]; ++ yl[i+16] = y4[i+64]; sumy[2] += yl[i+16]; ++ yl[i+24] = y4[i+96]; sumy[3] += yl[i+24]; ++ } ++ ++ device const uint8_t * sc_up = (device const uint8_t *)x_up[ib].scales + 8*iq + is; ++ device const uint16_t * qs_up = (device const uint16_t *)x_up[ib].qs + 16 * iq + 4 * ir; ++ device const half * dh_up = &x_up[ib].d; ++ device const uint8_t * sc_gate = (device const uint8_t *)x_gate[ib].scales + 8*iq + is; ++ device const uint16_t * qs_gate = (device const uint16_t *)x_gate[ib].qs + 16 * iq + 4 * ir; ++ device const half * dh_gate = &x_gate[ib].d; ++ ++ for (short row = 0; row < nr0; row++) { ++ float4 acc1_up = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2_up = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc1_gate = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2_gate = {0.f, 0.f, 0.f, 0.f}; ++ for (int i = 0; i < 8; i += 2) { ++ acc1_up[0] += yl[i+ 0] * (qs_up[i/2] & 0x0003); ++ acc2_up[0] += yl[i+ 1] * (qs_up[i/2] & 0x0300); ++ acc1_up[1] += yl[i+ 8] * (qs_up[i/2] & 0x000c); ++ acc2_up[1] += yl[i+ 9] * (qs_up[i/2] & 0x0c00); ++ acc1_up[2] += yl[i+16] * (qs_up[i/2] & 0x0030); ++ acc2_up[2] += yl[i+17] * (qs_up[i/2] & 0x3000); ++ acc1_up[3] += yl[i+24] * (qs_up[i/2] & 0x00c0); ++ acc2_up[3] += yl[i+25] * (qs_up[i/2] & 0xc000); ++ ++ acc1_gate[0] += yl[i+ 0] * (qs_gate[i/2] & 0x0003); ++ acc2_gate[0] += yl[i+ 1] * (qs_gate[i/2] & 0x0300); ++ acc1_gate[1] += yl[i+ 8] * (qs_gate[i/2] & 0x000c); ++ acc2_gate[1] += yl[i+ 9] * (qs_gate[i/2] & 0x0c00); ++ acc1_gate[2] += yl[i+16] * (qs_gate[i/2] & 0x0030); ++ acc2_gate[2] += yl[i+17] * (qs_gate[i/2] & 0x3000); ++ acc1_gate[3] += yl[i+24] * (qs_gate[i/2] & 0x00c0); ++ acc2_gate[3] += yl[i+25] * (qs_gate[i/2] & 0xc000); ++ } ++ ++ float dall = dh_up[0]; ++ float dmin = dh_up[1] * 1.f/16.f; ++ sum_up[row] += dall * ((acc1_up[0] + 1.f/256.f * acc2_up[0]) * (sc_up[0] & 0xF) * 1.f/ 1.f + ++ (acc1_up[1] + 1.f/256.f * acc2_up[1]) * (sc_up[2] & 0xF) * 1.f/ 4.f + ++ (acc1_up[2] + 1.f/256.f * acc2_up[2]) * (sc_up[4] & 0xF) * 1.f/16.f + ++ (acc1_up[3] + 1.f/256.f * acc2_up[3]) * (sc_up[6] & 0xF) * 1.f/64.f) - ++ dmin * (sumy[0] * (sc_up[0] & 0xF0) + sumy[1] * (sc_up[2] & 0xF0) + ++ sumy[2] * (sc_up[4] & 0xF0) + sumy[3] * (sc_up[6] & 0xF0)); ++ ++ dall = dh_gate[0]; ++ dmin = dh_gate[1] * 1.f/16.f; ++ sum_gate[row] += dall * ((acc1_gate[0] + 1.f/256.f * acc2_gate[0]) * (sc_gate[0] & 0xF) * 1.f/ 1.f + ++ (acc1_gate[1] + 1.f/256.f * acc2_gate[1]) * (sc_gate[2] & 0xF) * 1.f/ 4.f + ++ (acc1_gate[2] + 1.f/256.f * acc2_gate[2]) * (sc_gate[4] & 0xF) * 1.f/16.f + ++ (acc1_gate[3] + 1.f/256.f * acc2_gate[3]) * (sc_gate[6] & 0xF) * 1.f/64.f) - ++ dmin * (sumy[0] * (sc_gate[0] & 0xF0) + sumy[1] * (sc_gate[2] & 0xF0) + ++ sumy[2] * (sc_gate[4] & 0xF0) + sumy[3] * (sc_gate[6] & 0xF0)); ++ ++ qs_up += args.nb01/2; ++ sc_up += args.nb01; ++ dh_up += args.nb01/2; ++ qs_gate += args.nb01/2; ++ sc_gate += args.nb01; ++ dh_gate += args.nb01/2; ++ } ++ ++ y4 += 4 * QK_K; ++ } ++ } ++ ++ const bool valid_output = slot < args.nei0 && token < args.nei1; ++ const float route_weight = valid_output && args.weighted != 0 ? top_probs[slot] : 1.0f; ++ for (int row = 0; row < nr0; ++row) { ++ const float up_value = simd_sum(sum_up[row]); ++ const float gate_value = simd_sum(sum_gate[row]); ++ if (tiisg == 0 && valid_output && first_row + row < args.ne0) { ++ const float silu = gate_value / (1.0f + exp(-gate_value)); ++ ((device float *) dst)[uint64_t(slot)*args.ne0 + uint64_t(token)*args.ne1*args.ne0 + first_row + row] = ++ up_value * silu * route_weight; ++ } ++ } ++ ++ (void) partials; ++} ++ ++template ++void kernel_mul_mv_q3_K_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = fixed_nb > 0 ? fixed_nb : args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_q3_K * x = (device const block_q3_K *) (src0 + offset0); ++ device const float * yy = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ ++ //const uint16_t kmask1 = 0x3030; ++ //const uint16_t kmask2 = 0x0f0f; ++ ++ const short tid = tiisg/4; ++ const short ix = tiisg%4; ++ const short ip = tid/4; // 0 or 1 ++ const short il = 2*((tid%4)/2); // 0 or 2 ++ const short ir = tid%2; ++ const short l0 = 8*ir; ++ ++ // One would think that the Metal compiler would figure out that ip and il can only have ++ // 4 possible states, and optimize accordingly. Well, no. It needs help, and we do it ++ // with these two tales. ++ // ++ // Possible masks for the high bit ++ const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200}, // ip = 0, il = 0 ++ {0x0004, 0x0400, 0x0008, 0x0800}, // ip = 0, il = 2 ++ {0x0010, 0x1000, 0x0020, 0x2000}, // ip = 1, il = 0 ++ {0x0040, 0x4000, 0x0080, 0x8000}}; // ip = 1, il = 2 ++ ++ // Possible masks for the low 2 bits ++ const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}}; ++ ++ const ushort4 hm = mm[2*ip + il/2]; ++ ++ const short shift = 2*il; ++ ++ const float v1 = il == 0 ? 4.f : 64.f; ++ const float v2 = 4.f * v1; ++ ++ const uint16_t s_shift1 = 4*ip; ++ const uint16_t s_shift2 = s_shift1 + il; ++ ++ const short q_offset = 32*ip + l0; ++ const short y_offset = 128*ip + 32*il + l0; ++ ++ device const float * y1 = yy + ix*QK_K + y_offset; ++ ++ uint32_t scales32, aux32; ++ thread uint16_t * scales16 = (thread uint16_t *)&scales32; ++ thread const int8_t * scales = (thread const int8_t *)&scales32; ++ ++ float sumf1[nr0] = {0.f}; ++ float sumf2[nr0] = {0.f}; ++ ++ for (int i = ix; i < nb; i += 4) { ++ for (short l = 0; l < 8; ++l) { ++ yl[l+ 0] = y1[l+ 0]; ++ yl[l+ 8] = y1[l+16]; ++ yl[l+16] = y1[l+32]; ++ yl[l+24] = y1[l+48]; ++ } ++ ++ device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset); ++ device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0); ++ device const uint16_t * a = (device const uint16_t *)(x[i].scales); ++ device const half * dh = &x[i].d; ++ ++ for (short row = 0; row < nr0; ++row) { ++ const float d_all = (float)dh[0]; ++ ++ scales16[0] = a[4]; ++ scales16[1] = a[5]; ++ aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030; ++ scales16[0] = a[il+0]; ++ scales16[1] = a[il+1]; ++ scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32; ++ ++ float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2]; ++ s1 += yl[l+0] * (qs & qm[il/2][0]); ++ s2 += yl[l+1] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]); ++ s4 += yl[l+16] * (qs & qm[il/2][2]); ++ s5 += yl[l+17] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]); ++ } ++ float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[0] - 32); ++ sumf2[row] += d2 * (scales[2] - 32); ++ ++ s1 = s2 = s3 = s4 = s5 = s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2+8]; ++ s1 += yl[l+8] * (qs & qm[il/2][0]); ++ s2 += yl[l+9] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]); ++ s4 += yl[l+24] * (qs & qm[il/2][2]); ++ s5 += yl[l+25] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]); ++ } ++ d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[1] - 32); ++ sumf2[row] += d2 * (scales[3] - 32); ++ ++ q += args.nb01/2; ++ h += args.nb01/2; ++ a += args.nb01/2; ++ dh += args.nb01/2; ++ } ++ ++ y1 += 4 * QK_K; ++ } ++ ++ for (int row = 0; row < nr0; ++row) { ++ const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift); ++ sumf1[row] = simd_sum(sumf); ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ dst_f32[first_row + row] = sumf1[row]; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_q3_K_f32")]] ++kernel void kernel_mul_mv_q3_K_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_q3_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++} ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const short NSG = FC_mul_mv_nsg; ++ const int nr0 = N_R0_Q3_K; ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int token = tgpig.y; ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const short tid = tiisg/4; ++ const short ix = tiisg%4; ++ const short ip = tid/4; ++ const short il = 2*((tid%4)/2); ++ const short ir = tid%2; ++ const short l0 = 8*ir; ++ ++ const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200}, ++ {0x0004, 0x0400, 0x0008, 0x0800}, ++ {0x0010, 0x1000, 0x0020, 0x2000}, ++ {0x0040, 0x4000, 0x0080, 0x8000}}; ++ const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}}; ++ ++ const ushort4 hm = mm[2*ip + il/2]; ++ const short shift = 2*il; ++ const float v1 = il == 0 ? 4.f : 64.f; ++ const float v2 = 4.f * v1; ++ const uint16_t s_shift1 = 4*ip; ++ const uint16_t s_shift2 = s_shift1 + il; ++ const short q_offset = 32*ip + l0; ++ const short y_offset = 128*ip + 32*il + l0; ++ ++ float sum_out[nr0] = {0.f}; ++ float yl[32]; ++ uint32_t scales32, aux32; ++ thread uint16_t * scales16 = (thread uint16_t *)&scales32; ++ thread const int8_t * scales = (thread const int8_t *)&scales32; ++ ++ for (int slot = 0; slot < args.nei0; ++slot) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id < 0 || expert_id >= args.ne02) { ++ continue; ++ } ++ ++ const float route_weight = extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0]; ++ if (route_weight == 0.0f) { ++ continue; ++ } ++ ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot)*args.nb11 + uint64_t(token)*args.nb12; ++ device const block_q3_K * x = (device const block_q3_K *) (src0s + offset0); ++ device const float * yy = (device const float *) (src1 + offset1); ++ device const float * y1 = yy + ix*QK_K + y_offset; ++ ++ float sumf1[nr0] = {0.f}; ++ float sumf2[nr0] = {0.f}; ++ ++ for (int i = ix; i < nb; i += 4) { ++ for (short l = 0; l < 8; ++l) { ++ yl[l+ 0] = y1[l+ 0]; ++ yl[l+ 8] = y1[l+16]; ++ yl[l+16] = y1[l+32]; ++ yl[l+24] = y1[l+48]; ++ } ++ ++ device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset); ++ device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0); ++ device const uint16_t * a = (device const uint16_t *)(x[i].scales); ++ device const half * dh = &x[i].d; ++ ++ for (short row = 0; row < nr0; ++row) { ++ const float d_all = (float)dh[0]; ++ ++ scales16[0] = a[4]; ++ scales16[1] = a[5]; ++ aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030; ++ scales16[0] = a[il+0]; ++ scales16[1] = a[il+1]; ++ scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32; ++ ++ float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2]; ++ s1 += yl[l+0] * (qs & qm[il/2][0]); ++ s2 += yl[l+1] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]); ++ s4 += yl[l+16] * (qs & qm[il/2][2]); ++ s5 += yl[l+17] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]); ++ } ++ float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[0] - 32); ++ sumf2[row] += d2 * (scales[2] - 32); ++ ++ s1 = s2 = s3 = s4 = s5 = s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2+8]; ++ s1 += yl[l+8] * (qs & qm[il/2][0]); ++ s2 += yl[l+9] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]); ++ s4 += yl[l+24] * (qs & qm[il/2][2]); ++ s5 += yl[l+25] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]); ++ } ++ d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[1] - 32); ++ sumf2[row] += d2 * (scales[3] - 32); ++ ++ q += args.nb01/2; ++ h += args.nb01/2; ++ a += args.nb01/2; ++ dh += args.nb01/2; ++ } ++ ++ y1 += 4 * QK_K; ++ } ++ ++ for (int row = 0; row < nr0; ++row) { ++ const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift); ++ sum_out[row] += route_weight * simd_sum(sumf); ++ } ++ } ++ ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ ((device float *) (dst + uint64_t(first_row + row)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum_out[row]; ++ } ++ } ++} ++ ++template ++void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const int n_slots = 8; ++ const int nb = NB_FIXED == 0 ? args.ne00/QK_K : NB_FIXED; ++ ++ const int r0 = tgpig.x; ++ const int token = tgpig.y; ++ const int slot = sgitg; ++ const int first_row = r0 * nr0; ++ ++ const short tid = tiisg/4; ++ const short ix = tiisg%4; ++ const short ip = tid/4; ++ const short il = 2*((tid%4)/2); ++ const short ir = tid%2; ++ const short l0 = 8*ir; ++ ++ const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200}, ++ {0x0004, 0x0400, 0x0008, 0x0800}, ++ {0x0010, 0x1000, 0x0020, 0x2000}, ++ {0x0040, 0x4000, 0x0080, 0x8000}}; ++ const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}}; ++ ++ const ushort4 hm = mm[2*ip + il/2]; ++ const short shift = 2*il; ++ const float v1 = il == 0 ? 4.f : 64.f; ++ const float v2 = 4.f * v1; ++ const uint16_t s_shift1 = 4*ip; ++ const uint16_t s_shift2 = s_shift1 + il; ++ const short q_offset = 32*ip + l0; ++ const short y_offset = 128*ip + 32*il + l0; ++ ++ float sumf1[nr0] = {0.f}; ++ float sumf2[nr0] = {0.f}; ++ ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id >= 0 && expert_id < args.ne02) { ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot)*args.nb11 + uint64_t(token)*args.nb12; ++ device const block_q3_K * x = (device const block_q3_K *) (src0s + offset0); ++ device const act_t * yy = (device const act_t *) (src1 + offset1); ++ device const act_t * y1 = yy + ix*QK_K + y_offset; ++ ++ float yl[32]; ++ uint32_t scales32, aux32; ++ thread uint16_t * scales16 = (thread uint16_t *)&scales32; ++ thread const int8_t * scales = (thread const int8_t *)&scales32; ++ ++ for (int i = ix; i < nb; i += 4) { ++ for (short l = 0; l < 8; ++l) { ++ yl[l+ 0] = float(y1[l+ 0]); ++ yl[l+ 8] = float(y1[l+16]); ++ yl[l+16] = float(y1[l+32]); ++ yl[l+24] = float(y1[l+48]); ++ } ++ ++ device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset); ++ device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0); ++ device const uint16_t * a = (device const uint16_t *)(x[i].scales); ++ device const half * dh = &x[i].d; ++ ++ for (short row = 0; row < nr0; ++row) { ++ const float d_all = (float)dh[0]; ++ ++ scales16[0] = a[4]; ++ scales16[1] = a[5]; ++ aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030; ++ scales16[0] = a[il+0]; ++ scales16[1] = a[il+1]; ++ scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32; ++ ++ float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2]; ++ s1 += yl[l+0] * (qs & qm[il/2][0]); ++ s2 += yl[l+1] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]); ++ s4 += yl[l+16] * (qs & qm[il/2][2]); ++ s5 += yl[l+17] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]); ++ } ++ float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[0] - 32); ++ sumf2[row] += d2 * (scales[2] - 32); ++ ++ s1 = s2 = s3 = s4 = s5 = s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2+8]; ++ s1 += yl[l+8] * (qs & qm[il/2][0]); ++ s2 += yl[l+9] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]); ++ s4 += yl[l+24] * (qs & qm[il/2][2]); ++ s5 += yl[l+25] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]); ++ } ++ d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[1] - 32); ++ sumf2[row] += d2 * (scales[3] - 32); ++ ++ q += args.nb01/2; ++ h += args.nb01/2; ++ a += args.nb01/2; ++ dh += args.nb01/2; ++ } ++ ++ y1 += 4 * QK_K; ++ } ++ ++ const float route_weight = ALREADY_WEIGHTED_MODE == 1 ? 1.0f : ++ (ALREADY_WEIGHTED_MODE == 0 ? ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0] : ++ (extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0])); ++ for (int row = 0; row < nr0; ++row) { ++ const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift); ++ const float dot = simd_sum(sumf); ++ sumf1[row] = nr0 % 4 == 0 && first_row + nr0 <= args.ne0 ? dot : route_weight * dot; ++ } ++ } ++ } ++ ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0; ++row) { ++ partials[slot*nr0 + row] = first_row + row < args.ne0 ? sumf1[row] : 0.0f; ++ } ++ } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (nr0 % 4 == 0 && first_row + nr0 <= args.ne0) { ++ if (sgitg == 0 && tiisg < nr0/4) { ++ float4 sum = float4(0.0f); ++ for (int s = 0; s < n_slots; ++s) { ++ const float weight = ALREADY_WEIGHTED_MODE == 1 ? 1.0f : ++ (ALREADY_WEIGHTED_MODE == 0 ? ++ ((device const float *) (weights + s*extra.weights_nb1 + token*extra.weights_nb2))[0] : ++ (extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + s*extra.weights_nb1 + token*extra.weights_nb2))[0])); ++ sum += ((threadgroup const float4 *) (partials + s*nr0))[tiisg] * weight; ++ } ++ ((device float4 *) (dst + uint64_t(first_row + 4*tiisg)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum; ++ } ++ } else if (sgitg == 0 && tiisg < nr0 && first_row + tiisg < args.ne0) { ++ float sum = 0.0f; ++ for (int s = 0; s < n_slots; ++s) { ++ sum += partials[s*nr0 + tiisg]; ++ } ++ ((device float *) (dst + uint64_t(first_row + tiisg)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum; ++ } ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r8( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<8>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r8_nb8( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<8, 8>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r8_nb8_w0( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<8, 8, 0>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r8_nb8_w0_f16( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<8, 8, 0, half>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r6_nb8_w0( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<6, 8, 0>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r10_nb8_w0( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<10, 8, 0>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_glm52_w0( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const int nr0 = 8; ++ const int nb = 8; ++ ++ const int r0 = tgpig.x; ++ const int slot = sgitg; ++ const int first_row = r0 * nr0; ++ ++ const short tid = tiisg/4; ++ const short ix = tiisg%4; ++ const short ip = tid/4; ++ const short il = 2*((tid%4)/2); ++ const short ir = tid%2; ++ const short l0 = 8*ir; ++ ++ const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200}, ++ {0x0004, 0x0400, 0x0008, 0x0800}, ++ {0x0010, 0x1000, 0x0020, 0x2000}, ++ {0x0040, 0x4000, 0x0080, 0x8000}}; ++ const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}}; ++ ++ const ushort4 hm = mm[2*ip + il/2]; ++ const short shift = 2*il; ++ const float v1 = il == 0 ? 4.f : 64.f; ++ const float v2 = 4.f * v1; ++ const uint16_t s_shift1 = 4*ip; ++ const uint16_t s_shift2 = s_shift1 + il; ++ const short q_offset = 32*ip + l0; ++ const short y_offset = 128*ip + 32*il + l0; ++ ++ float sumf1[nr0] = {0.f}; ++ float sumf2[nr0] = {0.f}; ++ ++ const int32_t expert_id = ((device const int32_t *) ids)[slot]; ++ if (expert_id >= 0 && expert_id < 256) { ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot)*args.nb11; ++ device const block_q3_K * x = (device const block_q3_K *) (src0s + offset0); ++ device const float * yy = (device const float *) (src1 + offset1); ++ device const float * y1 = yy + ix*QK_K + y_offset; ++ ++ float yl[32]; ++ uint32_t scales32, aux32; ++ thread uint16_t * scales16 = (thread uint16_t *)&scales32; ++ thread const int8_t * scales = (thread const int8_t *)&scales32; ++ ++ for (int i = ix; i < nb; i += 4) { ++ for (short l = 0; l < 8; ++l) { ++ yl[l+ 0] = y1[l+ 0]; ++ yl[l+ 8] = y1[l+16]; ++ yl[l+16] = y1[l+32]; ++ yl[l+24] = y1[l+48]; ++ } ++ ++ device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset); ++ device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0); ++ device const uint16_t * a = (device const uint16_t *)(x[i].scales); ++ device const half * dh = &x[i].d; ++ ++ for (short row = 0; row < nr0; ++row) { ++ const float d_all = (float)dh[0]; ++ ++ scales16[0] = a[4]; ++ scales16[1] = a[5]; ++ aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030; ++ scales16[0] = a[il+0]; ++ scales16[1] = a[il+1]; ++ scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32; ++ ++ float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2]; ++ s1 += yl[l+0] * (qs & qm[il/2][0]); ++ s2 += yl[l+1] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]); ++ s4 += yl[l+16] * (qs & qm[il/2][2]); ++ s5 += yl[l+17] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]); ++ } ++ float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[0] - 32); ++ sumf2[row] += d2 * (scales[2] - 32); ++ ++ s1 = s2 = s3 = s4 = s5 = s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2+8]; ++ s1 += yl[l+8] * (qs & qm[il/2][0]); ++ s2 += yl[l+9] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]); ++ s4 += yl[l+24] * (qs & qm[il/2][2]); ++ s5 += yl[l+25] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]); ++ } ++ d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[1] - 32); ++ sumf2[row] += d2 * (scales[3] - 32); ++ ++ q += args.nb01/2; ++ h += args.nb01/2; ++ a += args.nb01/2; ++ dh += args.nb01/2; ++ } ++ ++ y1 += 4 * QK_K; ++ } ++ ++ for (int row = 0; row < nr0; ++row) { ++ const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift); ++ sumf1[row] = simd_sum(sumf); ++ } ++ } ++ ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0; ++row) { ++ partials[slot*nr0 + row] = sumf1[row]; ++ } ++ } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (sgitg == 0 && tiisg < nr0) { ++ float sum = 0.0f; ++ for (int s = 0; s < 8; ++s) { ++ const float route_weight = ((device const float *) (weights + uint64_t(s)*extra.weights_nb1))[0]; ++ sum += partials[s*nr0 + tiisg] * route_weight; ++ } ++ ((device float *) (dst + uint64_t(first_row + tiisg)*extra.dst_nb0))[0] = sum; ++ } ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r12_nb8_w0( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<12, 8, 0>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_r16( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<16>( ++ args, src0s, src1, dst, ids, weights, extra, partials, tgpig, tiisg, sgitg); ++} ++ ++kernel void kernel_glm_moe_swiglu_q3_K_down_weighted( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * gate, ++ device const char * up, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ constant ggml_metal_kargs_glm_moe_swiglu_q3_down & swiglu_args, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const int nr0 = 8; ++ const int n_slots = 8; ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int token = tgpig.y; ++ const int slot = sgitg; ++ const int first_row = r0 * nr0; ++ ++ const short tid = tiisg/4; ++ const short ix = tiisg%4; ++ const short ip = tid/4; ++ const short il = 2*((tid%4)/2); ++ const short ir = tid%2; ++ const short l0 = 8*ir; ++ ++ const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200}, ++ {0x0004, 0x0400, 0x0008, 0x0800}, ++ {0x0010, 0x1000, 0x0020, 0x2000}, ++ {0x0040, 0x4000, 0x0080, 0x8000}}; ++ const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}}; ++ ++ const ushort4 hm = mm[2*ip + il/2]; ++ const short shift = 2*il; ++ const float v1 = il == 0 ? 4.f : 64.f; ++ const float v2 = 4.f * v1; ++ const uint16_t s_shift1 = 4*ip; ++ const uint16_t s_shift2 = s_shift1 + il; ++ const short q_offset = 32*ip + l0; ++ const short y_offset = 128*ip + 32*il + l0; ++ ++ float sumf1[nr0] = {0.f}; ++ float sumf2[nr0] = {0.f}; ++ ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id >= 0 && expert_id < args.ne02) { ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ device const block_q3_K * x = (device const block_q3_K *) (src0s + offset0); ++ device const float * gate_base = (device const float *) (gate + ++ uint64_t(slot)*swiglu_args.gate_nb1 + uint64_t(token)*swiglu_args.gate_nb2); ++ device const float * up_base = (device const float *) (up + ++ uint64_t(slot)*swiglu_args.up_nb1 + uint64_t(token)*swiglu_args.up_nb2); ++ float yl[32]; ++ uint32_t scales32, aux32; ++ thread uint16_t * scales16 = (thread uint16_t *)&scales32; ++ thread const int8_t * scales = (thread const int8_t *)&scales32; ++ ++ for (int i = ix; i < nb; i += 4) { ++ const int y0 = i*QK_K + y_offset; ++ for (short l = 0; l < 8; ++l) { ++ const float g0 = gate_base[y0 + l + 0]; ++ const float g1 = gate_base[y0 + l + 16]; ++ const float g2 = gate_base[y0 + l + 32]; ++ const float g3 = gate_base[y0 + l + 48]; ++ yl[l+ 0] = (g0 / (1.0f + exp(-g0))) * up_base[y0 + l + 0]; ++ yl[l+ 8] = (g1 / (1.0f + exp(-g1))) * up_base[y0 + l + 16]; ++ yl[l+16] = (g2 / (1.0f + exp(-g2))) * up_base[y0 + l + 32]; ++ yl[l+24] = (g3 / (1.0f + exp(-g3))) * up_base[y0 + l + 48]; ++ } ++ ++ device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset); ++ device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0); ++ device const uint16_t * a = (device const uint16_t *)(x[i].scales); ++ device const half * dh = &x[i].d; ++ ++ for (short row = 0; row < nr0; ++row) { ++ const float d_all = (float)dh[0]; ++ ++ scales16[0] = a[4]; ++ scales16[1] = a[5]; ++ aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030; ++ scales16[0] = a[il+0]; ++ scales16[1] = a[il+1]; ++ scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32; ++ ++ float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2]; ++ s1 += yl[l+0] * (qs & qm[il/2][0]); ++ s2 += yl[l+1] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]); ++ s4 += yl[l+16] * (qs & qm[il/2][2]); ++ s5 += yl[l+17] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]); ++ } ++ float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[0] - 32); ++ sumf2[row] += d2 * (scales[2] - 32); ++ ++ s1 = s2 = s3 = s4 = s5 = s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2+8]; ++ s1 += yl[l+8] * (qs & qm[il/2][0]); ++ s2 += yl[l+9] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]); ++ s4 += yl[l+24] * (qs & qm[il/2][2]); ++ s5 += yl[l+25] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]); ++ } ++ d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[1] - 32); ++ sumf2[row] += d2 * (scales[3] - 32); ++ ++ q += args.nb01/2; ++ h += args.nb01/2; ++ a += args.nb01/2; ++ dh += args.nb01/2; ++ } ++ } ++ ++ const float route_weight = extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0]; ++ for (int row = 0; row < nr0; ++row) { ++ const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift); ++ sumf1[row] = route_weight * simd_sum(sumf); ++ } ++ } ++ } ++ ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0; ++row) { ++ partials[slot*nr0 + row] = first_row + row < args.ne0 ? sumf1[row] : 0.0f; ++ } ++ } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (sgitg == 0 && tiisg < nr0 && first_row + tiisg < args.ne0) { ++ float sum = 0.0f; ++ for (int s = 0; s < n_slots; ++s) { ++ sum += partials[s*nr0 + tiisg]; ++ } ++ ((device float *) (dst + uint64_t(first_row + tiisg)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum; ++ } ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_split2( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const int nr0 = N_R0_Q3_K; ++ const int n_slots = 8; ++ const int n_parts = 2; ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int token = tgpig.y; ++ const int slot = int(sgitg / n_parts); ++ const int part = int(sgitg % n_parts); ++ const int first_row = r0 * nr0; ++ ++ const short tid = tiisg/4; ++ const short ix = tiisg%4; ++ const short ip = tid/4; ++ const short il = 2*((tid%4)/2); ++ const short ir = tid%2; ++ const short l0 = 8*ir; ++ ++ const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200}, ++ {0x0004, 0x0400, 0x0008, 0x0800}, ++ {0x0010, 0x1000, 0x0020, 0x2000}, ++ {0x0040, 0x4000, 0x0080, 0x8000}}; ++ const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}}; ++ ++ const ushort4 hm = mm[2*ip + il/2]; ++ const short shift = 2*il; ++ const float v1 = il == 0 ? 4.f : 64.f; ++ const float v2 = 4.f * v1; ++ const uint16_t s_shift1 = 4*ip; ++ const uint16_t s_shift2 = s_shift1 + il; ++ const short q_offset = 32*ip + l0; ++ const short y_offset = 128*ip + 32*il + l0; ++ ++ float sumf1[nr0] = {0.f}; ++ float sumf2[nr0] = {0.f}; ++ ++ if (slot < args.nei0 && token < args.nei1) { ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id >= 0 && expert_id < args.ne02) { ++ const int first_block = ix + part * 4; ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot)*args.nb11 + uint64_t(token)*args.nb12; ++ device const block_q3_K * x = (device const block_q3_K *) (src0s + offset0); ++ device const float * yy = (device const float *) (src1 + offset1); ++ device const float * y1 = yy + first_block*QK_K + y_offset; ++ ++ float yl[32]; ++ uint32_t scales32, aux32; ++ thread uint16_t * scales16 = (thread uint16_t *)&scales32; ++ thread const int8_t * scales = (thread const int8_t *)&scales32; ++ ++ for (int i = first_block; i < nb; i += 8) { ++ for (short l = 0; l < 8; ++l) { ++ yl[l+ 0] = y1[l+ 0]; ++ yl[l+ 8] = y1[l+16]; ++ yl[l+16] = y1[l+32]; ++ yl[l+24] = y1[l+48]; ++ } ++ ++ device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset); ++ device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0); ++ device const uint16_t * a = (device const uint16_t *)(x[i].scales); ++ device const half * dh = &x[i].d; ++ ++ for (short row = 0; row < nr0; ++row) { ++ const float d_all = (float)dh[0]; ++ ++ scales16[0] = a[4]; ++ scales16[1] = a[5]; ++ aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030; ++ scales16[0] = a[il+0]; ++ scales16[1] = a[il+1]; ++ scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32; ++ ++ float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2]; ++ s1 += yl[l+0] * (qs & qm[il/2][0]); ++ s2 += yl[l+1] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]); ++ s4 += yl[l+16] * (qs & qm[il/2][2]); ++ s5 += yl[l+17] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]); ++ } ++ float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[0] - 32); ++ sumf2[row] += d2 * (scales[2] - 32); ++ ++ s1 = s2 = s3 = s4 = s5 = s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2+8]; ++ s1 += yl[l+8] * (qs & qm[il/2][0]); ++ s2 += yl[l+9] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]); ++ s4 += yl[l+24] * (qs & qm[il/2][2]); ++ s5 += yl[l+25] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]); ++ } ++ d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[1] - 32); ++ sumf2[row] += d2 * (scales[3] - 32); ++ ++ q += args.nb01/2; ++ h += args.nb01/2; ++ a += args.nb01/2; ++ dh += args.nb01/2; ++ } ++ ++ y1 += 8 * QK_K; ++ } ++ ++ const float route_weight = extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0]; ++ for (int row = 0; row < nr0; ++row) { ++ const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift); ++ sumf1[row] = route_weight * simd_sum(sumf); ++ } ++ } ++ } ++ ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0; ++row) { ++ partials[sgitg*nr0 + row] = first_row + row < args.ne0 ? sumf1[row] : 0.0f; ++ } ++ } ++ ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (sgitg == 0 && tiisg < nr0 && first_row + tiisg < args.ne0) { ++ float sum = 0.0f; ++ for (int s = 0; s < n_slots*n_parts; ++s) { ++ sum += partials[s*nr0 + tiisg]; ++ } ++ ((device float *) (dst + uint64_t(first_row + tiisg)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1))[0] = sum; ++ } ++} ++ ++kernel void kernel_mul_mv_id_q3_K_weighted_accum_atomic( ++ constant ggml_metal_kargs_mul_mv_id & args, ++ device const char * src0s, ++ device const char * src1, ++ device char * dst, ++ device const char * ids, ++ device const char * weights, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ const short NSG = FC_mul_mv_nsg; ++ const int nr0 = N_R0_Q3_K; ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int token = tgpig.y; ++ const int slot = tgpig.z; ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ if (slot >= args.nei0 || token >= args.nei1) { ++ return; ++ } ++ ++ const int32_t expert_id = ((device const int32_t *) (ids + token*args.nbi1))[slot]; ++ if (expert_id < 0 || expert_id >= args.ne02) { ++ return; ++ } ++ ++ const float route_weight = extra.already_weighted ? 1.0f : ++ ((device const float *) (weights + slot*extra.weights_nb1 + token*extra.weights_nb2))[0]; ++ if (route_weight == 0.0f) { ++ return; ++ } ++ ++ const short tid = tiisg/4; ++ const short ix = tiisg%4; ++ const short ip = tid/4; ++ const short il = 2*((tid%4)/2); ++ const short ir = tid%2; ++ const short l0 = 8*ir; ++ ++ const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200}, ++ {0x0004, 0x0400, 0x0008, 0x0800}, ++ {0x0010, 0x1000, 0x0020, 0x2000}, ++ {0x0040, 0x4000, 0x0080, 0x8000}}; ++ const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}}; ++ ++ const ushort4 hm = mm[2*ip + il/2]; ++ const short shift = 2*il; ++ const float v1 = il == 0 ? 4.f : 64.f; ++ const float v2 = 4.f * v1; ++ const uint16_t s_shift1 = 4*ip; ++ const uint16_t s_shift2 = s_shift1 + il; ++ const short q_offset = 32*ip + l0; ++ const short y_offset = 128*ip + 32*il + l0; ++ ++ const uint64_t offset0 = uint64_t(expert_id)*args.nb02 + uint64_t(first_row)*args.nb01; ++ const uint64_t offset1 = uint64_t(slot)*args.nb11 + uint64_t(token)*args.nb12; ++ device const block_q3_K * x = (device const block_q3_K *) (src0s + offset0); ++ device const float * yy = (device const float *) (src1 + offset1); ++ device const float * y1 = yy + ix*QK_K + y_offset; ++ ++ float yl[32]; ++ uint32_t scales32, aux32; ++ thread uint16_t * scales16 = (thread uint16_t *)&scales32; ++ thread const int8_t * scales = (thread const int8_t *)&scales32; ++ ++ float sumf1[nr0] = {0.f}; ++ float sumf2[nr0] = {0.f}; ++ ++ for (int i = ix; i < nb; i += 4) { ++ for (short l = 0; l < 8; ++l) { ++ yl[l+ 0] = y1[l+ 0]; ++ yl[l+ 8] = y1[l+16]; ++ yl[l+16] = y1[l+32]; ++ yl[l+24] = y1[l+48]; ++ } ++ ++ device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset); ++ device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0); ++ device const uint16_t * a = (device const uint16_t *)(x[i].scales); ++ device const half * dh = &x[i].d; ++ ++ for (short row = 0; row < nr0; ++row) { ++ const float d_all = (float)dh[0]; ++ ++ scales16[0] = a[4]; ++ scales16[1] = a[5]; ++ aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030; ++ scales16[0] = a[il+0]; ++ scales16[1] = a[il+1]; ++ scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32; ++ ++ float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2]; ++ s1 += yl[l+0] * (qs & qm[il/2][0]); ++ s2 += yl[l+1] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]); ++ s4 += yl[l+16] * (qs & qm[il/2][2]); ++ s5 += yl[l+17] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]); ++ } ++ float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[0] - 32); ++ sumf2[row] += d2 * (scales[2] - 32); ++ ++ s1 = s2 = s3 = s4 = s5 = s6 = 0; ++ for (short l = 0; l < 8; l += 2) { ++ const int32_t qs = q[l/2+8]; ++ s1 += yl[l+8] * (qs & qm[il/2][0]); ++ s2 += yl[l+9] * (qs & qm[il/2][1]); ++ s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]); ++ s4 += yl[l+24] * (qs & qm[il/2][2]); ++ s5 += yl[l+25] * (qs & qm[il/2][3]); ++ s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]); ++ } ++ d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1); ++ d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2); ++ sumf1[row] += d1 * (scales[1] - 32); ++ sumf2[row] += d2 * (scales[3] - 32); ++ ++ q += args.nb01/2; ++ h += args.nb01/2; ++ a += args.nb01/2; ++ dh += args.nb01/2; ++ } ++ ++ y1 += 4 * QK_K; ++ } ++ ++ for (int row = 0; row < nr0; ++row) { ++ const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift); ++ sumf1[row] = simd_sum(sumf); ++ } ++ ++ if (tiisg == 0) { ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ device char * dst_row = dst + uint64_t(first_row + row)*extra.dst_nb0 + uint64_t(token)*extra.dst_nb1; ++ const float value = route_weight * sumf1[row]; ++ device atomic_float * dst_atomic = (device atomic_float *) dst_row; ++ atomic_fetch_add_explicit(dst_atomic, value, memory_order_relaxed); ++ } ++ } ++} ++ ++static inline uint4 glm_moe_q2_scan_projection( ++ device const char * src, ++ uint64_t expert_stride, ++ uint64_t expert_bytes, ++ int32_t expert_id, ++ int32_t chunk, ++ int32_t chunks_per_expert, ++ ushort tid) { ++ constexpr uint64_t vector_bytes = 4*sizeof(uint32_t); ++ constexpr uint64_t threads_per_threadgroup = 256; ++ const uint64_t vector_count = expert_bytes/vector_bytes; ++ const uint64_t vector_begin = vector_count*uint64_t(chunk)/uint64_t(chunks_per_expert); ++ const uint64_t vector_end = vector_count*uint64_t(chunk + 1)/uint64_t(chunks_per_expert); ++ device const uint4 * vectors = ++ (device const uint4 *) (src + uint64_t(expert_id)*expert_stride); ++ ++ uint4 checksum = uint4(0); ++ for (uint64_t i = vector_begin + tid; i < vector_end; i += threads_per_threadgroup) { ++ checksum ^= vectors[i]; ++ } ++ return checksum; ++} ++ ++kernel void kernel_glm_moe_q2_selected_weight_scan( ++ constant ggml_metal_kargs_glm_moe_q2_weight_scan & args, ++ device const char * gate, ++ device const char * up, ++ device const char * down, ++ device const char * ids, ++ device float * dst, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tid [[thread_index_in_threadgroup]], ++ ushort tiisg [[thread_index_in_simdgroup]], ++ ushort sgitg [[simdgroup_index_in_threadgroup]]) { ++ const int32_t chunk = int32_t(tgpig.x); ++ const int32_t slot = int32_t(tgpig.y); ++ const int32_t token = int32_t(tgpig.z); ++ ++ uint4 checksum = uint4(0); ++ if (chunk < args.chunks_per_expert && slot < args.top_k && token < args.n_tokens) { ++ const int32_t routed_expert_id = ++ ((device const int32_t *) (ids + uint64_t(token)*args.ids_nb1))[slot]; ++ const int32_t expert_id = args.storage_block_bytes == int32_t(sizeof(block_q2_K)) ? ++ routed_expert_id : 3*slot; ++ if (expert_id >= 0 && expert_id < args.n_experts) { ++ checksum ^= glm_moe_q2_scan_projection( ++ gate, args.gate_expert_stride, args.gate_expert_bytes, ++ expert_id, chunk, args.chunks_per_expert, tid); ++ checksum ^= glm_moe_q2_scan_projection( ++ up, args.up_expert_stride, args.up_expert_bytes, ++ expert_id, chunk, args.chunks_per_expert, tid); ++ checksum ^= glm_moe_q2_scan_projection( ++ down, args.down_expert_stride, args.down_expert_bytes, ++ expert_id, chunk, args.chunks_per_expert, tid); ++ } ++ } ++ ++ const float subgroup_checksum = simd_sum(float( ++ (checksum.x ^ checksum.y ^ checksum.z ^ checksum.w) & 0xffffu)); ++ if (tiisg == 0) { ++ partials[sgitg] = subgroup_checksum; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (tid == 0) { ++ float total = 0.0f; ++ for (int32_t i = 0; i < 8; ++i) { ++ total += partials[i]; ++ } ++ dst[(token*args.top_k + slot)*args.chunks_per_expert + chunk] = total; ++ } ++} ++ ++kernel void kernel_zero_f32( ++ device float * dst, ++ uint gid[[thread_position_in_grid]]) { ++ dst[gid] = 0.0f; ++} ++ ++struct glm_decode_block_scan_args { ++ ulong offset; ++ ulong nbytes; ++}; ++ ++// Optimistic command-buffer scheduling ceiling for an eight-layer GLM-5.2 ++// block. Each dispatch streams one contiguous dependency phase and writes one ++// checksum per 64 KiB task. No useful model arithmetic is performed. ++[[host_name("kernel_glm_decode_block_phase_scan")]] ++kernel void kernel_glm_decode_block_phase_scan( ++ device const uint4 * src, ++ device uint * dst, ++ constant glm_decode_block_scan_args & args, ++ uint tgid[[threadgroup_position_in_grid]], ++ uint tid[[thread_index_in_threadgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]], ++ ushort tiisg[[thread_index_in_simdgroup]]) { ++ constexpr uint chunk_bytes = 64*1024; ++ threadgroup uint partials[4]; ++ ++ const ulong begin = args.offset + ulong(tgid)*chunk_bytes; ++ const ulong end = min(begin + ulong(chunk_bytes), args.offset + args.nbytes); ++ ++ uint value = 0u; ++ for (ulong byte = begin + ulong(tid)*sizeof(uint4); ++ byte + sizeof(uint4) <= end; ++ byte += 128ul*sizeof(uint4)) { ++ const uint4 words = src[byte/sizeof(uint4)]; ++ value += words.x + words.y + words.z + words.w; ++ } ++ value = simd_sum(value); ++ if (tiisg == 0) { ++ partials[sgitg] = value; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (tid == 0) { ++ dst[tgid] = partials[0] + partials[1] + partials[2] + partials[3]; ++ } ++} ++ ++template ++void kernel_mul_mv_q4_K_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ constexpr uint16_t kmask1 = 0x3f3f; ++ constexpr uint16_t kmask2 = 0x0f0f; ++ constexpr uint16_t kmask3 = 0xc0c0; ++ ++ const short ix = tiisg/8; // 0...3 ++ const short it = tiisg%8; // 0...7 ++ const short iq = it/4; // 0 or 1 ++ const short ir = it%4; // 0...3 ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_q4_K * x = (device const block_q4_K *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[16]; ++ float yh[16]; ++ ++ float sumf[nr0]={0.f}; ++ ++ device const float * y4 = y + ix * QK_K + 64 * iq + 8 * ir; ++ ++ uint16_t sc16[4]; ++ thread const uint8_t * sc8 = (thread const uint8_t *)sc16; ++ ++ for (int ib = ix; ib < nb; ib += 4) { ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ ++ for (short i = 0; i < 8; ++i) { ++ yl[i+0] = y4[i+ 0]; sumy[0] += yl[i+0]; ++ yl[i+8] = y4[i+ 32]; sumy[1] += yl[i+8]; ++ yh[i+0] = y4[i+128]; sumy[2] += yh[i+0]; ++ yh[i+8] = y4[i+160]; sumy[3] += yh[i+8]; ++ } ++ ++ device const uint16_t * sc = (device const uint16_t *)x[ib].scales + iq; ++ device const uint16_t * q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir; ++ device const half * dh = &x[ib].d; ++ ++ for (short row = 0; row < nr0; row++) { ++ sc16[0] = sc[0] & kmask1; ++ sc16[1] = sc[2] & kmask1; ++ sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); ++ sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); ++ ++ device const uint16_t * q2 = q1 + 32; ++ ++ float4 acc1 = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2 = {0.f, 0.f, 0.f, 0.f}; ++ ++ FOR_UNROLL (short i = 0; i < 4; ++i) { ++ acc1[0] += yl[2*i + 0] * (q1[i] & 0x000F); ++ acc1[1] += yl[2*i + 1] * (q1[i] & 0x0F00); ++ acc1[2] += yl[2*i + 8] * (q1[i] & 0x00F0); ++ acc1[3] += yl[2*i + 9] * (q1[i] & 0xF000); ++ acc2[0] += yh[2*i + 0] * (q2[i] & 0x000F); ++ acc2[1] += yh[2*i + 1] * (q2[i] & 0x0F00); ++ acc2[2] += yh[2*i + 8] * (q2[i] & 0x00F0); ++ acc2[3] += yh[2*i + 9] * (q2[i] & 0xF000); ++ } ++ ++ sumf[row] += dh[0] * ((acc1[0] + 1.f/256.f * acc1[1]) * sc8[0] + ++ (acc1[2] + 1.f/256.f * acc1[3]) * sc8[1] * 1.f/16.f + ++ (acc2[0] + 1.f/256.f * acc2[1]) * sc8[4] + ++ (acc2[2] + 1.f/256.f * acc2[3]) * sc8[5] * 1.f/16.f) - ++ dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]); ++ ++ q1 += args.nb01/2; ++ sc += args.nb01/2; ++ dh += args.nb01/2; ++ } ++ ++ y4 += 4 * QK_K; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (int64_t)im*args.ne0*args.ne1 + (int64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_q4_K_f32")]] ++kernel void kernel_mul_mv_q4_K_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_q4_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++} ++ ++static inline float glm_moe_q4_k_dot_row( ++ device const char * weights, ++ uint64_t row_stride, ++ device const char * activation, ++ int row, ++ int n_embd, ++ ushort tiisg) { ++ const int nb = n_embd/QK_K; ++ device const block_q4_K * x = ++ (device const block_q4_K *) (weights + uint64_t(row)*row_stride); ++ device const float * y = (device const float *) activation; ++ ++ constexpr uint16_t kmask1 = 0x3f3f; ++ constexpr uint16_t kmask2 = 0x0f0f; ++ constexpr uint16_t kmask3 = 0xc0c0; ++ ++ const short ix = tiisg/8; ++ const short it = tiisg%8; ++ const short iq = it/4; ++ const short ir = it%4; ++ ++ device const float * y4 = y + ix*QK_K + 64*iq + 8*ir; ++ float sumf = 0.0f; ++ float yl[16]; ++ float yh[16]; ++ uint16_t sc16[4]; ++ thread const uint8_t * sc8 = (thread const uint8_t *) sc16; ++ ++ for (int ib = ix; ib < nb; ib += 4) { ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ for (short i = 0; i < 8; ++i) { ++ yl[i + 0] = y4[i + 0]; sumy[0] += yl[i + 0]; ++ yl[i + 8] = y4[i + 32]; sumy[1] += yl[i + 8]; ++ yh[i + 0] = y4[i + 128]; sumy[2] += yh[i + 0]; ++ yh[i + 8] = y4[i + 160]; sumy[3] += yh[i + 8]; ++ } ++ ++ device const uint16_t * sc = (device const uint16_t *) x[ib].scales + iq; ++ device const uint16_t * q1 = (device const uint16_t *) x[ib].qs + 16*iq + 4*ir; ++ device const half * dh = &x[ib].d; ++ ++ sc16[0] = sc[0] & kmask1; ++ sc16[1] = sc[2] & kmask1; ++ sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2); ++ sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2); ++ ++ device const uint16_t * q2 = q1 + 32; ++ float4 acc1 = {0.f, 0.f, 0.f, 0.f}; ++ float4 acc2 = {0.f, 0.f, 0.f, 0.f}; ++ FOR_UNROLL (short i = 0; i < 4; ++i) { ++ acc1[0] += yl[2*i + 0] * (q1[i] & 0x000F); ++ acc1[1] += yl[2*i + 1] * (q1[i] & 0x0F00); ++ acc1[2] += yl[2*i + 8] * (q1[i] & 0x00F0); ++ acc1[3] += yl[2*i + 9] * (q1[i] & 0xF000); ++ acc2[0] += yh[2*i + 0] * (q2[i] & 0x000F); ++ acc2[1] += yh[2*i + 1] * (q2[i] & 0x0F00); ++ acc2[2] += yh[2*i + 8] * (q2[i] & 0x00F0); ++ acc2[3] += yh[2*i + 9] * (q2[i] & 0xF000); ++ } ++ ++ sumf += dh[0] * ((acc1[0] + 1.f/256.f*acc1[1]) * sc8[0] + ++ (acc1[2] + 1.f/256.f*acc1[3]) * sc8[1] * 1.f/16.f + ++ (acc2[0] + 1.f/256.f*acc2[1]) * sc8[4] + ++ (acc2[2] + 1.f/256.f*acc2[3]) * sc8[5] * 1.f/16.f) - ++ dh[1] * (sumy[0]*sc8[2] + sumy[1]*sc8[3] + ++ sumy[2]*sc8[6] + sumy[3]*sc8[7]); ++ y4 += 4*QK_K; ++ } ++ ++ return simd_sum(sumf); ++} ++ ++[[host_name("kernel_glm_moe_two_phase_gate")]] ++kernel void kernel_glm_moe_two_phase_gate( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & routed_args, ++ constant ggml_metal_kargs_glm_moe_two_phase & phase_args, ++ device const char * routed_up, ++ device const char * routed_gate, ++ device const char * shared_up, ++ device const char * shared_gate, ++ device const char * cur, ++ device char * routed_activation, ++ device char * shared_activation, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint tgid[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ if (tgid < uint(phase_args.routed_gate_groups)) { ++ const uint slot = tgid & 7u; ++ const uint row_group = tgid >> 3u; ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl<1, 8>( ++ routed_args, ++ routed_up, ++ routed_gate, ++ cur, ++ routed_activation, ++ ids, ++ weights, ++ cur, ++ partials, ++ uint3(row_group, 0, slot), ++ tiisg, ++ sgitg); ++ return; ++ } ++ ++ const int row = int(tgid) - phase_args.routed_gate_groups; ++ if (row >= phase_args.n_ff) { ++ return; ++ } ++ const float gate_value = glm_moe_q4_k_dot_row( ++ shared_gate, phase_args.shared_gate_nb1, cur, row, phase_args.n_embd, tiisg); ++ const float up_value = glm_moe_q4_k_dot_row( ++ shared_up, phase_args.shared_up_nb1, cur, row, phase_args.n_embd, tiisg); ++ if (tiisg == 0) { ++ const float silu = gate_value/(1.0f + exp(-gate_value)); ++ ((device float *) shared_activation)[row] = up_value*silu; ++ } ++} ++ ++[[host_name("kernel_glm_moe_dual_lane_swiglu")]] ++kernel void kernel_glm_moe_dual_lane_swiglu( ++ constant ggml_metal_kargs_glm_moe_two_phase & phase_args, ++ device const float * shared_gate, ++ device const float * shared_up, ++ device float * shared_activation, ++ uint row[[thread_position_in_grid]]) { ++ if (row < uint(phase_args.n_ff)) { ++ const float gate_value = shared_gate[row]; ++ const float silu = gate_value/(1.0f + exp(-gate_value)); ++ shared_activation[row] = silu*shared_up[row]; ++ } ++} ++ ++[[host_name("kernel_glm_moe_dual_lane_gate_slot4")]] ++kernel void kernel_glm_moe_dual_lane_gate_slot4( ++ constant ggml_metal_kargs_mul_mv_id_gate_up_swiglu & routed_args, ++ constant ggml_metal_kargs_glm_moe_two_phase & phase_args, ++ device const char * routed_up, ++ device const char * routed_gate, ++ device const char * shared_up, ++ device const char * shared_gate, ++ device const char * cur, ++ device char * routed_activation, ++ device char * shared_activation, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint tgid[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ constexpr uint slots_per_tg = 4; ++ const uint slot_groups = (uint(routed_args.nei0) + slots_per_tg - 1)/slots_per_tg; ++ const uint row_group = tgid/slot_groups; ++ const uint slot_group = tgid%slot_groups; ++ kernel_mul_mv_id_q2_K_gate_up_swiglu_pair_sg_dual_impl( ++ routed_args, ++ routed_up, ++ routed_gate, ++ cur, ++ routed_activation, ++ ids, ++ weights, ++ cur, ++ partials, ++ uint3(row_group, 0, slot_group), ++ tiisg, ++ sgitg); ++ ++ (void) phase_args; ++ (void) shared_up; ++ (void) shared_gate; ++ (void) shared_activation; ++} ++ ++[[host_name("kernel_glm_moe_two_phase_down")]] ++kernel void kernel_glm_moe_two_phase_down( ++ constant ggml_metal_kargs_mul_mv_id & routed_args, ++ constant ggml_metal_kargs_mul_mv_id_weighted_reduce_extra & extra, ++ constant ggml_metal_kargs_glm_moe_two_phase & phase_args, ++ device const char * routed_down, ++ device const char * routed_activation, ++ device const char * shared_down, ++ device const char * shared_activation, ++ device char * routed_out, ++ device char * final_out, ++ device const char * ids, ++ device const char * weights, ++ threadgroup float * partials [[threadgroup(0)]], ++ uint tgid[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_id_q3_K_weighted_reduce_slots_sg_impl<8, 8, 0>( ++ routed_args, ++ routed_down, ++ routed_activation, ++ routed_out, ++ ids, ++ weights, ++ extra, ++ partials, ++ uint3(tgid, 0, 0), ++ tiisg, ++ sgitg); ++ ++ threadgroup_barrier(mem_flags::mem_device); ++ const int row = int(tgid)*8 + int(sgitg); ++ if (row >= phase_args.n_out) { ++ return; ++ } ++ const float shared_value = glm_moe_q4_k_dot_row( ++ shared_down, ++ phase_args.shared_down_nb1, ++ shared_activation, ++ row, ++ phase_args.n_ff, ++ tiisg); ++ if (tiisg == 0) { ++ device const float * routed = ++ (device const float *) (routed_out + uint64_t(row)*extra.dst_nb0); ++ device float * out = ++ (device float *) (final_out + uint64_t(row)*extra.dst_nb0); ++ *out = *routed + shared_value; ++ } ++} ++ ++[[host_name("kernel_glm_moe_dual_lane_add")]] ++kernel void kernel_glm_moe_dual_lane_add( ++ constant ggml_metal_kargs_glm_moe_two_phase & phase_args, ++ device const float * routed_out, ++ device const float * shared_out, ++ device float * final_out, ++ uint row[[thread_position_in_grid]]) { ++ if (row < uint(phase_args.n_out)) { ++ final_out[row] = routed_out[row] + shared_out[row]; ++ } ++} ++ ++template ++void kernel_mul_mv_q5_K_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_q5_K * x = (device const block_q5_K *) (src0 + offset0); ++ device const float * yy = (device const float *) (src1 + offset1); ++ ++ float sumf[nr0]={0.f}; ++ ++ float yl[16], yh[16]; ++ ++ constexpr uint16_t kmask1 = 0x3f3f; ++ constexpr uint16_t kmask2 = 0x0f0f; ++ constexpr uint16_t kmask3 = 0xc0c0; ++ ++ const short tid = tiisg/4; ++ const short ix = tiisg%4; ++ const short iq = tid/4; ++ const short ir = tid%4; ++ ++ const short l0 = 8*ir; ++ const short q_offset = 32*iq + l0; ++ const short y_offset = 64*iq + l0; ++ ++ const uint8_t hm1 = 1u << (2*iq); ++ const uint8_t hm2 = hm1 << 1; ++ const uint8_t hm3 = hm1 << 4; ++ const uint8_t hm4 = hm2 << 4; ++ ++ uint16_t sc16[4]; ++ thread const uint8_t * sc8 = (thread const uint8_t *)sc16; ++ ++ device const float * y1 = yy + ix*QK_K + y_offset; ++ ++ for (int i = ix; i < nb; i += 4) { ++ device const uint8_t * q1 = x[i].qs + q_offset; ++ device const uint8_t * qh = x[i].qh + l0; ++ device const half * dh = &x[i].d; ++ device const uint16_t * a = (device const uint16_t *)x[i].scales + iq; ++ ++ device const float * y2 = y1 + 128; ++ float4 sumy = {0.f, 0.f, 0.f, 0.f}; ++ for (short l = 0; l < 8; ++l) { ++ yl[l+0] = y1[l+ 0]; sumy[0] += yl[l+0]; ++ yl[l+8] = y1[l+32]; sumy[1] += yl[l+8]; ++ yh[l+0] = y2[l+ 0]; sumy[2] += yh[l+0]; ++ yh[l+8] = y2[l+32]; sumy[3] += yh[l+8]; ++ } ++ ++ for (short row = 0; row < nr0; ++row) { ++ device const uint8_t * q2 = q1 + 64; ++ ++ sc16[0] = a[0] & kmask1; ++ sc16[1] = a[2] & kmask1; ++ sc16[2] = ((a[4] >> 0) & kmask2) | ((a[0] & kmask3) >> 2); ++ sc16[3] = ((a[4] >> 4) & kmask2) | ((a[2] & kmask3) >> 2); ++ ++ float4 acc1 = {0.f}; ++ float4 acc2 = {0.f}; ++ FOR_UNROLL (short l = 0; l < 8; ++l) { ++ uint8_t h = qh[l]; ++ acc1[0] += yl[l+0] * (q1[l] & 0x0F); ++ acc1[1] += yl[l+8] * (q1[l] & 0xF0); ++ acc1[2] += yh[l+0] * (q2[l] & 0x0F); ++ acc1[3] += yh[l+8] * (q2[l] & 0xF0); ++ acc2[0] += h & hm1 ? yl[l+0] : 0.f; ++ acc2[1] += h & hm2 ? yl[l+8] : 0.f; ++ acc2[2] += h & hm3 ? yh[l+0] : 0.f; ++ acc2[3] += h & hm4 ? yh[l+8] : 0.f; ++ } ++ ++ sumf[row] += dh[0] * (sc8[0] * (acc1[0] + 16.f*acc2[0]) + ++ sc8[1] * (acc1[1]/16.f + 16.f*acc2[1]) + ++ sc8[4] * (acc1[2] + 16.f*acc2[2]) + ++ sc8[5] * (acc1[3]/16.f + 16.f*acc2[3])) - ++ dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]); ++ ++ q1 += args.nb01; ++ qh += args.nb01; ++ dh += args.nb01/2; ++ a += args.nb01/2; ++ } ++ ++ y1 += 4 * QK_K; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ const float tot = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = tot; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_q5_K_f32")]] ++kernel void kernel_mul_mv_q5_K_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_q5_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_q6_K_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ constexpr uint8_t kmask1 = 0x03; ++ constexpr uint8_t kmask2 = 0x0C; ++ constexpr uint8_t kmask3 = 0x30; ++ constexpr uint8_t kmask4 = 0xC0; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_q6_K * x = (device const block_q6_K *) (src0 + offset0); ++ device const float * yy = (device const float *) (src1 + offset1); ++ ++ float sumf[nr0] = { 0.f }; ++ ++ float yl[16]; ++ ++ const short tid = tiisg/2; ++ const short ix = tiisg%2; ++ const short ip = tid/8; // 0 or 1 ++ const short il = tid%8; ++ const short l0 = 4*il; ++ const short is = 8*ip + l0/16; ++ ++ const short y_offset = 128*ip + l0; ++ const short q_offset_l = 64*ip + l0; ++ const short q_offset_h = 32*ip + l0; ++ ++ for (int i = ix; i < nb; i += 2) { ++ device const uint8_t * q1 = x[i].ql + q_offset_l; ++ device const uint8_t * q2 = q1 + 32; ++ device const uint8_t * qh = x[i].qh + q_offset_h; ++ device const int8_t * sc = x[i].scales + is; ++ device const half * dh = &x[i].d; ++ ++ device const float * y = yy + i * QK_K + y_offset; ++ ++ for (short l = 0; l < 4; ++l) { ++ yl[4*l + 0] = y[l + 0]; ++ yl[4*l + 1] = y[l + 32]; ++ yl[4*l + 2] = y[l + 64]; ++ yl[4*l + 3] = y[l + 96]; ++ } ++ ++ for (short row = 0; row < nr0; ++row) { ++ float4 sums = {0.f, 0.f, 0.f, 0.f}; ++ ++ FOR_UNROLL (short l = 0; l < 4; ++l) { ++ sums[0] += yl[4*l + 0] * ((int8_t)((q1[l] & 0xF) | ((qh[l] & kmask1) << 4)) - 32); ++ sums[1] += yl[4*l + 1] * ((int8_t)((q2[l] & 0xF) | ((qh[l] & kmask2) << 2)) - 32); ++ sums[2] += yl[4*l + 2] * ((int8_t)((q1[l] >> 4) | ((qh[l] & kmask3) << 0)) - 32); ++ sums[3] += yl[4*l + 3] * ((int8_t)((q2[l] >> 4) | ((qh[l] & kmask4) >> 2)) - 32); ++ } ++ ++ sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + sums[2] * sc[4] + sums[3] * sc[6]); ++ ++ q1 += args.nb01; ++ q2 += args.nb01; ++ qh += args.nb01; ++ sc += args.nb01; ++ dh += args.nb01/2; ++ } ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_q6_K_f32")]] ++kernel void kernel_mul_mv_q6_K_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_q6_K_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++} ++ ++// ======================= "True" 2-bit ++ ++template ++void kernel_mul_mv_iq2_xxs_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq2_xxs * x = (device const block_iq2_xxs *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ float sumf[nr0]={0.f}; ++ ++ const int nb32 = nb * (QK_K / 32); ++ ++ threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem); ++ threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256); ++ { ++ int nval = 4; ++ int pos = (32*sgitg + tiisg)*nval; ++ for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2xxs_grid[pos + i]; ++ nval = 2; ++ pos = (32*sgitg + tiisg)*nval; ++ for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ const int ix = tiisg; ++ ++ device const float * y4 = y + 32 * ix; ++ ++ for (int ib32 = ix; ib32 < nb32; ib32 += 32) { ++ for (short i = 0; i < 32; ++i) { ++ yl[i] = y4[i]; ++ } ++ ++ const int ibl = ib32 / (QK_K / 32); ++ const int ib = ib32 % (QK_K / 32); ++ ++ device const block_iq2_xxs * xr = x + ibl; ++ device const uint16_t * q2 = xr->qs + 4 * ib; ++ device const half * dh = &xr->d; ++ ++ for (short row = 0; row < nr0; row++) { ++ const float db = dh[0]; ++ device const uint8_t * aux8 = (device const uint8_t *)q2; ++ const uint32_t aux32 = q2[2] | (q2[3] << 16); ++ const float d = db * (0.5f + (aux32 >> 28)); ++ ++ float sum = 0; ++ for (short l = 0; l < 4; ++l) { ++ const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + aux8[l]); ++ const uint8_t signs = ssigns[(aux32 >> 7*l) & 127]; ++ for (short j = 0; j < 8; ++j) { ++ sum += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); ++ } ++ } ++ sumf[row] += d * sum; ++ ++ dh += args.nb01/2; ++ q2 += args.nb01/2; ++ } ++ ++ y4 += 32 * 32; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all * 0.25f; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq2_xxs_f32")]] ++kernel void kernel_mul_mv_iq2_xxs_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ kernel_mul_mv_iq2_xxs_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_iq2_xs_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq2_xs * x = (device const block_iq2_xs *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ float sumf[nr0]={0.f}; ++ ++ const int nb32 = nb * (QK_K / 32); ++ ++ threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem); ++ threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 512); ++ { ++ int nval = 8; ++ int pos = (32*sgitg + tiisg)*nval; ++ for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2xs_grid[pos + i]; ++ nval = 2; ++ pos = (32*sgitg + tiisg)*nval; ++ for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ const int ix = tiisg; ++ ++ device const float * y4 = y + 32 * ix; ++ ++ for (int ib32 = ix; ib32 < nb32; ib32 += 32) { ++ for (short i = 0; i < 32; ++i) { ++ yl[i] = y4[i]; ++ } ++ ++ const int ibl = ib32 / (QK_K / 32); ++ const int ib = ib32 % (QK_K / 32); ++ ++ device const block_iq2_xs * xr = x + ibl; ++ device const uint16_t * q2 = xr->qs + 4 * ib; ++ device const uint8_t * sc = xr->scales + ib; ++ device const half * dh = &xr->d; ++ ++ for (short row = 0; row < nr0; row++) { ++ const float db = dh[0]; ++ const uint8_t ls1 = sc[0] & 0xf; ++ const uint8_t ls2 = sc[0] >> 4; ++ const float d1 = db * (0.5f + ls1); ++ const float d2 = db * (0.5f + ls2); ++ ++ float sum1 = 0, sum2 = 0; ++ for (short l = 0; l < 2; ++l) { ++ const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + (q2[l] & 511)); ++ const uint8_t signs = ssigns[(q2[l] >> 9)]; ++ for (short j = 0; j < 8; ++j) { ++ sum1 += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); ++ } ++ } ++ for (short l = 2; l < 4; ++l) { ++ const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + (q2[l] & 511)); ++ const uint8_t signs = ssigns[(q2[l] >> 9)]; ++ for (short j = 0; j < 8; ++j) { ++ sum2 += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f); ++ } ++ } ++ sumf[row] += d1 * sum1 + d2 * sum2; ++ ++ dh += args.nb01/2; ++ q2 += args.nb01/2; ++ sc += args.nb01; ++ } ++ ++ y4 += 32 * 32; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all * 0.25f; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq2_xs_f32")]] ++kernel void kernel_mul_mv_iq2_xs_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_iq2_xs_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_iq3_xxs_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq3_xxs * x = (device const block_iq3_xxs *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ float sumf[nr0]={0.f}; ++ ++ const int nb32 = nb * (QK_K / 32); ++ ++ threadgroup uint32_t * svalues = (threadgroup uint32_t *)(shmem); ++ threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256); ++ { ++ int nval = 4; ++ int pos = (32*sgitg + tiisg)*nval; ++ for (int i = 0; i < nval; ++i) svalues[pos + i] = iq3xxs_grid[pos + i]; ++ nval = 2; ++ pos = (32*sgitg + tiisg)*nval; ++ for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ const int ix = tiisg; ++ ++ device const float * y4 = y + 32 * ix; ++ ++ for (int ib32 = ix; ib32 < nb32; ib32 += 32) { ++ for (short i = 0; i < 32; ++i) { ++ yl[i] = y4[i]; ++ } ++ ++ const int ibl = ib32 / (QK_K / 32); ++ const int ib = ib32 % (QK_K / 32); ++ ++ device const block_iq3_xxs * xr = x + ibl; ++ device const uint8_t * q3 = xr->qs + 8 * ib; ++ device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib; ++ device const half * dh = &xr->d; ++ ++ for (short row = 0; row < nr0; row++) { ++ const float db = dh[0]; ++ const uint32_t aux32 = gas[0] | (gas[1] << 16); ++ const float d = db * (0.5f + (aux32 >> 28)); ++ ++ float2 sum = {0}; ++ for (short l = 0; l < 4; ++l) { ++ const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(svalues + q3[2*l+0]); ++ const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(svalues + q3[2*l+1]); ++ const uint8_t signs = ssigns[(aux32 >> 7*l) & 127]; ++ for (short j = 0; j < 4; ++j) { ++ sum[0] += yl[8*l + j + 0] * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f); ++ sum[1] += yl[8*l + j + 4] * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f); ++ } ++ } ++ sumf[row] += d * (sum[0] + sum[1]); ++ ++ dh += args.nb01/2; ++ q3 += args.nb01; ++ gas += args.nb01/2; ++ } ++ ++ y4 += 32 * 32; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all * 0.5f; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq3_xxs_f32")]] ++kernel void kernel_mul_mv_iq3_xxs_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_iq3_xxs_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_iq3_s_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq3_s * x = (device const block_iq3_s *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ float sumf[nr0]={0.f}; ++ ++ const int nb32 = nb * (QK_K / 32); ++ ++ threadgroup uint32_t * svalues = (threadgroup uint32_t *) shmem; ++ { ++ int nval = 8; ++ int pos = (32*sgitg + tiisg)*nval; ++ for (int i = 0; i < nval; ++i) svalues[pos + i] = iq3s_grid[pos + i]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ const int ix = tiisg; ++ ++ device const float * y4 = y + 32 * ix; ++ ++ for (int ib32 = ix; ib32 < nb32; ib32 += 32) { ++ for (short i = 0; i < 32; ++i) { ++ yl[i] = y4[i]; ++ } ++ ++ const int ibl = ib32 / (QK_K / 32); ++ const int ib = ib32 % (QK_K / 32); ++ ++ device const block_iq3_s * xr = x + ibl; ++ device const uint8_t * qs = xr->qs + 8 * ib; ++ device const uint8_t * qh = xr->qh + ib; ++ device const uint8_t * sc = xr->scales + (ib/2); ++ device const uint8_t * signs = xr->signs + 4 * ib; ++ device const half * dh = &xr->d; ++ ++ for (short row = 0; row < nr0; row++) { ++ const float db = dh[0]; ++ const float d = db * (1 + 2*((sc[0] >> 4*(ib%2)) & 0xf)); ++ ++ float2 sum = {0}; ++ for (short l = 0; l < 4; ++l) { ++ const threadgroup uint32_t * table1 = qh[0] & kmask_iq2xs[2*l+0] ? svalues + 256 : svalues; ++ const threadgroup uint32_t * table2 = qh[0] & kmask_iq2xs[2*l+1] ? svalues + 256 : svalues; ++ const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(table1 + qs[2*l+0]); ++ const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(table2 + qs[2*l+1]); ++ for (short j = 0; j < 4; ++j) { ++ sum[0] += yl[8*l + j + 0] * grid1[j] * select(1, -1, signs[l] & kmask_iq2xs[j+0]); ++ sum[1] += yl[8*l + j + 4] * grid2[j] * select(1, -1, signs[l] & kmask_iq2xs[j+4]); ++ } ++ } ++ sumf[row] += d * (sum[0] + sum[1]); ++ ++ dh += args.nb01/2; ++ qs += args.nb01; ++ qh += args.nb01; ++ sc += args.nb01; ++ signs += args.nb01; ++ } ++ ++ y4 += 32 * 32; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq3_s_f32")]] ++kernel void kernel_mul_mv_iq3_s_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_iq3_s_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_iq2_s_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq2_s * x = (device const block_iq2_s *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ float sumf[nr0]={0.f}; ++ ++ const int nb32 = nb * (QK_K / 32); ++ ++ //threadgroup uint64_t * svalues = (threadgroup uint64_t *) shmem; ++ //{ ++ // int nval = 32; ++ // int pos = (32*sgitg + tiisg)*nval; ++ // for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2s_grid[pos + i]; ++ // threadgroup_barrier(mem_flags::mem_threadgroup); ++ //} ++ ++ const short ix = tiisg; ++ ++ device const float * y4 = y + 32 * ix; ++ ++ for (int ib32 = ix; ib32 < nb32; ib32 += 32) { ++ for (short i = 0; i < 32; ++i) { ++ yl[i] = y4[i]; ++ } ++ ++ const int ibl = ib32 / (QK_K / 32); ++ const int ib = ib32 % (QK_K / 32); ++ ++ device const block_iq2_s * xr = x + ibl; ++ device const uint8_t * qs = xr->qs + 4 * ib; ++ device const uint8_t * qh = xr->qh + ib; ++ device const uint8_t * sc = xr->scales + ib; ++ device const uint8_t * signs = qs + QK_K/8; ++ device const half * dh = &xr->d; ++ ++ for (short row = 0; row < nr0; row++) { ++ const float db = dh[0]; ++ const float d1 = db * (0.5f + (sc[0] & 0xf)); ++ const float d2 = db * (0.5f + (sc[0] >> 4)); ++ ++ float2 sum = {0}; ++ for (short l = 0; l < 2; ++l) { ++ //const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(svalues + (qs[l+0] | ((qh[0] << (8-2*l)) & 0x300))); ++ //const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(svalues + (qs[l+2] | ((qh[0] << (4-2*l)) & 0x300))); ++ constant uint8_t * grid1 = (constant uint8_t *)(iq2s_grid + (qs[l+0] | ((qh[0] << (8-2*l)) & 0x300))); ++ constant uint8_t * grid2 = (constant uint8_t *)(iq2s_grid + (qs[l+2] | ((qh[0] << (4-2*l)) & 0x300))); ++ for (short j = 0; j < 8; ++j) { ++ sum[0] += yl[8*l + j + 0] * grid1[j] * select(1, -1, signs[l+0] & kmask_iq2xs[j]); ++ sum[1] += yl[8*l + j + 16] * grid2[j] * select(1, -1, signs[l+2] & kmask_iq2xs[j]); ++ } ++ } ++ sumf[row] += d1 * sum[0] + d2 * sum[1]; ++ ++ dh += args.nb01/2; ++ qs += args.nb01; ++ qh += args.nb01; ++ sc += args.nb01; ++ signs += args.nb01; ++ } ++ ++ y4 += 32 * 32; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all * 0.25f; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq2_s_f32")]] ++kernel void kernel_mul_mv_iq2_s_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_iq2_s_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_iq1_s_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq1_s * x = (device const block_iq1_s *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ float sumf[nr0]={0.f}; ++ ++ const int nb32 = nb * (QK_K / 32); ++ ++ const short ix = tiisg; ++ ++ device const float * y4 = y + 32 * ix; ++ ++ for (int ib32 = ix; ib32 < nb32; ib32 += 32) { ++ float sumy = 0; ++ for (short i = 0; i < 32; ++i) { ++ yl[i] = y4[i]; ++ sumy += yl[i]; ++ } ++ ++ const int ibl = ib32 / (QK_K / 32); ++ const int ib = ib32 % (QK_K / 32); ++ ++ device const block_iq1_s * xr = x + ibl; ++ device const uint8_t * qs = xr->qs + 4 * ib; ++ device const uint16_t * qh = xr->qh + ib; ++ device const half * dh = &xr->d; ++ ++ for (short row = 0; row < nr0; row++) { ++ constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700))); ++ constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 5) & 0x700))); ++ constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[0] << 2) & 0x700))); ++ constant uint8_t * grid4 = (constant uint8_t *)(iq1s_grid_gpu + (qs[3] | ((qh[0] >> 1) & 0x700))); ++ ++ float sum = 0; ++ for (short j = 0; j < 4; ++j) { ++ sum += yl[j+ 0] * (grid1[j] & 0xf) + yl[j+ 4] * (grid1[j] >> 4) ++ + yl[j+ 8] * (grid2[j] & 0xf) + yl[j+12] * (grid2[j] >> 4) ++ + yl[j+16] * (grid3[j] & 0xf) + yl[j+20] * (grid3[j] >> 4) ++ + yl[j+24] * (grid4[j] & 0xf) + yl[j+28] * (grid4[j] >> 4); ++ } ++ sumf[row] += (float)dh[0] * (sum + sumy * (qh[0] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA)) * (2*((qh[0] >> 12) & 7) + 1); ++ ++ dh += args.nb01/2; ++ qs += args.nb01; ++ qh += args.nb01/2; ++ } ++ ++ y4 += 32 * 32; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq1_s_f32")]] ++kernel void kernel_mul_mv_iq1_s_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_iq1_s_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_iq1_m_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ const int nb = args.ne00/QK_K; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * nr0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq1_m * x = (device const block_iq1_m *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ float yl[32]; ++ float sumf[nr0]={0.f}; ++ ++ const int nb32 = nb * (QK_K / 32); ++ ++ const short ix = tiisg; ++ ++ device const float * y4 = y + 32 * ix; ++ ++ iq1m_scale_t scale; ++ ++ for (int ib32 = ix; ib32 < nb32; ib32 += 32) { ++ float4 sumy = {0.f}; ++ for (short i = 0; i < 8; ++i) { ++ yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0]; ++ yl[i+ 8] = y4[i+ 8]; sumy[1] += yl[i+ 8]; ++ yl[i+16] = y4[i+16]; sumy[2] += yl[i+16]; ++ yl[i+24] = y4[i+24]; sumy[3] += yl[i+24]; ++ } ++ ++ const int ibl = ib32 / (QK_K / 32); ++ const int ib = ib32 % (QK_K / 32); ++ ++ device const block_iq1_m * xr = x + ibl; ++ device const uint8_t * qs = xr->qs + 4 * ib; ++ device const uint8_t * qh = xr->qh + 2 * ib; ++ device const uint16_t * sc = (device const uint16_t *)xr->scales; ++ ++ for (short row = 0; row < nr0; row++) { ++ scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); ++ ++ constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700))); ++ constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 4) & 0x700))); ++ constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[1] << 8) & 0x700))); ++ constant uint8_t * grid4 = (constant uint8_t *)(iq1s_grid_gpu + (qs[3] | ((qh[1] << 4) & 0x700))); ++ ++ float2 sum = {0.f}; ++ for (short j = 0; j < 4; ++j) { ++ sum[0] += yl[j+ 0] * (grid1[j] & 0xf) + yl[j+ 4] * (grid1[j] >> 4) ++ + yl[j+ 8] * (grid2[j] & 0xf) + yl[j+12] * (grid2[j] >> 4); ++ sum[1] += yl[j+16] * (grid3[j] & 0xf) + yl[j+20] * (grid3[j] >> 4) ++ + yl[j+24] * (grid4[j] & 0xf) + yl[j+28] * (grid4[j] >> 4); ++ } ++ const float delta1 = sumy[0] * (qh[0] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA) + sumy[1] * (qh[0] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA); ++ const float delta2 = sumy[2] * (qh[1] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA) + sumy[3] * (qh[1] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA); ++ ++ sumf[row] += (float)scale.f16 * ((sum[0] + delta1) * (2*((sc[ib/2] >> (6*(ib%2)+0)) & 7) + 1) + ++ (sum[1] + delta2) * (2*((sc[ib/2] >> (6*(ib%2)+3)) & 7) + 1)); ++ ++ sc += args.nb01/2; ++ qs += args.nb01; ++ qh += args.nb01; ++ } ++ ++ y4 += 32 * 32; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq1_m_f32")]] ++kernel void kernel_mul_mv_iq1_m_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_iq1_m_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_iq4_nl_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ threadgroup float * shmem_f32 = (threadgroup float *) shmem; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * NR0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq4_nl * x = (device const block_iq4_nl *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ const int nb = args.ne00/QK4_NL; ++ const int ns01 = args.nb01/args.nb00; ++ ++ const short ix = tiisg/2; // 0...15 ++ const short it = tiisg%2; // 0 or 1 ++ ++ shmem_f32[tiisg] = kvalues_iq4nl_f[tiisg%16]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ float4 yl[4]; ++ float sumf[NR0]={0.f}; ++ ++ device const float * yb = y + ix*QK4_NL + it*8; ++ ++ uint32_t aux32[2]; ++ thread const uint8_t * q8 = (thread const uint8_t *)aux32; ++ ++ float4 qf1, qf2; ++ ++ // [TAG_MUL_MV_WEIRD] ++ for (int ib = ix; ib < nb && ib < ns01; ib += 16) { ++ device const float4 * y4 = (device const float4 *)yb; ++ yl[0] = y4[0]; ++ yl[1] = y4[4]; ++ yl[2] = y4[1]; ++ yl[3] = y4[5]; ++ ++ for (short row = 0; row < NR0; row++) { ++ device const block_iq4_nl & xb = x[row*ns01 + ib]; ++ device const uint16_t * q4 = (device const uint16_t *)(xb.qs + 8*it); ++ ++ float4 acc1 = {0.f}, acc2 = {0.f}; ++ ++ aux32[0] = q4[0] | (q4[1] << 16); ++ aux32[1] = (aux32[0] >> 4) & 0x0f0f0f0f; ++ aux32[0] &= 0x0f0f0f0f; ++ qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]}; ++ qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]}; ++ acc1 += yl[0] * qf1; ++ acc2 += yl[1] * qf2; ++ ++ aux32[0] = q4[2] | (q4[3] << 16); ++ aux32[1] = (aux32[0] >> 4) & 0x0f0f0f0f; ++ aux32[0] &= 0x0f0f0f0f; ++ qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]}; ++ qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]}; ++ acc1 += yl[2] * qf1; ++ acc2 += yl[3] * qf2; ++ ++ acc1 += acc2; ++ ++ sumf[row] += (float)xb.d * (acc1[0] + acc1[1] + acc1[2] + acc1[3]); ++ } ++ ++ yb += 16 * QK4_NL; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq4_nl_f32")]] ++kernel void kernel_mul_mv_iq4_nl_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_iq4_nl_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_iq4_xs_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ threadgroup float * shmem_f32 = (threadgroup float *) shmem; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ const int first_row = (r0 * NSG + sgitg) * NR0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_iq4_xs * x = (device const block_iq4_xs *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ const int nb = args.ne00/QK_K; ++ const int ns01 = args.nb01/args.nb00; ++ ++ const short ix = tiisg/16; // 0 or 1 ++ const short it = tiisg%16; // 0...15 ++ const short ib = it/2; ++ const short il = it%2; ++ ++ shmem_f32[tiisg] = kvalues_iq4nl_f[tiisg%16]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ float4 yl[4]; ++ float sumf[NR0]={0.f}; ++ ++ device const float * yb = y + ix * QK_K + ib * 32 + il * 8; ++ ++ uint32_t aux32[2]; ++ thread const uint8_t * q8 = (thread const uint8_t *)aux32; ++ ++ float4 qf1, qf2; ++ ++ // [TAG_MUL_MV_WEIRD] ++ for (int ibl = ix; ibl < nb && ibl < ns01; ibl += 2) { ++ device const float4 * y4 = (device const float4 *)yb; ++ yl[0] = y4[0]; ++ yl[1] = y4[4]; ++ yl[2] = y4[1]; ++ yl[3] = y4[5]; ++ ++ for (short row = 0; row < NR0; ++row) { ++ device const block_iq4_xs & xb = x[row*ns01 + ibl]; ++ device const uint32_t * q4 = (device const uint32_t *)(xb.qs + 16*ib + 8*il); ++ ++ float4 acc1 = {0.f}, acc2 = {0.f}; ++ ++ aux32[0] = (q4[0] ) & 0x0f0f0f0f; ++ aux32[1] = (q4[0] >> 4) & 0x0f0f0f0f; ++ qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]}; ++ qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]}; ++ acc1 += yl[0] * qf1; ++ acc2 += yl[1] * qf2; ++ ++ aux32[0] = (q4[1] ) & 0x0f0f0f0f; ++ aux32[1] = (q4[1] >> 4) & 0x0f0f0f0f; ++ qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]}; ++ qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]}; ++ acc1 += yl[2] * qf1; ++ acc2 += yl[3] * qf2; ++ ++ acc1 += acc2; ++ ++ const int ls = (((xb.scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((xb.scales_h >> 2*ib) & 3) << 4)) - 32; ++ sumf[row] += (float)xb.d * ls * (acc1[0] + acc1[1] + acc1[2] + acc1[3]); ++ } ++ ++ yb += 2 * QK_K; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_iq4_xs_f32")]] ++kernel void kernel_mul_mv_iq4_xs_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_iq4_xs_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} ++ ++template ++void kernel_mul_mv_mxfp4_f32_impl( ++ args_t args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem, ++ uint3 tgpig, ++ ushort tiisg, ++ ushort sgitg) { ++ const short NSG = FC_mul_mv_nsg; ++ ++ threadgroup float * shmem_f32 = (threadgroup float *) shmem; ++ ++ const int r0 = tgpig.x; ++ const int r1 = tgpig.y; ++ const int im = tgpig.z; ++ ++ const int first_row = (r0 * NSG + sgitg) * NR0; ++ ++ const uint i12 = im%FC_mul_mv_ne12; ++ const uint i13 = im/FC_mul_mv_ne12; ++ ++ const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; ++ const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ ++ device const block_mxfp4 * x = (device const block_mxfp4 *) (src0 + offset0); ++ device const float * y = (device const float *) (src1 + offset1); ++ ++ const int nb = args.ne00/QK_MXFP4; ++ const int ns01 = args.nb01/args.nb00; // this can be larger than nb for permuted src0 tensors ++ ++ const short ix = tiisg/2; // 0...15 ++ const short it = tiisg%2; // 0 or 1 ++ ++ shmem_f32[tiisg] = kvalues_mxfp4_f[tiisg%16]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ float4 yl[4]; ++ float sumf[NR0]={0.f}; ++ ++ device const float * yb = y + ix*QK_MXFP4 + it*8; ++ ++ // note: just the check `ib < nb` is enough, but adding the redundant `&& ib < ns01` check makes the kernel a bit faster ++ // no idea why that is - needs some deeper investigation [TAG_MUL_MV_WEIRD] ++ for (int ib = ix; ib < nb && ib < ns01; ib += 16) { ++ device const float4 * y4 = (device const float4 *) yb; ++ ++ yl[0] = y4[0]; ++ yl[1] = y4[4]; ++ yl[2] = y4[1]; ++ yl[3] = y4[5]; ++ ++ FOR_UNROLL (short row = 0; row < NR0; row++) { ++ device const block_mxfp4 & xb = x[row*ns01 + ib]; ++ device const uint8_t * q2 = (device const uint8_t *)(xb.qs + 8*it); ++ ++ float4 acc1 = yl[0]*float4(shmem_f32[q2[0] & 0x0F], shmem_f32[q2[1] & 0x0F], shmem_f32[q2[2] & 0x0F], shmem_f32[q2[3] & 0x0F]); ++ float4 acc2 = yl[1]*float4(shmem_f32[q2[0] >> 4 ], shmem_f32[q2[1] >> 4 ], shmem_f32[q2[2] >> 4 ], shmem_f32[q2[3] >> 4 ]); ++ float4 acc3 = yl[2]*float4(shmem_f32[q2[4] & 0x0F], shmem_f32[q2[5] & 0x0F], shmem_f32[q2[6] & 0x0F], shmem_f32[q2[7] & 0x0F]); ++ float4 acc4 = yl[3]*float4(shmem_f32[q2[4] >> 4 ], shmem_f32[q2[5] >> 4 ], shmem_f32[q2[6] >> 4 ], shmem_f32[q2[7] >> 4 ]); ++ ++ acc1 = (acc1 + acc3) + (acc2 + acc4); ++ ++ sumf[row] += e8m0_to_fp32(xb.e) * ((acc1[0] + acc1[1]) + (acc1[2] + acc1[3])); ++ } ++ ++ yb += 16 * QK_MXFP4; ++ } ++ ++ device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ ++ for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) { ++ float sum_all = simd_sum(sumf[row]); ++ if (tiisg == 0) { ++ dst_f32[first_row + row] = sum_all; ++ } ++ } ++} ++ ++[[host_name("kernel_mul_mv_mxfp4_f32")]] ++kernel void kernel_mul_mv_mxfp4_f32( ++ constant ggml_metal_kargs_mul_mv & args, ++ device const char * src0, ++ device const char * src1, ++ device char * dst, ++ threadgroup char * shmem [[threadgroup(0)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiisg[[thread_index_in_simdgroup]], ++ ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ ++ kernel_mul_mv_mxfp4_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++} ++ ++template ++kernel void kernel_get_rows_q( ++ constant ggml_metal_kargs_get_rows & args, ++ device const void * src0, ++ device const void * src1, ++ device void * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiitg[[thread_index_in_threadgroup]], ++ ushort3 ntg [[threads_per_threadgroup]]) { ++ const int32_t iw0 = tgpig.x/args.ne10; ++ const int32_t i10 = tgpig.x%args.ne10; ++ const int32_t i11 = tgpig.y; ++ const int32_t i12 = tgpig.z; ++ ++ const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0]; ++ ++ const int32_t i02 = i11; ++ const int32_t i03 = i12; ++ ++ auto psrc = (device const block_q *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01); ++ auto pdst = (device float4x4 *) (( device char *) dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1); ++ ++ for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) { ++ float4x4 temp; ++ dequantize_func(psrc + ind/nl, ind%nl, temp); ++ pdst[ind] = temp; ++ ++ break; ++ } ++} ++ ++template ++kernel void kernel_get_rows_f( ++ constant ggml_metal_kargs_get_rows & args, ++ device const void * src0, ++ device const void * src1, ++ device void * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiitg[[thread_index_in_threadgroup]], ++ ushort3 ntg [[threads_per_threadgroup]]) { ++ const int32_t iw0 = tgpig.x/args.ne10; ++ const int32_t i10 = tgpig.x%args.ne10; ++ const int32_t i11 = tgpig.y; ++ const int32_t i12 = tgpig.z; ++ ++ const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0]; ++ ++ const int32_t i02 = i11; ++ const int32_t i03 = i12; ++ ++ auto psrc = (const device T0 *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01); ++ auto pdst = ( device T *) (( device char *) dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1); ++ ++ for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) { ++ pdst[ind] = psrc[ind]; ++ ++ break; ++ } ++} ++ ++kernel void kernel_get_rows_f16_f16_vec4( ++ constant ggml_metal_kargs_get_rows & args, ++ device const void * src0, ++ device const void * src1, ++ device void * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiitg[[thread_index_in_threadgroup]], ++ ushort3 ntg [[threads_per_threadgroup]]) { ++ const int32_t iw0 = tgpig.x/args.ne10; ++ const int32_t i10 = tgpig.x%args.ne10; ++ const int32_t i11 = tgpig.y; ++ const int32_t i12 = tgpig.z; ++ ++ const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0]; ++ ++ const int32_t i02 = i11; ++ const int32_t i03 = i12; ++ ++ auto psrc = (const device half *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01); ++ auto pdst = ( device half *) (( device char *) dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1); ++ ++ for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) { ++ const int base = ind*4; ++ if (base + 3 < args.ne00) { ++ ((device half4 *) pdst)[ind] = ((const device half4 *) psrc)[ind]; ++ } else { ++ for (int i = base; i < args.ne00; ++i) { ++ pdst[i] = psrc[i]; ++ } ++ } ++ ++ break; ++ } ++} ++ ++kernel void kernel_get_rows_f16_f16_packed_rows( ++ constant ggml_metal_kargs_get_rows & args, ++ device const void * src0, ++ device const void * src1, ++ device void * dst, ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort3 tpitg [[thread_position_in_threadgroup]], ++ ushort3 tptg [[threads_per_threadgroup]]) { ++ const int32_t i10 = tgpig.x*tptg.y + tpitg.y; ++ if (i10 >= args.ne10) { ++ return; ++ } ++ ++ const int32_t i11 = tgpig.y; ++ const int32_t i12 = tgpig.z; ++ const int32_t r = ((const device int32_t *) ((const device char *) src1 + ++ i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0]; ++ ++ const int32_t i02 = i11; ++ const int32_t i03 = i12; ++ ++ const device half4 * psrc = (const device half4 *) ((const device char *) src0 + ++ i03*args.nb03 + i02*args.nb02 + r*args.nb01); ++ device half4 * pdst = (device half4 *) ((device char *) dst + ++ i12*args.nb3 + i11*args.nb2 + i10*args.nb1); ++ ++ const int32_t vec_count = args.ne00/4; ++ for (int32_t i = tpitg.x; i < vec_count; i += tptg.x) { ++ pdst[i] = psrc[i]; ++ } ++} ++ ++template ++kernel void kernel_set_rows_q32( ++ constant ggml_metal_kargs_set_rows & args, ++ device const void * src0, ++ device const void * src1, ++ device float * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ uint tiitg[[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ const int32_t i03 = tgpig.z; ++ const int32_t i02 = tgpig.y; ++ ++ const int32_t i12 = i03%args.ne12; ++ const int32_t i11 = i02%args.ne11; ++ ++ const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x; ++ if (i01 >= args.ne01) { ++ return; ++ } ++ ++ const int32_t i10 = i01; ++ const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0]; ++ ++ device block_q * dst_row = ( device block_q *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3); ++ const device float * src_row = (const device float *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); ++ ++ for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) { ++ quantize_func(src_row + 32*ind, dst_row[ind]); ++ } ++} ++ ++template ++kernel void kernel_set_rows_f( ++ constant ggml_metal_kargs_set_rows & args, ++ device const void * src0, ++ device const void * src1, ++ device float * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ uint tiitg[[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ const int32_t i03 = tgpig.z; ++ const int32_t i02 = tgpig.y; ++ ++ const int32_t i12 = i03%args.ne12; ++ const int32_t i11 = i02%args.ne11; ++ ++ const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x; ++ if (i01 >= args.ne01) { ++ return; ++ } ++ ++ const int32_t i10 = i01; ++ const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0]; + +- dh += args.nb01/2; +- q2 += args.nb01/2; +- sc += args.nb01; +- } ++ device T * dst_row = ( device T *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3); ++ const device float * src_row = (const device float *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); + +- y4 += 32 * 32; ++ for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) { ++ dst_row[ind] = (T) src_row[ind]; + } ++} + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++kernel void kernel_dsa_sparse_mask_fill( ++ constant ggml_metal_kargs_dsa_sparse_mask & args, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ uint tiitg[[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ const int32_t i_kv = tgpig.x*tptg.x + tiitg; ++ if (i_kv >= args.n_kv) { ++ return; ++ } + +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all * 0.25f; +- } ++ const int32_t i_batch = tgpig.y; ++ const int32_t i_stream = tgpig.z; ++ ++ device char * dst_ptr = dst + i_kv*args.nb1 + i_batch*args.nb2 + i_stream*args.nb3; ++ if (args.elem_size == 2) { ++ ((device half *) dst_ptr)[0] = (half) -INFINITY; ++ } else { ++ ((device float *) dst_ptr)[0] = -INFINITY; + } + } + +-[[host_name("kernel_mul_mv_iq2_xs_f32")]] +-kernel void kernel_mul_mv_iq2_xs_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem [[threadgroup(0)]], +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++kernel void kernel_dsa_sparse_mask_set( ++ constant ggml_metal_kargs_dsa_sparse_mask & args, ++ device const char * kq_mask, ++ device const int32_t * top_k, ++ device char * dst, ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ uint tiitg[[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ const int32_t i_top = tgpig.x*tptg.x + tiitg; ++ if (i_top >= args.n_top_k) { ++ return; ++ } + +- kernel_mul_mv_iq2_xs_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +-} ++ const int32_t i_batch = tgpig.y; ++ const int32_t i_stream = tgpig.z; ++ const int32_t i12 = i_stream%args.n_top_stream; ++ const int32_t i_kv = ((device const int32_t *) ((device const char *) top_k + i_top*args.nb10 + i_batch*args.nb11 + i12*args.nb12))[0]; + +-template +-void kernel_mul_mv_iq3_xxs_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ if (i_kv < 0 || i_kv >= args.n_kv) { ++ return; ++ } + +- const int nb = args.ne00/QK_K; ++ const device char * src_ptr = kq_mask + i_kv*args.nb01 + i_batch*args.nb02 + i_stream*args.nb03; ++ device char * dst_ptr = dst + i_kv*args.nb1 + i_batch*args.nb2 + i_stream*args.nb3; + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; ++ if (args.elem_size == 2) { ++ ((device half *) dst_ptr)[0] = ((const device half *) src_ptr)[0]; ++ } else { ++ ((device float *) dst_ptr)[0] = ((const device float *) src_ptr)[0]; ++ } ++} + +- const int first_row = (r0 * NSG + sgitg) * nr0; ++kernel void kernel_selected_row_flash_vec_f16_dk576_dv512( ++ constant ggml_metal_kargs_selected_row_flash & args, ++ device const char * q, ++ device const char * k, ++ device const char * v, ++ device const int32_t * top_k, ++ device char * dst, ++ device const char * kq_mask, ++ threadgroup half * shmem_f16 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]], ++ ushort sgitg [[simdgroup_index_in_threadgroup]]) { ++#define NWG (FC_flash_attn_ext_vec_nwg) ++#define HPTG (FC_flash_attn_ext_vec_nsg) ++ ++ constexpr short DK = 576; ++ constexpr short DV = 512; ++ constexpr short C = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ constexpr short NE = 2; ++ constexpr short NW = N_SIMDWIDTH; ++ constexpr short NL = NW/NE; + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++ constexpr short DK4 = DK/4; ++ constexpr short DV4 = DV/4; ++ constexpr short PK = PAD2(DK, 128); ++ constexpr short PK4 = PK/4; ++ constexpr short PV = PAD2(DV, 128); ++ constexpr short PV4 = PV/4; ++ constexpr short SH = 4*C; + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ const short iwg = tgpig.z%NWG; + +- device const block_iq3_xxs * x = (device const block_iq3_xxs *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++ const ushort iq3 = tgpig.z/NWG; ++ const ushort iq2 = tgpig.y*HPTG + sgitg; ++ const ushort iq1 = tgpig.x; + +- float yl[32]; +- float sumf[nr0]={0.f}; ++ if (iq1 >= args.ne01 || iq2 >= args.ne02 || iq3 >= args.ne03) { ++ return; ++ } + +- const int nb32 = nb * (QK_K / 32); ++ constexpr short T = PK + SH + 2*PV; ++ threadgroup half * head_smem = shmem_f16 + sgitg*T; ++ threadgroup half4 * sq4 = (threadgroup half4 *) head_smem; ++ threadgroup float * ss = (threadgroup float *) (head_smem + PK); ++ threadgroup float4 * ss4 = (threadgroup float4 *) (head_smem + PK); ++ threadgroup float4 * so4 = (threadgroup float4 *) (head_smem + PK + SH); ++ ++ so4 += tiisg; + +- threadgroup uint32_t * svalues = (threadgroup uint32_t *)(shmem); +- threadgroup uint8_t * ssigns = (threadgroup uint8_t *)(svalues + 256); + { +- int nval = 4; +- int pos = (32*sgitg + tiisg)*nval; +- for (int i = 0; i < nval; ++i) svalues[pos + i] = iq3xxs_grid[pos + i]; +- nval = 2; +- pos = (32*sgitg + tiisg)*nval; +- for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i]; +- threadgroup_barrier(mem_flags::mem_threadgroup); +- } ++ q += iq1*args.nb01 + iq2*args.nb02 + iq3*args.nb03; + +- const int ix = tiisg; ++ const short i_k_head = iq2/(args.ne02/args.ne12); ++ const short i_v_head = iq2/(args.ne02/args.ne22); ++ const short i_k_stream = iq3%args.ne13; ++ const short i_v_stream = iq3%args.ne23; + +- device const float * y4 = y + 32 * ix; ++ k += i_k_head*args.nb12 + i_k_stream*args.nb13; ++ v += i_v_head*args.nb22 + i_v_stream*args.nb23; ++ } + +- for (int ib32 = ix; ib32 < nb32; ib32 += 32) { +- for (short i = 0; i < 32; ++i) { +- yl[i] = y4[i]; ++ device const float4 * q4 = (device const float4 *) ((device const char *) q); ++ for (short i = tiisg; i < PK4; i += NW) { ++ if (i < DK4) { ++ sq4[i] = half4(q4[i]); ++ } else { ++ sq4[i] = half4(0.0h); + } ++ } + +- const int ibl = ib32 / (QK_K / 32); +- const int ib = ib32 % (QK_K / 32); ++ FOR_UNROLL (short i = 0; i < DV4/NL; ++i) { ++ so4[i*NL] = float4(0.0f); ++ } + +- device const block_iq3_xxs * xr = x + ibl; +- device const uint8_t * q3 = xr->qs + 8 * ib; +- device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib; +- device const half * dh = &xr->d; ++ for (short i = tiisg; i < SH/4; i += NW) { ++ ss4[i] = float4(0.0f); ++ } + +- for (short row = 0; row < nr0; row++) { +- const float db = dh[0]; +- const uint32_t aux32 = gas[0] | (gas[1] << 16); +- const float d = db * (0.5f + (aux32 >> 28)); ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- float2 sum = {0}; +- for (short l = 0; l < 4; ++l) { +- const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(svalues + q3[2*l+0]); +- const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(svalues + q3[2*l+1]); +- const uint8_t signs = ssigns[(aux32 >> 7*l) & 127]; +- for (short j = 0; j < 4; ++j) { +- sum[0] += yl[8*l + j + 0] * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f); +- sum[1] += yl[8*l + j + 4] * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f); ++ float S = 0.0f; ++ float M = -FLT_MAX/2; ++ ++ const short tx = tiisg%NL; ++ const short ty = tiisg/NL; ++ const short top_stream = iq3%args.ne42; ++ ++ for (int ic0 = iwg; ; ic0 += NWG) { ++ const int ic = ic0*C; ++ if (ic >= args.ne40) { ++ break; ++ } ++ ++ float mqk[C/NE] = { [0 ... C/NE - 1] = 0.0f }; ++ ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ const int row_slot = ic + NE*cc + ty; ++ int32_t i_kv = -1; ++ if (row_slot < args.ne40) { ++ i_kv = ((device const int32_t *) ((device const char *) top_k + ++ row_slot*args.nb40 + iq1*args.nb41 + top_stream*args.nb42))[0]; ++ } ++ ++ if (i_kv >= 0 && i_kv < args.ne11) { ++ device const half4 * pk4 = (device const half4 *) (k + i_kv*args.nb11); ++ FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { ++ const short i = ii*NL + tx; ++ mqk[cc] += dot((float4) pk4[i], (float4) sq4[i]); + } + } +- sumf[row] += d * (sum[0] + sum[1]); + +- dh += args.nb01/2; +- q3 += args.nb01; +- gas += args.nb01/2; +- } ++ if (NE <= 1) { ++ mqk[cc] += simd_shuffle_down(mqk[cc], 16); ++ } ++ if (NE <= 2) { ++ mqk[cc] += simd_shuffle_down(mqk[cc], 8); ++ } ++ if (NE <= 4) { ++ mqk[cc] += simd_shuffle_down(mqk[cc], 4); ++ } ++ if (NE <= 8) { ++ mqk[cc] += simd_shuffle_down(mqk[cc], 2); ++ } ++ if (NE <= 16) { ++ mqk[cc] += simd_shuffle_down(mqk[cc], 1); ++ } + +- y4 += 32 * 32; +- } ++ mqk[cc] = simd_shuffle(mqk[cc], NL*ty); ++ } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ { ++ const int row_slot = ic + NE*tx + ty; ++ int32_t i_kv = -1; ++ if (row_slot < args.ne40) { ++ i_kv = ((device const int32_t *) ((device const char *) top_k + ++ row_slot*args.nb40 + iq1*args.nb41 + top_stream*args.nb42))[0]; ++ } + +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all * 0.5f; ++ float mask = 0.0f; ++ if (i_kv >= 0 && i_kv < args.ne11 && args.ne30 > 0) { ++ mask = float(((device const half *) (kq_mask + ++ i_kv*args.nb31 + iq1*args.nb32 + iq3*args.nb33))[0]); ++ } ++ if (i_kv >= 0 && i_kv < args.ne11 && isfinite(mask)) { ++ mqk[tx] *= args.scale; ++ mqk[tx] += mask; ++ ss[NE*tx + ty] = mqk[tx]; ++ } else { ++ ss[NE*tx + ty] = -MAXHALF; ++ } + } +- } +-} + +-[[host_name("kernel_mul_mv_iq3_xxs_f32")]] +-kernel void kernel_mul_mv_iq3_xxs_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem [[threadgroup(0)]], +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- kernel_mul_mv_iq3_xxs_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +-} ++ { ++ const float m = M; ++ const float s = ss[tiisg]; + +-template +-void kernel_mul_mv_iq3_s_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ M = simd_max(max(M, s)); + +- const int nb = args.ne00/QK_K; ++ const float ms = exp(m - M); ++ const float vs = exp(s - M); + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; ++ S = S*ms + simd_sum(vs); ++ ss[tiisg] = vs; + +- const int first_row = (r0 * NSG + sgitg) * nr0; ++ if ((DV4/NL % NW == 0) || ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] *= ms; ++ } ++ } ++ } + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ { ++ float4 lo[DV4/NL]; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii] = float4(0.0f); ++ } + +- device const block_iq3_s * x = (device const block_iq3_s *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ const int row_slot = ic + NE*cc + ty; ++ int32_t i_kv = -1; ++ if (row_slot < args.ne40) { ++ i_kv = ((device const int32_t *) ((device const char *) top_k + ++ row_slot*args.nb40 + iq1*args.nb41 + top_stream*args.nb42))[0]; ++ } + +- float yl[32]; +- float sumf[nr0]={0.f}; ++ if (i_kv >= 0 && i_kv < args.ne11) { ++ device const half4 * pv4 = (device const half4 *) (v + i_kv*args.nb21); ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ const short i = ii*NL + tx; ++ lo[ii] += float4(pv4[i])*float4(ss[NE*cc + ty]); ++ } ++ } ++ } + +- const int nb32 = nb * (QK_K / 32); ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ if (NE > 1) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 16); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 16); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 16); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 16); ++ } ++ if (NE > 2) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 8); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 8); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 8); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 8); ++ } ++ if (NE > 4) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 4); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 4); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 4); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 4); ++ } ++ if (NE > 8) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 2); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 2); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 2); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 2); ++ } ++ if (NE > 16) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 1); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 1); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 1); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 1); ++ } ++ } + +- threadgroup uint32_t * svalues = (threadgroup uint32_t *) shmem; +- { +- int nval = 8; +- int pos = (32*sgitg + tiisg)*nval; +- for (int i = 0; i < nval; ++i) svalues[pos + i] = iq3s_grid[pos + i]; +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ if ((DV4/NL % NW == 0) || ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] += lo[ii]; ++ } ++ } ++ } + } + +- const int ix = tiisg; ++ if (tiisg == 0) { ++ ss[0] = S; ++ ss[1] = M; ++ } + +- device const float * y4 = y + 32 * ix; ++ so4 -= tiisg; + +- for (int ib32 = ix; ib32 < nb32; ib32 += 32) { +- for (short i = 0; i < 32; ++i) { +- yl[i] = y4[i]; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ { ++ const int64_t nrows = args.ne3*args.ne2*args.ne1; ++ const int64_t rid = iq3*args.ne2*args.ne1 + (args.ne30 > 0 ? ++ iq1 + iq2*args.ne1 : iq2 + iq1*args.ne1); ++ ++ device float4 * dst4 = (device float4 *) dst; ++ device float * dst1 = (device float *) dst + nrows*DV*NWG; ++ ++ const float inv_s = NWG == 1 ? (ss[0] == 0.0f ? 0.0f : 1.0f/ss[0]) : 1.0f; ++ ++ for (short i = tiisg; i < DV4; i += NW) { ++ dst4[rid*DV4*NWG + NWG*i + iwg] = (float4) so4[i]*inv_s; ++ } ++ ++ if (NWG > 1 && tiisg == 0) { ++ dst1[rid*(2*NWG) + 2*iwg + 0] = ss[0]; ++ dst1[rid*(2*NWG) + 2*iwg + 1] = ss[1]; + } ++ } ++ ++#undef NWG ++#undef HPTG ++} ++ ++kernel void kernel_selected_row_flash_pair_f16_dk576_dv512( ++ constant ggml_metal_kargs_selected_row_flash & args, ++ device const char * q, ++ device const char * k, ++ device const char * v, ++ device const int32_t * top_k, ++ device char * dst, ++ device const char * kq_mask, ++ threadgroup half * shmem_f16 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]]) { ++#define NWG (FC_flash_attn_ext_vec_nwg) ++ ++ constexpr short DK = 576; ++ constexpr short DV = 512; ++ constexpr short C = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ constexpr short NE = 2; ++ constexpr short NW = N_SIMDWIDTH; ++ constexpr short NL = NW/NE; ++ ++ constexpr short DK4 = DK/4; ++ constexpr short DV4 = DV/4; ++ constexpr short PK = PAD2(DK, 128); ++ constexpr short PK4 = PK/4; ++ constexpr short PV = PAD2(DV, 128); ++ constexpr short SH = 4*C; ++ constexpr short T = PK + SH + 2*PV; ++ ++ const short iwg = tgpig.z%NWG; ++ const ushort iq3 = tgpig.z/NWG; ++ const ushort iq1 = tgpig.x; ++ const ushort head0 = 2*tgpig.y; ++ const ushort head1 = head0 + 1; ++ ++ if (iq1 >= args.ne01 || head1 >= args.ne02 || iq3 >= args.ne03) { ++ return; ++ } + +- const int ibl = ib32 / (QK_K / 32); +- const int ib = ib32 % (QK_K / 32); ++ threadgroup half * smem0 = shmem_f16; ++ threadgroup half * smem1 = shmem_f16 + T; + +- device const block_iq3_s * xr = x + ibl; +- device const uint8_t * qs = xr->qs + 8 * ib; +- device const uint8_t * qh = xr->qh + ib; +- device const uint8_t * sc = xr->scales + (ib/2); +- device const uint8_t * signs = xr->signs + 4 * ib; +- device const half * dh = &xr->d; ++ threadgroup half4 * sq40 = (threadgroup half4 *) smem0; ++ threadgroup half4 * sq41 = (threadgroup half4 *) smem1; ++ threadgroup float * ss0 = (threadgroup float *) (smem0 + PK); ++ threadgroup float * ss1 = (threadgroup float *) (smem1 + PK); ++ threadgroup float4 * ss40 = (threadgroup float4 *) (smem0 + PK); ++ threadgroup float4 * ss41 = (threadgroup float4 *) (smem1 + PK); ++ threadgroup float4 * so40 = (threadgroup float4 *) (smem0 + PK + SH); ++ threadgroup float4 * so41 = (threadgroup float4 *) (smem1 + PK + SH); + +- for (short row = 0; row < nr0; row++) { +- const float db = dh[0]; +- const float d = db * (1 + 2*((sc[0] >> 4*(ib%2)) & 0xf)); ++ so40 += tiisg; ++ so41 += tiisg; + +- float2 sum = {0}; +- for (short l = 0; l < 4; ++l) { +- const threadgroup uint32_t * table1 = qh[0] & kmask_iq2xs[2*l+0] ? svalues + 256 : svalues; +- const threadgroup uint32_t * table2 = qh[0] & kmask_iq2xs[2*l+1] ? svalues + 256 : svalues; +- const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(table1 + qs[2*l+0]); +- const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(table2 + qs[2*l+1]); +- for (short j = 0; j < 4; ++j) { +- sum[0] += yl[8*l + j + 0] * grid1[j] * select(1, -1, signs[l] & kmask_iq2xs[j+0]); +- sum[1] += yl[8*l + j + 4] * grid2[j] * select(1, -1, signs[l] & kmask_iq2xs[j+4]); +- } +- } +- sumf[row] += d * (sum[0] + sum[1]); ++ device const char * q0 = q + iq1*args.nb01 + head0*args.nb02 + iq3*args.nb03; ++ device const char * q1 = q + iq1*args.nb01 + head1*args.nb02 + iq3*args.nb03; + +- dh += args.nb01/2; +- qs += args.nb01; +- qh += args.nb01; +- sc += args.nb01; +- signs += args.nb01; +- } ++ const short k_stream = iq3%args.ne13; ++ const short v_stream = iq3%args.ne23; ++ k += k_stream*args.nb13; ++ v += v_stream*args.nb23; + +- y4 += 32 * 32; ++ device const float4 * q40 = (device const float4 *) q0; ++ device const float4 * q41 = (device const float4 *) q1; ++ for (short i = tiisg; i < PK4; i += NW) { ++ if (i < DK4) { ++ sq40[i] = half4(q40[i]); ++ sq41[i] = half4(q41[i]); ++ } else { ++ sq40[i] = half4(0.0h); ++ sq41[i] = half4(0.0h); ++ } + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ FOR_UNROLL (short i = 0; i < DV4/NL; ++i) { ++ so40[i*NL] = float4(0.0f); ++ so41[i*NL] = float4(0.0f); ++ } + +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all; +- } ++ for (short i = tiisg; i < SH/4; i += NW) { ++ ss40[i] = float4(0.0f); ++ ss41[i] = float4(0.0f); + } +-} + +-[[host_name("kernel_mul_mv_iq3_s_f32")]] +-kernel void kernel_mul_mv_iq3_s_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem [[threadgroup(0)]], +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- kernel_mul_mv_iq3_s_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +-} ++ float S0 = 0.0f; ++ float S1 = 0.0f; ++ float M0 = -FLT_MAX/2; ++ float M1 = -FLT_MAX/2; + +-template +-void kernel_mul_mv_iq2_s_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ const short tx = tiisg%NL; ++ const short ty = tiisg/NL; ++ const short top_stream = iq3%args.ne42; + +- const int nb = args.ne00/QK_K; ++ for (int ic0 = iwg; ; ic0 += NWG) { ++ const int ic = ic0*C; ++ if (ic >= args.ne40) { ++ break; ++ } + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; ++ float mqk0[C/NE] = { [0 ... C/NE - 1] = 0.0f }; ++ float mqk1[C/NE] = { [0 ... C/NE - 1] = 0.0f }; + +- const int first_row = (r0 * NSG + sgitg) * nr0; ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ const int row_slot = ic + NE*cc + ty; ++ int32_t i_kv = -1; ++ if (row_slot < args.ne40) { ++ i_kv = ((device const int32_t *) ((device const char *) top_k + ++ row_slot*args.nb40 + iq1*args.nb41 + top_stream*args.nb42))[0]; ++ } + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++ if (i_kv >= 0 && i_kv < args.ne11) { ++ device const half4 * pk4 = (device const half4 *) (k + i_kv*args.nb11); ++ FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { ++ const short i = ii*NL + tx; ++ const float4 kval = (float4) pk4[i]; ++ mqk0[cc] += dot(kval, (float4) sq40[i]); ++ mqk1[cc] += dot(kval, (float4) sq41[i]); ++ } ++ } + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ if (NE <= 1) { ++ mqk0[cc] += simd_shuffle_down(mqk0[cc], 16); ++ mqk1[cc] += simd_shuffle_down(mqk1[cc], 16); ++ } ++ if (NE <= 2) { ++ mqk0[cc] += simd_shuffle_down(mqk0[cc], 8); ++ mqk1[cc] += simd_shuffle_down(mqk1[cc], 8); ++ } ++ if (NE <= 4) { ++ mqk0[cc] += simd_shuffle_down(mqk0[cc], 4); ++ mqk1[cc] += simd_shuffle_down(mqk1[cc], 4); ++ } ++ if (NE <= 8) { ++ mqk0[cc] += simd_shuffle_down(mqk0[cc], 2); ++ mqk1[cc] += simd_shuffle_down(mqk1[cc], 2); ++ } ++ if (NE <= 16) { ++ mqk0[cc] += simd_shuffle_down(mqk0[cc], 1); ++ mqk1[cc] += simd_shuffle_down(mqk1[cc], 1); ++ } + +- device const block_iq2_s * x = (device const block_iq2_s *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++ mqk0[cc] = simd_shuffle(mqk0[cc], NL*ty); ++ mqk1[cc] = simd_shuffle(mqk1[cc], NL*ty); ++ } + +- float yl[32]; +- float sumf[nr0]={0.f}; ++ { ++ const int row_slot = ic + NE*tx + ty; ++ int32_t i_kv = -1; ++ if (row_slot < args.ne40) { ++ i_kv = ((device const int32_t *) ((device const char *) top_k + ++ row_slot*args.nb40 + iq1*args.nb41 + top_stream*args.nb42))[0]; ++ } + +- const int nb32 = nb * (QK_K / 32); ++ float mask = 0.0f; ++ if (i_kv >= 0 && i_kv < args.ne11 && args.ne30 > 0) { ++ mask = float(((device const half *) (kq_mask + ++ i_kv*args.nb31 + iq1*args.nb32 + iq3*args.nb33))[0]); ++ } ++ if (i_kv >= 0 && i_kv < args.ne11 && isfinite(mask)) { ++ mqk0[tx] *= args.scale; ++ mqk0[tx] += mask; ++ mqk1[tx] *= args.scale; ++ mqk1[tx] += mask; ++ ss0[NE*tx + ty] = mqk0[tx]; ++ ss1[NE*tx + ty] = mqk1[tx]; ++ } else { ++ ss0[NE*tx + ty] = -MAXHALF; ++ ss1[NE*tx + ty] = -MAXHALF; ++ } ++ } + +- //threadgroup uint64_t * svalues = (threadgroup uint64_t *) shmem; +- //{ +- // int nval = 32; +- // int pos = (32*sgitg + tiisg)*nval; +- // for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2s_grid[pos + i]; +- // threadgroup_barrier(mem_flags::mem_threadgroup); +- //} ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- const short ix = tiisg; ++ { ++ const float m0 = M0; ++ const float m1 = M1; ++ const float s0 = ss0[tiisg]; ++ const float s1 = ss1[tiisg]; + +- device const float * y4 = y + 32 * ix; ++ M0 = simd_max(max(M0, s0)); ++ M1 = simd_max(max(M1, s1)); + +- for (int ib32 = ix; ib32 < nb32; ib32 += 32) { +- for (short i = 0; i < 32; ++i) { +- yl[i] = y4[i]; ++ const float ms0 = exp(m0 - M0); ++ const float ms1 = exp(m1 - M1); ++ const float vs0 = exp(s0 - M0); ++ const float vs1 = exp(s1 - M1); ++ ++ S0 = S0*ms0 + simd_sum(vs0); ++ S1 = S1*ms1 + simd_sum(vs1); ++ ss0[tiisg] = vs0; ++ ss1[tiisg] = vs1; ++ ++ if ((DV4/NL % NW == 0) || ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so40[ii*NL] *= ms0; ++ so41[ii*NL] *= ms1; ++ } ++ } + } + +- const int ibl = ib32 / (QK_K / 32); +- const int ib = ib32 % (QK_K / 32); ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- device const block_iq2_s * xr = x + ibl; +- device const uint8_t * qs = xr->qs + 4 * ib; +- device const uint8_t * qh = xr->qh + ib; +- device const uint8_t * sc = xr->scales + ib; +- device const uint8_t * signs = qs + QK_K/8; +- device const half * dh = &xr->d; ++ { ++ float4 lo0[DV4/NL]; ++ float4 lo1[DV4/NL]; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo0[ii] = float4(0.0f); ++ lo1[ii] = float4(0.0f); ++ } + +- for (short row = 0; row < nr0; row++) { +- const float db = dh[0]; +- const float d1 = db * (0.5f + (sc[0] & 0xf)); +- const float d2 = db * (0.5f + (sc[0] >> 4)); ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ const int row_slot = ic + NE*cc + ty; ++ int32_t i_kv = -1; ++ if (row_slot < args.ne40) { ++ i_kv = ((device const int32_t *) ((device const char *) top_k + ++ row_slot*args.nb40 + iq1*args.nb41 + top_stream*args.nb42))[0]; ++ } + +- float2 sum = {0}; +- for (short l = 0; l < 2; ++l) { +- //const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(svalues + (qs[l+0] | ((qh[0] << (8-2*l)) & 0x300))); +- //const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(svalues + (qs[l+2] | ((qh[0] << (4-2*l)) & 0x300))); +- constant uint8_t * grid1 = (constant uint8_t *)(iq2s_grid + (qs[l+0] | ((qh[0] << (8-2*l)) & 0x300))); +- constant uint8_t * grid2 = (constant uint8_t *)(iq2s_grid + (qs[l+2] | ((qh[0] << (4-2*l)) & 0x300))); +- for (short j = 0; j < 8; ++j) { +- sum[0] += yl[8*l + j + 0] * grid1[j] * select(1, -1, signs[l+0] & kmask_iq2xs[j]); +- sum[1] += yl[8*l + j + 16] * grid2[j] * select(1, -1, signs[l+2] & kmask_iq2xs[j]); ++ if (i_kv >= 0 && i_kv < args.ne11) { ++ device const half4 * pv4 = (device const half4 *) (v + i_kv*args.nb21); ++ const float p0 = ss0[NE*cc + ty]; ++ const float p1 = ss1[NE*cc + ty]; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ const short i = ii*NL + tx; ++ const float4 value = (float4) pv4[i]; ++ lo0[ii] += value*float4(p0); ++ lo1[ii] += value*float4(p1); ++ } + } + } +- sumf[row] += d1 * sum[0] + d2 * sum[1]; + +- dh += args.nb01/2; +- qs += args.nb01; +- qh += args.nb01; +- sc += args.nb01; +- signs += args.nb01; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ if (NE > 1) { ++ lo0[ii][0] += simd_shuffle_down(lo0[ii][0], 16); ++ lo0[ii][1] += simd_shuffle_down(lo0[ii][1], 16); ++ lo0[ii][2] += simd_shuffle_down(lo0[ii][2], 16); ++ lo0[ii][3] += simd_shuffle_down(lo0[ii][3], 16); ++ lo1[ii][0] += simd_shuffle_down(lo1[ii][0], 16); ++ lo1[ii][1] += simd_shuffle_down(lo1[ii][1], 16); ++ lo1[ii][2] += simd_shuffle_down(lo1[ii][2], 16); ++ lo1[ii][3] += simd_shuffle_down(lo1[ii][3], 16); ++ } ++ if (NE > 2) { ++ lo0[ii][0] += simd_shuffle_down(lo0[ii][0], 8); ++ lo0[ii][1] += simd_shuffle_down(lo0[ii][1], 8); ++ lo0[ii][2] += simd_shuffle_down(lo0[ii][2], 8); ++ lo0[ii][3] += simd_shuffle_down(lo0[ii][3], 8); ++ lo1[ii][0] += simd_shuffle_down(lo1[ii][0], 8); ++ lo1[ii][1] += simd_shuffle_down(lo1[ii][1], 8); ++ lo1[ii][2] += simd_shuffle_down(lo1[ii][2], 8); ++ lo1[ii][3] += simd_shuffle_down(lo1[ii][3], 8); ++ } ++ if (NE > 4) { ++ lo0[ii][0] += simd_shuffle_down(lo0[ii][0], 4); ++ lo0[ii][1] += simd_shuffle_down(lo0[ii][1], 4); ++ lo0[ii][2] += simd_shuffle_down(lo0[ii][2], 4); ++ lo0[ii][3] += simd_shuffle_down(lo0[ii][3], 4); ++ lo1[ii][0] += simd_shuffle_down(lo1[ii][0], 4); ++ lo1[ii][1] += simd_shuffle_down(lo1[ii][1], 4); ++ lo1[ii][2] += simd_shuffle_down(lo1[ii][2], 4); ++ lo1[ii][3] += simd_shuffle_down(lo1[ii][3], 4); ++ } ++ if (NE > 8) { ++ lo0[ii][0] += simd_shuffle_down(lo0[ii][0], 2); ++ lo0[ii][1] += simd_shuffle_down(lo0[ii][1], 2); ++ lo0[ii][2] += simd_shuffle_down(lo0[ii][2], 2); ++ lo0[ii][3] += simd_shuffle_down(lo0[ii][3], 2); ++ lo1[ii][0] += simd_shuffle_down(lo1[ii][0], 2); ++ lo1[ii][1] += simd_shuffle_down(lo1[ii][1], 2); ++ lo1[ii][2] += simd_shuffle_down(lo1[ii][2], 2); ++ lo1[ii][3] += simd_shuffle_down(lo1[ii][3], 2); ++ } ++ if (NE > 16) { ++ lo0[ii][0] += simd_shuffle_down(lo0[ii][0], 1); ++ lo0[ii][1] += simd_shuffle_down(lo0[ii][1], 1); ++ lo0[ii][2] += simd_shuffle_down(lo0[ii][2], 1); ++ lo0[ii][3] += simd_shuffle_down(lo0[ii][3], 1); ++ lo1[ii][0] += simd_shuffle_down(lo1[ii][0], 1); ++ lo1[ii][1] += simd_shuffle_down(lo1[ii][1], 1); ++ lo1[ii][2] += simd_shuffle_down(lo1[ii][2], 1); ++ lo1[ii][3] += simd_shuffle_down(lo1[ii][3], 1); ++ } ++ } ++ ++ if ((DV4/NL % NW == 0) || ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so40[ii*NL] += lo0[ii]; ++ so41[ii*NL] += lo1[ii]; ++ } ++ } + } ++ } + +- y4 += 32 * 32; ++ if (tiisg == 0) { ++ ss0[0] = S0; ++ ss0[1] = M0; ++ ss1[0] = S1; ++ ss1[1] = M1; + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ so40 -= tiisg; ++ so41 -= tiisg; + +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all * 0.25f; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ { ++ const int64_t nrows = args.ne3*args.ne2*args.ne1; ++ const int64_t rid0 = iq3*args.ne2*args.ne1 + (args.ne30 > 0 ? ++ iq1 + head0*args.ne1 : head0 + iq1*args.ne1); ++ const int64_t rid1 = iq3*args.ne2*args.ne1 + (args.ne30 > 0 ? ++ iq1 + head1*args.ne1 : head1 + iq1*args.ne1); ++ ++ device float4 * dst4 = (device float4 *) dst; ++ device float * dst1 = (device float *) dst + nrows*DV*NWG; ++ ++ const float inv_s0 = NWG == 1 ? (ss0[0] == 0.0f ? 0.0f : 1.0f/ss0[0]) : 1.0f; ++ const float inv_s1 = NWG == 1 ? (ss1[0] == 0.0f ? 0.0f : 1.0f/ss1[0]) : 1.0f; ++ ++ for (short i = tiisg; i < DV4; i += NW) { ++ dst4[rid0*DV4*NWG + NWG*i + iwg] = (float4) so40[i]*inv_s0; ++ dst4[rid1*DV4*NWG + NWG*i + iwg] = (float4) so41[i]*inv_s1; ++ } ++ ++ if (NWG > 1 && tiisg == 0) { ++ dst1[rid0*(2*NWG) + 2*iwg + 0] = ss0[0]; ++ dst1[rid0*(2*NWG) + 2*iwg + 1] = ss0[1]; ++ dst1[rid1*(2*NWG) + 2*iwg + 0] = ss1[0]; ++ dst1[rid1*(2*NWG) + 2*iwg + 1] = ss1[1]; + } + } ++ ++#undef NWG + } + +-[[host_name("kernel_mul_mv_iq2_s_f32")]] +-kernel void kernel_mul_mv_iq2_s_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem [[threadgroup(0)]], +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++kernel void kernel_selected_row_flash_tiled_f16_dk576_dv512( ++ constant ggml_metal_kargs_selected_row_flash & args, ++ device const char * q, ++ device const char * k, ++ device const char * v, ++ device const int32_t * top_k, ++ device char * dst, ++ device const char * kq_mask, ++ threadgroup half * shmem_f16 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]], ++ ushort sgitg [[simdgroup_index_in_threadgroup]]) { ++#define NWG (FC_flash_attn_ext_vec_nwg) ++ ++ constexpr short DK = 576; ++ constexpr short DV = 512; ++ constexpr short Q = 8; ++ constexpr short C = 64; ++ constexpr short NSG = 4; ++ constexpr short NQ = Q/NSG; ++ constexpr short DK4 = DK/4; ++ constexpr short DK16 = DK/16; ++ constexpr short DV4 = DV/4; ++ constexpr short DV16 = DV/16; ++ constexpr short SH = C; ++ constexpr short NW = N_SIMDWIDTH; ++ ++ const short iwg = tgpig.z%NWG; ++ const short iq3 = tgpig.z/NWG; ++ const short iq1 = tgpig.y; ++ const short head_first = tgpig.x*Q; ++ ++ threadgroup half * sq = shmem_f16; ++ threadgroup float * so = (threadgroup float *) (shmem_f16 + Q*DK); ++ threadgroup float * ss = (threadgroup float *) (shmem_f16 + Q*DK + 2*Q*DV); ++ threadgroup half * st = shmem_f16 + Q*DK + 2*Q*DV + 2*Q*SH; ++ ++ threadgroup half4 * sq4 = (threadgroup half4 *) sq; ++ threadgroup float4 * so4 = (threadgroup float4 *) so; ++ threadgroup float2 * ss2 = (threadgroup float2 *) ss; ++ threadgroup half * sk = st + sgitg*(4*16*8); ++ threadgroup half4x4 * sk4x4 = (threadgroup half4x4 *) sk; ++ ++ const short k_stream = iq3%args.ne13; ++ const short v_stream = iq3%args.ne23; ++ const short top_stream = iq3%args.ne43; ++ k += k_stream*args.nb13; ++ v += v_stream*args.nb23; + +- kernel_mul_mv_iq2_s_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +-} ++ FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { ++ const short j = jj*NSG + sgitg; ++ const short head = head_first + j; ++ device const char * q_head = q + iq1*args.nb01 + head*args.nb02 + iq3*args.nb03; ++ device const float4 * q4 = (device const float4 *) q_head; ++ for (short i = tiisg; i < DK4; i += NW) { ++ sq4[j*DK4 + i] = head < args.ne02 ? half4(q4[i]) : half4(0.0h); ++ } + +-template +-void kernel_mul_mv_iq1_s_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ for (short i = tiisg; i < DV4; i += NW) { ++ so4[j*DV4 + i] = float4(0.0f); ++ } ++ for (short i = tiisg; i < SH/2; i += NW) { ++ ss2[j*SH/2 + i] = float2(0.0f); ++ } ++ } + +- const int nb = args.ne00/QK_K; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; ++ float S[NQ] = { [0 ... NQ - 1] = 0.0f }; ++ float M[NQ] = { [0 ... NQ - 1] = -FLT_MAX/2 }; + +- const int first_row = (r0 * NSG + sgitg) * nr0; ++ const short tx = tiisg%4; ++ const short ty = tiisg/4; + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++ for (int ic0 = iwg; ; ic0 += NWG) { ++ const int ic = ic0*C; ++ if (ic >= args.ne40) { ++ break; ++ } + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ FOR_UNROLL (short ccc = 0; ccc < (C/8)/NSG; ++ccc) { ++ const short cc = ccc*NSG + sgitg; ++ simdgroup_float8x8 mqk = make_filled_simdgroup_matrix(0.0f); ++ const int k_slot = ic + 8*cc + ty; ++ int32_t i_kv = -1; ++ if (k_slot < args.ne40) { ++ i_kv = ((device const int32_t *) ((device const char *) top_k + ++ k_slot*args.nb40 + iq1*args.nb41 + top_stream*args.nb43))[0]; ++ } + +- device const block_iq1_s * x = (device const block_iq1_s *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++ for (short ii = 0; ii < DK16; ii += 4) { ++ half4x4 tmp = half4x4(0.0h); ++ if (i_kv >= 0 && i_kv < args.ne11) { ++ device const half4x4 * pk4x4 = ++ (device const half4x4 *) (k + i_kv*args.nb11); ++ tmp = pk4x4[ii + tx]; ++ } ++ sk4x4[4*ty + tx] = tmp; + +- float yl[32]; +- float sumf[nr0]={0.f}; ++ simdgroup_barrier(mem_flags::mem_threadgroup); + +- const int nb32 = nb * (QK_K / 32); ++ FOR_UNROLL (short kk = 0; kk < 4; ++kk) { ++ simdgroup_half8x8 mk; ++ simdgroup_half8x8 mq[2]; + +- const short ix = tiisg; ++ simdgroup_load(mk, sk + 16*kk + 0*8, 4*16, 0, true); ++ simdgroup_load(mq[0], sq + (2*(ii + kk) + 0)*8, DK); ++ simdgroup_multiply_accumulate(mqk, mq[0], mk, mqk); + +- device const float * y4 = y + 32 * ix; ++ simdgroup_load(mk, sk + 16*kk + 1*8, 4*16, 0, true); ++ simdgroup_load(mq[1], sq + (2*(ii + kk) + 1)*8, DK); ++ simdgroup_multiply_accumulate(mqk, mq[1], mk, mqk); ++ } ++ } + +- for (int ib32 = ix; ib32 < nb32; ib32 += 32) { +- float sumy = 0; +- for (short i = 0; i < 32; ++i) { +- yl[i] = y4[i]; +- sumy += yl[i]; ++ simdgroup_store(mqk, ss + 8*cc, SH, 0, false); + } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { ++ const short j = jj*NSG + sgitg; ++ float2 s2 = ss2[j*SH/2 + tiisg]*args.scale; ++ const int row0 = ic + 2*tiisg; ++ const int row1 = row0 + 1; ++ int32_t idx0 = -1; ++ int32_t idx1 = -1; ++ if (row0 < args.ne40) { ++ idx0 = ((device const int32_t *) ((device const char *) top_k + ++ row0*args.nb40 + iq1*args.nb41 + top_stream*args.nb43))[0]; ++ } ++ if (row1 < args.ne40) { ++ idx1 = ((device const int32_t *) ((device const char *) top_k + ++ row1*args.nb40 + iq1*args.nb41 + top_stream*args.nb43))[0]; ++ } ++ float mask0 = 0.0f; ++ float mask1 = 0.0f; ++ if (idx0 >= 0 && idx0 < args.ne11 && args.ne30 > 0) { ++ mask0 = float(((device const half *) (kq_mask + ++ idx0*args.nb31 + iq1*args.nb32 + iq3*args.nb33))[0]); ++ } ++ if (idx1 >= 0 && idx1 < args.ne11 && args.ne30 > 0) { ++ mask1 = float(((device const half *) (kq_mask + ++ idx1*args.nb31 + iq1*args.nb32 + iq3*args.nb33))[0]); ++ } ++ if (idx0 < 0 || idx0 >= args.ne11 || !isfinite(mask0)) { ++ s2[0] = -MAXHALF; ++ } else { ++ s2[0] += mask0; ++ } ++ if (idx1 < 0 || idx1 >= args.ne11 || !isfinite(mask1)) { ++ s2[1] = -MAXHALF; ++ } else { ++ s2[1] += mask1; ++ } ++ ++ const float m = M[jj]; ++ M[jj] = simd_max(max(M[jj], max(s2[0], s2[1]))); ++ const float ms = exp(m - M[jj]); ++ const float2 vs2 = exp(s2 - M[jj]); ++ S[jj] = S[jj]*ms + simd_sum(vs2[0] + vs2[1]); ++ ss2[j*SH/2 + tiisg] = vs2; + +- const int ibl = ib32 / (QK_K / 32); +- const int ib = ib32 % (QK_K / 32); ++ for (short i = tiisg; i < DV4; i += NW) { ++ so4[j*DV4 + i] *= ms; ++ } ++ } + +- device const block_iq1_s * xr = x + ibl; +- device const uint8_t * qs = xr->qs + 4 * ib; +- device const uint16_t * qh = xr->qh + ib; +- device const half * dh = &xr->d; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- for (short row = 0; row < nr0; row++) { +- constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700))); +- constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 5) & 0x700))); +- constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[0] << 2) & 0x700))); +- constant uint8_t * grid4 = (constant uint8_t *)(iq1s_grid_gpu + (qs[3] | ((qh[0] >> 1) & 0x700))); ++ FOR_UNROLL (short cc = 0; cc < C/8; ++cc) { ++ simdgroup_float8x8 ps; ++ simdgroup_load(ps, ss + 8*cc, SH, 0, false); + +- float sum = 0; +- for (short j = 0; j < 4; ++j) { +- sum += yl[j+ 0] * (grid1[j] & 0xf) + yl[j+ 4] * (grid1[j] >> 4) +- + yl[j+ 8] * (grid2[j] & 0xf) + yl[j+12] * (grid2[j] >> 4) +- + yl[j+16] * (grid3[j] & 0xf) + yl[j+20] * (grid3[j] >> 4) +- + yl[j+24] * (grid4[j] & 0xf) + yl[j+28] * (grid4[j] >> 4); ++ const int v_slot = ic + 8*cc + ty; ++ int32_t i_v = -1; ++ if (v_slot < args.ne40) { ++ i_v = ((device const int32_t *) ((device const char *) top_k + ++ v_slot*args.nb40 + iq1*args.nb41 + top_stream*args.nb43))[0]; + } +- sumf[row] += (float)dh[0] * (sum + sumy * (qh[0] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA)) * (2*((qh[0] >> 12) & 7) + 1); + +- dh += args.nb01/2; +- qs += args.nb01; +- qh += args.nb01/2; ++ for (short ii = 4*sgitg; ii < DV16; ii += 4*NSG) { ++ half4x4 tmp = half4x4(0.0h); ++ if (i_v >= 0 && i_v < args.ne11) { ++ device const half4x4 * pv4x4 = ++ (device const half4x4 *) (v + i_v*args.nb21); ++ tmp = pv4x4[ii + tx]; ++ } ++ sk4x4[4*ty + tx] = tmp; ++ ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ ++ FOR_UNROLL (short kk = 0; kk < 4; ++kk) { ++ simdgroup_half8x8 mv[2]; ++ simdgroup_float8x8 lo[2]; ++ simdgroup_load(mv[0], sk + 16*kk + 0*8, 4*16, 0, false); ++ simdgroup_load(mv[1], sk + 16*kk + 1*8, 4*16, 0, false); ++ simdgroup_load(lo[0], so + 8*(2*(ii + kk) + 0), DV, 0, false); ++ simdgroup_load(lo[1], so + 8*(2*(ii + kk) + 1), DV, 0, false); ++ simdgroup_multiply_accumulate(lo[0], ps, mv[0], lo[0]); ++ simdgroup_multiply_accumulate(lo[1], ps, mv[1], lo[1]); ++ simdgroup_store(lo[0], so + 8*(2*(ii + kk) + 0), DV, 0, false); ++ simdgroup_store(lo[1], so + 8*(2*(ii + kk) + 1), DV, 0, false); ++ } ++ } + } + +- y4 += 32 * 32; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ FOR_UNROLL (short jj = 0; jj < NQ; ++jj) { ++ const short j = jj*NSG + sgitg; ++ const short head = head_first + j; ++ if (head < args.ne02) { ++ const int64_t nrows = args.ne3*args.ne2*args.ne1; ++ const int64_t rid = iq3*args.ne2*args.ne1 + (args.ne30 > 0 ? ++ iq1 + head*args.ne1 : head + iq1*args.ne1); + +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all; ++ device float4 * dst4 = (device float4 *) dst; ++ device float * dst1 = (device float *) dst + nrows*DV*NWG; ++ const float inv_s = NWG == 1 ? (S[jj] == 0.0f ? 0.0f : 1.0f/S[jj]) : 1.0f; ++ ++ for (short i = tiisg; i < DV4; i += NW) { ++ dst4[rid*DV4*NWG + NWG*i + iwg] = so4[j*DV4 + i]*inv_s; ++ } ++ ++ if (NWG > 1 && tiisg == 0) { ++ dst1[rid*(2*NWG) + 2*iwg + 0] = S[jj]; ++ dst1[rid*(2*NWG) + 2*iwg + 1] = M[jj]; ++ } + } + } +-} +- +-[[host_name("kernel_mul_mv_iq1_s_f32")]] +-kernel void kernel_mul_mv_iq1_s_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { + +- kernel_mul_mv_iq1_s_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); ++#undef NWG + } + +-template +-void kernel_mul_mv_iq1_m_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, ++// Group GLM query heads while preserving the stock vector Flash arithmetic. ++// Each SIMD group owns one head and cooperatively staged K/V rows are reused ++// by all four heads in the threadgroup. ++kernel void kernel_glm_compact_multihead_flash_staged_group4_f16_dk576_dv512( ++ constant ggml_metal_kargs_flash_attn_ext_vec & args, ++ device const char * q, ++ device const char * k, ++ device const char * v, + device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ threadgroup half * shmem_f16 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]], ++ ushort sgitg [[simdgroup_index_in_threadgroup]]) { ++#define NWG (FC_flash_attn_ext_vec_nwg) ++ ++ constexpr short DK = 576; ++ constexpr short DV = 512; ++ constexpr short Q = 4; ++ constexpr short C = 32; ++ constexpr short NE = 2; ++ constexpr short NSG = Q; ++ constexpr short SCORE_BATCH = 4; ++ constexpr short STAGE_ROWS = SCORE_BATCH*NE; ++ constexpr short DK4 = DK/4; ++ constexpr short DV4 = DV/4; ++ constexpr short NW = N_SIMDWIDTH; ++ constexpr short NL = NW/NE; ++ ++ const short iwg = tgpig.z%NWG; ++ const ushort iq3 = tgpig.z/NWG; ++ const ushort iq1 = tgpig.y; ++ const ushort head_first = tgpig.x*Q; ++ const ushort head_slot = sgitg; ++ const ushort head = head_first + head_slot; ++ const ushort tid = sgitg*NW + tiisg; ++ const short tx = tiisg%NL; ++ const short ty = tiisg/NL; ++ ++ threadgroup half * sq = shmem_f16; ++ threadgroup half * st = sq + Q*DK; ++ threadgroup float * ss = (threadgroup float *) (st + STAGE_ROWS*DK); ++ threadgroup float * so = ss + Q*C; ++ ++ threadgroup half4 * sq4 = (threadgroup half4 *) sq; ++ threadgroup half4 * st4 = (threadgroup half4 *) st; ++ threadgroup float4 * so4 = (threadgroup float4 *) so + head_slot*DV4 + tiisg; ++ threadgroup float * sh = ss + head_slot*C; ++ ++ const short kv_stream = iq3%args.ne_12_3; ++ k += kv_stream*args.nb13; ++ v += kv_stream*args.nb23; ++ ++ device const float4 * q4 = (device const float4 *) ++ (q + iq1*args.nb01 + head*args.nb02 + iq3*args.nb03); ++ for (short i = tiisg; i < DK4; i += NW) { ++ sq4[head_slot*DK4 + i] = head < args.ne02 ? half4(q4[i]) : half4(0.0h); ++ } ++ for (short i = 0; i < DV4/NL; ++i) { ++ so4[i*NL] = float4(0.0f); ++ } + +- const int nb = args.ne00/QK_K; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; ++ float S = 0.0f; ++ float M = -FLT_MAX/2; + +- const int first_row = (r0 * NSG + sgitg) * nr0; ++ for (int ic0 = iwg; ; ic0 += NWG) { ++ const int ic = ic0*C; ++ if (ic >= args.ne11) { ++ break; ++ } + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++ FOR_UNROLL (short cc0 = 0; cc0 < C/NE; cc0 += SCORE_BATCH) { ++ const int row_first = ic + NE*cc0; ++ for (short i = tid; i < STAGE_ROWS*DK4; i += NSG*NW) { ++ const short row = i/DK4; ++ const short col = i%DK4; ++ const int source_row = row_first + row; ++ st4[i] = source_row < args.ne11 ? ++ ((device const half4 *) (k + source_row*args.nb11))[col] : half4(0.0h); ++ } + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- device const block_iq1_m * x = (device const block_iq1_m *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++ FOR_UNROLL (short ccc = 0; ccc < SCORE_BATCH; ++ccc) { ++ float mqk = 0.0f; ++ threadgroup const half4 * pk4 = st4 + (NE*ccc + ty)*DK4 + tx; ++ threadgroup const half4 * pq4 = sq4 + head_slot*DK4 + tx; + +- float yl[32]; +- float sumf[nr0]={0.f}; ++ FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { ++ mqk += dot(float4(pk4[ii*NL]), float4(pq4[ii*NL])); ++ } + +- const int nb32 = nb * (QK_K / 32); ++ mqk += simd_shuffle_down(mqk, 8); ++ mqk += simd_shuffle_down(mqk, 4); ++ mqk += simd_shuffle_down(mqk, 2); ++ mqk += simd_shuffle_down(mqk, 1); ++ mqk = simd_shuffle(mqk, NL*ty); + +- const short ix = tiisg; ++ if (tx == cc0 + ccc) { ++ sh[NE*(cc0 + ccc) + ty] = mqk*args.scale; ++ } ++ } + +- device const float * y4 = y + 32 * ix; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +- iq1m_scale_t scale; ++ { ++ const float m = M; ++ const float s = sh[tiisg]; ++ M = simd_max(max(M, s)); ++ const float ms = exp(m - M); ++ const float vs = exp(s - M); ++ S = S*ms + simd_sum(vs); ++ sh[tiisg] = vs; + +- for (int ib32 = ix; ib32 < nb32; ib32 += 32) { +- float4 sumy = {0.f}; +- for (short i = 0; i < 8; ++i) { +- yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0]; +- yl[i+ 8] = y4[i+ 8]; sumy[1] += yl[i+ 8]; +- yl[i+16] = y4[i+16]; sumy[2] += yl[i+16]; +- yl[i+24] = y4[i+24]; sumy[3] += yl[i+24]; ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] *= ms; ++ } ++ } + } + +- const int ibl = ib32 / (QK_K / 32); +- const int ib = ib32 % (QK_K / 32); +- +- device const block_iq1_m * xr = x + ibl; +- device const uint8_t * qs = xr->qs + 4 * ib; +- device const uint8_t * qh = xr->qh + 2 * ib; +- device const uint16_t * sc = (device const uint16_t *)xr->scales; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- for (short row = 0; row < nr0; row++) { +- scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); ++ float4 lo[DV4/NL] = { [0 ... DV4/NL - 1] = float4(0.0f) }; ++ FOR_UNROLL (short cc0 = 0; cc0 < C/NE; cc0 += SCORE_BATCH) { ++ const int row_first = ic + NE*cc0; ++ for (short i = tid; i < STAGE_ROWS*DV4; i += NSG*NW) { ++ const short row = i/DV4; ++ const short col = i%DV4; ++ const int source_row = row_first + row; ++ st4[i] = source_row < args.ne11 ? ++ ((device const half4 *) (v + source_row*args.nb21))[col] : half4(0.0h); ++ } + +- constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700))); +- constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 4) & 0x700))); +- constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[1] << 8) & 0x700))); +- constant uint8_t * grid4 = (constant uint8_t *)(iq1s_grid_gpu + (qs[3] | ((qh[1] << 4) & 0x700))); ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- float2 sum = {0.f}; +- for (short j = 0; j < 4; ++j) { +- sum[0] += yl[j+ 0] * (grid1[j] & 0xf) + yl[j+ 4] * (grid1[j] >> 4) +- + yl[j+ 8] * (grid2[j] & 0xf) + yl[j+12] * (grid2[j] >> 4); +- sum[1] += yl[j+16] * (grid3[j] & 0xf) + yl[j+20] * (grid3[j] >> 4) +- + yl[j+24] * (grid4[j] & 0xf) + yl[j+28] * (grid4[j] >> 4); ++ FOR_UNROLL (short ccc = 0; ccc < SCORE_BATCH; ++ccc) { ++ const float weight = sh[NE*(cc0 + ccc) + ty]; ++ threadgroup const half4 * pv4 = st4 + (NE*ccc + ty)*DV4 + tx; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii] += float4(pv4[ii*NL])*float4(weight); ++ } + } +- const float delta1 = sumy[0] * (qh[0] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA) + sumy[1] * (qh[0] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA); +- const float delta2 = sumy[2] * (qh[1] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA) + sumy[3] * (qh[1] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA); + +- sumf[row] += (float)scale.f16 * ((sum[0] + delta1) * (2*((sc[ib/2] >> (6*(ib%2)+0)) & 7) + 1) + +- (sum[1] + delta2) * (2*((sc[ib/2] >> (6*(ib%2)+3)) & 7) + 1)); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +- sc += args.nb01/2; +- qs += args.nb01; +- qh += args.nb01; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 16); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 16); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 16); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 16); ++ } ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] += lo[ii]; ++ } + } + +- y4 += 32 * 32; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ if (head < args.ne02) { ++ const int64_t rid = iq3*args.ne02*args.ne01 + head + iq1*args.ne01; ++ device float4 * dst4 = (device float4 *) dst; ++ device float * dst1 = (device float *) dst + args.ne1*args.ne2*args.ne3*DV*NWG; ++ const float inv_s = NWG == 1 ? (S == 0.0f ? 0.0f : 1.0f/S) : 1.0f; + +- for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all; ++ threadgroup float4 * out4 = (threadgroup float4 *) so + head_slot*DV4; ++ for (short i = tiisg; i < DV4; i += NW) { ++ dst4[rid*DV4*NWG + NWG*i + iwg] = out4[i]*inv_s; ++ } ++ if (NWG > 1 && tiisg == 0) { ++ dst1[rid*(2*NWG) + 2*iwg + 0] = S; ++ dst1[rid*(2*NWG) + 2*iwg + 1] = M; + } + } ++ ++#undef NWG + } + +-[[host_name("kernel_mul_mv_iq1_m_f32")]] +-kernel void kernel_mul_mv_iq1_m_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, ++// Keep stock GLM compact-Flash chunking and arithmetic while a second SIMD ++// group computes the upper half of the 512-wide value projection. ++kernel void kernel_glm_compact_multihead_flash_f16_dk576_dv512( ++ constant ggml_metal_kargs_flash_attn_ext_vec & args, ++ device const char * q, ++ device const char * k, ++ device const char * v, + device char * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ threadgroup half * shmem_f16 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]], ++ ushort sgitg [[simdgroup_index_in_threadgroup]]) { ++#define NWG (FC_flash_attn_ext_vec_nwg) ++ ++ constexpr short DK = 576; ++ constexpr short DV = 512; ++ constexpr short C = 32; ++ constexpr short NE = 2; ++ constexpr short NSG = 2; ++ constexpr short DK4 = DK/4; ++ constexpr short DV4 = DV/4; ++ constexpr short NW = N_SIMDWIDTH; ++ constexpr short NL = NW/NE; ++ constexpr short NVO = DV4/NL; ++ constexpr short NVP = NVO/NSG; + +- kernel_mul_mv_iq1_m_f32_impl(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg); +-} ++ const short iwg = tgpig.z%NWG; ++ const ushort iq3 = tgpig.z/NWG; ++ const ushort head = tgpig.y; ++ const ushort iq1 = tgpig.x; ++ const ushort tid = sgitg*NW + tiisg; ++ const short tx = tiisg%NL; ++ const short ty = tiisg/NL; + +-template +-void kernel_mul_mv_iq4_nl_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ const short kv_stream = iq3%args.ne_12_3; ++ k += kv_stream*args.nb13; ++ v += kv_stream*args.nb23; + +- threadgroup float * shmem_f32 = (threadgroup float *) shmem; ++ threadgroup half4 * sq4 = (threadgroup half4 *) shmem_f16; ++ threadgroup float * ss = (threadgroup float *) (shmem_f16 + DK); ++ threadgroup float * state = ss + C; ++ threadgroup float4 * out4 = (threadgroup float4 *) (state + 4); + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; ++ device const float4 * q4 = (device const float4 *) ++ (q + iq1*args.nb01 + head*args.nb02 + iq3*args.nb03); ++ for (short i = tid; i < DK4; i += NSG*NW) { ++ sq4[i] = half4(q4[i]); ++ } ++ for (short i = tid; i < DV4; i += NSG*NW) { ++ out4[i] = float4(0.0f); ++ } + +- const int first_row = (r0 * NSG + sgitg) * NR0; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++ float S = 0.0f; ++ float M = -FLT_MAX/2; + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ for (int ic0 = iwg; ; ic0 += NWG) { ++ const int ic = ic0*C; ++ if (ic >= args.ne11) { ++ break; ++ } + +- device const block_iq4_nl * x = (device const block_iq4_nl *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++ if (sgitg == 0) { ++ device const half4 * pk4 = (device const half4 *) (k + ic*args.nb11); ++ pk4 += ty*DK4 + tx; ++ threadgroup const half4 * pq4 = sq4 + tx; ++ ++ float mqk[C/NE] = { [0 ... C/NE - 1] = 0.0f }; ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { ++ mqk[cc] += dot( ++ float4(pk4[cc*NE*DK4 + ii*NL]), ++ float4(pq4[ii*NL])); ++ } + +- const int nb = args.ne00/QK4_NL; +- const int ns01 = args.nb01/args.nb00; ++ mqk[cc] += simd_shuffle_down(mqk[cc], 8); ++ mqk[cc] += simd_shuffle_down(mqk[cc], 4); ++ mqk[cc] += simd_shuffle_down(mqk[cc], 2); ++ mqk[cc] += simd_shuffle_down(mqk[cc], 1); ++ mqk[cc] = simd_shuffle(mqk[cc], NL*ty); ++ } ++ ss[NE*tx + ty] = mqk[tx]*args.scale; + +- const short ix = tiisg/2; // 0...15 +- const short it = tiisg%2; // 0 or 1 ++ simdgroup_barrier(mem_flags::mem_threadgroup); + +- shmem_f32[tiisg] = kvalues_iq4nl_f[tiisg%16]; +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ const float m = M; ++ const float s = ss[tiisg]; ++ M = simd_max(max(M, s)); ++ const float ms = exp(m - M); ++ const float vs = exp(s - M); ++ S = S*ms + simd_sum(vs); ++ ss[tiisg] = vs; ++ if (tiisg == 0) { ++ state[0] = ms; ++ state[1] = S; ++ state[2] = M; ++ } ++ } + +- float4 yl[4]; +- float sumf[NR0]={0.f}; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- device const float * yb = y + ix*QK4_NL + it*8; ++ const short ii0 = sgitg*NVP; ++ threadgroup float4 * so4 = out4 + tiisg; ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < NVP; ++ii) { ++ so4[(ii0 + ii)*NL] *= state[0]; ++ } ++ } + +- uint32_t aux32[2]; +- thread const uint8_t * q8 = (thread const uint8_t *)aux32; ++ float4 lo[NVP] = { [0 ... NVP - 1] = float4(0.0f) }; ++ device const half4 * pv4 = (device const half4 *) (v + ic*args.nb21); ++ pv4 += ty*DV4 + tx; ++ threadgroup const float * sst = ss + ty; + +- float4 qf1, qf2; ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ FOR_UNROLL (short ii = 0; ii < NVP; ++ii) { ++ lo[ii] += float4(pv4[cc*NE*DV4 + (ii0 + ii)*NL])*float4(sst[cc*NE]); ++ } ++ } + +- // [TAG_MUL_MV_WEIRD] +- for (int ib = ix; ib < nb && ib < ns01; ib += 16) { +- device const float4 * y4 = (device const float4 *)yb; +- yl[0] = y4[0]; +- yl[1] = y4[4]; +- yl[2] = y4[1]; +- yl[3] = y4[5]; ++ FOR_UNROLL (short ii = 0; ii < NVP; ++ii) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 16); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 16); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 16); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 16); ++ } ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < NVP; ++ii) { ++ so4[(ii0 + ii)*NL] += lo[ii]; ++ } ++ } + +- for (short row = 0; row < NR0; row++) { +- device const block_iq4_nl & xb = x[row*ns01 + ib]; +- device const uint16_t * q4 = (device const uint16_t *)(xb.qs + 8*it); ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ const int64_t rid = iq3*args.ne02*args.ne01 + head + iq1*args.ne01; ++ device float4 * dst4 = (device float4 *) dst; ++ device float * dst1 = (device float *) dst + args.ne1*args.ne2*args.ne3*DV*NWG; ++ const float inv_s = NWG == 1 ? (state[1] == 0.0f ? 0.0f : 1.0f/state[1]) : 1.0f; ++ ++ for (short i = tid; i < DV4; i += NSG*NW) { ++ dst4[rid*DV4*NWG + NWG*i + iwg] = out4[i]*inv_s; ++ } ++ if (NWG > 1 && sgitg == 0 && tiisg == 0) { ++ dst1[rid*(2*NWG) + 2*iwg + 0] = state[1]; ++ dst1[rid*(2*NWG) + 2*iwg + 1] = state[2]; ++ } + +- float4 acc1 = {0.f}, acc2 = {0.f}; ++#undef NWG ++} + +- aux32[0] = q4[0] | (q4[1] << 16); +- aux32[1] = (aux32[0] >> 4) & 0x0f0f0f0f; +- aux32[0] &= 0x0f0f0f0f; +- qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]}; +- qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]}; +- acc1 += yl[0] * qf1; +- acc2 += yl[1] * qf2; ++kernel void kernel_glm_compact_qk_scores_f16_dk576( ++ constant ggml_metal_kargs_flash_attn_ext_vec & args, ++ device const char * q, ++ device const char * k, ++ device float * scores, ++ device float * debug_scores, ++ threadgroup half * shmem_f16 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]]) { ++ constexpr short DK = 576; ++ constexpr short C = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ constexpr short NE = 2; ++ constexpr short NL = N_SIMDWIDTH/NE; ++ constexpr short DK4 = DK/4; ++ constexpr short PK = PAD2(DK, 128); ++ constexpr short PK4 = PK/4; + +- aux32[0] = q4[2] | (q4[3] << 16); +- aux32[1] = (aux32[0] >> 4) & 0x0f0f0f0f; +- aux32[0] &= 0x0f0f0f0f; +- qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]}; +- qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]}; +- acc1 += yl[2] * qf1; +- acc2 += yl[3] * qf2; ++ const int chunk = tgpig.x; ++ const int head = tgpig.y; ++ const int stream = tgpig.z; ++ if (head >= args.ne02) { ++ return; ++ } + +- acc1 += acc2; ++ threadgroup half4 * sq4 = (threadgroup half4 *) shmem_f16; ++ device const float4 * q4 = (device const float4 *) ++ (q + head*args.nb02 + stream*args.nb03); ++ for (short i = tiisg; i < PK4; i += N_SIMDWIDTH) { ++ sq4[i] = i < DK4 ? half4(q4[i]) : half4(0.0h); ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); + +- sumf[row] += (float)xb.d * (acc1[0] + acc1[1] + acc1[2] + acc1[3]); ++ const int k_stream = stream%args.ne_12_3; ++ device const half4 * pk4 = (device const half4 *) ++ (k + k_stream*args.nb13 + chunk*C*args.nb11); ++ const short tx = tiisg%NL; ++ const short ty = tiisg/NL; ++ pk4 += ty*args.ns10/4 + tx; ++ ++ float mqk[C/NE] = { [0 ... C/NE - 1] = 0.0f }; ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ FOR_UNROLL (short ii = 0; ii < DK4/NL; ++ii) { ++ mqk[cc] += dot( ++ float4(pk4[cc*NE*args.ns10/4 + ii*NL]), ++ float4(sq4[ii*NL + tx])); + } + +- yb += 16 * QK4_NL; ++ mqk[cc] += simd_shuffle_down(mqk[cc], 8); ++ mqk[cc] += simd_shuffle_down(mqk[cc], 4); ++ mqk[cc] += simd_shuffle_down(mqk[cc], 2); ++ mqk[cc] += simd_shuffle_down(mqk[cc], 1); ++ mqk[cc] = simd_shuffle(mqk[cc], NL*ty); + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; +- +- for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all; ++ const int row = chunk*C + NE*tx + ty; ++ if (row < args.ne11) { ++ const int64_t score_row = (int64_t(stream)*args.ne02 + head)*args.ne11 + row; ++ scores[score_row] = mqk[tx]*args.scale; ++ if (stream == 0 && head == 0) { ++ debug_scores[row] = scores[score_row]; + } + } + } + +-[[host_name("kernel_mul_mv_iq4_nl_f32")]] +-kernel void kernel_mul_mv_iq4_nl_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem [[threadgroup(0)]], +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { +- +- kernel_mul_mv_iq4_nl_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +-} ++kernel void kernel_glm_compact_scores_v_f16_dv512_nwg4( ++ constant ggml_metal_kargs_flash_attn_ext_vec & args, ++ device const char * v, ++ device const float * scores, ++ device char * dst, ++ threadgroup half * shmem_f16 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]]) { ++ constexpr short DV = 512; ++ constexpr short C = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ constexpr short NE = 2; ++ constexpr short NL = N_SIMDWIDTH/NE; ++ constexpr short DV4 = DV/4; ++ constexpr short NWG = 4; + +-template +-void kernel_mul_mv_iq4_xs_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++ const short iwg = tgpig.z%NWG; ++ const short stream = tgpig.z/NWG; ++ const short head = tgpig.y; ++ const short query = tgpig.x; + +- threadgroup float * shmem_f32 = (threadgroup float *) shmem; ++ const int v_stream = stream%args.ne_12_3; ++ v += v_stream*args.nb23; ++ scores += (int64_t(stream)*args.ne02 + head)*args.ne11; + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; +- const int first_row = (r0 * NSG + sgitg) * NR0; ++ threadgroup float * ss = (threadgroup float *) shmem_f16; ++ threadgroup float4 * so4 = (threadgroup float4 *) (ss + C); ++ so4 += tiisg; + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++ const short tx = tiisg%NL; ++ const short ty = tiisg/NL; + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ if (ty == 0) { ++ FOR_UNROLL (short i = 0; i < DV4/NL; ++i) { ++ so4[i*NL] = float4(0.0f); ++ } ++ } ++ for (short i = tiisg; i < C; i += N_SIMDWIDTH) { ++ ss[i] = 0.0f; ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); + +- device const block_iq4_xs * x = (device const block_iq4_xs *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++ float S = 0.0f; ++ float M = -FLT_MAX/2; + +- const int nb = args.ne00/QK_K; +- const int ns01 = args.nb01/args.nb00; ++ for (int ic0 = iwg; ; ic0 += NWG) { ++ const int ic = ic0*C; ++ if (ic >= args.ne11) { ++ break; ++ } + +- const short ix = tiisg/16; // 0 or 1 +- const short it = tiisg%16; // 0...15 +- const short ib = it/2; +- const short il = it%2; ++ const int row = ic + NE*tx + ty; ++ ss[NE*tx + ty] = row < args.ne11 ? scores[row] : -MAXHALF; ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ ++ const float m = M; ++ const float s = ss[tiisg]; ++ M = simd_max(max(M, s)); ++ const float ms = exp(m - M); ++ const float vs = exp(s - M); ++ S = S*ms + simd_sum(vs); ++ ss[tiisg] = vs; ++ ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] *= ms; ++ } ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); + +- shmem_f32[tiisg] = kvalues_iq4nl_f[tiisg%16]; +- threadgroup_barrier(mem_flags::mem_threadgroup); ++ float4 lo[DV4/NL]; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii] = float4(0.0f); ++ } + +- float4 yl[4]; +- float sumf[NR0]={0.f}; ++ device const half4 * pv4 = (device const half4 *) (v + ic*args.nb21); ++ pv4 += ty*args.ns20/4 + tx; ++ const threadgroup float * sst = ss + ty; ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii] += float4(pv4[cc*NE*args.ns20/4 + ii*NL])*float4(sst[cc*NE]); ++ } ++ } + +- device const float * yb = y + ix * QK_K + ib * 32 + il * 8; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 16); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 16); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 16); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 16); ++ } + +- uint32_t aux32[2]; +- thread const uint8_t * q8 = (thread const uint8_t *)aux32; ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] += lo[ii]; ++ } ++ } ++ } + +- float4 qf1, qf2; ++ if (tiisg == 0) { ++ ss[0] = S; ++ ss[1] = M; ++ } ++ so4 -= tiisg; ++ simdgroup_barrier(mem_flags::mem_threadgroup); + +- // [TAG_MUL_MV_WEIRD] +- for (int ibl = ix; ibl < nb && ibl < ns01; ibl += 2) { +- device const float4 * y4 = (device const float4 *)yb; +- yl[0] = y4[0]; +- yl[1] = y4[4]; +- yl[2] = y4[1]; +- yl[3] = y4[5]; ++ const int64_t nrows = int64_t(args.ne1)*args.ne2*args.ne3; ++ const int64_t rid = int64_t(stream)*args.ne02*args.ne01 + head + query*args.ne01; ++ device float4 * dst4 = (device float4 *) dst; ++ device float * dst1 = (device float *) dst + nrows*DV*NWG; + +- for (short row = 0; row < NR0; ++row) { +- device const block_iq4_xs & xb = x[row*ns01 + ibl]; +- device const uint32_t * q4 = (device const uint32_t *)(xb.qs + 16*ib + 8*il); ++ for (short i = tiisg; i < DV4; i += N_SIMDWIDTH) { ++ dst4[rid*DV4*NWG + NWG*i + iwg] = so4[i]; ++ } ++ if (tiisg == 0) { ++ dst1[rid*(2*NWG) + 2*iwg + 0] = S; ++ dst1[rid*(2*NWG) + 2*iwg + 1] = M; ++ } ++} + +- float4 acc1 = {0.f}, acc2 = {0.f}; ++kernel void kernel_glm_compact_softmax_prefix_f32_nwg4( ++ constant ggml_metal_kargs_flash_attn_ext_vec & args, ++ device float * scores_probs, ++ device float * chunk_ms, ++ device char * partials, ++ threadgroup float * ss [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]]) { ++ constexpr short C = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ constexpr short NE = 2; ++ constexpr short NL = N_SIMDWIDTH/NE; ++ constexpr short NWG = 4; ++ constexpr short DV = 512; ++ ++ const short iwg = tgpig.x; ++ const short head = tgpig.y; ++ const short stream = tgpig.z; ++ const int chunk_count = (args.ne11 + C - 1)/C; ++ const int row_in_chunk = NE*(tiisg%NL) + tiisg/NL; ++ const int64_t head_base = (int64_t(stream)*args.ne02 + head)*args.ne11; ++ const int64_t chunk_base = (int64_t(stream)*args.ne02 + head)*chunk_count; ++ ++ float S = 0.0f; ++ float M = -FLT_MAX/2; ++ FOR_UNROLL (short chunk = iwg; chunk < 64; chunk += NWG) { ++ if (chunk >= chunk_count) { ++ break; ++ } + +- aux32[0] = (q4[0] ) & 0x0f0f0f0f; +- aux32[1] = (q4[0] >> 4) & 0x0f0f0f0f; +- qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]}; +- qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]}; +- acc1 += yl[0] * qf1; +- acc2 += yl[1] * qf2; ++ const int row = chunk*C + row_in_chunk; ++ ss[tiisg] = row < args.ne11 ? scores_probs[head_base + row] : -MAXHALF; ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ ++ const float m = M; ++ const float s = ss[tiisg]; ++ M = simd_max(max(M, s)); ++ const float ms = exp(m - M); ++ const float vs = exp(s - M); ++ const float vs_sum = simd_sum(vs); ++ ss[C + tiisg] = S*ms; ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ S = ss[C + tiisg] + vs_sum; ++ ss[tiisg] = vs; ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (row < args.ne11) { ++ scores_probs[head_base + row] = ss[tiisg]; ++ } ++ if (tiisg == 0) { ++ chunk_ms[chunk_base + chunk] = ms; ++ } ++ } + +- aux32[0] = (q4[1] ) & 0x0f0f0f0f; +- aux32[1] = (q4[1] >> 4) & 0x0f0f0f0f; +- qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]}; +- qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]}; +- acc1 += yl[2] * qf1; +- acc2 += yl[3] * qf2; ++ if (tiisg == 0) { ++ const int64_t nrows = int64_t(args.ne1)*args.ne2*args.ne3; ++ const int64_t rid = int64_t(stream)*args.ne02*args.ne01 + head; ++ device float * dst1 = (device float *) partials + nrows*DV*NWG; ++ dst1[rid*(2*NWG) + 2*iwg + 0] = S; ++ dst1[rid*(2*NWG) + 2*iwg + 1] = M; ++ } ++} + +- acc1 += acc2; ++kernel void kernel_glm_compact_chunk_v_f16_dv512( ++ constant ggml_metal_kargs_flash_attn_ext_vec & args, ++ device const char * v, ++ device const float * probs, ++ device float4 * chunk_v, ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]]) { ++ constexpr short DV = 512; ++ constexpr short C = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ constexpr short NE = 2; ++ constexpr short NL = N_SIMDWIDTH/NE; ++ constexpr short DV4 = DV/4; + +- const int ls = (((xb.scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((xb.scales_h >> 2*ib) & 3) << 4)) - 32; +- sumf[row] += (float)xb.d * ls * (acc1[0] + acc1[1] + acc1[2] + acc1[3]); ++ const int chunk = tgpig.x; ++ const int head = tgpig.y; ++ const int stream = tgpig.z; ++ const int chunk_count = (args.ne11 + C - 1)/C; ++ const short tx = tiisg%NL; ++ const short ty = tiisg/NL; ++ const int v_stream = stream%args.ne_12_3; ++ v += v_stream*args.nb23 + chunk*C*args.nb21; ++ probs += (int64_t(stream)*args.ne02 + head)*args.ne11 + chunk*C; ++ ++ float4 lo[DV4/NL]; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii] = float4(0.0f); ++ } ++ ++ device const half4 * pv4 = (device const half4 *) v; ++ pv4 += ty*args.ns20/4 + tx; ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ const float p = probs[NE*cc + ty]; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii] += float4(pv4[cc*NE*args.ns20/4 + ii*NL])*float4(p); + } ++ } + +- yb += 2 * QK_K; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 16); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 16); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 16); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 16); + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ if (ty == 0) { ++ const int64_t chunk_base = ++ ((int64_t(stream)*args.ne02 + head)*chunk_count + chunk)*DV4; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ chunk_v[chunk_base + ii*NL + tx] = lo[ii]; ++ } ++ } ++} + +- for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all; ++kernel void kernel_glm_compact_chunk_fold_f32_dv512_nwg4( ++ constant ggml_metal_kargs_flash_attn_ext_vec & args, ++ device const float * chunk_ms, ++ device const float4 * chunk_v, ++ device char * partials, ++ threadgroup float4 * shmem_f32x4 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]]) { ++ constexpr short DV = 512; ++ constexpr short DV4 = DV/4; ++ constexpr short C = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ constexpr short NWG = 4; ++ ++ const short iwg = tgpig.x; ++ const short head = tgpig.y; ++ const short stream = tgpig.z; ++ const int chunk_count = (args.ne11 + C - 1)/C; ++ const int64_t chunk_base = (int64_t(stream)*args.ne02 + head)*chunk_count; ++ const int64_t rid = int64_t(stream)*args.ne02*args.ne01 + head; ++ device float4 * dst4 = (device float4 *) partials; ++ threadgroup float4 * so4 = shmem_f32x4 + tiisg; ++ ++ FOR_UNROLL (short ii = 0; ii < DV4/N_SIMDWIDTH; ++ii) { ++ so4[ii*N_SIMDWIDTH] = float4(0.0f); ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ ++ FOR_UNROLL (short chunk = iwg; chunk < 64; chunk += NWG) { ++ if (chunk >= chunk_count) { ++ break; ++ } ++ const float ms = chunk_ms[chunk_base + chunk]; ++ FOR_UNROLL (short ii = 0; ii < DV4/N_SIMDWIDTH; ++ii) { ++ so4[ii*N_SIMDWIDTH] *= ms; + } ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ FOR_UNROLL (short ii = 0; ii < DV4/N_SIMDWIDTH; ++ii) { ++ const short i = ii*N_SIMDWIDTH + tiisg; ++ so4[ii*N_SIMDWIDTH] += chunk_v[(chunk_base + chunk)*DV4 + i]; ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ FOR_UNROLL (short ii = 0; ii < DV4/N_SIMDWIDTH; ++ii) { ++ const short i = ii*N_SIMDWIDTH + tiisg; ++ dst4[rid*DV4*NWG + NWG*i + iwg] = so4[ii*N_SIMDWIDTH]; + } + } + +-[[host_name("kernel_mul_mv_iq4_xs_f32")]] +-kernel void kernel_mul_mv_iq4_xs_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem [[threadgroup(0)]], +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++kernel void kernel_glm_compact_probs_v_sequential_f16_dv512_nwg4( ++ constant ggml_metal_kargs_flash_attn_ext_vec & args, ++ device const char * v, ++ device const float * probs, ++ device const float * chunk_ms, ++ device char * partials, ++ threadgroup float4 * shmem_f32x4 [[threadgroup(0)]], ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ ushort tiisg [[thread_index_in_simdgroup]]) { ++ constexpr short DV = 512; ++ constexpr short C = OP_FLASH_ATTN_EXT_VEC_NCPSG; ++ constexpr short NE = 2; ++ constexpr short NL = N_SIMDWIDTH/NE; ++ constexpr short DV4 = DV/4; ++ constexpr short NWG = 4; ++ ++ const short iwg = tgpig.x; ++ const short head = tgpig.y; ++ const short stream = tgpig.z; ++ const short tx = tiisg%NL; ++ const short ty = tiisg/NL; ++ const int chunk_count = (args.ne11 + C - 1)/C; ++ const int64_t chunk_base = (int64_t(stream)*args.ne02 + head)*chunk_count; ++ const int64_t score_base = (int64_t(stream)*args.ne02 + head)*args.ne11; ++ const int v_stream = stream%args.ne_12_3; ++ v += v_stream*args.nb23; ++ ++ threadgroup float4 * so4 = shmem_f32x4 + tiisg; ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] = float4(0.0f); ++ } ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); + +- kernel_mul_mv_iq4_xs_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); ++ FOR_UNROLL (short chunk = iwg; chunk < 64; chunk += NWG) { ++ if (chunk >= chunk_count) { ++ break; ++ } ++ const float ms = chunk_ms[chunk_base + chunk]; ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] *= ms; ++ } ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ ++ float4 lo[DV4/NL]; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii] = float4(0.0f); ++ } ++ device const half4 * pv4 = ++ (device const half4 *) (v + chunk*C*args.nb21); ++ pv4 += ty*args.ns20/4 + tx; ++ device const float * chunk_probs = probs + score_base + chunk*C; ++ FOR_UNROLL (short cc = 0; cc < C/NE; ++cc) { ++ const float p = chunk_probs[NE*cc + ty]; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii] += float4(pv4[cc*NE*args.ns20/4 + ii*NL])*float4(p); ++ } ++ } ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ lo[ii][0] += simd_shuffle_down(lo[ii][0], 16); ++ lo[ii][1] += simd_shuffle_down(lo[ii][1], 16); ++ lo[ii][2] += simd_shuffle_down(lo[ii][2], 16); ++ lo[ii][3] += simd_shuffle_down(lo[ii][3], 16); ++ } ++ if (ty == 0) { ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ so4[ii*NL] += lo[ii]; ++ } ++ } ++ simdgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ if (ty == 0) { ++ const int64_t rid = int64_t(stream)*args.ne02*args.ne01 + head; ++ device float4 * dst4 = (device float4 *) partials; ++ FOR_UNROLL (short ii = 0; ii < DV4/NL; ++ii) { ++ const short i = ii*NL + tx; ++ dst4[rid*DV4*NWG + NWG*i + iwg] = so4[ii*NL]; ++ } ++ } + } + +-template +-void kernel_mul_mv_mxfp4_f32_impl( +- args_t args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem, +- uint3 tgpig, +- ushort tiisg, +- ushort sgitg) { +- const short NSG = FC_mul_mv_nsg; ++template ++kernel void kernel_dsa_sparse_attn_impl( ++ constant ggml_metal_kargs_dsa_sparse_attn & args, ++ device const char * q, ++ device const char * k, ++ device const char * v, ++ device const char * kq_mask, ++ device const int32_t * top_k, ++ device char * dst, ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ uint tiitg [[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ constexpr int32_t DSA_SPARSE_ATTN_MAX_TOP_K = 4096; ++ constexpr int32_t DSA_SPARSE_ATTN_MAX_THREADS = 256; ++ ++ threadgroup float scores[DSA_SPARSE_ATTN_MAX_TOP_K]; ++ threadgroup float reduce[DSA_SPARSE_ATTN_MAX_THREADS]; ++ ++ const int32_t i_batch = tgpig.x; ++ const int32_t i_head = tgpig.y; ++ const int32_t i_stream = tgpig.z; ++ const int32_t tid = tiitg; ++ const int32_t nth = tptg.x; ++ ++ if (i_batch >= args.ne1 || i_head >= args.ne2 || i_stream >= args.ne3 || args.ne40 > DSA_SPARSE_ATTN_MAX_TOP_K) { ++ return; ++ } + +- threadgroup float * shmem_f32 = (threadgroup float *) shmem; ++ const int32_t n_head_per_kv = args.ne2/args.ne12; ++ const int32_t n_head_per_v = args.ne2/args.ne22; ++ const int32_t i_kv_head = i_head/n_head_per_kv; ++ const int32_t i_v_head = i_head/n_head_per_v; ++ const int32_t i_top_stream = i_stream%args.ne42; ++ ++ float local_max = -FLT_MAX; ++ float local_active_top_end = 0.0f; ++ ++ for (int32_t i_top = tid; i_top < args.ne40; i_top += nth) { ++ const int32_t i_kv = ((device const int32_t *) ((device const char *) top_k + ++ i_top*args.nb40 + i_batch*args.nb41 + i_top_stream*args.nb42))[0]; ++ ++ float score = -INFINITY; ++ if (i_kv >= 0 && i_kv < args.ne11) { ++ device const M * mask_ptr = (device const M *) (kq_mask + ++ i_kv*args.nb31 + i_batch*args.nb32 + i_stream*args.nb33); ++ const float mask = float(*mask_ptr); ++ ++ if (isfinite(mask)) { ++ local_active_top_end = max(local_active_top_end, float(i_top + 1)); ++ ++ float qk = 0.0f; ++ ++ int32_t i_dk = 0; ++ for (; i_dk + 3 < args.ne00; i_dk += 4) { ++ device const float * q0_ptr = (device const float *) (q + ++ (i_dk + 0)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q1_ptr = (device const float *) (q + ++ (i_dk + 1)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q2_ptr = (device const float *) (q + ++ (i_dk + 2)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q3_ptr = (device const float *) (q + ++ (i_dk + 3)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const K * k0_ptr = (device const K *) (k + ++ (i_dk + 0)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k1_ptr = (device const K *) (k + ++ (i_dk + 1)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k2_ptr = (device const K *) (k + ++ (i_dk + 2)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k3_ptr = (device const K *) (k + ++ (i_dk + 3)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ qk += (*q0_ptr) * float(*k0_ptr); ++ qk += (*q1_ptr) * float(*k1_ptr); ++ qk += (*q2_ptr) * float(*k2_ptr); ++ qk += (*q3_ptr) * float(*k3_ptr); ++ } ++ for (; i_dk < args.ne00; ++i_dk) { ++ device const float * q_ptr = (device const float *) (q + ++ i_dk*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const K * k_ptr = (device const K *) (k + ++ i_dk*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ qk += (*q_ptr) * float(*k_ptr); ++ } + +- const int r0 = tgpig.x; +- const int r1 = tgpig.y; +- const int im = tgpig.z; ++ score = qk*args.scale + mask; ++ } ++ } + +- const int first_row = (r0 * NSG + sgitg) * NR0; ++ scores[i_top] = score; ++ local_max = max(local_max, score); ++ } + +- const uint i12 = im%FC_mul_mv_ne12; +- const uint i13 = im/FC_mul_mv_ne12; ++ (void) local_max; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03; +- const uint64_t offset1 = r1*args.nb11 + (i12 )*args.nb12 + (i13 )*args.nb13; ++ if (tid < 32) { ++ float scan_max = -FLT_MAX; ++ float scan_active_top_end = 0.0f; ++ for (int32_t i_top = tid; i_top < args.ne40; i_top += 32) { ++ const float score = scores[i_top]; ++ scan_max = max(scan_max, score); ++ if (isfinite(score)) { ++ scan_active_top_end = max(scan_active_top_end, float(i_top + 1)); ++ } ++ } ++ reduce[tid] = scan_max; ++ reduce[32 + tid] = scan_active_top_end; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- device const block_mxfp4 * x = (device const block_mxfp4 *) (src0 + offset0); +- device const float * y = (device const float *) (src1 + offset1); ++ for (int32_t stride = 16; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ reduce[tid] = max(reduce[tid], reduce[tid + stride]); ++ reduce[32 + tid] = max(reduce[32 + tid], reduce[32 + tid + stride]); ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +- const int nb = args.ne00/QK_MXFP4; +- const int ns01 = args.nb01/args.nb00; // this can be larger than nb for permuted src0 tensors ++ const float max_score = reduce[0]; ++ const int32_t active_top_end = int32_t(reduce[32]); ++ float local_sum = 0.0f; + +- const short ix = tiisg/2; // 0...15 +- const short it = tiisg%2; // 0 or 1 ++ if (max_score > -FLT_MAX/2) { ++ for (int32_t i_top = tid; i_top < active_top_end; i_top += nth) { ++ scores[i_top] = exp(scores[i_top] - max_score); ++ local_sum += scores[i_top]; ++ } ++ } else { ++ for (int32_t i_top = tid; i_top < args.ne40; i_top += nth) { ++ scores[i_top] = 0.0f; ++ } ++ } + +- shmem_f32[tiisg] = kvalues_mxfp4_f[tiisg%16]; ++ (void) local_sum; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ if (tid < 32) { ++ float scan_sum = 0.0f; ++ for (int32_t i_top = tid; i_top < active_top_end; i_top += 32) { ++ scan_sum += scores[i_top]; ++ } ++ reduce[tid] = scan_sum; ++ } + threadgroup_barrier(mem_flags::mem_threadgroup); + +- float4 yl[4]; +- float sumf[NR0]={0.f}; ++ for (int32_t stride = 16; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ reduce[tid] += reduce[tid + stride]; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ const float sum_score = reduce[0]; ++ ++ for (int32_t i_dv = tid; i_dv < args.ne20; i_dv += nth) { ++ float acc = 0.0f; ++ ++ if (sum_score > 0.0f && isfinite(sum_score)) { ++ for (int32_t i_top = 0; i_top < active_top_end; ++i_top) { ++ const int32_t i_kv = ((device const int32_t *) ((device const char *) top_k + ++ i_top*args.nb40 + i_batch*args.nb41 + i_top_stream*args.nb42))[0]; ++ if (i_kv < 0 || i_kv >= args.ne11) { ++ continue; ++ } ++ ++ const float p = scores[i_top]/sum_score; ++ if (p == 0.0f) { ++ continue; ++ } ++ ++ device const V * v_ptr = (device const V *) (v + ++ i_dv*args.nb20 + i_kv*args.nb21 + i_v_head*args.nb22 + i_stream*args.nb23); ++ acc += p * float(*v_ptr); ++ } ++ } + +- device const float * yb = y + ix*QK_MXFP4 + it*8; ++ device float * dst_ptr = (device float *) (dst + ++ i_dv*args.nb0 + i_batch*args.nb1 + i_head*args.nb2 + i_stream*args.nb3); ++ *dst_ptr = acc; ++ } ++} + +- // note: just the check `ib < nb` is enough, but adding the redundant `&& ib < ns01` check makes the kernel a bit faster +- // no idea why that is - needs some deeper investigation [TAG_MUL_MV_WEIRD] +- for (int ib = ix; ib < nb && ib < ns01; ib += 16) { +- device const float4 * y4 = (device const float4 *) yb; ++template ++kernel void kernel_dsa_sparse_attn_cached_topk_impl( ++ constant ggml_metal_kargs_dsa_sparse_attn & args, ++ device const char * q, ++ device const char * k, ++ device const char * v, ++ device const char * kq_mask, ++ device const int32_t * top_k, ++ device char * dst, ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ uint tiitg [[thread_index_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ constexpr int32_t DSA_SPARSE_ATTN_CACHED_TOPK_MAX_TOP_K = 1024; ++ constexpr int32_t DSA_SPARSE_ATTN_MAX_THREADS = 256; ++ ++ threadgroup float scores[DSA_SPARSE_ATTN_CACHED_TOPK_MAX_TOP_K]; ++ threadgroup int32_t top_indices[DSA_SPARSE_ATTN_CACHED_TOPK_MAX_TOP_K]; ++ threadgroup float reduce[DSA_SPARSE_ATTN_MAX_THREADS]; ++ ++ const int32_t i_batch = tgpig.x; ++ const int32_t i_head = tgpig.y; ++ const int32_t i_stream = tgpig.z; ++ const int32_t tid = tiitg; ++ const int32_t nth = tptg.x; ++ ++ if (i_batch >= args.ne1 || i_head >= args.ne2 || i_stream >= args.ne3 || args.ne40 > DSA_SPARSE_ATTN_CACHED_TOPK_MAX_TOP_K) { ++ return; ++ } + +- yl[0] = y4[0]; +- yl[1] = y4[4]; +- yl[2] = y4[1]; +- yl[3] = y4[5]; ++ const int32_t n_head_per_kv = args.ne2/args.ne12; ++ const int32_t n_head_per_v = args.ne2/args.ne22; ++ const int32_t i_kv_head = i_head/n_head_per_kv; ++ const int32_t i_v_head = i_head/n_head_per_v; ++ const int32_t i_top_stream = i_stream%args.ne42; ++ ++ float local_max = -FLT_MAX; ++ float local_active_top_end = 0.0f; ++ ++ for (int32_t i_top = tid; i_top < args.ne40; i_top += nth) { ++ const int32_t i_kv = ((device const int32_t *) ((device const char *) top_k + ++ i_top*args.nb40 + i_batch*args.nb41 + i_top_stream*args.nb42))[0]; ++ top_indices[i_top] = i_kv; ++ ++ float score = -INFINITY; ++ if (i_kv >= 0 && i_kv < args.ne11) { ++ device const M * mask_ptr = (device const M *) (kq_mask + ++ i_kv*args.nb31 + i_batch*args.nb32 + i_stream*args.nb33); ++ const float mask = float(*mask_ptr); ++ ++ if (isfinite(mask)) { ++ local_active_top_end = max(local_active_top_end, float(i_top + 1)); ++ ++ float qk = 0.0f; ++ ++ int32_t i_dk = 0; ++ for (; i_dk + 3 < args.ne00; i_dk += 4) { ++ device const float * q0_ptr = (device const float *) (q + ++ (i_dk + 0)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q1_ptr = (device const float *) (q + ++ (i_dk + 1)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q2_ptr = (device const float *) (q + ++ (i_dk + 2)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q3_ptr = (device const float *) (q + ++ (i_dk + 3)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const K * k0_ptr = (device const K *) (k + ++ (i_dk + 0)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k1_ptr = (device const K *) (k + ++ (i_dk + 1)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k2_ptr = (device const K *) (k + ++ (i_dk + 2)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k3_ptr = (device const K *) (k + ++ (i_dk + 3)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ qk += (*q0_ptr) * float(*k0_ptr); ++ qk += (*q1_ptr) * float(*k1_ptr); ++ qk += (*q2_ptr) * float(*k2_ptr); ++ qk += (*q3_ptr) * float(*k3_ptr); ++ } ++ for (; i_dk < args.ne00; ++i_dk) { ++ device const float * q_ptr = (device const float *) (q + ++ i_dk*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const K * k_ptr = (device const K *) (k + ++ i_dk*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ qk += (*q_ptr) * float(*k_ptr); ++ } + +- FOR_UNROLL (short row = 0; row < NR0; row++) { +- device const block_mxfp4 & xb = x[row*ns01 + ib]; +- device const uint8_t * q2 = (device const uint8_t *)(xb.qs + 8*it); ++ score = qk*args.scale + mask; ++ } ++ } + +- float4 acc1 = yl[0]*float4(shmem_f32[q2[0] & 0x0F], shmem_f32[q2[1] & 0x0F], shmem_f32[q2[2] & 0x0F], shmem_f32[q2[3] & 0x0F]); +- float4 acc2 = yl[1]*float4(shmem_f32[q2[0] >> 4 ], shmem_f32[q2[1] >> 4 ], shmem_f32[q2[2] >> 4 ], shmem_f32[q2[3] >> 4 ]); +- float4 acc3 = yl[2]*float4(shmem_f32[q2[4] & 0x0F], shmem_f32[q2[5] & 0x0F], shmem_f32[q2[6] & 0x0F], shmem_f32[q2[7] & 0x0F]); +- float4 acc4 = yl[3]*float4(shmem_f32[q2[4] >> 4 ], shmem_f32[q2[5] >> 4 ], shmem_f32[q2[6] >> 4 ], shmem_f32[q2[7] >> 4 ]); ++ scores[i_top] = score; ++ local_max = max(local_max, score); ++ } + +- acc1 = (acc1 + acc3) + (acc2 + acc4); ++ reduce[tid] = local_max; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- sumf[row] += e8m0_to_fp32(xb.e) * ((acc1[0] + acc1[1]) + (acc1[2] + acc1[3])); ++ for (int32_t stride = nth/2; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ reduce[tid] = max(reduce[tid], reduce[tid + stride]); + } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +- yb += 16 * QK_MXFP4; ++ const float max_score = reduce[0]; ++ reduce[tid] = local_active_top_end; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ for (int32_t stride = nth/2; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ reduce[tid] = max(reduce[tid], reduce[tid + stride]); ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); + } + +- device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0; ++ const int32_t active_top_end = int32_t(reduce[0]); ++ float local_sum = 0.0f; + +- for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) { +- float sum_all = simd_sum(sumf[row]); +- if (tiisg == 0) { +- dst_f32[first_row + row] = sum_all; ++ if (max_score > -FLT_MAX/2) { ++ for (int32_t i_top = tid; i_top < active_top_end; i_top += nth) { ++ scores[i_top] = exp(scores[i_top] - max_score); ++ local_sum += scores[i_top]; ++ } ++ } else { ++ for (int32_t i_top = tid; i_top < args.ne40; i_top += nth) { ++ scores[i_top] = 0.0f; + } + } +-} + +-[[host_name("kernel_mul_mv_mxfp4_f32")]] +-kernel void kernel_mul_mv_mxfp4_f32( +- constant ggml_metal_kargs_mul_mv & args, +- device const char * src0, +- device const char * src1, +- device char * dst, +- threadgroup char * shmem [[threadgroup(0)]], +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiisg[[thread_index_in_simdgroup]], +- ushort sgitg[[simdgroup_index_in_threadgroup]]) { ++ reduce[tid] = local_sum; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- kernel_mul_mv_mxfp4_f32_impl(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); +-} ++ for (int32_t stride = nth/2; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ reduce[tid] += reduce[tid + stride]; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +-template +-kernel void kernel_get_rows_q( +- constant ggml_metal_kargs_get_rows & args, +- device const void * src0, +- device const void * src1, +- device void * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiitg[[thread_index_in_threadgroup]], +- ushort3 ntg [[threads_per_threadgroup]]) { +- const int32_t iw0 = tgpig.x/args.ne10; +- const int32_t i10 = tgpig.x%args.ne10; +- const int32_t i11 = tgpig.y; +- const int32_t i12 = tgpig.z; ++ const float sum_score = reduce[0]; + +- const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0]; ++ for (int32_t i_dv = tid; i_dv < args.ne20; i_dv += nth) { ++ float acc = 0.0f; + +- const int32_t i02 = i11; +- const int32_t i03 = i12; ++ if (sum_score > 0.0f && isfinite(sum_score)) { ++ for (int32_t i_top = 0; i_top < active_top_end; ++i_top) { ++ const int32_t i_kv = top_indices[i_top]; ++ if (i_kv < 0 || i_kv >= args.ne11) { ++ continue; ++ } + +- auto psrc = (device const block_q *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01); +- auto pdst = (device float4x4 *) (( device char *) dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1); ++ const float p = scores[i_top]/sum_score; ++ if (p == 0.0f) { ++ continue; ++ } + +- for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) { +- float4x4 temp; +- dequantize_func(psrc + ind/nl, ind%nl, temp); +- pdst[ind] = temp; ++ device const V * v_ptr = (device const V *) (v + ++ i_dv*args.nb20 + i_kv*args.nb21 + i_v_head*args.nb22 + i_stream*args.nb23); ++ acc += p * float(*v_ptr); ++ } ++ } + +- break; ++ device float * dst_ptr = (device float *) (dst + ++ i_dv*args.nb0 + i_batch*args.nb1 + i_head*args.nb2 + i_stream*args.nb3); ++ *dst_ptr = acc; + } + } + +-template +-kernel void kernel_get_rows_f( +- constant ggml_metal_kargs_get_rows & args, +- device const void * src0, +- device const void * src1, +- device void * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- ushort tiitg[[thread_index_in_threadgroup]], +- ushort3 ntg [[threads_per_threadgroup]]) { +- const int32_t iw0 = tgpig.x/args.ne10; +- const int32_t i10 = tgpig.x%args.ne10; +- const int32_t i11 = tgpig.y; +- const int32_t i12 = tgpig.z; ++typedef decltype(kernel_dsa_sparse_attn_impl) kernel_dsa_sparse_attn_t; ++typedef decltype(kernel_dsa_sparse_attn_cached_topk_impl) kernel_dsa_sparse_attn_cached_topk_t; ++ ++template [[host_name("kernel_dsa_sparse_attn_f32_f32_f32")]] kernel kernel_dsa_sparse_attn_t kernel_dsa_sparse_attn_impl; ++template [[host_name("kernel_dsa_sparse_attn_f32_f32_f16")]] kernel kernel_dsa_sparse_attn_t kernel_dsa_sparse_attn_impl; ++template [[host_name("kernel_dsa_sparse_attn_f32_f16_f32")]] kernel kernel_dsa_sparse_attn_t kernel_dsa_sparse_attn_impl; ++template [[host_name("kernel_dsa_sparse_attn_f32_f16_f16")]] kernel kernel_dsa_sparse_attn_t kernel_dsa_sparse_attn_impl; ++template [[host_name("kernel_dsa_sparse_attn_f16_f32_f32")]] kernel kernel_dsa_sparse_attn_t kernel_dsa_sparse_attn_impl; ++template [[host_name("kernel_dsa_sparse_attn_f16_f32_f16")]] kernel kernel_dsa_sparse_attn_t kernel_dsa_sparse_attn_impl; ++template [[host_name("kernel_dsa_sparse_attn_f16_f16_f32")]] kernel kernel_dsa_sparse_attn_t kernel_dsa_sparse_attn_impl; ++template [[host_name("kernel_dsa_sparse_attn_f16_f16_f16")]] kernel kernel_dsa_sparse_attn_t kernel_dsa_sparse_attn_impl; ++ ++template [[host_name("kernel_dsa_sparse_attn_cached_topk_f32_f32_f32")]] kernel kernel_dsa_sparse_attn_cached_topk_t kernel_dsa_sparse_attn_cached_topk_impl; ++template [[host_name("kernel_dsa_sparse_attn_cached_topk_f32_f32_f16")]] kernel kernel_dsa_sparse_attn_cached_topk_t kernel_dsa_sparse_attn_cached_topk_impl; ++template [[host_name("kernel_dsa_sparse_attn_cached_topk_f32_f16_f32")]] kernel kernel_dsa_sparse_attn_cached_topk_t kernel_dsa_sparse_attn_cached_topk_impl; ++template [[host_name("kernel_dsa_sparse_attn_cached_topk_f32_f16_f16")]] kernel kernel_dsa_sparse_attn_cached_topk_t kernel_dsa_sparse_attn_cached_topk_impl; ++template [[host_name("kernel_dsa_sparse_attn_cached_topk_f16_f32_f32")]] kernel kernel_dsa_sparse_attn_cached_topk_t kernel_dsa_sparse_attn_cached_topk_impl; ++template [[host_name("kernel_dsa_sparse_attn_cached_topk_f16_f32_f16")]] kernel kernel_dsa_sparse_attn_cached_topk_t kernel_dsa_sparse_attn_cached_topk_impl; ++template [[host_name("kernel_dsa_sparse_attn_cached_topk_f16_f16_f32")]] kernel kernel_dsa_sparse_attn_cached_topk_t kernel_dsa_sparse_attn_cached_topk_impl; ++template [[host_name("kernel_dsa_sparse_attn_cached_topk_f16_f16_f16")]] kernel kernel_dsa_sparse_attn_cached_topk_t kernel_dsa_sparse_attn_cached_topk_impl; ++ ++template ++kernel void kernel_dsa_sparse_attn_decode_grouped_impl( ++ constant ggml_metal_kargs_dsa_sparse_attn & args, ++ device const char * q, ++ device const char * k, ++ device const char * v, ++ device const char * kq_mask, ++ device const int32_t * top_k, ++ device char * dst, ++ uint3 tgpig [[threadgroup_position_in_grid]], ++ uint3 tpitg [[thread_position_in_threadgroup]], ++ uint3 tptg [[threads_per_threadgroup]]) { ++ constexpr int32_t DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_TOP_K = 1024; ++ constexpr int32_t DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_HEADS = 4; ++ constexpr int32_t DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_THREADS = 256; ++ ++ threadgroup int32_t top_indices[DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_TOP_K]; ++ threadgroup float scores[ ++ DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_HEADS* ++ DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_TOP_K ++ ]; ++ threadgroup float reduce[ ++ DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_HEADS* ++ DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_THREADS ++ ]; ++ ++ const int32_t i_batch = tgpig.x; ++ const int32_t i_stream = tgpig.z; ++ const int32_t tid = tpitg.x; ++ const int32_t head_lane = tpitg.y; ++ const int32_t nth = tptg.x; ++ const int32_t n_heads_per_group = tptg.y; ++ const int32_t i_head = tgpig.y*n_heads_per_group + head_lane; ++ ++ if (i_batch >= args.ne1 || ++ i_stream >= args.ne3 || ++ args.ne40 > DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_TOP_K || ++ n_heads_per_group > DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_HEADS || ++ nth > DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_THREADS) { ++ return; ++ } + +- const int32_t r = ((const device int32_t *) ((const device char *) src1 + i12*args.nb12 + i11*args.nb11 + i10*args.nb10))[0]; ++ const bool valid_head = i_head < args.ne2; ++ const int32_t n_head_per_kv = args.ne2/args.ne12; ++ const int32_t n_head_per_v = args.ne2/args.ne22; ++ const int32_t i_kv_head = valid_head ? i_head/n_head_per_kv : 0; ++ const int32_t i_v_head = valid_head ? i_head/n_head_per_v : 0; ++ const int32_t i_top_stream = i_stream%args.ne42; ++ const int32_t score_offset = head_lane*DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_TOP_K; ++ const int32_t reduce_offset = head_lane*DSA_SPARSE_ATTN_DECODE_GROUPED_MAX_THREADS; + +- const int32_t i02 = i11; +- const int32_t i03 = i12; ++ if (head_lane == 0) { ++ for (int32_t i_top = tid; i_top < args.ne40; i_top += nth) { ++ top_indices[i_top] = ((device const int32_t *) ((device const char *) top_k + ++ i_top*args.nb40 + i_batch*args.nb41 + i_top_stream*args.nb42))[0]; ++ } ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +- auto psrc = (const device T0 *) ((const device char *) src0 + i03*args.nb03 + i02*args.nb02 + r*args.nb01); +- auto pdst = ( device T *) (( device char *) dst + i12*args.nb3 + i11*args.nb2 + i10*args.nb1); ++ float local_max = -FLT_MAX; ++ float local_active_top_end = 0.0f; ++ ++ for (int32_t i_top = tid; i_top < args.ne40; i_top += nth) { ++ const int32_t i_kv = top_indices[i_top]; ++ ++ float score = -INFINITY; ++ if (valid_head && i_kv >= 0 && i_kv < args.ne11) { ++ device const M * mask_ptr = (device const M *) (kq_mask + ++ i_kv*args.nb31 + i_batch*args.nb32 + i_stream*args.nb33); ++ const float mask = float(*mask_ptr); ++ ++ if (isfinite(mask)) { ++ local_active_top_end = max(local_active_top_end, float(i_top + 1)); ++ ++ float qk = 0.0f; ++ int32_t i_dk = 0; ++ for (; i_dk + 3 < args.ne00; i_dk += 4) { ++ device const float * q0_ptr = (device const float *) (q + ++ (i_dk + 0)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q1_ptr = (device const float *) (q + ++ (i_dk + 1)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q2_ptr = (device const float *) (q + ++ (i_dk + 2)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const float * q3_ptr = (device const float *) (q + ++ (i_dk + 3)*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const K * k0_ptr = (device const K *) (k + ++ (i_dk + 0)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k1_ptr = (device const K *) (k + ++ (i_dk + 1)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k2_ptr = (device const K *) (k + ++ (i_dk + 2)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ device const K * k3_ptr = (device const K *) (k + ++ (i_dk + 3)*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ qk += (*q0_ptr) * float(*k0_ptr); ++ qk += (*q1_ptr) * float(*k1_ptr); ++ qk += (*q2_ptr) * float(*k2_ptr); ++ qk += (*q3_ptr) * float(*k3_ptr); ++ } ++ for (; i_dk < args.ne00; ++i_dk) { ++ device const float * q_ptr = (device const float *) (q + ++ i_dk*args.nb00 + i_batch*args.nb01 + i_head*args.nb02 + i_stream*args.nb03); ++ device const K * k_ptr = (device const K *) (k + ++ i_dk*args.nb10 + i_kv*args.nb11 + i_kv_head*args.nb12 + i_stream*args.nb13); ++ qk += (*q_ptr) * float(*k_ptr); ++ } + +- for (int ind = iw0*ntg.x + tiitg; ind < args.ne00t;) { +- pdst[ind] = psrc[ind]; ++ score = qk*args.scale + mask; ++ } ++ } + +- break; ++ scores[score_offset + i_top] = score; ++ local_max = max(local_max, score); + } +-} + +-typedef decltype(kernel_get_rows_f) get_rows_f_t; ++ reduce[reduce_offset + tid] = local_max; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + +-template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f; +-template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f; +-template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f; +-#if defined(GGML_METAL_HAS_BF16) +-template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f; +-#endif ++ for (int32_t stride = nth/2; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ reduce[reduce_offset + tid] = max( ++ reduce[reduce_offset + tid], ++ reduce[reduce_offset + tid + stride]); ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +-typedef decltype(kernel_get_rows_q) get_rows_q_t; ++ const float max_score = reduce[reduce_offset]; ++ reduce[reduce_offset + tid] = local_active_top_end; ++ threadgroup_barrier(mem_flags::mem_threadgroup); + + template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q; + template [[host_name("kernel_get_rows_q2_0")]] kernel get_rows_q_t kernel_get_rows_q; +@@ -9837,101 +16837,126 @@ template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get + template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q; + template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q; + template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q; ++ for (int32_t stride = nth/2; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ reduce[reduce_offset + tid] = max( ++ reduce[reduce_offset + tid], ++ reduce[reduce_offset + tid + stride]); ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } + +-template +-kernel void kernel_set_rows_q32( +- constant ggml_metal_kargs_set_rows & args, +- device const void * src0, +- device const void * src1, +- device float * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- uint tiitg[[thread_index_in_threadgroup]], +- uint3 tptg [[threads_per_threadgroup]]) { +- const int32_t i03 = tgpig.z; +- const int32_t i02 = tgpig.y; ++ const int32_t active_top_end = int32_t(reduce[reduce_offset]); ++ float local_sum = 0.0f; + +- const int32_t i12 = i03%args.ne12; +- const int32_t i11 = i02%args.ne11; ++ if (max_score > -FLT_MAX/2) { ++ for (int32_t i_top = tid; i_top < active_top_end; i_top += nth) { ++ scores[score_offset + i_top] = exp(scores[score_offset + i_top] - max_score); ++ local_sum += scores[score_offset + i_top]; ++ } ++ } else { ++ for (int32_t i_top = tid; i_top < args.ne40; i_top += nth) { ++ scores[score_offset + i_top] = 0.0f; ++ } ++ } + +- const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x; +- if (i01 >= args.ne01) { ++ reduce[reduce_offset + tid] = local_sum; ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ ++ for (int32_t stride = nth/2; stride > 0; stride >>= 1) { ++ if (tid < stride) { ++ reduce[reduce_offset + tid] += reduce[reduce_offset + tid + stride]; ++ } ++ threadgroup_barrier(mem_flags::mem_threadgroup); ++ } ++ ++ const float sum_score = reduce[reduce_offset]; ++ if (!valid_head) { + return; + } + +- const int32_t i10 = i01; +- const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0]; ++ for (int32_t i_dv = tid; i_dv < args.ne20; i_dv += nth) { ++ float acc = 0.0f; + +- device block_q * dst_row = ( device block_q *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3); +- const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); ++ if (sum_score > 0.0f && isfinite(sum_score)) { ++ for (int32_t i_top = 0; i_top < active_top_end; ++i_top) { ++ const int32_t i_kv = top_indices[i_top]; ++ if (i_kv < 0 || i_kv >= args.ne11) { ++ continue; ++ } + +- for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) { +- quantize_func(src_row + 32*ind, dst_row[ind]); ++ const float p = scores[score_offset + i_top]/sum_score; ++ if (p == 0.0f) { ++ continue; ++ } ++ ++ device const V * v_ptr = (device const V *) (v + ++ i_dv*args.nb20 + i_kv*args.nb21 + i_v_head*args.nb22 + i_stream*args.nb23); ++ acc += p * float(*v_ptr); ++ } ++ } ++ ++ device float * dst_ptr = (device float *) (dst + ++ i_dv*args.nb0 + i_batch*args.nb1 + i_head*args.nb2 + i_stream*args.nb3); ++ *dst_ptr = acc; + } + } + +-template +-kernel void kernel_set_rows_f( +- constant ggml_metal_kargs_set_rows & args, +- device const void * src0, +- device const void * src1, +- device float * dst, +- uint3 tgpig[[threadgroup_position_in_grid]], +- uint tiitg[[thread_index_in_threadgroup]], +- uint3 tptg [[threads_per_threadgroup]]) { +- const int32_t i03 = tgpig.z; +- const int32_t i02 = tgpig.y; ++typedef decltype(kernel_dsa_sparse_attn_decode_grouped_impl) kernel_dsa_sparse_attn_decode_grouped_t; + +- const int32_t i12 = i03%args.ne12; +- const int32_t i11 = i02%args.ne11; ++template [[host_name("kernel_dsa_sparse_attn_decode_grouped_f32_f32_f32")]] kernel kernel_dsa_sparse_attn_decode_grouped_t kernel_dsa_sparse_attn_decode_grouped_impl; ++template [[host_name("kernel_dsa_sparse_attn_decode_grouped_f32_f32_f16")]] kernel kernel_dsa_sparse_attn_decode_grouped_t kernel_dsa_sparse_attn_decode_grouped_impl; ++template [[host_name("kernel_dsa_sparse_attn_decode_grouped_f32_f16_f32")]] kernel kernel_dsa_sparse_attn_decode_grouped_t kernel_dsa_sparse_attn_decode_grouped_impl; ++template [[host_name("kernel_dsa_sparse_attn_decode_grouped_f32_f16_f16")]] kernel kernel_dsa_sparse_attn_decode_grouped_t kernel_dsa_sparse_attn_decode_grouped_impl; ++template [[host_name("kernel_dsa_sparse_attn_decode_grouped_f16_f32_f32")]] kernel kernel_dsa_sparse_attn_decode_grouped_t kernel_dsa_sparse_attn_decode_grouped_impl; ++template [[host_name("kernel_dsa_sparse_attn_decode_grouped_f16_f32_f16")]] kernel kernel_dsa_sparse_attn_decode_grouped_t kernel_dsa_sparse_attn_decode_grouped_impl; ++template [[host_name("kernel_dsa_sparse_attn_decode_grouped_f16_f16_f32")]] kernel kernel_dsa_sparse_attn_decode_grouped_t kernel_dsa_sparse_attn_decode_grouped_impl; ++template [[host_name("kernel_dsa_sparse_attn_decode_grouped_f16_f16_f16")]] kernel kernel_dsa_sparse_attn_decode_grouped_t kernel_dsa_sparse_attn_decode_grouped_impl; + +- const int32_t i01 = tgpig.x*tptg.y + tiitg/tptg.x; +- if (i01 >= args.ne01) { ++template ++kernel void kernel_dsa_top1_attn_impl( ++ constant ggml_metal_kargs_dsa_top1_attn & args, ++ device const char * q [[buffer(1)]], ++ device const char * v [[buffer(2)]], ++ device const char * top_k [[buffer(3)]], ++ device char * dst [[buffer(4)]], ++ uint3 tgpig[[threadgroup_position_in_grid]], ++ ushort tiitg[[thread_index_in_threadgroup]], ++ ushort3 ntg[[threads_per_threadgroup]]) { ++ (void) q; ++ ++ const int32_t i0 = tgpig.x*ntg.x + tiitg; ++ const int32_t i_batch = tgpig.y % args.ne1; ++ const int32_t i_head = tgpig.y / args.ne1; ++ const int32_t i_stream = tgpig.z; ++ ++ if (i0 >= args.ne0 || i_batch >= args.ne1 || i_head >= args.ne2 || i_stream >= args.ne3) { + return; + } + +- const int32_t i10 = i01; +- const TI i1 = ((const device TI *) ((const device char *) src1 + i10*args.nb10 + i11*args.nb11 + i12*args.nb12))[0]; +- +- device TD * dst_row = ( device TD *) (( device char *) dst + i1*args.nb1 + i02*args.nb2 + i03*args.nb3); +- const device TS * src_row = (const device TS *) ((const device char *) src0 + i01*args.nb01 + i02*args.nb02 + i03*args.nb03); ++ const int32_t heads_per_v_head = args.ne2 / args.ne12; ++ const int32_t i_v_head = i_head / heads_per_v_head; ++ const int32_t i_top_stream = i_stream % args.ne22; + +- for (int ind = tiitg%tptg.x; ind < args.nk0; ind += tptg.x) { +- dst_row[ind] = (TD) src_row[ind]; ++ device const int32_t * top_k_ptr = (device const int32_t *) (top_k + ++ i_batch*args.nb21 + i_top_stream*args.nb22); ++ const int32_t i_kv = *top_k_ptr; ++ if (i_kv < 0 || i_kv >= args.ne11) { ++ return; + } +-} + +-typedef decltype(kernel_set_rows_f) set_rows_f_t; ++ device const T * v_row = (device const T *) (v + ++ i_kv*args.nb11 + i_v_head*args.nb12 + i_stream*args.nb13); ++ device float * dst_row = (device float *) (dst + ++ i_batch*args.nb1 + i_head*args.nb2 + i_stream*args.nb3); + +-template [[host_name("kernel_set_rows_f32_i64_f32")]] kernel set_rows_f_t kernel_set_rows_f; +-template [[host_name("kernel_set_rows_f32_i32_f32")]] kernel set_rows_f_t kernel_set_rows_f; +-template [[host_name("kernel_set_rows_f32_i64_f16")]] kernel set_rows_f_t kernel_set_rows_f; +-template [[host_name("kernel_set_rows_f32_i32_f16")]] kernel set_rows_f_t kernel_set_rows_f; +-#if defined(GGML_METAL_HAS_BF16) +-template [[host_name("kernel_set_rows_f32_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f; +-template [[host_name("kernel_set_rows_f32_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f; +-#endif ++ dst_row[i0] = float(v_row[i0]); ++} + +-template [[host_name("kernel_set_rows_f16_i64_f16")]] kernel set_rows_f_t kernel_set_rows_f; +-template [[host_name("kernel_set_rows_f16_i32_f16")]] kernel set_rows_f_t kernel_set_rows_f; +-#if defined(GGML_METAL_HAS_BF16) +-template [[host_name("kernel_set_rows_bf16_i64_bf16")]] kernel set_rows_f_t kernel_set_rows_f; +-template [[host_name("kernel_set_rows_bf16_i32_bf16")]] kernel set_rows_f_t kernel_set_rows_f; +-#endif ++typedef decltype(kernel_dsa_top1_attn_impl) kernel_dsa_top1_attn_t; + +-typedef decltype(kernel_set_rows_q32) set_rows_q32_t; +- +-template [[host_name("kernel_set_rows_f32_i64_q8_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i32_q8_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i64_q4_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i32_q4_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i64_q4_1")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i32_q4_1")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i64_q5_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i32_q5_0")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i64_q5_1")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i32_q5_1")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i64_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32; +-template [[host_name("kernel_set_rows_f32_i32_iq4_nl")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_dsa_top1_attn_f32")]] kernel kernel_dsa_top1_attn_t kernel_dsa_top1_attn_impl; ++template [[host_name("kernel_dsa_top1_attn_f16")]] kernel kernel_dsa_top1_attn_t kernel_dsa_top1_attn_impl; + + kernel void kernel_diag_f32( + constant ggml_metal_kargs_diag & args, +@@ -10675,6 +17700,76 @@ kernel void kernel_mul_mm_id( + } + } + ++#define QK_NL 16 ++ ++// ++// get rows ++// ++ ++typedef decltype(kernel_get_rows_f) get_rows_f_t; ++ ++template [[host_name("kernel_get_rows_f32")]] kernel get_rows_f_t kernel_get_rows_f; ++template [[host_name("kernel_get_rows_f16")]] kernel get_rows_f_t kernel_get_rows_f; ++template [[host_name("kernel_get_rows_f16_f16")]] kernel get_rows_f_t kernel_get_rows_f; ++template [[host_name("kernel_get_rows_i32")]] kernel get_rows_f_t kernel_get_rows_f; ++#if defined(GGML_METAL_HAS_BF16) ++template [[host_name("kernel_get_rows_bf16")]] kernel get_rows_f_t kernel_get_rows_f; ++#endif ++ ++typedef decltype(kernel_get_rows_q) get_rows_q_t; ++ ++template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q4_0")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q4_1")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q5_1")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q8_0")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_mxfp4")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q2_K")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q3_K")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q4_K")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q5_K")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_q6_K")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq2_xs")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq3_s")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq2_s")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q; ++template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q; ++ ++// ++// set rows ++// ++ ++typedef decltype(kernel_set_rows_f) set_rows_f_t; ++ ++template [[host_name("kernel_set_rows_f32_i64")]] kernel set_rows_f_t kernel_set_rows_f; ++template [[host_name("kernel_set_rows_f32_i32")]] kernel set_rows_f_t kernel_set_rows_f; ++template [[host_name("kernel_set_rows_f16_i64")]] kernel set_rows_f_t kernel_set_rows_f; ++template [[host_name("kernel_set_rows_f16_i32")]] kernel set_rows_f_t kernel_set_rows_f; ++#if defined(GGML_METAL_HAS_BF16) ++template [[host_name("kernel_set_rows_bf16_i64")]] kernel set_rows_f_t kernel_set_rows_f; ++template [[host_name("kernel_set_rows_bf16_i32")]] kernel set_rows_f_t kernel_set_rows_f; ++#endif ++ ++typedef decltype(kernel_set_rows_q32) set_rows_q32_t; ++ ++template [[host_name("kernel_set_rows_q8_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q8_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q4_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q4_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q4_1_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q4_1_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q5_0_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q5_0_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q5_1_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_q5_1_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_iq4_nl_i64")]] kernel set_rows_q32_t kernel_set_rows_q32; ++template [[host_name("kernel_set_rows_iq4_nl_i32")]] kernel set_rows_q32_t kernel_set_rows_q32; ++ + // + // matrix-matrix multiplication + // +@@ -10928,6 +18023,7 @@ template [[host_name("kernel_mul_mv_id_q8_0_f32")]] kernel kernel_mul_mv_id_t + + template [[host_name("kernel_mul_mv_id_q1_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + template [[host_name("kernel_mul_mv_id_q2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_tq2_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + template [[host_name("kernel_mul_mv_id_q4_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + template [[host_name("kernel_mul_mv_id_q4_1_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + template [[host_name("kernel_mul_mv_id_q5_0_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +@@ -10936,7 +18032,15 @@ template [[host_name("kernel_mul_mv_id_q5_1_f32")]] kernel kernel_mul_mv_id_t + template [[host_name("kernel_mul_mv_id_mxfp4_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + + template [[host_name("kernel_mul_mv_id_q2_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_q2_K_f32_glm_down")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_q2_K_f32_r1")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_q2_K_f32_r2")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_q2_K_f32_r8")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + template [[host_name("kernel_mul_mv_id_q3_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_q3_K_f32_glm_down")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_q3_K_f32_r1")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_q3_K_f32_r2")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; ++template [[host_name("kernel_mul_mv_id_q3_K_f32_r8")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + template [[host_name("kernel_mul_mv_id_q4_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + template [[host_name("kernel_mul_mv_id_q5_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; + template [[host_name("kernel_mul_mv_id_q6_K_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id>>; +diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c +index a7d1fe7d..cdc85b08 100644 +--- a/ggml/src/ggml.c ++++ b/ggml/src/ggml.c +@@ -1083,6 +1083,11 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + "DSV4_HC_COMB", + "DSV4_HC_PRE", + "DSV4_HC_POST", ++ "DSA_SPARSE_MASK", ++ "DSA_SPARSE_ATTN", ++ "DSA_TOP1_ATTN", ++ "MOE_ROUTE_WEIGHTS", ++ "MOE_WEIGHTED_SUM", + + "UNARY", + +@@ -1098,9 +1103,11 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + "OPT_STEP_SGD", + + "GLU", ++ ++ "MOE_MUL_MAT_ID", + }; + +-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); ++static_assert(GGML_OP_COUNT == 107, "GGML_OP_COUNT != 107"); + + static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "none", +@@ -1198,6 +1205,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "dsv4_hc_comb(mixes, scale, base)", + "dsv4_hc_pre(x, weights)", + "dsv4_hc_post(x, residual, post, comb)", ++ "dsa_sparse_mask(kq_mask, top_k)", ++ "dsa_sparse_attn(q, k, v, kq_mask, top_k)", ++ "dsa_top1_attn(q, v, top_k)", ++ "moe_route_weights(probs, ids)", ++ "moe_weighted_sum(experts, weights)", + + "unary(x)", + +@@ -1213,9 +1225,11 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "sgd(x)", + + "glu(x)", ++ ++ "moe_mul_mat_id(experts, input, ids, weights)", + }; + +-static_assert(GGML_OP_COUNT == 101, "GGML_OP_COUNT != 101"); ++static_assert(GGML_OP_COUNT == 107, "GGML_OP_COUNT != 107"); + + static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); + +@@ -3911,6 +3925,29 @@ struct ggml_tensor * ggml_get_rows( + return result; + } + ++struct ggml_tensor * ggml_get_rows_typed( ++ struct ggml_context * ctx, ++ struct ggml_tensor * a, ++ struct ggml_tensor * b) { ++ GGML_ASSERT(a->ne[2] == b->ne[1]); ++ GGML_ASSERT(a->ne[3] == b->ne[2]); ++ GGML_ASSERT(b->ne[3] == 1); ++ GGML_ASSERT(b->type == GGML_TYPE_I32); ++ ++ enum ggml_type type = a->type; ++ if (ggml_is_quantized(type) || type == GGML_TYPE_BF16) { ++ type = GGML_TYPE_F32; ++ } ++ ++ struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, a->ne[0], b->ne[0], b->ne[1], b->ne[2]); ++ ++ result->op = GGML_OP_GET_ROWS; ++ result->src[0] = a; ++ result->src[1] = b; ++ ++ return result; ++} ++ + // ggml_get_rows_back + + struct ggml_tensor * ggml_get_rows_back( +@@ -6470,6 +6507,200 @@ struct ggml_tensor * ggml_dsv4_hc_post( + return result; + } + ++// ggml_dsa_sparse_mask ++ ++struct ggml_tensor * ggml_dsa_sparse_mask( ++ struct ggml_context * ctx, ++ struct ggml_tensor * kq_mask_rows, ++ struct ggml_tensor * top_k) { ++ ++ GGML_ASSERT(kq_mask_rows->type == GGML_TYPE_F32 || kq_mask_rows->type == GGML_TYPE_F16); ++ GGML_ASSERT(top_k->type == GGML_TYPE_I32); ++ GGML_ASSERT(kq_mask_rows->ne[0] == 1); ++ GGML_ASSERT(kq_mask_rows->ne[2] == top_k->ne[1]); ++ GGML_ASSERT(kq_mask_rows->ne[3] % top_k->ne[2] == 0); ++ GGML_ASSERT(top_k->ne[3] == 1); ++ ++ int64_t ne[4] = { kq_mask_rows->ne[0], kq_mask_rows->ne[1], kq_mask_rows->ne[2], kq_mask_rows->ne[3] }; ++ struct ggml_tensor * result = ggml_new_tensor(ctx, kq_mask_rows->type, 4, ne); ++ ++ result->op = GGML_OP_DSA_SPARSE_MASK; ++ result->src[0] = kq_mask_rows; ++ result->src[1] = top_k; ++ ++ return result; ++} ++ ++// ggml_dsa_sparse_attn ++ ++struct ggml_tensor * ggml_dsa_sparse_attn( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * k, ++ struct ggml_tensor * v, ++ struct ggml_tensor * kq_mask_rows, ++ struct ggml_tensor * top_k, ++ float scale) { ++ ++ GGML_ASSERT(q->type == GGML_TYPE_F32); ++ GGML_ASSERT(k->type == GGML_TYPE_F32 || ggml_get_type_traits(k->type)->to_float); ++ GGML_ASSERT(v->type == GGML_TYPE_F32 || ggml_get_type_traits(v->type)->to_float); ++ GGML_ASSERT(kq_mask_rows->type == GGML_TYPE_F32 || kq_mask_rows->type == GGML_TYPE_F16); ++ GGML_ASSERT(top_k->type == GGML_TYPE_I32); ++ ++ GGML_ASSERT(q->ne[0] == k->ne[0]); ++ GGML_ASSERT(k->ne[1] == v->ne[1]); ++ GGML_ASSERT(q->ne[3] == k->ne[3]); ++ GGML_ASSERT(q->ne[3] == v->ne[3]); ++ ++ GGML_ASSERT(kq_mask_rows->ne[0] == 1); ++ GGML_ASSERT(kq_mask_rows->ne[1] == k->ne[1]); ++ GGML_ASSERT(kq_mask_rows->ne[2] == q->ne[1]); ++ GGML_ASSERT(kq_mask_rows->ne[3] == q->ne[3]); ++ ++ GGML_ASSERT(top_k->ne[1] == q->ne[1]); ++ GGML_ASSERT(kq_mask_rows->ne[3] % top_k->ne[2] == 0); ++ GGML_ASSERT(top_k->ne[3] == 1); ++ ++ GGML_ASSERT(q->ne[2] % k->ne[2] == 0); ++ GGML_ASSERT(q->ne[2] % v->ne[2] == 0); ++ ++ int64_t ne[4] = { v->ne[0], q->ne[1], q->ne[2], q->ne[3] }; ++ struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); ++ ++ result->op = GGML_OP_DSA_SPARSE_ATTN; ++ result->src[0] = q; ++ result->src[1] = k; ++ result->src[2] = v; ++ result->src[3] = kq_mask_rows; ++ result->src[4] = top_k; ++ ++ ggml_set_op_params_f32(result, 0, scale); ++ ++ return result; ++} ++ ++// ggml_dsa_top1_attn ++ ++struct ggml_tensor * ggml_dsa_top1_attn( ++ struct ggml_context * ctx, ++ struct ggml_tensor * q, ++ struct ggml_tensor * v, ++ struct ggml_tensor * top_k) { ++ ++ GGML_ASSERT(q->type == GGML_TYPE_F32); ++ GGML_ASSERT(v->type == GGML_TYPE_F32 || v->type == GGML_TYPE_F16); ++ GGML_ASSERT(top_k->type == GGML_TYPE_I32); ++ ++ GGML_ASSERT(q->ne[3] == v->ne[3]); ++ GGML_ASSERT(top_k->ne[0] == 1); ++ GGML_ASSERT(top_k->ne[1] == q->ne[1]); ++ GGML_ASSERT(q->ne[3] % top_k->ne[2] == 0); ++ GGML_ASSERT(top_k->ne[3] == 1); ++ GGML_ASSERT(q->ne[2] % v->ne[2] == 0); ++ ++ int64_t ne[4] = { v->ne[0], q->ne[1], q->ne[2], q->ne[3] }; ++ struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); ++ ++ result->op = GGML_OP_DSA_TOP1_ATTN; ++ result->src[0] = q; ++ result->src[1] = v; ++ result->src[2] = top_k; ++ ++ return result; ++} ++ ++// ggml_moe_weighted_sum ++ ++struct ggml_tensor * ggml_moe_route_weights( ++ struct ggml_context * ctx, ++ struct ggml_tensor * probs, ++ struct ggml_tensor * ids, ++ bool norm, ++ float clamp_min, ++ float scale) { ++ ++ GGML_ASSERT(probs->type == GGML_TYPE_F32); ++ GGML_ASSERT(ids->type == GGML_TYPE_I32); ++ GGML_ASSERT(probs->ne[0] == 1); ++ GGML_ASSERT(ids->ne[0] > 0); ++ GGML_ASSERT(ids->ne[1] == probs->ne[2]); ++ GGML_ASSERT(ids->ne[2] == 1); ++ GGML_ASSERT(ids->ne[3] == 1); ++ GGML_ASSERT(probs->ne[3] == 1); ++ ++ int64_t ne[3] = { 1, ids->ne[0], ids->ne[1] }; ++ struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 3, ne); ++ ++ ggml_set_op_params_f32(result, 0, clamp_min); ++ ggml_set_op_params_f32(result, 1, scale); ++ ggml_set_op_params_i32(result, 2, norm ? 1 : 0); ++ ++ result->op = GGML_OP_MOE_ROUTE_WEIGHTS; ++ result->src[0] = probs; ++ result->src[1] = ids; ++ ++ return result; ++} ++ ++struct ggml_tensor * ggml_moe_weighted_sum( ++ struct ggml_context * ctx, ++ struct ggml_tensor * experts, ++ struct ggml_tensor * weights) { ++ ++ GGML_ASSERT(experts->type == GGML_TYPE_F32); ++ GGML_ASSERT(weights->type == GGML_TYPE_F32); ++ GGML_ASSERT(weights->ne[0] == 1); ++ GGML_ASSERT(experts->ne[1] == weights->ne[1]); ++ GGML_ASSERT(experts->ne[2] == weights->ne[2]); ++ GGML_ASSERT(experts->ne[3] == 1); ++ GGML_ASSERT(weights->ne[3] == 1); ++ ++ int64_t ne[2] = { experts->ne[0], experts->ne[2] }; ++ struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne); ++ ++ result->op = GGML_OP_MOE_WEIGHTED_SUM; ++ result->src[0] = experts; ++ result->src[1] = weights; ++ ++ return result; ++} ++ ++struct ggml_tensor * ggml_moe_mul_mat_id( ++ struct ggml_context * ctx, ++ struct ggml_tensor * experts, ++ struct ggml_tensor * input, ++ struct ggml_tensor * ids, ++ struct ggml_tensor * weights) { ++ ++ GGML_ASSERT(experts->type == GGML_TYPE_F32 || ggml_get_type_traits(experts->type)->to_float); ++ GGML_ASSERT(input->type == GGML_TYPE_F32 || input->type == GGML_TYPE_F16); ++ GGML_ASSERT(ids->type == GGML_TYPE_I32); ++ GGML_ASSERT(weights->type == GGML_TYPE_F32); ++ GGML_ASSERT(experts->ne[0] == input->ne[0]); ++ GGML_ASSERT(experts->ne[2] > 0); ++ GGML_ASSERT(experts->ne[3] == 1); ++ GGML_ASSERT(ids->ne[0] == input->ne[1]); ++ GGML_ASSERT(ids->ne[1] == input->ne[2]); ++ GGML_ASSERT(ids->ne[2] == 1 && ids->ne[3] == 1); ++ GGML_ASSERT(weights->ne[0] == 1); ++ GGML_ASSERT(weights->ne[1] == ids->ne[0]); ++ GGML_ASSERT(weights->ne[2] == ids->ne[1]); ++ GGML_ASSERT(weights->ne[3] == 1); ++ GGML_ASSERT(input->ne[3] == 1); ++ ++ int64_t ne[2] = { experts->ne[1], input->ne[2] }; ++ struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne); ++ ++ result->op = GGML_OP_MOE_MUL_MAT_ID; ++ result->src[0] = experts; ++ result->src[1] = input; ++ result->src[2] = ids; ++ result->src[3] = weights; ++ ++ return result; ++} ++ + //////////////////////////////////////////////////////////////////////////////// + + struct ggml_hash_set ggml_hash_set_new(size_t size) { +diff --git a/src/llama-hparams.h b/src/llama-hparams.h +index 8be5f28f..92626b4d 100644 +--- a/src/llama-hparams.h ++++ b/src/llama-hparams.h +@@ -226,6 +226,10 @@ struct llama_hparams { + uint32_t indexer_n_head = 0; + uint32_t indexer_head_size = 0; + uint32_t indexer_top_k = 0; ++ uint32_t indexer_top_k_freq = 0; ++ uint32_t indexer_skip_top_k_offset = 0; ++ bool indexer_types_present = false; ++ std::array indexer_types; // -1 unknown, 0 shared, 1 full + + // DeepSeek-V4 + uint32_t dsv4_o_group_count = 0; +-- +2.54.0 (Apple Git-157) diff --git a/third_party/llama.cpp/patches/0023-skippy-expose-GLM-DSA-staged-runtime-controls.patch b/third_party/llama.cpp/patches/0023-skippy-expose-GLM-DSA-staged-runtime-controls.patch new file mode 100644 index 000000000..1ed508cf7 --- /dev/null +++ b/third_party/llama.cpp/patches/0023-skippy-expose-GLM-DSA-staged-runtime-controls.patch @@ -0,0 +1,2612 @@ +From e87de13ba5a2bb71fb5b1b874ca61d7338ec3f62 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Tue, 14 Jul 2026 09:16:19 +1000 +Subject: [PATCH 20/24] skippy: expose GLM-DSA staged runtime controls + +--- + include/skippy.h | 20 + + src/llama-graph.cpp | 45 + + src/llama-graph.h | 43 + + src/skippy.cpp | 2006 ++++++++++++++++++++++++++++++++++++++++--- + 4 files changed, 1986 insertions(+), 128 deletions(-) + +diff --git a/include/skippy.h b/include/skippy.h +index c9ecc349..cd0fe21a 100644 +--- a/include/skippy.h ++++ b/include/skippy.h +@@ -43,6 +43,15 @@ enum skippy_activation_layout { + + #define SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST (UINT64_C(1) << 0) + #define SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP (UINT64_C(1) << 1) ++#define SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K (UINT64_C(1) << 3) ++ ++#define SKIPPY_GLM_DSA_POLICY_PROFILE_NONE 0 ++#define SKIPPY_GLM_DSA_POLICY_PROFILE_V1 1 ++ ++#define SKIPPY_GLM_DSA_POLICY_DIRECT_SPARSE_ATTN (UINT32_C(1) << 0) ++#define SKIPPY_GLM_DSA_POLICY_DIRECT_SPARSE_PREFILL (UINT32_C(1) << 1) ++#define SKIPPY_GLM_DSA_POLICY_DISABLE_COMPACT_FLASH_ATTN (UINT32_C(1) << 2) ++#define SKIPPY_GLM_DSA_POLICY_UNPROVEN_LARGE_DIRECT_SPARSE_PREFILL (UINT32_C(1) << 3) + + struct skippy_model; + struct skippy_session; +@@ -71,6 +80,8 @@ struct skippy_runtime_config { + enum skippy_load_mode load_mode; + + bool disable_repack; ++ bool use_mmap_prefetch; ++ bool use_mmap_buffer; + bool filter_tensors_on_load; + bool include_embeddings; + bool include_output; +@@ -79,6 +90,15 @@ struct skippy_runtime_config { + // "Vulkan1", or "CPU". When set, skippy loads the model only on that + // backend device instead of relying on llama.cpp default device ordering. + const char * selected_backend_device; ++ ++ // Optional GLM-DSA execution policy resolved from package generation ++ // metadata. Zero values preserve llama.cpp built-in/env behavior. ++ int32_t glm_dsa_policy_profile; ++ uint32_t glm_dsa_policy_flags; ++ int32_t glm_dsa_short_prefill_max_tokens; ++ int32_t glm_dsa_direct_sparse_decode_max_top_k; ++ uint64_t glm_dsa_dense_sparse_mask_max_bytes; ++ int32_t glm_dsa_compact_flash_min_kv; + }; + + struct skippy_tensor_info { +diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp +index 6613747a..244a400a 100644 +--- a/src/llama-graph.cpp ++++ b/src/llama-graph.cpp +@@ -25,6 +25,7 @@ static thread_local skippy_graph_filter g_skippy_graph_filter; + static thread_local skippy_activation_tokens g_skippy_activation_tokens; + static thread_local skippy_activation_rwkv7_v_first g_skippy_rwkv7_v_first; + static thread_local skippy_activation_gemma3n_altup g_skippy_gemma3n_altup; ++static thread_local skippy_activation_glm_dsa_top_k g_skippy_glm_dsa_top_k; + + void skippy_graph_set_filter(const skippy_graph_filter & filter) { + g_skippy_graph_filter = filter; +@@ -74,6 +75,18 @@ const skippy_activation_gemma3n_altup & skippy_graph_get_gemma3n_altup() { + return g_skippy_gemma3n_altup; + } + ++void skippy_graph_set_glm_dsa_top_k(const skippy_activation_glm_dsa_top_k & values) { ++ g_skippy_glm_dsa_top_k = values; ++} ++ ++void skippy_graph_clear_glm_dsa_top_k() { ++ g_skippy_glm_dsa_top_k = {}; ++} ++ ++const skippy_activation_glm_dsa_top_k & skippy_graph_get_glm_dsa_top_k() { ++ return g_skippy_glm_dsa_top_k; ++} ++ + // dedup helpers + + static ggml_tensor * build_attn_inp_kq_mask( +@@ -213,6 +226,34 @@ bool llm_graph_input_gemma3n_altup::can_reuse(const llm_graph_params & params) { + return values && values->ne[0] == n_embd && values->ne[1] == params.ubatch.n_tokens && values->ne[2] == n_altup; + } + ++void llm_graph_input_glm_dsa_top_k::set_input(const llama_ubatch * ubatch) { ++ const skippy_activation_glm_dsa_top_k & sideband = skippy_graph_get_glm_dsa_top_k(); ++ GGML_ASSERT(sideband.values != nullptr); ++ GGML_ASSERT(sideband.token_count == ubatch->n_tokens); ++ GGML_ASSERT(sideband.n_top_k == n_top_k); ++ GGML_ASSERT(sideband.n_stream == n_stream); ++ ++ ggml_backend_tensor_set(values, sideband.values, 0, sideband.token_count*n_top_k*ggml_element_size(values)); ++} ++ ++bool llm_graph_input_glm_dsa_top_k::can_reuse(const llm_graph_params & params) { ++ const auto * mctx = static_cast(params.mctx); ++ const int64_t params_n_stream = params.cparams.kv_unified ? 1 : params.ubatch.n_seqs_unq; ++ if (mctx == nullptr || params_n_stream <= 0) { ++ return false; ++ } ++ ++ const int64_t params_n_top_k = std::min( ++ mctx->get_lid()->get_n_kv(), ++ params.hparams.indexer_top_k); ++ ++ return values && ++ values->ne[0] == params_n_top_k && ++ values->ne[1] == params.ubatch.n_tokens/params_n_stream && ++ values->ne[2] == 1 && ++ values->ne[3] == params_n_stream; ++} ++ + void llm_graph_input_pos::set_input(const llama_ubatch * ubatch) { + // A filtered stage can omit positional work entirely. In that case the + // scheduler leaves this otherwise declared graph input unallocated. +@@ -1325,6 +1366,7 @@ void llm_graph_result::reset() { + t_h_nextn = nullptr; + t_skippy_rwkv7_v_first = nullptr; + t_skippy_gemma3n_altup = nullptr; ++ t_skippy_glm_dsa_top_k = nullptr; + + t_layer_inp.resize(LLAMA_MAX_LAYERS); + std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr); +@@ -1377,6 +1419,9 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) { + if (t_skippy_gemma3n_altup != nullptr) { + ggml_set_output(t_skippy_gemma3n_altup); + } ++ if (t_skippy_glm_dsa_top_k != nullptr) { ++ ggml_set_output(t_skippy_glm_dsa_top_k); ++ } + { + const auto & embeddings_layer_inp = params.cparams.embeddings_layer_inp; + for (size_t il = 0; il < embeddings_layer_inp.size(); ++il) { +diff --git a/src/llama-graph.h b/src/llama-graph.h +index 2e0d9c3b..ab6e111e 100644 +--- a/src/llama-graph.h ++++ b/src/llama-graph.h +@@ -30,12 +30,26 @@ class llama_memory_recurrent_context; + class llama_memory_hybrid_context; + class llama_memory_hybrid_iswa_context; + ++enum skippy_glm_dsa_phase_hint : int32_t { ++ SKIPPY_GLM_DSA_PHASE_HINT_AUTO = 0, ++ SKIPPY_GLM_DSA_PHASE_HINT_PREFILL = 1, ++ SKIPPY_GLM_DSA_PHASE_HINT_DECODE = 2, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY = 3, ++}; ++ + struct skippy_graph_filter { + bool enabled = false; + int32_t layer_start = 0; + int32_t layer_end = 0; + bool include_embeddings = false; + bool include_output = false; ++ int32_t glm_dsa_policy_profile = 0; ++ uint32_t glm_dsa_policy_flags = 0; ++ int32_t glm_dsa_short_prefill_max_tokens = 0; ++ int32_t glm_dsa_direct_sparse_decode_max_top_k = 0; ++ uint64_t glm_dsa_dense_sparse_mask_max_bytes = 0; ++ int32_t glm_dsa_compact_flash_min_kv = 0; ++ int32_t glm_dsa_phase_hint = SKIPPY_GLM_DSA_PHASE_HINT_AUTO; + }; + + struct skippy_activation_tokens { +@@ -56,6 +70,14 @@ struct skippy_activation_gemma3n_altup { + uint32_t n_altup = 0; + }; + ++struct skippy_activation_glm_dsa_top_k { ++ const int32_t * values = nullptr; ++ uint32_t token_count = 0; ++ uint32_t n_top_k = 0; ++ uint32_t n_stream = 0; ++ llama_pos pos_start = 0; ++}; ++ + void skippy_graph_set_filter(const skippy_graph_filter & filter); + void skippy_graph_clear_filter(); + const skippy_graph_filter & skippy_graph_get_filter(); +@@ -68,6 +90,9 @@ const skippy_activation_rwkv7_v_first & skippy_graph_get_rwkv7_v_first(); + void skippy_graph_set_gemma3n_altup(const skippy_activation_gemma3n_altup & values); + void skippy_graph_clear_gemma3n_altup(); + const skippy_activation_gemma3n_altup & skippy_graph_get_gemma3n_altup(); ++void skippy_graph_set_glm_dsa_top_k(const skippy_activation_glm_dsa_top_k & values); ++void skippy_graph_clear_glm_dsa_top_k(); ++const skippy_activation_glm_dsa_top_k & skippy_graph_get_glm_dsa_top_k(); + + // certain models (typically multi-modal) can produce different types of graphs + enum llm_graph_type { +@@ -234,6 +259,22 @@ private: + int64_t n_altup; + }; + ++class llm_graph_input_glm_dsa_top_k : public llm_graph_input_i { ++public: ++ llm_graph_input_glm_dsa_top_k(int64_t n_top_k, int64_t n_stream) : n_top_k(n_top_k), n_stream(n_stream) {} ++ virtual ~llm_graph_input_glm_dsa_top_k() = default; ++ ++ void set_input(const llama_ubatch * ubatch) override; ++ ++ bool can_reuse(const llm_graph_params & params) override; ++ ++ ggml_tensor * values = nullptr; // runtime layout: I32 [n_top_k, n_batch, 1, n_stream] ++ ++private: ++ int64_t n_top_k; ++ int64_t n_stream; ++}; ++ + class llm_graph_input_pos : public llm_graph_input_i { + public: + llm_graph_input_pos(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {} +@@ -886,6 +927,7 @@ public: + ggml_tensor * get_h_nextn() const { return t_h_nextn; } + ggml_tensor * get_skippy_rwkv7_v_first() const { return t_skippy_rwkv7_v_first; } + ggml_tensor * get_skippy_gemma3n_altup() const { return t_skippy_gemma3n_altup; } ++ ggml_tensor * get_skippy_glm_dsa_top_k() const { return t_skippy_glm_dsa_top_k; } + + ggml_tensor * get_layer_inp(int il) const { return t_layer_inp[il]; } + +@@ -923,6 +965,7 @@ public: + ggml_tensor * t_h_nextn = nullptr; // [n_embd, n_outputs] hidden state before final output norm + ggml_tensor * t_skippy_rwkv7_v_first = nullptr; + ggml_tensor * t_skippy_gemma3n_altup = nullptr; ++ ggml_tensor * t_skippy_glm_dsa_top_k = nullptr; + + std::vector t_layer_inp; + +diff --git a/src/skippy.cpp b/src/skippy.cpp +index 46da7007..7e11d92b 100644 +--- a/src/skippy.cpp ++++ b/src/skippy.cpp +@@ -9,6 +9,7 @@ + #include "llama-ext.h" + #include "llama-graph.h" + #include "llama-kv-cache.h" ++#include "llama-kv-cache-dsa.h" + #include "llama-memory-hybrid.h" + #include "llama-memory-hybrid-iswa.h" + #include "llama-memory-recurrent.h" +@@ -18,6 +19,7 @@ + + #include + #include ++#include + #include + #include + #include +@@ -43,12 +45,79 @@ + + using json = nlohmann::ordered_json; + ++enum skippy_glm_dsa_op_kind { ++ SKIPPY_GLM_DSA_OP_INDEXER = 0, ++ SKIPPY_GLM_DSA_OP_TOP_K = 1, ++ SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL = 2, ++ SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK = 3, ++ SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD = 4, ++ SKIPPY_GLM_DSA_OP_DSA_SPARSE_ATTN = 5, ++ SKIPPY_GLM_DSA_OP_COMPACT_GET_ROWS = 6, ++ SKIPPY_GLM_DSA_OP_MLA_ATTENTION = 7, ++ SKIPPY_GLM_DSA_OP_ROUTED_MOE_ROUTE = 8, ++ SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE_UP = 9, ++ SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE = 10, ++ SKIPPY_GLM_DSA_OP_ROUTED_MOE_UP = 11, ++ SKIPPY_GLM_DSA_OP_ROUTED_MOE_ACT = 12, ++ SKIPPY_GLM_DSA_OP_ROUTED_MOE_DOWN = 13, ++ SKIPPY_GLM_DSA_OP_ROUTED_MOE_WEIGHTED = 14, ++ SKIPPY_GLM_DSA_OP_ROUTED_MOE_AGGREGATE = 15, ++ SKIPPY_GLM_DSA_OP_SHARED_EXPERT = 16, ++ SKIPPY_GLM_DSA_OP_COUNT = 17, ++ SKIPPY_GLM_DSA_OP_UNKNOWN = 255, ++}; ++ ++struct skippy_glm_dsa_op_stat { ++ uint64_t nodes = 0; ++ int64_t elapsed_us = 0; ++}; ++ ++static constexpr size_t SKIPPY_GLM_DSA_TIMING_GROUP_LIMIT = 96; ++static constexpr size_t SKIPPY_GLM_DSA_TIMING_GROUP_NAME_SIZE = 32; ++static constexpr size_t SKIPPY_GLM_DSA_HOT_TENSOR_LIMIT = 24; ++static constexpr size_t SKIPPY_GLM_DSA_HOT_TENSOR_NAME_SIZE = 64; ++ ++struct skippy_glm_dsa_op_group { ++ bool used = false; ++ char name[SKIPPY_GLM_DSA_TIMING_GROUP_NAME_SIZE] = {}; ++ skippy_glm_dsa_op_stat stats[SKIPPY_GLM_DSA_OP_COUNT] = {}; ++}; ++ ++struct skippy_glm_dsa_hot_tensor { ++ bool used = false; ++ skippy_glm_dsa_op_kind kind = SKIPPY_GLM_DSA_OP_UNKNOWN; ++ enum ggml_op op = GGML_OP_NONE; ++ int64_t elapsed_us = 0; ++ int64_t ne[4] = {}; ++ char name[SKIPPY_GLM_DSA_HOT_TENSOR_NAME_SIZE] = {}; ++}; ++ ++struct skippy_glm_dsa_op_timing { ++ bool enabled = false; ++ bool print_timing = false; ++ bool trace_tensors = false; ++ bool trace_stats = false; ++ int32_t stage_index = -1; ++ size_t token_count = 0; ++ skippy_glm_dsa_op_kind pending_kind = SKIPPY_GLM_DSA_OP_UNKNOWN; ++ int64_t pending_start_us = 0; ++ uint32_t trace_value_limit = 8; ++ uint32_t trace_node_limit = 32; ++ uint32_t trace_stats_max_bytes = 32 * 1024 * 1024; ++ uint32_t trace_nodes = 0; ++ std::string trace_filter; ++ skippy_glm_dsa_op_stat stats[SKIPPY_GLM_DSA_OP_COUNT] = {}; ++ skippy_glm_dsa_op_group groups[SKIPPY_GLM_DSA_TIMING_GROUP_LIMIT] = {}; ++ skippy_glm_dsa_hot_tensor hot_tensors[SKIPPY_GLM_DSA_HOT_TENSOR_LIMIT] = {}; ++}; ++ + struct skippy_model { + llama_model * model = nullptr; + llama_context * ctx = nullptr; + llama_model * mtp_model = nullptr; + llama_context * mtp_ctx = nullptr; + skippy_runtime_config config = {}; ++ skippy_glm_dsa_op_timing glm_dsa_timing = {}; + bool executable = true; + uint32_t lane_count = 1; + std::vector lane_in_use; +@@ -936,14 +1005,27 @@ struct skippy_filter_scope { + }; + + struct skippy_graph_filter_scope { +- explicit skippy_graph_filter_scope(const skippy_runtime_config * config) { +- if (config != nullptr && config->filter_tensors_on_load) { ++ explicit skippy_graph_filter_scope( ++ const skippy_runtime_config * config, ++ int32_t glm_dsa_phase_hint = SKIPPY_GLM_DSA_PHASE_HINT_AUTO) { ++ const bool has_phase_hint = glm_dsa_phase_hint != SKIPPY_GLM_DSA_PHASE_HINT_AUTO; ++ if ((config != nullptr && (config->filter_tensors_on_load || config->glm_dsa_policy_profile != 0)) || ++ has_phase_hint) { + skippy_graph_filter filter; +- filter.enabled = true; +- filter.layer_start = config->layer_start; +- filter.layer_end = config->layer_end; +- filter.include_embeddings = config->include_embeddings; +- filter.include_output = config->include_output; ++ if (config != nullptr) { ++ filter.enabled = config->filter_tensors_on_load; ++ filter.layer_start = config->layer_start; ++ filter.layer_end = config->layer_end; ++ filter.include_embeddings = config->include_embeddings; ++ filter.include_output = config->include_output; ++ filter.glm_dsa_policy_profile = config->glm_dsa_policy_profile; ++ filter.glm_dsa_policy_flags = config->glm_dsa_policy_flags; ++ filter.glm_dsa_short_prefill_max_tokens = config->glm_dsa_short_prefill_max_tokens; ++ filter.glm_dsa_direct_sparse_decode_max_top_k = config->glm_dsa_direct_sparse_decode_max_top_k; ++ filter.glm_dsa_dense_sparse_mask_max_bytes = config->glm_dsa_dense_sparse_mask_max_bytes; ++ filter.glm_dsa_compact_flash_min_kv = config->glm_dsa_compact_flash_min_kv; ++ } ++ filter.glm_dsa_phase_hint = glm_dsa_phase_hint; + skippy_graph_set_filter(filter); + enabled = true; + } +@@ -1070,6 +1152,55 @@ struct skippy_gemma3n_altup_scope { + bool enabled = false; + }; + ++struct skippy_glm_dsa_top_k_scope { ++ skippy_glm_dsa_top_k_scope( ++ const skippy_activation_desc * desc, ++ const void * payload, ++ size_t hidden_bytes, ++ uint32_t n_top_k, ++ uint32_t n_stream, ++ llama_pos pos_start) { ++ if (desc != nullptr && ++ payload != nullptr && ++ (desc->flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { ++ skippy_activation_glm_dsa_top_k sideband; ++ sideband.values = reinterpret_cast(static_cast(payload) + hidden_bytes); ++ sideband.token_count = desc->token_count; ++ sideband.n_top_k = n_top_k; ++ sideband.n_stream = n_stream; ++ sideband.pos_start = pos_start; ++ skippy_graph_set_glm_dsa_top_k(sideband); ++ enabled = true; ++ } ++ } ++ ++ skippy_glm_dsa_top_k_scope( ++ const int32_t * values, ++ uint32_t token_count, ++ uint32_t n_top_k, ++ uint32_t n_stream, ++ llama_pos pos_start) { ++ if (values != nullptr && token_count > 0 && n_top_k > 0 && n_stream > 0) { ++ skippy_activation_glm_dsa_top_k sideband; ++ sideband.values = values; ++ sideband.token_count = token_count; ++ sideband.n_top_k = n_top_k; ++ sideband.n_stream = n_stream; ++ sideband.pos_start = pos_start; ++ skippy_graph_set_glm_dsa_top_k(sideband); ++ enabled = true; ++ } ++ } ++ ++ ~skippy_glm_dsa_top_k_scope() { ++ if (enabled) { ++ skippy_graph_clear_glm_dsa_top_k(); ++ } ++ } ++ ++ bool enabled = false; ++}; ++ + static bool skippy_is_filtered(const skippy_session * session) { + return session != nullptr && + session->stage_model != nullptr && +@@ -1095,6 +1226,124 @@ static bool skippy_is_gemma3n_activation_model(const skippy_session * session) { + return session->stage_model->model->arch == LLM_ARCH_GEMMA3N; + } + ++static bool skippy_is_glm_dsa_activation_model(const skippy_session * session) { ++ if (session == nullptr || session->stage_model == nullptr || session->stage_model->model == nullptr) { ++ return false; ++ } ++ return session->stage_model->model->arch == LLM_ARCH_GLM_DSA; ++} ++ ++static bool skippy_glm_dsa_layer_has_indexer(const llama_layer & layer) { ++ const bool has_any = ++ layer.indexer_k_norm || ++ layer.indexer_k_norm_b || ++ layer.indexer_proj || ++ layer.indexer_attn_k || ++ layer.indexer_attn_q_b; ++ ++ const bool has_all = ++ layer.indexer_k_norm && ++ layer.indexer_k_norm_b && ++ layer.indexer_proj && ++ layer.indexer_attn_k && ++ layer.indexer_attn_q_b; ++ ++ return has_any == has_all && has_all; ++} ++ ++static int skippy_glm_dsa_indexshare_freq() { ++ const char * value = getenv("LLAMA_GLM_DSA_INDEXSHARE_FREQ"); ++ if (value == nullptr) { ++ return 1; ++ } ++ ++ const int freq = atoi(value); ++ return freq > 0 ? freq : 1; ++} ++ ++static bool skippy_glm_dsa_indexshare_pattern_layer_is_full(const char * pattern, int layer, bool * matched) { ++ *matched = false; ++ if (pattern == nullptr || pattern[0] == '\0') { ++ return false; ++ } ++ ++ int layer_index = 0; ++ for (const char * p = pattern; *p != '\0'; ++p) { ++ const char value = static_cast(std::toupper(static_cast(*p))); ++ if (value != 'F' && value != 'S') { ++ continue; ++ } ++ if (layer_index == layer) { ++ *matched = true; ++ return value == 'F'; ++ } ++ ++layer_index; ++ } ++ ++ return false; ++} ++ ++static bool skippy_glm_dsa_metadata_layer_uses_indexer(const llama_hparams & hparams, int32_t layer_index, bool * matched) { ++ *matched = false; ++ if (layer_index < 0 || static_cast(layer_index) >= hparams.n_layer()) { ++ return false; ++ } ++ ++ if (hparams.indexer_types_present) { ++ *matched = true; ++ return hparams.indexer_types[layer_index] == 1; ++ } ++ ++ if (hparams.indexer_top_k_freq > 0) { ++ const uint32_t layer = static_cast(layer_index); ++ const uint32_t offset = hparams.indexer_skip_top_k_offset; ++ *matched = true; ++ return layer < offset || (layer >= offset && ((layer - offset + 1) % hparams.indexer_top_k_freq) == 0); ++ } ++ ++ return false; ++} ++ ++static bool skippy_glm_dsa_layer_uses_indexer(const llama_hparams & hparams, int32_t layer_index, const llama_layer & layer) { ++ bool metadata_matched = false; ++ const bool metadata_full = skippy_glm_dsa_metadata_layer_uses_indexer(hparams, layer_index, &metadata_matched); ++ if (metadata_matched) { ++ return metadata_full; ++ } ++ ++ bool pattern_matched = false; ++ const bool pattern_full = skippy_glm_dsa_indexshare_pattern_layer_is_full( ++ getenv("LLAMA_GLM_DSA_INDEXSHARE_PATTERN"), ++ layer_index, ++ &pattern_matched); ++ if (pattern_matched) { ++ return pattern_full; ++ } ++ ++ if (!skippy_glm_dsa_layer_has_indexer(layer)) { ++ return false; ++ } ++ ++ const int freq = skippy_glm_dsa_indexshare_freq(); ++ return freq <= 1 || (layer_index % freq) == 0; ++} ++ ++static bool skippy_glm_dsa_layer_starts_consumer_group(const skippy_session * session, int32_t layer_start) { ++ if (!skippy_is_glm_dsa_activation_model(session)) { ++ return false; ++ } ++ const skippy_runtime_config & config = session->stage_model->config; ++ const llama_model * model = session->stage_model->model; ++ return config.filter_tensors_on_load && ++ layer_start > 0 && ++ layer_start < static_cast(model->hparams.n_layer()) && ++ !skippy_glm_dsa_layer_uses_indexer(model->hparams, layer_start, model->layers[layer_start]); ++} ++ ++static bool skippy_glm_dsa_stage_starts_in_consumer_group(const skippy_session * session) { ++ return skippy_glm_dsa_layer_starts_consumer_group(session, session->stage_model->config.layer_start); ++} ++ + static size_t skippy_activation_hidden_bytes(const skippy_session * session, size_t token_count) { + if (session == nullptr || session->stage_model == nullptr || session->stage_model->model == nullptr) { + return 0; +@@ -1105,6 +1354,113 @@ static size_t skippy_activation_hidden_bytes(const skippy_session * session, siz + sizeof(float); + } + ++static uint32_t skippy_glm_dsa_top_k_count_for_n_kv(const skippy_session * session, uint64_t n_kv) { ++ if (!skippy_is_glm_dsa_activation_model(session)) { ++ return 0; ++ } ++ ++ // GLM-DSA top-k tensors are shaped from llama_kv_cache::get_n_kv(), not ++ // the logical sequence length. That graph width is padded to at least 256 ++ // so the scheduler can reuse stable KV-cache graph shapes. ++ const uint64_t padded_n_kv = std::max(256, GGML_PAD(n_kv, 256)); ++ return static_cast( ++ std::min(padded_n_kv, session->stage_model->model->hparams.indexer_top_k)); ++} ++ ++static uint32_t skippy_glm_dsa_top_k_count(const skippy_session * session, size_t token_count) { ++ const uint64_t n_kv = static_cast(std::max(session->n_past, 0)) + token_count; ++ return skippy_glm_dsa_top_k_count_for_n_kv(session, n_kv); ++} ++ ++// Activation sidebands carry the actual top-k tensor width produced by ++// glm-dsa.cpp, not the padded graph-capacity width used for output buffers. ++static uint32_t skippy_glm_dsa_sideband_top_k_count_for_n_kv( ++ const skippy_session * session, ++ uint64_t n_kv) { ++ if (!skippy_is_glm_dsa_activation_model(session)) { ++ return 0; ++ } ++ ++ return static_cast( ++ std::min(n_kv, session->stage_model->model->hparams.indexer_top_k)); ++} ++ ++static uint32_t skippy_glm_dsa_sideband_top_k_count(const skippy_session * session, size_t token_count) { ++ const uint64_t n_kv = static_cast(std::max(session->n_past, 0)) + token_count; ++ return skippy_glm_dsa_sideband_top_k_count_for_n_kv(session, n_kv); ++} ++ ++static uint32_t skippy_glm_dsa_n_stream(const skippy_session * session) { ++ GGML_UNUSED(session); ++ return 1; ++} ++ ++static size_t skippy_glm_dsa_top_k_bytes_for_count(size_t token_count, uint32_t n_top_k) { ++ return static_cast(n_top_k)*token_count*sizeof(int32_t); ++} ++ ++static uint32_t skippy_glm_dsa_top_k_count_for_sideband_bytes(size_t token_count, size_t sideband_bytes) { ++ if (token_count == 0 || sideband_bytes == 0 || (sideband_bytes % sizeof(int32_t)) != 0) { ++ return 0; ++ } ++ ++ const size_t sideband_i32 = sideband_bytes/sizeof(int32_t); ++ if ((sideband_i32 % token_count) != 0) { ++ return 0; ++ } ++ ++ return static_cast(std::min( ++ sideband_i32/token_count, ++ std::numeric_limits::max())); ++} ++ ++static bool skippy_glm_dsa_sideband_top_k_count_is_valid( ++ size_t token_count, ++ uint32_t actual_top_k, ++ uint32_t expected_top_k) { ++ if (actual_top_k == 0 || expected_top_k == 0) { ++ return false; ++ } ++ ++ // Prefill sidebands can be chunk-local: llama.cpp emits the actual ++ // per-chunk top-k width, while the receiving staged session may have a ++ // larger accumulated KV prefix. ++ if (token_count > 1) { ++ return actual_top_k <= expected_top_k; ++ } ++ ++ // Decode producers export the actual top-k tensor width from their graph. ++ // The receiving session can have a larger accumulated KV prefix than the ++ // producer's current sideband width, especially when an OpenAI generation ++ // loop advances through cached or fused decode steps. The sideband is still ++ // valid as long as it is non-empty, token-major, and does not claim more ++ // rows than the receiver can make visible. ++ return actual_top_k <= expected_top_k; ++} ++ ++static size_t skippy_glm_dsa_top_k_bytes(const skippy_session * session, size_t token_count) { ++ const uint32_t n_top_k = skippy_glm_dsa_top_k_count(session, token_count); ++ return skippy_glm_dsa_top_k_bytes_for_count(token_count, n_top_k); ++} ++ ++static uint32_t skippy_glm_dsa_top_k_count_from_desc( ++ const skippy_session * session, ++ const skippy_activation_desc * desc) { ++ if (desc == nullptr || (desc->flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) == 0) { ++ return 0; ++ } ++ ++ const size_t token_count = desc->token_count; ++ const size_t hidden_bytes = skippy_activation_hidden_bytes(session, token_count); ++ if (desc->payload_bytes < hidden_bytes) { ++ return 0; ++ } ++ ++ return skippy_glm_dsa_top_k_count_for_sideband_bytes( ++ token_count, ++ static_cast(desc->payload_bytes - hidden_bytes)); ++} ++ + static size_t skippy_gemma3n_altup_bytes(const skippy_session * session, size_t token_count) { + if (session == nullptr || session->stage_model == nullptr || session->stage_model->model == nullptr) { + return 0; +@@ -1122,6 +1478,11 @@ static uint64_t skippy_output_activation_flags( + if (skippy_emits_activation_frame(session) && skippy_is_gemma3n_activation_model(session)) { + return SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP; + } ++ if (skippy_emits_activation_frame(session) && ++ skippy_is_glm_dsa_activation_model(session) && ++ skippy_glm_dsa_layer_starts_consumer_group(session, session->stage_model->config.layer_end)) { ++ return SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K; ++ } + if (!skippy_emits_activation_frame(session) || !skippy_is_rwkv7_activation_model(session)) { + return 0; + } +@@ -1145,6 +1506,9 @@ static size_t skippy_activation_payload_bytes( + return skippy_gemma3n_altup_bytes(session, token_count); + } + size_t payload_bytes = hidden_bytes; ++ if ((flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { ++ payload_bytes += skippy_glm_dsa_top_k_bytes(session, token_count); ++ } + if ((flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0) { + payload_bytes += hidden_bytes; + } +@@ -1155,11 +1519,12 @@ static bool skippy_has_activation_payload(const skippy_activation_desc * desc, c + return desc != nullptr && desc->payload_bytes > 0 && payload != nullptr; + } + +-static enum skippy_status skippy_validate_frame_input( ++static enum skippy_status skippy_validate_frame_input_sequences( + skippy_session * session, + const skippy_activation_desc * input_desc, + const void * input_payload, + size_t expected_token_count, ++ uint32_t expected_sequence_count, + struct skippy_error ** out_error) { + if (session == nullptr || session->stage_model == nullptr) { + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session is required"); +@@ -1185,7 +1550,7 @@ static enum skippy_status skippy_validate_frame_input( + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "activation frame must be version 1 F32 token-major"); + return SKIPPY_STATUS_INVALID_ARGUMENT; + } +- if (input_desc->sequence_count != 1 || input_desc->token_count != expected_token_count) { ++ if (input_desc->sequence_count != expected_sequence_count || input_desc->token_count != expected_token_count) { + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "activation frame token or sequence count does not match request"); + return SKIPPY_STATUS_INVALID_ARGUMENT; + } +@@ -1194,131 +1559,1177 @@ static enum skippy_status skippy_validate_frame_input( + return SKIPPY_STATUS_INVALID_ARGUMENT; + } + +- const uint64_t supported_flags = SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST | SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP; +- if ((input_desc->flags & ~supported_flags) != 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "activation frame has unsupported sideband flags"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ const uint64_t supported_flags = ++ SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST | ++ SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP | ++ SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K; ++ if ((input_desc->flags & ~supported_flags) != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "activation frame has unsupported sideband flags"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0 && !skippy_is_glm_dsa_activation_model(session)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "GLM-DSA top-k sideband is only valid for GLM-DSA stages"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (skippy_glm_dsa_stage_starts_in_consumer_group(session) && ++ (input_desc->flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) == 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "GLM-DSA consumer slices require top-k sideband input"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0 && ++ (input_desc->flags & (SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP | SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST)) != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "GLM-DSA top-k sideband cannot be combined with other activation sidebands"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0 && !skippy_is_gemma3n_activation_model(session)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "Gemma3n AltUp activation payload is only valid for Gemma3n stages"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (skippy_is_gemma3n_activation_model(session) && (input_desc->flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) == 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "non-first Gemma3n runtime slices require AltUp activation payload"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0 && ++ (input_desc->flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "Gemma3n AltUp and RWKV7 v_first activation flags cannot be combined"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0 && !skippy_is_rwkv7_activation_model(session)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "RWKV7 v_first sideband is only valid for RWKV7 stages"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (skippy_is_rwkv7_activation_model(session) && (input_desc->flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) == 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "non-first RWKV7 runtime slices require v_first activation sideband"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ const size_t hidden_bytes = skippy_activation_hidden_bytes(session, expected_token_count); ++ if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { ++ if (input_desc->payload_bytes < hidden_bytes) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "GLM-DSA activation frame is smaller than model hidden size"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ const size_t sideband_bytes = static_cast(input_desc->payload_bytes - hidden_bytes); ++ const uint32_t actual_top_k = skippy_glm_dsa_top_k_count_for_sideband_bytes(expected_token_count, sideband_bytes); ++ if (actual_top_k == 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "GLM-DSA top-k sideband payload size is not token-major i32"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ const uint32_t expected_top_k = skippy_glm_dsa_sideband_top_k_count(session, expected_token_count); ++ if (!skippy_glm_dsa_sideband_top_k_count_is_valid(expected_token_count, actual_top_k, expected_top_k)) { ++ char message[256]; ++ std::snprintf( ++ message, ++ sizeof(message), ++ "GLM-DSA top-k sideband width does not match expected IndexShare width: expected=%u actual=%u tokens=%zu sideband_bytes=%zu", ++ expected_top_k, ++ actual_top_k, ++ expected_token_count, ++ sideband_bytes); ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, message); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ } else { ++ const size_t expected_bytes = skippy_activation_payload_bytes(session, expected_token_count, input_desc->flags); ++ if (input_desc->payload_bytes != expected_bytes) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "activation frame payload size does not match model hidden size"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ } ++ ++ return SKIPPY_STATUS_OK; ++} ++ ++static enum skippy_status skippy_validate_frame_input( ++ skippy_session * session, ++ const skippy_activation_desc * input_desc, ++ const void * input_payload, ++ size_t expected_token_count, ++ struct skippy_error ** out_error) { ++ return skippy_validate_frame_input_sequences( ++ session, ++ input_desc, ++ input_payload, ++ expected_token_count, ++ 1, ++ out_error); ++} ++ ++static enum skippy_status skippy_prepare_output_activation_frame( ++ skippy_session * session, ++ size_t token_count, ++ void * output_payload, ++ size_t output_payload_capacity, ++ size_t * out_output_payload_bytes, ++ skippy_activation_desc * output_desc, ++ const skippy_activation_desc * input_desc, ++ struct skippy_error ** out_error) { ++ const uint64_t output_flags = skippy_output_activation_flags(session, input_desc); ++ const size_t payload_bytes = skippy_emits_activation_frame(session) ? ++ skippy_activation_payload_bytes(session, token_count, output_flags) : 0; ++ ++ if (out_output_payload_bytes != nullptr) { ++ *out_output_payload_bytes = payload_bytes; ++ } ++ ++ if (payload_bytes > 0) { ++ if (output_payload == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "output_payload is required for runtime-slice activation output"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ if (output_payload_capacity < payload_bytes) { ++ skippy_set_error(out_error, SKIPPY_STATUS_BUFFER_TOO_SMALL, "output activation buffer is too small"); ++ return SKIPPY_STATUS_BUFFER_TOO_SMALL; ++ } ++ } else if (output_payload_capacity > 0 && output_payload == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "output_payload is required when output capacity is non-zero"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ if (output_desc != nullptr) { ++ *output_desc = {}; ++ output_desc->version = 1; ++ output_desc->dtype = payload_bytes > 0 ? SKIPPY_ACTIVATION_DTYPE_F32 : SKIPPY_ACTIVATION_DTYPE_UNKNOWN; ++ output_desc->layout = payload_bytes > 0 ? SKIPPY_ACTIVATION_LAYOUT_TOKEN_MAJOR : SKIPPY_ACTIVATION_LAYOUT_OPAQUE; ++ output_desc->producer_stage_index = session != nullptr ? session->stage_model->config.stage_index : -1; ++ output_desc->layer_start = session != nullptr ? session->stage_model->config.layer_start : 0; ++ output_desc->layer_end = session != nullptr ? session->stage_model->config.layer_end : 0; ++ output_desc->token_count = static_cast(std::min(token_count, std::numeric_limits::max())); ++ output_desc->sequence_count = token_count > 0 ? 1 : 0; ++ output_desc->payload_bytes = payload_bytes; ++ output_desc->flags = output_flags; ++ } ++ ++ return SKIPPY_STATUS_OK; ++} ++ ++static void skippy_glm_dsa_op_timing_begin(skippy_session * session, size_t token_count); ++static void skippy_glm_dsa_op_timing_end(skippy_session * session); ++ ++static enum skippy_status skippy_decode_batch( ++ skippy_session * session, ++ llama_batch batch, ++ size_t token_count, ++ int32_t glm_dsa_phase_hint, ++ struct skippy_error ** out_error) { ++ if (session == nullptr || session->ctx == nullptr || token_count == 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session and at least one token are required"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ ++ skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config, glm_dsa_phase_hint); ++ const char * phase_log = getenv("SKIPPY_GLM_DSA_LOG_DIRECT_SPARSE_DECISIONS"); ++ if (phase_log != nullptr && atoi(phase_log) != 0) { ++ LLAMA_LOG_INFO( ++ "skippy: glm_dsa_phase_hint requested=%d active=%d tokens=%zu\n", ++ glm_dsa_phase_hint, ++ skippy_graph_get_filter().glm_dsa_phase_hint, ++ token_count); ++ } ++ skippy_glm_dsa_op_timing_begin(session, token_count); ++ const int32_t rc = llama_decode(session->ctx, batch); ++ skippy_glm_dsa_op_timing_end(session); ++ if (rc != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ ++ session->n_past += static_cast(token_count); ++ return skippy_success(out_error); ++} ++ ++static bool skippy_mtp_available(const skippy_session * session) { ++ return session != nullptr && ++ session->stage_model != nullptr && ++ session->stage_model->mtp_ctx != nullptr && ++ session->stage_model->config.include_output; ++} ++ ++static bool skippy_env_enabled(const char * name) { ++ const char * value = std::getenv(name); ++ if (value == nullptr || value[0] == '\0') { ++ return false; ++ } ++ return std::strcmp(value, "0") != 0 && ++ std::strcmp(value, "false") != 0 && ++ std::strcmp(value, "off") != 0 && ++ std::strcmp(value, "no") != 0; ++} ++ ++static bool skippy_env_disabled(const char * name) { ++ const char * value = std::getenv(name); ++ if (value == nullptr || value[0] == '\0') { ++ return false; ++ } ++ return std::strcmp(value, "0") == 0 || ++ std::strcmp(value, "false") == 0 || ++ std::strcmp(value, "FALSE") == 0 || ++ std::strcmp(value, "off") == 0 || ++ std::strcmp(value, "OFF") == 0 || ++ std::strcmp(value, "no") == 0 || ++ std::strcmp(value, "NO") == 0; ++} ++ ++static uint32_t skippy_env_u32(const char * name, uint32_t fallback, uint32_t min_value, uint32_t max_value) { ++ const char * value = std::getenv(name); ++ if (value == nullptr || value[0] == '\0') { ++ return fallback; ++ } ++ ++ char * end = nullptr; ++ errno = 0; ++ const unsigned long parsed = std::strtoul(value, &end, 10); ++ if (errno != 0 || end == value || *end != '\0') { ++ return fallback; ++ } ++ ++ return std::min(max_value, std::max(min_value, static_cast(parsed))); ++} ++ ++static bool skippy_mtp_greedy_sampling_fastpath_enabled() { ++ return skippy_env_enabled("SKIPPY_NATIVE_MTP_GREEDY_SAMPLING_FASTPATH"); ++} ++ ++static bool skippy_glm_dsa_op_timing_enabled() { ++ return skippy_env_enabled("SKIPPY_GLM_DSA_OP_TIMING"); ++} ++ ++static bool skippy_glm_dsa_tensor_trace_enabled() { ++ return skippy_env_enabled("SKIPPY_GLM_DSA_TENSOR_TRACE"); ++} ++ ++static bool skippy_glm_dsa_tensor_trace_stats_enabled() { ++ return skippy_env_enabled("SKIPPY_GLM_DSA_TENSOR_TRACE_STATS"); ++} ++ ++static bool skippy_glm_dsa_direct_sparse_decision_log_enabled() { ++ return skippy_env_enabled("SKIPPY_GLM_DSA_LOG_DIRECT_SPARSE_DECISIONS"); ++} ++ ++static bool skippy_glm_dsa_direct_sparse_attn_enabled() { ++ if (skippy_env_disabled("SKIPPY_GLM_DSA_ENABLE_DIRECT_SPARSE_ATTN") || ++ skippy_env_enabled("SKIPPY_GLM_DSA_DISABLE_DIRECT_SPARSE_ATTN")) { ++ return false; ++ } ++ ++ return true; ++} ++ ++static bool skippy_glm_dsa_direct_sparse_prefill_enabled() { ++ if (skippy_env_disabled("SKIPPY_GLM_DSA_ENABLE_DIRECT_SPARSE_PREFILL") || ++ skippy_env_enabled("SKIPPY_GLM_DSA_DISABLE_DIRECT_SPARSE_PREFILL")) { ++ return false; ++ } ++ ++ return true; ++} ++ ++static bool skippy_glm_dsa_unproven_large_direct_sparse_prefill_enabled() { ++ return skippy_env_enabled("SKIPPY_GLM_DSA_ENABLE_UNPROVEN_LARGE_DIRECT_SPARSE_PREFILL"); ++} ++ ++static uint32_t skippy_glm_dsa_direct_sparse_prefill_max_tokens() { ++ return skippy_env_u32("SKIPPY_GLM_DSA_DIRECT_SPARSE_PREFILL_MAX_TOKENS", 8, 1, 4096); ++} ++ ++static uint64_t skippy_glm_dsa_dense_sparse_mask_max_bytes() { ++ const char * value = std::getenv("SKIPPY_GLM_DSA_DENSE_SPARSE_MASK_MAX_BYTES"); ++ if (value == nullptr || value[0] == '\0') { ++ return 512ULL * 1024ULL * 1024ULL; ++ } ++ ++ char * end = nullptr; ++ errno = 0; ++ const unsigned long long parsed = std::strtoull(value, &end, 10); ++ if (errno != 0 || end == value || *end != '\0' || parsed == 0) { ++ return UINT64_MAX; ++ } ++ return parsed; ++} ++ ++static constexpr int64_t SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_DEFAULT_MAX_TOP_K = 256; ++ ++static int64_t skippy_glm_dsa_direct_sparse_decode_max_top_k() { ++ const char * value = std::getenv("SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_MAX_TOP_K"); ++ if (value == nullptr || value[0] == '\0') { ++ return SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_DEFAULT_MAX_TOP_K; ++ } ++ ++ char * end = nullptr; ++ errno = 0; ++ const long long parsed = std::strtoll(value, &end, 10); ++ if (errno != 0 || end == value || *end != '\0') { ++ return SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_DEFAULT_MAX_TOP_K; ++ } ++ ++ return parsed <= 0 ? INT64_MAX : parsed; ++} ++ ++ ++static bool skippy_name_starts_with(const char * name, const char * prefix) { ++ return name != nullptr && std::strncmp(name, prefix, std::strlen(prefix)) == 0; ++} ++ ++static const char * skippy_glm_dsa_op_name(skippy_glm_dsa_op_kind kind) { ++ switch (kind) { ++ case SKIPPY_GLM_DSA_OP_INDEXER: ++ return "indexer"; ++ case SKIPPY_GLM_DSA_OP_TOP_K: ++ return "top_k"; ++ case SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL: ++ return "sparse_mask_fill"; ++ case SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK: ++ return "sparse_mask_topk"; ++ case SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD: ++ return "sparse_mask_add"; ++ case SKIPPY_GLM_DSA_OP_DSA_SPARSE_ATTN: ++ return "dsa_sparse_attn"; ++ case SKIPPY_GLM_DSA_OP_COMPACT_GET_ROWS: ++ return "compact_get_rows"; ++ case SKIPPY_GLM_DSA_OP_MLA_ATTENTION: ++ return "mla_attention"; ++ case SKIPPY_GLM_DSA_OP_ROUTED_MOE_ROUTE: ++ return "routed_moe_route"; ++ case SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE_UP: ++ return "routed_moe_gate_up"; ++ case SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE: ++ return "routed_moe_gate"; ++ case SKIPPY_GLM_DSA_OP_ROUTED_MOE_UP: ++ return "routed_moe_up"; ++ case SKIPPY_GLM_DSA_OP_ROUTED_MOE_ACT: ++ return "routed_moe_act"; ++ case SKIPPY_GLM_DSA_OP_ROUTED_MOE_DOWN: ++ return "routed_moe_down"; ++ case SKIPPY_GLM_DSA_OP_ROUTED_MOE_WEIGHTED: ++ return "routed_moe_weighted"; ++ case SKIPPY_GLM_DSA_OP_ROUTED_MOE_AGGREGATE: ++ return "routed_moe_aggregate"; ++ case SKIPPY_GLM_DSA_OP_SHARED_EXPERT: ++ return "shared_expert"; ++ case SKIPPY_GLM_DSA_OP_COUNT: ++ case SKIPPY_GLM_DSA_OP_UNKNOWN: ++ break; ++ } ++ return "unknown"; ++} ++ ++static skippy_glm_dsa_op_kind skippy_glm_dsa_op_kind_for_tensor(const char * name) { ++ if (skippy_name_starts_with(name, "top_k")) { ++ return SKIPPY_GLM_DSA_OP_TOP_K; ++ } ++ if (skippy_name_starts_with(name, "indexer_")) { ++ return SKIPPY_GLM_DSA_OP_INDEXER; ++ } ++ if (skippy_name_starts_with(name, "dsa_sparse_mask_fill")) { ++ return SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL; ++ } ++ if (skippy_name_starts_with(name, "dsa_sparse_mask_topk")) { ++ return SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK; ++ } ++ if (skippy_name_starts_with(name, "dsa_sparse_mask")) { ++ return SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD; ++ } ++ if (skippy_name_starts_with(name, "dsa_sparse_attn")) { ++ return SKIPPY_GLM_DSA_OP_DSA_SPARSE_ATTN; ++ } ++ if (skippy_name_starts_with(name, "dsa_compact_")) { ++ return SKIPPY_GLM_DSA_OP_COMPACT_GET_ROWS; ++ } ++ if (skippy_name_starts_with(name, "kqv_out")) { ++ return SKIPPY_GLM_DSA_OP_MLA_ATTENTION; ++ } ++ if (skippy_name_starts_with(name, "ffn_moe_gate_up")) { ++ return SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE_UP; ++ } ++ if (skippy_name_starts_with(name, "ffn_moe_gate")) { ++ return SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE; ++ } ++ if (skippy_name_starts_with(name, "ffn_moe_up")) { ++ return SKIPPY_GLM_DSA_OP_ROUTED_MOE_UP; ++ } ++ if (skippy_name_starts_with(name, "ffn_moe_silu") || ++ skippy_name_starts_with(name, "ffn_moe_swiglu") || ++ skippy_name_starts_with(name, "ffn_moe_geglu") || ++ skippy_name_starts_with(name, "ffn_moe_gelu") || ++ skippy_name_starts_with(name, "ffn_moe_reglu") || ++ skippy_name_starts_with(name, "ffn_moe_relu")) { ++ return SKIPPY_GLM_DSA_OP_ROUTED_MOE_ACT; ++ } ++ if (skippy_name_starts_with(name, "ffn_moe_down")) { ++ return SKIPPY_GLM_DSA_OP_ROUTED_MOE_DOWN; ++ } ++ if (skippy_name_starts_with(name, "ffn_moe_weighted")) { ++ return SKIPPY_GLM_DSA_OP_ROUTED_MOE_WEIGHTED; ++ } ++ if (skippy_name_starts_with(name, "ffn_moe_out")) { ++ return SKIPPY_GLM_DSA_OP_ROUTED_MOE_AGGREGATE; ++ } ++ if (skippy_name_starts_with(name, "ffn_moe_")) { ++ return SKIPPY_GLM_DSA_OP_ROUTED_MOE_ROUTE; ++ } ++ if (skippy_name_starts_with(name, "ffn_shexp") || ++ skippy_name_starts_with(name, "ffn_up") || ++ skippy_name_starts_with(name, "ffn_gate") || ++ skippy_name_starts_with(name, "ffn_silu") || ++ skippy_name_starts_with(name, "ffn_swiglu") || ++ skippy_name_starts_with(name, "ffn_down")) { ++ return SKIPPY_GLM_DSA_OP_SHARED_EXPERT; ++ } ++ return SKIPPY_GLM_DSA_OP_UNKNOWN; ++} ++ ++static void skippy_glm_dsa_group_name_for_tensor( ++ const char * tensor_name, ++ char * group_name, ++ size_t group_name_size) { ++ if (group_name == nullptr || group_name_size == 0) { ++ return; ++ } ++ ++ group_name[0] = '\0'; ++ if (tensor_name == nullptr || tensor_name[0] == '\0') { ++ std::snprintf(group_name, group_name_size, "global"); ++ return; ++ } ++ ++ const char * suffix = std::strrchr(tensor_name, '-'); ++ if (suffix == nullptr || suffix[1] < '0' || suffix[1] > '9') { ++ std::snprintf(group_name, group_name_size, "global"); ++ return; ++ } ++ ++ const char * layer = suffix + 1; ++ size_t layer_len = 0; ++ while (layer[layer_len] >= '0' && layer[layer_len] <= '9') { ++ ++layer_len; ++ } ++ std::snprintf(group_name, group_name_size, "layer_%.*s", static_cast(layer_len), layer); ++} ++ ++static int skippy_glm_dsa_layer_from_group_name(const char * group_name) { ++ const char * prefix = "layer_"; ++ if (!skippy_name_starts_with(group_name, prefix)) { ++ return -1; ++ } ++ ++ const char * layer = group_name + std::strlen(prefix); ++ if (layer[0] == '\0') { ++ return -1; ++ } ++ ++ char * end = nullptr; ++ errno = 0; ++ const long parsed = std::strtol(layer, &end, 10); ++ if (errno != 0 || end == layer || *end != '\0' || parsed < 0 || parsed > std::numeric_limits::max()) { ++ return -1; ++ } ++ ++ return static_cast(parsed); ++} ++ ++static void skippy_glm_dsa_log_direct_sparse_decision_for_tensor( ++ const skippy_glm_dsa_op_timing & timing, ++ const ggml_tensor * tensor, ++ skippy_glm_dsa_op_kind kind) { ++ if (!skippy_glm_dsa_direct_sparse_decision_log_enabled() || tensor == nullptr) { ++ return; ++ } ++ ++ const ggml_tensor * top_k = nullptr; ++ bool use_direct = false; ++ if (kind == SKIPPY_GLM_DSA_OP_DSA_SPARSE_ATTN && tensor->op == GGML_OP_DSA_SPARSE_ATTN) { ++ top_k = tensor->src[4]; ++ use_direct = true; ++ } else if ((kind == SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK || kind == SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD) && ++ tensor->op == GGML_OP_DSA_SPARSE_MASK) { ++ top_k = tensor->src[1]; ++ } else { ++ return; ++ } ++ ++ char group_name[SKIPPY_GLM_DSA_TIMING_GROUP_NAME_SIZE] = {}; ++ skippy_glm_dsa_group_name_for_tensor(tensor->name, group_name, sizeof(group_name)); ++ const int layer = skippy_glm_dsa_layer_from_group_name(group_name); ++ const int64_t sparse_batch = top_k != nullptr ? top_k->ne[1] : -1; ++ const int64_t sparse_streams = top_k != nullptr ? top_k->ne[2] : -1; ++ const uint32_t prefill_cap = skippy_glm_dsa_direct_sparse_prefill_max_tokens(); ++ const bool prefill_enabled = skippy_glm_dsa_direct_sparse_prefill_enabled(); ++ const ggml_tensor * dense_mask_shape = nullptr; ++ if (tensor->op == GGML_OP_DSA_SPARSE_ATTN) { ++ dense_mask_shape = tensor->src[3]; ++ } else if (tensor->op == GGML_OP_DSA_SPARSE_MASK) { ++ dense_mask_shape = tensor; ++ } ++ const uint64_t dense_mask_bytes = dense_mask_shape != nullptr ? ggml_nbytes(dense_mask_shape) : 0; ++ const uint64_t dense_mask_limit = skippy_glm_dsa_dense_sparse_mask_max_bytes(); ++ const int64_t sparse_kv = dense_mask_shape != nullptr ? dense_mask_shape->ne[1] : -1; ++ const int64_t sparse_top_k = top_k != nullptr ? top_k->ne[0] : -1; ++ const int64_t decode_max_top_k = skippy_glm_dsa_direct_sparse_decode_max_top_k(); ++ const int64_t kv_topk_ratio = sparse_top_k > 0 ? sparse_kv / sparse_top_k : 0; ++ const bool direct_enabled = use_direct || skippy_glm_dsa_direct_sparse_attn_enabled(); ++ const bool decode_shape = ++ sparse_streams > 0 && ++ sparse_batch == 1 && ++ timing.token_count == static_cast(sparse_streams); ++ const bool prefill_shape = ++ prefill_enabled && ++ sparse_batch >= 1 && ++ sparse_batch <= static_cast(prefill_cap); ++ const bool large_prefill_shape = ++ prefill_enabled && ++ skippy_glm_dsa_unproven_large_direct_sparse_prefill_enabled() && ++ sparse_batch > static_cast(prefill_cap) && ++ dense_mask_bytes > dense_mask_limit; ++ const bool token_shape_allowed = decode_shape || prefill_shape || large_prefill_shape; ++ const bool direct_top_k_too_large = ++ decode_shape && ++ sparse_top_k > decode_max_top_k; ++ const char * phase = ++ decode_shape ? "decode" : ++ (prefill_shape || large_prefill_shape) ? "prefill" : ++ "other"; ++ const char * selector_reason = ++ use_direct && decode_shape ? "decode" : ++ use_direct && prefill_shape ? "short_prefill" : ++ use_direct && large_prefill_shape ? "dense_mask_guard_large_prefill" : ++ !direct_enabled ? "direct_sparse_disabled" : ++ !token_shape_allowed ? "token_shape_not_allowed" : ++ direct_top_k_too_large ? "direct_sparse_top_k_too_large" : ++ !use_direct ? "dense_mask_fallback" : ++ "unknown"; ++ ++ LLAMA_LOG_INFO( ++ "skippy: glm_dsa_direct_sparse_decision layer=%d ubatch_tokens=%lld sparse_batch=%lld sparse_streams=%lld prefill_cap=%lld decode_max_top_k=%lld sparse_kv=%lld sparse_top_k=%lld min_kv_topk_ratio=%lld kv_topk_ratio=%lld dense_mask_bytes=%llu dense_mask_limit=%llu phase=%s selector_reason=%s direct_enabled=%d prefill_enabled=%d decode_shape=%d prefill_shape=%d large_prefill_shape=%d token_shape_allowed=%d backend_sparse_supported=%d kq_b_ok=%d sinks_ok=%d alibi_ok=%d soft_cap_ok=%d use_direct=%d\n", ++ layer, ++ static_cast(timing.token_count), ++ static_cast(sparse_batch), ++ static_cast(sparse_streams), ++ static_cast(prefill_cap), ++ static_cast(decode_max_top_k), ++ static_cast(sparse_kv), ++ static_cast(sparse_top_k), ++ 0LL, ++ static_cast(kv_topk_ratio), ++ static_cast(dense_mask_bytes), ++ static_cast(dense_mask_limit), ++ phase, ++ selector_reason, ++ direct_enabled ? 1 : 0, ++ prefill_enabled ? 1 : 0, ++ decode_shape ? 1 : 0, ++ prefill_shape ? 1 : 0, ++ large_prefill_shape ? 1 : 0, ++ token_shape_allowed ? 1 : 0, ++ 1, ++ 1, ++ 1, ++ 1, ++ 1, ++ use_direct ? 1 : 0); ++} ++ ++static skippy_glm_dsa_op_group * skippy_glm_dsa_op_timing_group( ++ skippy_glm_dsa_op_timing & timing, ++ const char * group_name) { ++ if (group_name == nullptr || group_name[0] == '\0') { ++ return nullptr; ++ } ++ ++ skippy_glm_dsa_op_group * unused = nullptr; ++ for (size_t i = 0; i < SKIPPY_GLM_DSA_TIMING_GROUP_LIMIT; ++i) { ++ skippy_glm_dsa_op_group & group = timing.groups[i]; ++ if (!group.used) { ++ if (unused == nullptr) { ++ unused = &group; ++ } ++ continue; ++ } ++ if (std::strncmp(group.name, group_name, sizeof(group.name)) == 0) { ++ return &group; ++ } ++ } ++ ++ if (unused == nullptr) { ++ return nullptr; ++ } ++ ++ unused->used = true; ++ std::snprintf(unused->name, sizeof(unused->name), "%s", group_name); ++ return unused; ++} ++ ++static void skippy_glm_dsa_record_group_timing( ++ skippy_glm_dsa_op_timing & timing, ++ skippy_glm_dsa_op_kind kind, ++ const char * tensor_name, ++ int64_t elapsed_us) { ++ char group_name[SKIPPY_GLM_DSA_TIMING_GROUP_NAME_SIZE] = {}; ++ skippy_glm_dsa_group_name_for_tensor(tensor_name, group_name, sizeof(group_name)); ++ skippy_glm_dsa_op_group * group = skippy_glm_dsa_op_timing_group(timing, group_name); ++ if (group == nullptr) { ++ return; ++ } ++ ++ group->stats[kind].nodes += 1; ++ group->stats[kind].elapsed_us += elapsed_us; ++} ++ ++static void skippy_glm_dsa_record_hot_tensor( ++ skippy_glm_dsa_op_timing & timing, ++ skippy_glm_dsa_op_kind kind, ++ const ggml_tensor * tensor, ++ int64_t elapsed_us) { ++ if (tensor == nullptr || elapsed_us <= 0) { ++ return; ++ } ++ ++ size_t slot = SKIPPY_GLM_DSA_HOT_TENSOR_LIMIT; ++ for (size_t i = 0; i < SKIPPY_GLM_DSA_HOT_TENSOR_LIMIT; ++i) { ++ if (!timing.hot_tensors[i].used || elapsed_us > timing.hot_tensors[i].elapsed_us) { ++ slot = i; ++ break; ++ } ++ } ++ if (slot == SKIPPY_GLM_DSA_HOT_TENSOR_LIMIT) { ++ return; ++ } ++ ++ for (size_t i = SKIPPY_GLM_DSA_HOT_TENSOR_LIMIT - 1; i > slot; --i) { ++ timing.hot_tensors[i] = timing.hot_tensors[i - 1]; ++ } ++ ++ skippy_glm_dsa_hot_tensor & record = timing.hot_tensors[slot]; ++ record.used = true; ++ record.kind = kind; ++ record.op = tensor->op; ++ record.elapsed_us = elapsed_us; ++ for (int i = 0; i < 4; ++i) { ++ record.ne[i] = tensor->ne[i]; ++ } ++ std::snprintf(record.name, sizeof(record.name), "%s", tensor->name); ++} ++ ++static bool skippy_glm_dsa_tensor_has_buffer(const ggml_tensor * tensor) { ++ return tensor != nullptr && ++ tensor->data != nullptr && ++ (tensor->view_src != nullptr ? tensor->view_src->buffer : tensor->buffer) != nullptr; ++} ++ ++static bool skippy_glm_dsa_trace_filter_matches(const std::string & filter, const char * name) { ++ if (filter.empty()) { ++ return true; ++ } ++ if (name == nullptr || name[0] == '\0') { ++ return false; ++ } ++ ++ size_t start = 0; ++ while (start <= filter.size()) { ++ const size_t end = filter.find(',', start); ++ const size_t len = (end == std::string::npos ? filter.size() : end) - start; ++ if (len > 0 && std::string(name).find(filter.substr(start, len)) != std::string::npos) { ++ return true; ++ } ++ if (end == std::string::npos) { ++ break; ++ } ++ start = end + 1; ++ } ++ return false; ++} ++ ++template ++static bool skippy_glm_dsa_get_tensor_prefix( ++ const ggml_tensor * tensor, ++ std::vector & values, ++ size_t value_count) { ++ if (tensor->nb[0] != sizeof(T)) { ++ return false; ++ } ++ ++ values.resize(value_count); ++ size_t copied = 0; ++ for (int64_t i3 = 0; i3 < tensor->ne[3] && copied < value_count; ++i3) { ++ for (int64_t i2 = 0; i2 < tensor->ne[2] && copied < value_count; ++i2) { ++ for (int64_t i1 = 0; i1 < tensor->ne[1] && copied < value_count; ++i1) { ++ const size_t row_values = std::min( ++ static_cast(tensor->ne[0]), ++ value_count - copied); ++ const size_t offset = ++ static_cast(i1) * tensor->nb[1] + ++ static_cast(i2) * tensor->nb[2] + ++ static_cast(i3) * tensor->nb[3]; ++ ggml_backend_tensor_get( ++ tensor, ++ values.data() + copied, ++ offset, ++ row_values * sizeof(T)); ++ copied += row_values; ++ } ++ } ++ } ++ return copied == value_count; ++} ++ ++static void skippy_glm_dsa_trace_tensor_values(std::string & line, const ggml_tensor * tensor, uint32_t limit) { ++ const size_t value_count = static_cast( ++ std::min(static_cast(limit), std::max(0, ggml_nelements(tensor)))); ++ ++ line += " values=["; ++ if (value_count == 0) { ++ line += "]"; ++ return; ++ } ++ ++ if (!skippy_glm_dsa_tensor_has_buffer(tensor)) { ++ line += "unallocated]"; ++ return; ++ } ++ ++ switch (tensor->type) { ++ case GGML_TYPE_F32: { ++ std::vector values; ++ if (!skippy_glm_dsa_get_tensor_prefix(tensor, values, value_count)) { ++ line += "unsupported-stride"; ++ break; ++ } ++ for (size_t i = 0; i < value_count; ++i) { ++ line += format("%s%.9g", i == 0 ? "" : ",", static_cast(values[i])); ++ } ++ break; ++ } ++ case GGML_TYPE_F16: ++ case GGML_TYPE_BF16: { ++ const ggml_type_traits * traits = ggml_get_type_traits(tensor->type); ++ if (traits == nullptr || traits->to_float == nullptr) { ++ line += format("unsupported:%s", ggml_type_name(tensor->type)); ++ break; ++ } ++ std::vector raw; ++ std::vector values(value_count); ++ if (!skippy_glm_dsa_get_tensor_prefix(tensor, raw, value_count)) { ++ line += "unsupported-stride"; ++ break; ++ } ++ traits->to_float(raw.data(), values.data(), value_count); ++ for (size_t i = 0; i < value_count; ++i) { ++ line += format("%s%.9g", i == 0 ? "" : ",", static_cast(values[i])); ++ } ++ break; ++ } ++ case GGML_TYPE_I32: { ++ std::vector values; ++ if (!skippy_glm_dsa_get_tensor_prefix(tensor, values, value_count)) { ++ line += "unsupported-stride"; ++ break; ++ } ++ for (size_t i = 0; i < value_count; ++i) { ++ line += format("%s%d", i == 0 ? "" : ",", values[i]); ++ } ++ break; ++ } ++ default: ++ line += format("unsupported:%s", ggml_type_name(tensor->type)); ++ break; ++ } ++ ++ line += "]"; ++} ++ ++static uint64_t skippy_glm_dsa_fnv1a64(const uint8_t * data, size_t size) { ++ uint64_t hash = 1469598103934665603ull; ++ for (size_t i = 0; i < size; ++i) { ++ hash ^= static_cast(data[i]); ++ hash *= 1099511628211ull; + } +- if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0 && !skippy_is_gemma3n_activation_model(session)) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "Gemma3n AltUp activation payload is only valid for Gemma3n stages"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ return hash; ++} ++ ++static void skippy_glm_dsa_trace_numeric_stats_from_f32(std::string & line, const float * values, size_t count) { ++ double sum = 0.0; ++ double max_abs = 0.0; ++ for (size_t i = 0; i < count; ++i) { ++ const double value = static_cast(values[i]); ++ sum += value; ++ max_abs = std::max(max_abs, std::fabs(value)); + } +- if (skippy_is_gemma3n_activation_model(session) && (input_desc->flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) == 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "non-first Gemma3n runtime slices require AltUp activation payload"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ const double mean = count == 0 ? 0.0 : sum / static_cast(count); ++ line += format( ++ " count=%zu sum=%.17g mean=%.17g max_abs=%.17g", ++ count, ++ sum, ++ mean, ++ max_abs); ++} ++ ++static void skippy_glm_dsa_trace_tensor_stats(std::string & line, const ggml_tensor * tensor, uint32_t max_bytes) { ++ if (!skippy_glm_dsa_tensor_has_buffer(tensor)) { ++ line += " stats=unallocated"; ++ return; + } +- if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0 && +- (input_desc->flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "Gemma3n AltUp and RWKV7 v_first activation flags cannot be combined"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ ++ const size_t nbytes = ggml_nbytes(tensor); ++ if (nbytes > static_cast(max_bytes)) { ++ line += format(" stats=skipped:nbytes>%u", max_bytes); ++ return; + } +- if ((input_desc->flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0 && !skippy_is_rwkv7_activation_model(session)) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "RWKV7 v_first sideband is only valid for RWKV7 stages"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ ++ std::vector raw(nbytes); ++ ggml_backend_tensor_get(tensor, raw.data(), 0, nbytes); ++ const uint64_t hash = skippy_glm_dsa_fnv1a64(raw.data(), raw.size()); ++ line += format(" stats=fnv64:%016llx", static_cast(hash)); ++ ++ const size_t count = static_cast(std::max(0, ggml_nelements(tensor))); ++ switch (tensor->type) { ++ case GGML_TYPE_F32: ++ skippy_glm_dsa_trace_numeric_stats_from_f32(line, reinterpret_cast(raw.data()), count); ++ break; ++ case GGML_TYPE_F16: ++ case GGML_TYPE_BF16: { ++ const ggml_type_traits * traits = ggml_get_type_traits(tensor->type); ++ if (traits == nullptr || traits->to_float == nullptr) { ++ break; ++ } ++ std::vector values(count); ++ traits->to_float(raw.data(), values.data(), count); ++ skippy_glm_dsa_trace_numeric_stats_from_f32(line, values.data(), count); ++ break; ++ } ++ case GGML_TYPE_I32: { ++ double sum = 0.0; ++ double max_abs = 0.0; ++ const int32_t * values = reinterpret_cast(raw.data()); ++ for (size_t i = 0; i < count; ++i) { ++ const double value = static_cast(values[i]); ++ sum += value; ++ max_abs = std::max(max_abs, std::fabs(value)); ++ } ++ const double mean = count == 0 ? 0.0 : sum / static_cast(count); ++ line += format( ++ " count=%zu sum=%.17g mean=%.17g max_abs=%.17g", ++ count, ++ sum, ++ mean, ++ max_abs); ++ break; ++ } ++ default: ++ break; + } +- if (skippy_is_rwkv7_activation_model(session) && (input_desc->flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) == 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "non-first RWKV7 runtime slices require v_first activation sideband"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++} ++ ++static void skippy_glm_dsa_dump_trace_tensor( ++ const skippy_glm_dsa_op_timing & timing, ++ const ggml_tensor * tensor) { ++ const char * dump_dir = std::getenv("SKIPPY_GLM_DSA_TENSOR_DUMP_DIR"); ++ if (dump_dir == nullptr || dump_dir[0] == '\0' || timing.token_count != 1 || ++ !skippy_glm_dsa_tensor_has_buffer(tensor)) { ++ return; + } + +- const size_t expected_bytes = skippy_activation_payload_bytes(session, expected_token_count, input_desc->flags); +- if (input_desc->payload_bytes != expected_bytes) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "activation frame payload size does not match model hidden size"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ std::string filename = tensor->name; ++ for (char & c : filename) { ++ if (!std::isalnum(static_cast(c)) && c != '-' && c != '_') { ++ c = '_'; ++ } + } ++ filename += format( ++ "_%lld_%lld_%lld_%lld_%s.bin", ++ static_cast(tensor->ne[0]), ++ static_cast(tensor->ne[1]), ++ static_cast(tensor->ne[2]), ++ static_cast(tensor->ne[3]), ++ ggml_type_name(tensor->type)); + +- return SKIPPY_STATUS_OK; ++ const size_t nbytes = ggml_nbytes(tensor); ++ std::vector raw(nbytes); ++ ggml_backend_tensor_get(tensor, raw.data(), 0, nbytes); ++ std::ofstream output(std::string(dump_dir) + "/" + filename, std::ios::binary | std::ios::trunc); ++ if (output) { ++ output.write(reinterpret_cast(raw.data()), static_cast(raw.size())); ++ } + } + +-static enum skippy_status skippy_prepare_output_activation_frame( +- skippy_session * session, +- size_t token_count, +- void * output_payload, +- size_t output_payload_capacity, +- size_t * out_output_payload_bytes, +- skippy_activation_desc * output_desc, +- const skippy_activation_desc * input_desc, +- struct skippy_error ** out_error) { +- const uint64_t output_flags = skippy_output_activation_flags(session, input_desc); +- const size_t payload_bytes = skippy_emits_activation_frame(session) ? +- skippy_activation_payload_bytes(session, token_count, output_flags) : 0; ++static void skippy_glm_dsa_trace_tensor( ++ skippy_glm_dsa_op_timing & timing, ++ skippy_glm_dsa_op_kind kind, ++ const ggml_tensor * tensor) { ++ if (!timing.trace_tensors || ++ timing.trace_nodes >= timing.trace_node_limit || ++ !skippy_glm_dsa_trace_filter_matches(timing.trace_filter, tensor->name)) { ++ return; ++ } + +- if (out_output_payload_bytes != nullptr) { +- *out_output_payload_bytes = payload_bytes; ++ timing.trace_nodes += 1; ++ skippy_glm_dsa_dump_trace_tensor(timing, tensor); ++ std::string line = format( ++ "skippy: glm_dsa_tensor_trace stage=%d tokens=%zu op=%s node=%u name=%s type=%s ne=[%lld,%lld,%lld,%lld] nb=[%zu,%zu,%zu,%zu] contiguous=%d nbytes=%zu", ++ timing.stage_index, ++ timing.token_count, ++ skippy_glm_dsa_op_name(kind), ++ timing.trace_nodes, ++ tensor->name, ++ ggml_type_name(tensor->type), ++ static_cast(tensor->ne[0]), ++ static_cast(tensor->ne[1]), ++ static_cast(tensor->ne[2]), ++ static_cast(tensor->ne[3]), ++ static_cast(tensor->nb[0]), ++ static_cast(tensor->nb[1]), ++ static_cast(tensor->nb[2]), ++ static_cast(tensor->nb[3]), ++ ggml_is_contiguous(tensor) ? 1 : 0, ++ ggml_nbytes(tensor)); ++ skippy_glm_dsa_trace_tensor_values(line, tensor, timing.trace_value_limit); ++ if (timing.trace_stats) { ++ skippy_glm_dsa_trace_tensor_stats(line, tensor, timing.trace_stats_max_bytes); ++ } ++ LLAMA_LOG_INFO("%s\n", line.c_str()); ++} ++ ++static bool skippy_glm_dsa_op_timing_cb(ggml_tensor * tensor, bool ask, void * user_data) { ++ skippy_model * model = static_cast(user_data); ++ if (model == nullptr || !model->glm_dsa_timing.enabled || tensor == nullptr) { ++ return false; + } + +- if (payload_bytes > 0) { +- if (output_payload == nullptr) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "output_payload is required for runtime-slice activation output"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; +- } +- if (output_payload_capacity < payload_bytes) { +- skippy_set_error(out_error, SKIPPY_STATUS_BUFFER_TOO_SMALL, "output activation buffer is too small"); +- return SKIPPY_STATUS_BUFFER_TOO_SMALL; +- } +- } else if (output_payload_capacity > 0 && output_payload == nullptr) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "output_payload is required when output capacity is non-zero"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ skippy_glm_dsa_op_timing & timing = model->glm_dsa_timing; ++ const skippy_glm_dsa_op_kind kind = skippy_glm_dsa_op_kind_for_tensor(tensor->name); ++ const bool trace_match = timing.trace_tensors && ++ skippy_glm_dsa_trace_filter_matches(timing.trace_filter, tensor->name); ++ if (kind == SKIPPY_GLM_DSA_OP_UNKNOWN && !trace_match) { ++ return false; + } + +- if (output_desc != nullptr) { +- *output_desc = {}; +- output_desc->version = 1; +- output_desc->dtype = payload_bytes > 0 ? SKIPPY_ACTIVATION_DTYPE_F32 : SKIPPY_ACTIVATION_DTYPE_UNKNOWN; +- output_desc->layout = payload_bytes > 0 ? SKIPPY_ACTIVATION_LAYOUT_TOKEN_MAJOR : SKIPPY_ACTIVATION_LAYOUT_OPAQUE; +- output_desc->producer_stage_index = session != nullptr ? session->stage_model->config.stage_index : -1; +- output_desc->layer_start = session != nullptr ? session->stage_model->config.layer_start : 0; +- output_desc->layer_end = session != nullptr ? session->stage_model->config.layer_end : 0; +- output_desc->token_count = static_cast(std::min(token_count, std::numeric_limits::max())); +- output_desc->sequence_count = token_count > 0 ? 1 : 0; +- output_desc->payload_bytes = payload_bytes; +- output_desc->flags = output_flags; ++ // Tensor tracing can be used as a narrow synchronization probe. Keep the ++ // normal op-timing callback broad, but only split the graph at tensors the ++ // trace filter actually requests when timing is disabled. ++ if (timing.trace_tensors && !timing.print_timing && !trace_match) { ++ return false; + } + +- return SKIPPY_STATUS_OK; +-} ++ if (ask) { ++ timing.pending_kind = kind; ++ timing.pending_start_us = ggml_time_us(); ++ return true; ++ } + +-static enum skippy_status skippy_decode_batch( +- skippy_session * session, +- llama_batch batch, +- size_t token_count, +- struct skippy_error ** out_error) { +- if (session == nullptr || session->ctx == nullptr || token_count == 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session and at least one token are required"); +- return SKIPPY_STATUS_INVALID_ARGUMENT; ++ const int64_t elapsed_us = std::max(0, ggml_time_us() - timing.pending_start_us); ++ if (kind != SKIPPY_GLM_DSA_OP_UNKNOWN) { ++ timing.stats[kind].nodes += 1; ++ timing.stats[kind].elapsed_us += elapsed_us; ++ skippy_glm_dsa_record_group_timing(timing, kind, tensor->name, elapsed_us); ++ skippy_glm_dsa_record_hot_tensor(timing, kind, tensor, elapsed_us); ++ skippy_glm_dsa_log_direct_sparse_decision_for_tensor(timing, tensor, kind); + } ++ skippy_glm_dsa_trace_tensor(timing, kind, tensor); ++ timing.pending_kind = SKIPPY_GLM_DSA_OP_UNKNOWN; ++ timing.pending_start_us = 0; ++ return true; ++} + +- skippy_graph_filter_scope graph_filter_scope(&session->stage_model->config); +- const int32_t rc = llama_decode(session->ctx, batch); +- if (rc != 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed"); +- return SKIPPY_STATUS_RUNTIME_ERROR; ++static skippy_glm_dsa_op_stat skippy_glm_dsa_sparse_mask_stat_from_stats(const skippy_glm_dsa_op_stat stats[]) { ++ skippy_glm_dsa_op_stat stat = {}; ++ stat.nodes += stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL].nodes; ++ stat.nodes += stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK].nodes; ++ stat.nodes += stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD].nodes; ++ stat.elapsed_us += stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL].elapsed_us; ++ stat.elapsed_us += stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK].elapsed_us; ++ stat.elapsed_us += stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD].elapsed_us; ++ return stat; ++} ++ ++static skippy_glm_dsa_op_stat skippy_glm_dsa_sparse_mask_stat(const skippy_glm_dsa_op_timing & timing) { ++ return skippy_glm_dsa_sparse_mask_stat_from_stats(timing.stats); ++} ++ ++static skippy_glm_dsa_op_stat skippy_glm_dsa_indexer_topk_stat_from_stats(const skippy_glm_dsa_op_stat stats[]) { ++ skippy_glm_dsa_op_stat stat = {}; ++ stat.nodes += stats[SKIPPY_GLM_DSA_OP_INDEXER].nodes; ++ stat.nodes += stats[SKIPPY_GLM_DSA_OP_TOP_K].nodes; ++ stat.elapsed_us += stats[SKIPPY_GLM_DSA_OP_INDEXER].elapsed_us; ++ stat.elapsed_us += stats[SKIPPY_GLM_DSA_OP_TOP_K].elapsed_us; ++ return stat; ++} ++ ++static skippy_glm_dsa_op_stat skippy_glm_dsa_indexer_topk_stat(const skippy_glm_dsa_op_timing & timing) { ++ return skippy_glm_dsa_indexer_topk_stat_from_stats(timing.stats); ++} ++ ++static skippy_glm_dsa_op_stat skippy_glm_dsa_routed_moe_stat_from_stats(const skippy_glm_dsa_op_stat stats[]) { ++ skippy_glm_dsa_op_stat stat = {}; ++ for (int i = SKIPPY_GLM_DSA_OP_ROUTED_MOE_ROUTE; i <= SKIPPY_GLM_DSA_OP_ROUTED_MOE_AGGREGATE; ++i) { ++ stat.nodes += stats[i].nodes; ++ stat.elapsed_us += stats[i].elapsed_us; + } ++ return stat; ++} + +- session->n_past += static_cast(token_count); +- return skippy_success(out_error); ++static skippy_glm_dsa_op_stat skippy_glm_dsa_routed_moe_stat(const skippy_glm_dsa_op_timing & timing) { ++ return skippy_glm_dsa_routed_moe_stat_from_stats(timing.stats); + } + +-static bool skippy_mtp_available(const skippy_session * session) { +- return session != nullptr && +- session->stage_model != nullptr && +- session->stage_model->mtp_ctx != nullptr && +- session->stage_model->config.include_output; ++static void skippy_glm_dsa_append_op_stat(std::string & line, const char * name, const skippy_glm_dsa_op_stat & stat) { ++ line += format( ++ " %s_nodes=%llu %s_us=%lld", ++ name, ++ static_cast(stat.nodes), ++ name, ++ static_cast(stat.elapsed_us)); + } + +-static bool skippy_env_enabled(const char * name) { +- const char * value = std::getenv(name); +- if (value == nullptr || value[0] == '\0') { +- return false; ++static void skippy_glm_dsa_append_op_kind( ++ std::string & line, ++ skippy_glm_dsa_op_kind kind, ++ const skippy_glm_dsa_op_stat & stat) { ++ skippy_glm_dsa_append_op_stat(line, skippy_glm_dsa_op_name(kind), stat); ++} ++ ++static void skippy_glm_dsa_print_op_group( ++ const skippy_glm_dsa_op_timing & timing, ++ const skippy_glm_dsa_op_group & group) { ++ uint64_t total_nodes = 0; ++ int64_t total_us = 0; ++ for (int i = 0; i < SKIPPY_GLM_DSA_OP_COUNT; ++i) { ++ total_nodes += group.stats[i].nodes; ++ total_us += group.stats[i].elapsed_us; ++ } ++ if (total_nodes == 0) { ++ return; ++ } ++ ++ std::string line = format( ++ "skippy: glm_dsa_group_timing stage=%d tokens=%zu group=%s total_us=%lld", ++ timing.stage_index, ++ timing.token_count, ++ group.name, ++ static_cast(total_us)); ++ skippy_glm_dsa_append_op_stat(line, "indexer_topk", skippy_glm_dsa_indexer_topk_stat_from_stats(group.stats)); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_INDEXER, group.stats[SKIPPY_GLM_DSA_OP_INDEXER]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_TOP_K, group.stats[SKIPPY_GLM_DSA_OP_TOP_K]); ++ skippy_glm_dsa_append_op_stat(line, "sparse_mask", skippy_glm_dsa_sparse_mask_stat_from_stats(group.stats)); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL, group.stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK, group.stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD, group.stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_DSA_SPARSE_ATTN, group.stats[SKIPPY_GLM_DSA_OP_DSA_SPARSE_ATTN]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_COMPACT_GET_ROWS, group.stats[SKIPPY_GLM_DSA_OP_COMPACT_GET_ROWS]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_MLA_ATTENTION, group.stats[SKIPPY_GLM_DSA_OP_MLA_ATTENTION]); ++ skippy_glm_dsa_append_op_stat(line, "routed_moe", skippy_glm_dsa_routed_moe_stat_from_stats(group.stats)); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_ROUTE, group.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_ROUTE]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE_UP, group.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE_UP]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE, group.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_UP, group.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_UP]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_ACT, group.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_ACT]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_DOWN, group.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_DOWN]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_WEIGHTED, group.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_WEIGHTED]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_AGGREGATE, group.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_AGGREGATE]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_SHARED_EXPERT, group.stats[SKIPPY_GLM_DSA_OP_SHARED_EXPERT]); ++ LLAMA_LOG_INFO("%s\n", line.c_str()); ++} ++ ++static void skippy_glm_dsa_print_hot_tensors(const skippy_glm_dsa_op_timing & timing) { ++ for (size_t i = 0; i < SKIPPY_GLM_DSA_HOT_TENSOR_LIMIT; ++i) { ++ const skippy_glm_dsa_hot_tensor & record = timing.hot_tensors[i]; ++ if (!record.used) { ++ return; ++ } ++ ++ LLAMA_LOG_INFO( ++ "skippy: glm_dsa_hot_tensor stage=%d tokens=%zu rank=%zu op=%s kind=%s elapsed_us=%lld name=%s ne0=%lld ne1=%lld ne2=%lld ne3=%lld\n", ++ timing.stage_index, ++ timing.token_count, ++ i + 1, ++ ggml_op_name(record.op), ++ skippy_glm_dsa_op_name(record.kind), ++ static_cast(record.elapsed_us), ++ record.name[0] == '\0' ? "" : record.name, ++ static_cast(record.ne[0]), ++ static_cast(record.ne[1]), ++ static_cast(record.ne[2]), ++ static_cast(record.ne[3])); + } +- return std::strcmp(value, "0") != 0 && +- std::strcmp(value, "false") != 0 && +- std::strcmp(value, "off") != 0 && +- std::strcmp(value, "no") != 0; + } + +-static bool skippy_mtp_greedy_sampling_fastpath_enabled() { +- return skippy_env_enabled("SKIPPY_NATIVE_MTP_GREEDY_SAMPLING_FASTPATH"); ++static void skippy_glm_dsa_op_timing_begin(skippy_session * session, size_t token_count) { ++ if (session == nullptr || session->stage_model == nullptr || !session->stage_model->glm_dsa_timing.enabled) { ++ return; ++ } ++ ++ skippy_glm_dsa_op_timing & timing = session->stage_model->glm_dsa_timing; ++ const bool enabled = timing.enabled; ++ const bool print_timing = timing.print_timing; ++ const bool trace_tensors = timing.trace_tensors; ++ const bool trace_stats = timing.trace_stats; ++ const int32_t stage_index = timing.stage_index; ++ const uint32_t trace_value_limit = timing.trace_value_limit; ++ const uint32_t trace_node_limit = timing.trace_node_limit; ++ const uint32_t trace_stats_max_bytes = timing.trace_stats_max_bytes; ++ const std::string trace_filter = timing.trace_filter; ++ timing = skippy_glm_dsa_op_timing{}; ++ timing.enabled = enabled; ++ timing.print_timing = print_timing; ++ timing.trace_tensors = trace_tensors; ++ timing.trace_stats = trace_stats; ++ timing.stage_index = stage_index; ++ timing.token_count = token_count; ++ timing.trace_value_limit = trace_value_limit; ++ timing.trace_node_limit = trace_node_limit; ++ timing.trace_stats_max_bytes = trace_stats_max_bytes; ++ timing.trace_filter = trace_filter; ++} ++ ++static void skippy_glm_dsa_op_timing_end(skippy_session * session) { ++ if (session == nullptr || session->stage_model == nullptr || !session->stage_model->glm_dsa_timing.enabled) { ++ return; ++ } ++ ++ const skippy_glm_dsa_op_timing & timing = session->stage_model->glm_dsa_timing; ++ uint64_t total_nodes = 0; ++ int64_t total_us = 0; ++ for (int i = 0; i < SKIPPY_GLM_DSA_OP_COUNT; ++i) { ++ total_nodes += timing.stats[i].nodes; ++ total_us += timing.stats[i].elapsed_us; ++ } ++ if (total_nodes == 0) { ++ return; ++ } ++ if (!timing.print_timing) { ++ return; ++ } ++ ++ std::string line = format( ++ "skippy: glm_dsa_op_timing stage=%d tokens=%zu total_us=%lld", ++ timing.stage_index, ++ timing.token_count, ++ static_cast(total_us)); ++ skippy_glm_dsa_append_op_stat(line, "indexer_topk", skippy_glm_dsa_indexer_topk_stat(timing)); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_INDEXER, timing.stats[SKIPPY_GLM_DSA_OP_INDEXER]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_TOP_K, timing.stats[SKIPPY_GLM_DSA_OP_TOP_K]); ++ skippy_glm_dsa_append_op_stat(line, "sparse_mask", skippy_glm_dsa_sparse_mask_stat(timing)); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL, timing.stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_FILL]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK, timing.stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_TOPK]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD, timing.stats[SKIPPY_GLM_DSA_OP_SPARSE_MASK_ADD]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_DSA_SPARSE_ATTN, timing.stats[SKIPPY_GLM_DSA_OP_DSA_SPARSE_ATTN]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_COMPACT_GET_ROWS, timing.stats[SKIPPY_GLM_DSA_OP_COMPACT_GET_ROWS]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_MLA_ATTENTION, timing.stats[SKIPPY_GLM_DSA_OP_MLA_ATTENTION]); ++ skippy_glm_dsa_append_op_stat(line, "routed_moe", skippy_glm_dsa_routed_moe_stat(timing)); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_ROUTE, timing.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_ROUTE]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE_UP, timing.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE_UP]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE, timing.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_GATE]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_UP, timing.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_UP]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_ACT, timing.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_ACT]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_DOWN, timing.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_DOWN]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_WEIGHTED, timing.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_WEIGHTED]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_ROUTED_MOE_AGGREGATE, timing.stats[SKIPPY_GLM_DSA_OP_ROUTED_MOE_AGGREGATE]); ++ skippy_glm_dsa_append_op_kind(line, SKIPPY_GLM_DSA_OP_SHARED_EXPERT, timing.stats[SKIPPY_GLM_DSA_OP_SHARED_EXPERT]); ++ LLAMA_LOG_INFO("%s\n", line.c_str()); ++ ++ for (size_t i = 0; i < SKIPPY_GLM_DSA_TIMING_GROUP_LIMIT; ++i) { ++ if (timing.groups[i].used) { ++ skippy_glm_dsa_print_op_group(timing, timing.groups[i]); ++ } ++ } ++ skippy_glm_dsa_print_hot_tensors(timing); + } + + static void skippy_mtp_clear_session_state(skippy_session * session) { +@@ -1460,6 +2871,7 @@ static enum skippy_status skippy_decode_tokens( + const llama_token * token_ids, + size_t token_count, + bool request_logits, ++ int32_t glm_dsa_phase_hint, + struct skippy_error ** out_error) { + (void) request_logits; + +@@ -1493,7 +2905,7 @@ static enum skippy_status skippy_decode_tokens( + /*seq_id =*/ &seq_ids, + /*logits =*/ &logits, + }; +- enum skippy_status status = skippy_decode_batch(session, batch, 1, out_error); ++ enum skippy_status status = skippy_decode_batch(session, batch, 1, glm_dsa_phase_hint, out_error); + if (status == SKIPPY_STATUS_OK) { + status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, pos, out_error); + } +@@ -1514,7 +2929,7 @@ static enum skippy_status skippy_decode_tokens( + batch.logits[i] = request_logits && i == n_tokens - 1 ? 1 : 0; + } + +- enum skippy_status status = skippy_decode_batch(session, batch, token_count, out_error); ++ enum skippy_status status = skippy_decode_batch(session, batch, token_count, glm_dsa_phase_hint, out_error); + llama_batch_free(batch); + if (status == SKIPPY_STATUS_OK) { + status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); +@@ -1554,7 +2969,12 @@ static enum skippy_status skippy_verify_token_batch( + batch.logits[i] = 1; + } + +- enum skippy_status status = skippy_decode_batch(session, batch, token_count, out_error); ++ enum skippy_status status = skippy_decode_batch( ++ session, ++ batch, ++ token_count, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, ++ out_error); + llama_batch_free(batch); + return status; + } +@@ -2201,6 +3621,8 @@ static enum skippy_status skippy_copy_output_activation_frame( + void * output_payload, + const skippy_activation_desc * input_desc, + const void * input_payload, ++ size_t * out_output_payload_bytes, ++ skippy_activation_desc * output_desc, + struct skippy_error ** out_error) { + if (!skippy_emits_activation_frame(session)) { + return skippy_success(out_error); +@@ -2208,6 +3630,7 @@ static enum skippy_status skippy_copy_output_activation_frame( + + const uint64_t output_flags = skippy_output_activation_flags(session, input_desc); + const size_t hidden_bytes = skippy_activation_hidden_bytes(session, token_count); ++ size_t actual_payload_bytes = hidden_bytes; + if ((output_flags & SKIPPY_ACTIVATION_FLAG_GEMMA3N_ALTUP) != 0) { + llm_graph_result * res = session->ctx->get_gf_res_prev(); + ggml_tensor * altup = res != nullptr ? res->get_skippy_gemma3n_altup() : nullptr; +@@ -2221,6 +3644,12 @@ static enum skippy_status skippy_copy_output_activation_frame( + return SKIPPY_STATUS_RUNTIME_ERROR; + } + ggml_backend_tensor_get(altup, output_payload, 0, altup_bytes); ++ if (out_output_payload_bytes != nullptr) { ++ *out_output_payload_bytes = altup_bytes; ++ } ++ if (output_desc != nullptr) { ++ output_desc->payload_bytes = altup_bytes; ++ } + return skippy_success(out_error); + } + +@@ -2232,6 +3661,66 @@ static enum skippy_status skippy_copy_output_activation_frame( + + std::memcpy(output_payload, embeddings, hidden_bytes); + ++ if ((output_flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { ++ uint8_t * sideband_output = static_cast(output_payload) + hidden_bytes; ++ llm_graph_result * res = session->ctx->get_gf_res_prev(); ++ ggml_tensor * top_k = res != nullptr ? res->get_skippy_glm_dsa_top_k() : nullptr; ++ if (top_k != nullptr) { ++ const size_t top_k_rows = top_k->ne[1]*top_k->ne[2]*top_k->ne[3]; ++ if (top_k->type != GGML_TYPE_I32 || top_k->ne[0] <= 0 || ++ static_cast(top_k->ne[0]) > std::numeric_limits::max() || ++ top_k_rows != token_count) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "GLM-DSA top-k sideband tensor shape is not token-major i32"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ const uint32_t visible_top_k = skippy_glm_dsa_sideband_top_k_count_for_n_kv( ++ session, ++ static_cast(std::max(session->n_past, 0))); ++ const uint32_t output_top_k = static_cast( ++ std::min(static_cast(top_k->ne[0]), visible_top_k)); ++ if (output_top_k == 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "GLM-DSA top-k sideband output has no visible entries"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ const size_t top_k_bytes_per_row = static_cast(output_top_k)*sizeof(int32_t); ++ const size_t top_k_bytes = token_count*top_k_bytes_per_row; ++ if (static_cast(output_top_k) == top_k->ne[0] && ggml_is_contiguous(top_k)) { ++ ggml_backend_tensor_get(top_k, sideband_output, 0, top_k_bytes); ++ } else { ++ for (size_t row = 0; row < token_count; ++row) { ++ size_t rem = row; ++ const size_t i1 = rem % static_cast(top_k->ne[1]); ++ rem /= static_cast(top_k->ne[1]); ++ const size_t i2 = rem % static_cast(top_k->ne[2]); ++ const size_t i3 = rem / static_cast(top_k->ne[2]); ++ const size_t src_offset = i1*top_k->nb[1] + i2*top_k->nb[2] + i3*top_k->nb[3]; ++ ggml_backend_tensor_get( ++ top_k, ++ sideband_output + row*top_k_bytes_per_row, ++ src_offset, ++ top_k_bytes_per_row); ++ } ++ } ++ actual_payload_bytes += top_k_bytes; ++ } else { ++ if (input_desc != nullptr && input_desc->payload_bytes < hidden_bytes) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "GLM-DSA downstream slice input is smaller than hidden payload"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ const size_t top_k_bytes = input_desc != nullptr ? ++ static_cast(input_desc->payload_bytes - hidden_bytes) : 0; ++ if (input_desc == nullptr || ++ input_payload == nullptr || ++ (input_desc->flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) == 0 || ++ input_desc->payload_bytes < hidden_bytes + top_k_bytes) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "GLM-DSA downstream slice cannot forward missing top-k sideband"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ std::memcpy(sideband_output, static_cast(input_payload) + hidden_bytes, top_k_bytes); ++ actual_payload_bytes += top_k_bytes; ++ } ++ } ++ + if ((output_flags & SKIPPY_ACTIVATION_FLAG_RWKV7_V_FIRST) != 0) { + uint8_t * sideband_output = static_cast(output_payload) + hidden_bytes; + const skippy_runtime_config & config = session->stage_model->config; +@@ -2257,6 +3746,14 @@ static enum skippy_status skippy_copy_output_activation_frame( + } + std::memcpy(sideband_output, static_cast(input_payload) + hidden_bytes, hidden_bytes); + } ++ actual_payload_bytes += hidden_bytes; ++ } ++ ++ if (out_output_payload_bytes != nullptr) { ++ *out_output_payload_bytes = actual_payload_bytes; ++ } ++ if (output_desc != nullptr) { ++ output_desc->payload_bytes = actual_payload_bytes; + } + return skippy_success(out_error); + } +@@ -2270,6 +3767,7 @@ static enum skippy_status skippy_decode_activation_frame( + size_t position_count, + size_t token_count, + bool request_logits, ++ int32_t glm_dsa_phase_hint, + struct skippy_error ** out_error) { + if (session == nullptr || session->ctx == nullptr || input_desc == nullptr || input_payload == nullptr || token_count == 0) { + skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "session and activation frame input are required"); +@@ -2350,7 +3848,19 @@ static enum skippy_status skippy_decode_activation_frame( + skippy_rwkv7_v_first_scope rwkv7_v_first_scope(input_desc, input_payload, hidden_bytes, n_embd); + skippy_gemma3n_altup_scope gemma3n_altup_scope(input_desc, input_payload, n_embd, n_altup); + const llama_pos token_start = pos_storage.empty() ? session->n_past : pos_storage[0]; +- enum skippy_status status = skippy_decode_batch(session, batch, token_count, out_error); ++ skippy_glm_dsa_top_k_scope glm_dsa_top_k_scope( ++ input_desc, ++ input_payload, ++ hidden_bytes, ++ skippy_glm_dsa_top_k_count_from_desc(session, input_desc), ++ skippy_glm_dsa_n_stream(session), ++ token_start); ++ enum skippy_status status = skippy_decode_batch( ++ session, ++ batch, ++ token_count, ++ glm_dsa_phase_hint, ++ out_error); + if (status == SKIPPY_STATUS_OK && token_ids != nullptr) { + status = skippy_mtp_sync_target_tokens(session, token_ids, token_count, token_start, out_error); + } +@@ -2416,7 +3926,19 @@ static enum skippy_status skippy_verify_activation_frame( + skippy_rwkv7_v_first_scope rwkv7_v_first_scope(input_desc, input_payload, hidden_bytes, n_embd); + skippy_gemma3n_altup_scope gemma3n_altup_scope(input_desc, input_payload, n_embd, n_altup); + const llama_pos token_start = session->n_past; +- enum skippy_status status = skippy_decode_batch(session, batch, token_count, out_error); ++ skippy_glm_dsa_top_k_scope glm_dsa_top_k_scope( ++ input_desc, ++ input_payload, ++ hidden_bytes, ++ skippy_glm_dsa_top_k_count_from_desc(session, input_desc), ++ skippy_glm_dsa_n_stream(session), ++ token_start); ++ enum skippy_status status = skippy_decode_batch( ++ session, ++ batch, ++ token_count, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, ++ out_error); + if (!alias_input_payload) { + llama_batch_free(batch); + } +@@ -2719,8 +4241,28 @@ static enum skippy_status skippy_finish_model_open( + params.type_v = config != nullptr && config->cache_type_v > 0 ? static_cast(config->cache_type_v) : GGML_TYPE_F16; + params.flash_attn_type = config != nullptr ? static_cast(config->flash_attn_type) : LLAMA_FLASH_ATTN_TYPE_AUTO; + params.embeddings = config != nullptr && config->filter_tensors_on_load && !config->include_output; ++ const bool glm_dsa_op_timing_enabled = skippy_glm_dsa_op_timing_enabled(); ++ const bool glm_dsa_tensor_trace_enabled = skippy_glm_dsa_tensor_trace_enabled(); ++ if (model->arch == LLM_ARCH_GLM_DSA && (glm_dsa_op_timing_enabled || glm_dsa_tensor_trace_enabled)) { ++ stage_model->glm_dsa_timing.enabled = true; ++ stage_model->glm_dsa_timing.print_timing = glm_dsa_op_timing_enabled; ++ stage_model->glm_dsa_timing.trace_tensors = glm_dsa_tensor_trace_enabled; ++ stage_model->glm_dsa_timing.trace_stats = skippy_glm_dsa_tensor_trace_stats_enabled(); ++ stage_model->glm_dsa_timing.trace_value_limit = ++ skippy_env_u32("SKIPPY_GLM_DSA_TENSOR_TRACE_VALUES", 8, 0, 4096); ++ stage_model->glm_dsa_timing.trace_node_limit = ++ skippy_env_u32("SKIPPY_GLM_DSA_TENSOR_TRACE_NODES", 32, 1, 1024); ++ stage_model->glm_dsa_timing.trace_stats_max_bytes = ++ skippy_env_u32("SKIPPY_GLM_DSA_TENSOR_TRACE_STATS_MAX_BYTES", 32 * 1024 * 1024, 0, 1024 * 1024 * 1024); ++ if (const char * trace_filter = std::getenv("SKIPPY_GLM_DSA_TENSOR_TRACE_FILTER")) { ++ stage_model->glm_dsa_timing.trace_filter = trace_filter; ++ } ++ stage_model->glm_dsa_timing.stage_index = config != nullptr ? config->stage_index : -1; ++ params.cb_eval = skippy_glm_dsa_op_timing_cb; ++ params.cb_eval_user_data = stage_model; ++ } + if (llm_arch_is_recurrent(model->arch) || llm_arch_is_hybrid(model->arch)) { +- params.n_seq_max = std::max(2, stage_model->lane_count * 2); ++ params.n_seq_max = std::max(params.n_seq_max, std::max(2, stage_model->lane_count * 2)); + params.n_rs_seq = std::max(params.n_rs_seq, 2); + params.kv_unified = true; + } +@@ -3414,7 +4956,13 @@ enum skippy_status skippy_prefill_chunk( + *out_output_activation_bytes = 0; + } + +- return skippy_decode_tokens(session, token_ids, token_count, false, out_error); ++ return skippy_decode_tokens( ++ session, ++ token_ids, ++ token_count, ++ false, ++ SKIPPY_GLM_DSA_PHASE_HINT_PREFILL, ++ out_error); + } + + enum skippy_status skippy_decode_step_sampled( +@@ -3437,7 +4985,13 @@ enum skippy_status skippy_decode_step_sampled( + *out_output_activation_bytes = 0; + } + +- enum skippy_status status = skippy_decode_tokens(session, &token_id, 1, true, out_error); ++ enum skippy_status status = skippy_decode_tokens( ++ session, ++ &token_id, ++ 1, ++ true, ++ SKIPPY_GLM_DSA_PHASE_HINT_DECODE, ++ out_error); + if (status != SKIPPY_STATUS_OK) { + return status; + } +@@ -3540,7 +5094,9 @@ enum skippy_status skippy_decode_batch_sampled( + batch.logits[i] = 1; + } + +- skippy_graph_filter_scope graph_filter_scope(&first->stage_model->config); ++ skippy_graph_filter_scope graph_filter_scope( ++ &first->stage_model->config, ++ SKIPPY_GLM_DSA_PHASE_HINT_DECODE); + const int32_t rc = llama_decode(first->ctx, batch); + if (rc != 0) { + llama_batch_free(batch); +@@ -3651,9 +5207,25 @@ static enum skippy_status skippy_prefill_chunk_frame_impl( + } + + if (skippy_is_filtered(session) && session->stage_model->config.layer_start > 0) { +- status = skippy_decode_activation_frame(session, input_desc, input_payload, token_ids, positions, position_count, token_count, request_logits, out_error); ++ status = skippy_decode_activation_frame( ++ session, ++ input_desc, ++ input_payload, ++ token_ids, ++ positions, ++ position_count, ++ token_count, ++ request_logits, ++ SKIPPY_GLM_DSA_PHASE_HINT_PREFILL, ++ out_error); + } else { +- status = skippy_decode_tokens(session, token_ids, token_count, request_logits, out_error); ++ status = skippy_decode_tokens( ++ session, ++ token_ids, ++ token_count, ++ request_logits, ++ SKIPPY_GLM_DSA_PHASE_HINT_PREFILL, ++ out_error); + } + if (status != SKIPPY_STATUS_OK) { + return status; +@@ -3664,7 +5236,8 @@ static enum skippy_status skippy_prefill_chunk_frame_impl( + skippy_sample_token(session, sampling) : -1; + } + +- return skippy_copy_output_activation_frame(session, token_count, output_payload, input_desc, input_payload, out_error); ++ return skippy_copy_output_activation_frame( ++ session, token_count, output_payload, input_desc, input_payload, out_output_payload_bytes, output_desc, out_error); + } + + enum skippy_status skippy_prefill_chunk_frame( +@@ -3827,9 +5400,25 @@ enum skippy_status skippy_decode_step_frame_sampled( + } + + if (skippy_is_filtered(session) && session->stage_model->config.layer_start > 0) { +- status = skippy_decode_activation_frame(session, input_desc, input_payload, &token_id, nullptr, 0, 1, true, out_error); ++ status = skippy_decode_activation_frame( ++ session, ++ input_desc, ++ input_payload, ++ &token_id, ++ nullptr, ++ 0, ++ 1, ++ true, ++ SKIPPY_GLM_DSA_PHASE_HINT_DECODE, ++ out_error); + } else { +- status = skippy_decode_tokens(session, &token_id, 1, true, out_error); ++ status = skippy_decode_tokens( ++ session, ++ &token_id, ++ 1, ++ true, ++ SKIPPY_GLM_DSA_PHASE_HINT_DECODE, ++ out_error); + } + if (status != SKIPPY_STATUS_OK) { + return status; +@@ -3839,7 +5428,8 @@ enum skippy_status skippy_decode_step_frame_sampled( + *out_predicted_token = session->stage_model->config.include_output ? skippy_sample_token(session, sampling) : -1; + } + +- return skippy_copy_output_activation_frame(session, 1, output_payload, input_desc, input_payload, out_error); ++ return skippy_copy_output_activation_frame( ++ session, 1, output_payload, input_desc, input_payload, out_output_payload_bytes, output_desc, out_error); + } + + enum skippy_status skippy_decode_step_frame_sampled_mtp( +@@ -3932,6 +5522,9 @@ enum skippy_status skippy_decode_step_frame_batch_sampled( + } + + const size_t hidden_bytes_per_request = skippy_activation_hidden_bytes(first, 1); ++ uint64_t batch_input_flags = 0; ++ uint32_t batch_input_top_k = 0; ++ uint64_t batch_output_flags = 0; + for (int32_t i = 0; i < n_tokens; ++i) { + skippy_session * session = sessions[i]; + if (session == nullptr || session->ctx != first->ctx || session->stage_model != first->stage_model) { +@@ -3944,14 +5537,30 @@ enum skippy_status skippy_decode_step_frame_batch_sampled( + if (status != SKIPPY_STATUS_OK) { + return status; + } +- if (input_desc != nullptr && input_desc->flags != 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "batched activation decode does not support activation sidebands yet"); ++ const uint64_t input_flags = input_desc != nullptr ? input_desc->flags : 0; ++ if ((input_flags & ~SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "batched activation decode only supports GLM-DSA input sidebands"); + return SKIPPY_STATUS_UNSUPPORTED; + } +- if (skippy_output_activation_flags(session, input_desc) != 0) { +- skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "batched activation decode does not support output activation sidebands yet"); ++ const uint32_t input_top_k = skippy_glm_dsa_top_k_count_from_desc(session, input_desc); ++ if (i == 0) { ++ batch_input_flags = input_flags; ++ batch_input_top_k = input_top_k; ++ } else if (input_flags != batch_input_flags || input_top_k != batch_input_top_k) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "batched activation inputs must use the same GLM-DSA sideband shape"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } ++ const uint64_t output_flags = skippy_output_activation_flags(session, input_desc); ++ if ((output_flags & ~SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "batched activation decode only supports GLM-DSA output sidebands"); + return SKIPPY_STATUS_UNSUPPORTED; + } ++ if (i == 0) { ++ batch_output_flags = output_flags; ++ } else if (output_flags != batch_output_flags) { ++ skippy_set_error(out_error, SKIPPY_STATUS_INVALID_ARGUMENT, "batched activation outputs must use the same sideband flags"); ++ return SKIPPY_STATUS_INVALID_ARGUMENT; ++ } + status = skippy_prepare_output_activation_frame( + session, + 1, +@@ -3967,9 +5576,13 @@ enum skippy_status skippy_decode_step_frame_batch_sampled( + } + + std::vector embd_storage; ++ std::vector glm_dsa_input_top_k_storage; + std::vector token_storage; + if (activation_input) { + embd_storage.resize(static_cast(n_tokens)*n_embd_inp); ++ if ((batch_input_flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { ++ glm_dsa_input_top_k_storage.resize(static_cast(n_tokens)*batch_input_top_k); ++ } + for (int32_t i = 0; i < n_tokens; ++i) { + const skippy_activation_desc * input_desc = input_descs[i]; + const void * input_payload = input_payloads[i]; +@@ -3982,6 +5595,14 @@ enum skippy_status skippy_decode_step_frame_batch_sampled( + std::memcpy(dst, input_payload, static_cast(n_embd)*sizeof(float)); + std::memset(dst + n_embd, 0, static_cast(n_embd_inp - n_embd)*sizeof(float)); + } ++ if (!glm_dsa_input_top_k_storage.empty()) { ++ const int32_t * top_k_src = reinterpret_cast( ++ static_cast(input_payload) + hidden_bytes_per_request); ++ std::memcpy( ++ glm_dsa_input_top_k_storage.data() + static_cast(i)*batch_input_top_k, ++ top_k_src, ++ static_cast(batch_input_top_k)*sizeof(int32_t)); ++ } + } + } else { + token_storage.assign(token_ids, token_ids + request_count); +@@ -4019,7 +5640,16 @@ enum skippy_status skippy_decode_step_frame_batch_sampled( + }; + + skippy_activation_tokens_scope activation_tokens_scope(token_ids, request_count); +- skippy_graph_filter_scope graph_filter_scope(&first->stage_model->config); ++ const llama_pos batch_pos_start = pos_storage.empty() ? first->n_past : pos_storage[0]; ++ skippy_glm_dsa_top_k_scope glm_dsa_top_k_scope( ++ glm_dsa_input_top_k_storage.empty() ? nullptr : glm_dsa_input_top_k_storage.data(), ++ static_cast(n_tokens), ++ batch_input_top_k, ++ 1, ++ batch_pos_start); ++ skippy_graph_filter_scope graph_filter_scope( ++ &first->stage_model->config, ++ SKIPPY_GLM_DSA_PHASE_HINT_DECODE); + const int32_t rc = llama_decode(first->ctx, batch); + if (rc != 0) { + skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "llama_decode failed"); +@@ -4035,6 +5665,75 @@ enum skippy_status skippy_decode_step_frame_batch_sampled( + } + } + ++ std::vector glm_dsa_top_k_storage; ++ std::vector glm_dsa_top_k_offsets(request_count + 1, 0); ++ std::vector glm_dsa_top_k_bytes_per_request(request_count, 0); ++ if ((batch_output_flags & SKIPPY_ACTIVATION_FLAG_GLM_DSA_TOP_K) != 0) { ++ llm_graph_result * res = first->ctx->get_gf_res_prev(); ++ ggml_tensor * top_k = res != nullptr ? res->get_skippy_glm_dsa_top_k() : nullptr; ++ if (top_k != nullptr) { ++ const size_t top_k_rows = top_k->ne[1]*top_k->ne[2]*top_k->ne[3]; ++ if (top_k->type != GGML_TYPE_I32 || top_k->ne[0] <= 0 || ++ static_cast(top_k->ne[0]) > std::numeric_limits::max() || ++ top_k_rows != request_count) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "batched GLM-DSA top-k output is not request-major i32"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ const size_t top_k_row_bytes = static_cast(top_k->ne[0])*sizeof(int32_t); ++ bool full_rows = true; ++ for (size_t row = 0; row < request_count; ++row) { ++ const uint32_t visible_top_k = skippy_glm_dsa_sideband_top_k_count(sessions[row], 1); ++ const uint32_t request_top_k = static_cast(std::min(top_k->ne[0], visible_top_k)); ++ const size_t request_bytes = static_cast(request_top_k)*sizeof(int32_t); ++ if (request_bytes == 0) { ++ const std::string message = format( ++ "batched GLM-DSA top-k output has no visible entries: request=%zu n_past=%d tensor_top_k=%lld visible_top_k=%u", ++ row, ++ sessions[row]->n_past, ++ static_cast(top_k->ne[0]), ++ visible_top_k); ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, message.c_str()); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ glm_dsa_top_k_bytes_per_request[row] = request_bytes; ++ glm_dsa_top_k_offsets[row + 1] = glm_dsa_top_k_offsets[row] + request_bytes; ++ full_rows = full_rows && request_bytes == top_k_row_bytes; ++ } ++ glm_dsa_top_k_storage.resize(glm_dsa_top_k_offsets.back()); ++ if (full_rows && ggml_is_contiguous(top_k)) { ++ ggml_backend_tensor_get(top_k, glm_dsa_top_k_storage.data(), 0, glm_dsa_top_k_storage.size()); ++ } else { ++ for (size_t row = 0; row < request_count; ++row) { ++ size_t rem = row; ++ const size_t i1 = rem % static_cast(top_k->ne[1]); ++ rem /= static_cast(top_k->ne[1]); ++ const size_t i2 = rem % static_cast(top_k->ne[2]); ++ const size_t i3 = rem / static_cast(top_k->ne[2]); ++ const size_t src_offset = i1*top_k->nb[1] + i2*top_k->nb[2] + i3*top_k->nb[3]; ++ ggml_backend_tensor_get( ++ top_k, ++ glm_dsa_top_k_storage.data() + glm_dsa_top_k_offsets[row], ++ src_offset, ++ glm_dsa_top_k_bytes_per_request[row]); ++ } ++ } ++ } else if (!glm_dsa_input_top_k_storage.empty()) { ++ const size_t request_bytes = static_cast(batch_input_top_k)*sizeof(int32_t); ++ for (size_t row = 0; row < request_count; ++row) { ++ glm_dsa_top_k_bytes_per_request[row] = request_bytes; ++ glm_dsa_top_k_offsets[row + 1] = glm_dsa_top_k_offsets[row] + request_bytes; ++ } ++ glm_dsa_top_k_storage.resize(glm_dsa_top_k_offsets.back()); ++ std::memcpy( ++ glm_dsa_top_k_storage.data(), ++ glm_dsa_input_top_k_storage.data(), ++ glm_dsa_top_k_storage.size()); ++ } else { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "batched GLM-DSA top-k output was not available"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } ++ } ++ + for (int32_t i = 0; i < n_tokens; ++i) { + skippy_session * session = sessions[i]; + session->n_past += 1; +@@ -4054,6 +5753,23 @@ enum skippy_status skippy_decode_step_frame_batch_sampled( + output_payloads[i], + embeddings + static_cast(i)*n_embd, + hidden_bytes_per_request); ++ size_t actual_output_bytes = hidden_bytes_per_request; ++ const size_t request_top_k_bytes = glm_dsa_top_k_bytes_per_request[static_cast(i)]; ++ if (request_top_k_bytes > 0) { ++ uint8_t * output = static_cast(output_payloads[i]); ++ const size_t sideband_offset = glm_dsa_top_k_offsets[static_cast(i)]; ++ std::memcpy( ++ output + hidden_bytes_per_request, ++ glm_dsa_top_k_storage.data() + sideband_offset, ++ request_top_k_bytes); ++ actual_output_bytes += request_top_k_bytes; ++ } ++ if (out_output_payload_bytes != nullptr) { ++ out_output_payload_bytes[i] = actual_output_bytes; ++ } ++ if (output_descs != nullptr) { ++ output_descs[i].payload_bytes = actual_output_bytes; ++ } + } + } + +@@ -4129,17 +5845,34 @@ enum skippy_status skippy_verify_tokens_frame_sampled( + if (skippy_is_filtered(session) && session->stage_model->config.layer_start > 0) { + status = session->stage_model->config.include_output ? + skippy_verify_activation_frame(session, input_desc, input_payload, token_ids, token_count, out_error) : +- skippy_decode_activation_frame(session, input_desc, input_payload, token_ids, nullptr, 0, token_count, false, out_error); ++ skippy_decode_activation_frame( ++ session, ++ input_desc, ++ input_payload, ++ token_ids, ++ nullptr, ++ 0, ++ token_count, ++ false, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, ++ out_error); + } else { + status = session->stage_model->config.include_output ? + skippy_verify_token_batch(session, token_ids, token_count, out_error) : +- skippy_decode_tokens(session, token_ids, token_count, false, out_error); ++ skippy_decode_tokens( ++ session, ++ token_ids, ++ token_count, ++ false, ++ SKIPPY_GLM_DSA_PHASE_HINT_VERIFY, ++ out_error); + } + if (status != SKIPPY_STATUS_OK) { + return status; + } + +- status = skippy_copy_output_activation_frame(session, token_count, output_payload, input_desc, input_payload, out_error); ++ status = skippy_copy_output_activation_frame( ++ session, token_count, output_payload, input_desc, input_payload, out_output_payload_bytes, output_desc, out_error); + if (status != SKIPPY_STATUS_OK) { + return status; + } +@@ -4198,7 +5931,8 @@ enum skippy_status skippy_session_copy_output_activation_frame( + if (status != SKIPPY_STATUS_OK) { + return status; + } +- return skippy_copy_output_activation_frame(session, token_count, output_payload, nullptr, nullptr, out_error); ++ return skippy_copy_output_activation_frame( ++ session, token_count, output_payload, nullptr, nullptr, out_output_payload_bytes, output_desc, out_error); + } + + static enum skippy_status skippy_validate_state_range( +@@ -4566,6 +6300,13 @@ static llama_kv_cache * skippy_get_kv_cache( + if (auto * kv = dynamic_cast(memory)) { + return kv; + } ++ if (auto * dsa = dynamic_cast(memory)) { ++ llama_kv_cache * kv = dsa->get_mla(); ++ if (kv == nullptr) { ++ skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "runtime has no GLM-DSA MLA KV cache"); ++ } ++ return kv; ++ } + + skippy_set_error(out_error, SKIPPY_STATUS_UNSUPPORTED, "runtime memory type is not supported for native KV pages"); + return nullptr; +@@ -4635,6 +6376,10 @@ enum skippy_status skippy_import_kv_page( + skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, error.c_str()); + return SKIPPY_STATUS_RUNTIME_ERROR; + } ++ // KV residency determines graph topology for paths such as GLM-DSA, where ++ // the selected-index width is capped by the visible KV length. Do not reuse ++ // a graph that was built before this page was imported. ++ session->ctx->get_gf_res_prev()->reset(); + session->n_past = std::max( + session->n_past, + static_cast(std::min( +@@ -4679,6 +6424,11 @@ enum skippy_status skippy_trim_session( + skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim hybrid ISWA memory suffix"); + return SKIPPY_STATUS_RUNTIME_ERROR; + } ++ } else if (auto * dsa = dynamic_cast(memory)) { ++ if (!dsa->seq_rm(session->seq_id, p0, -1)) { ++ skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim GLM-DSA KV suffix"); ++ return SKIPPY_STATUS_RUNTIME_ERROR; ++ } + } else if (auto * kv = dynamic_cast(memory)) { + if (!kv->seq_rm(session->seq_id, p0, -1)) { + skippy_set_error(out_error, SKIPPY_STATUS_RUNTIME_ERROR, "failed to trim native KV suffix"); +-- +2.54.0 (Apple Git-157) diff --git a/third_party/llama.cpp/patches/0024-tests-cover-native-GLM-DSA-execution-paths.patch b/third_party/llama.cpp/patches/0024-tests-cover-native-GLM-DSA-execution-paths.patch new file mode 100644 index 000000000..7fc0e78fe --- /dev/null +++ b/third_party/llama.cpp/patches/0024-tests-cover-native-GLM-DSA-execution-paths.patch @@ -0,0 +1,8502 @@ +From 356609cab13dccdfb2069bcd739a3fff39e55e66 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Tue, 14 Jul 2026 09:16:32 +1000 +Subject: [PATCH 21/24] tests: cover native GLM-DSA execution paths + +--- + scripts/glm52-phase-d-policy-report.sh | 287 ++ + scripts/glm52-phase-e-moe-report.sh | 950 ++++ + tests/test-backend-ops.cpp | 5988 +++++++++++++++++++++++- + tests/test-llama-archs.cpp | 985 +++- + 4 files changed, 8085 insertions(+), 125 deletions(-) + create mode 100755 scripts/glm52-phase-d-policy-report.sh + create mode 100755 scripts/glm52-phase-e-moe-report.sh + +diff --git a/scripts/glm52-phase-d-policy-report.sh b/scripts/glm52-phase-d-policy-report.sh +new file mode 100755 +index 00000000..5b475d81 +--- /dev/null ++++ b/scripts/glm52-phase-d-policy-report.sh +@@ -0,0 +1,287 @@ ++#!/usr/bin/env bash ++set -euo pipefail ++ ++script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" ++llama_root="$(cd "${script_dir}/.." && pwd)" ++ ++find_binary() { ++ local name="$1" ++ shift ++ for candidate in "$@"; do ++ if [[ -x "${candidate}/${name}" ]]; then ++ printf '%s\n' "${candidate}/${name}" ++ return 0 ++ fi ++ done ++ return 1 ++} ++ ++build_dir="${LLAMA_BUILD_DIR:-}" ++binary_dirs=() ++if [[ -n "${build_dir}" ]]; then ++ binary_dirs+=("${build_dir}/bin" "${build_dir}") ++fi ++binary_dirs+=( ++ "${llama_root}/build/bin" ++ "${llama_root}/build" ++ "${llama_root}/../llama-build/build-stage-abi-static-metal-tests/bin" ++ "${PWD}/.deps/llama-build/build-stage-abi-static-metal-tests/bin" ++) ++ ++test_llama_archs="${TEST_LLAMA_ARCHS:-}" ++if [[ -z "${test_llama_archs}" ]]; then ++ if ! test_llama_archs="$(find_binary test-llama-archs "${binary_dirs[@]}")"; then ++ echo "error: could not find test-llama-archs; set LLAMA_BUILD_DIR or TEST_LLAMA_ARCHS" >&2 ++ exit 1 ++ fi ++fi ++test_backend_ops="${TEST_BACKEND_OPS:-}" ++if [[ -z "${test_backend_ops}" ]]; then ++ if ! test_backend_ops="$(find_binary test-backend-ops "${binary_dirs[@]}")"; then ++ echo "error: could not find test-backend-ops; set LLAMA_BUILD_DIR or TEST_BACKEND_OPS" >&2 ++ exit 1 ++ fi ++fi ++ ++timestamp="$(date -u +%Y%m%dT%H%M%SZ)" ++out_dir="${GLM52_PHASE_D_REPORT_DIR:-${PWD}/glm52-phase-d-policy-report-${timestamp}}" ++mkdir -p "${out_dir}" ++ ++phase_log="${out_dir}/phase-gates.log" ++perf_sparse_log="${out_dir}/sparse-policy-perf.console.log" ++perf_sparse_csv="${out_dir}/sparse-policy-perf.csv" ++perf_compact_fallback_log="${out_dir}/compact-gather-flash-fallback.console.log" ++perf_compact_fallback_csv="${out_dir}/compact-gather-flash-fallback.csv" ++perf_compact_components_log="${out_dir}/compact-decode-components.console.log" ++perf_compact_components_csv="${out_dir}/compact-decode-components.csv" ++perf_dense_log="${out_dir}/dense-flash-perf.console.log" ++perf_dense_csv="${out_dir}/dense-flash-perf.csv" ++summary="${out_dir}/summary.md" ++ ++echo "== GLM-DSA phase policy graph gates ==" ++"${test_llama_archs}" --arch glm-dsa 2>&1 | tee "${phase_log}" ++ ++echo "== GLM-DSA sparse policy perf rows ==" ++sparse_filter='kv=1024|n_kv=1024|n_kv=257,n_batch=4|n_kv=257,n_batch=8|n_kv=257,n_batch=16' ++"${test_backend_ops}" perf \ ++ -o GLM_DSA_SELECTED_ROW_FLASH,DSA_SPARSE_ATTN \ ++ -p "${sparse_filter}" \ ++ --output console 2>&1 | tee "${perf_sparse_log}" ++"${test_backend_ops}" perf \ ++ -o GLM_DSA_SELECTED_ROW_FLASH,DSA_SPARSE_ATTN \ ++ -p "${sparse_filter}" \ ++ --output csv > "${perf_sparse_csv}" ++ ++echo "== GLM-DSA compact gather + flash fallback row ==" ++compact_fallback_filter='kv=1024,top_k=768' ++SKIPPY_GLM_DSA_DISABLE_SELECTED_ROW_FLASH=1 \ ++"${test_backend_ops}" perf \ ++ -o GLM_DSA_SELECTED_ROW_FLASH \ ++ -p "${compact_fallback_filter}" \ ++ --output console 2>&1 | tee "${perf_compact_fallback_log}" ++SKIPPY_GLM_DSA_DISABLE_SELECTED_ROW_FLASH=1 \ ++"${test_backend_ops}" perf \ ++ -o GLM_DSA_SELECTED_ROW_FLASH \ ++ -p "${compact_fallback_filter}" \ ++ --output csv > "${perf_compact_fallback_csv}" ++ ++echo "== GLM-DSA compact decode component rows ==" ++compact_component_filter='kv=1024,top_k=768|hsk=576.*hsv=512.*nh=1.*kv=768.*nb=1.*mask=0.*prec=f32.*type_K=f16.*type_V=f16' ++"${test_backend_ops}" perf \ ++ -o GLM_DSA_COMPACT_K_GATHER,FLASH_ATTN_EXT \ ++ -p "${compact_component_filter}" \ ++ --output console 2>&1 | tee "${perf_compact_components_log}" ++"${test_backend_ops}" perf \ ++ -o GLM_DSA_COMPACT_K_GATHER,FLASH_ATTN_EXT \ ++ -p "${compact_component_filter}" \ ++ --output csv > "${perf_compact_components_csv}" ++ ++echo "== GLM-DSA dense flash comparison rows ==" ++dense_filter='kv=257,nb=4|kv=257,nb=8|kv=257,nb=16' ++"${test_backend_ops}" perf \ ++ -o FLASH_ATTN_EXT \ ++ -p "${dense_filter}" \ ++ --output console 2>&1 | tee "${perf_dense_log}" ++"${test_backend_ops}" perf \ ++ -o FLASH_ATTN_EXT \ ++ -p "${dense_filter}" \ ++ --output csv > "${perf_dense_csv}" ++ ++python3 - "${phase_log}" "${perf_sparse_log}" "${perf_compact_fallback_log}" "${perf_compact_components_log}" "${perf_dense_log}" "${summary}" <<'PY' ++import re ++import sys ++from pathlib import Path ++ ++phase_log, sparse_log, compact_fallback_log, compact_components_log, dense_log, summary = map(Path, sys.argv[1:]) ++ ++line_start_re = re.compile(r"^\s+[A-Z0-9_]+\([^)]*\):") ++line_re = re.compile( ++ r"^\s+(?P[A-Z0-9_]+)\((?P[^)]*)\):.*?" ++ r"(?P\d+) runs -\s+(?P[0-9.]+) us/run", ++ re.DOTALL, ++) ++ ++def parse(path): ++ rows = [] ++ pending = None ++ for line in path.read_text(errors="replace").splitlines(): ++ if line_start_re.match(line): ++ pending = line ++ elif pending is not None: ++ pending = f"{pending}\n{line}" ++ else: ++ continue ++ ++ if "not supported" in pending: ++ pending = None ++ continue ++ ++ match = line_re.search(pending) ++ if not match: ++ continue ++ ++ rows.append({ ++ "name": match.group("name"), ++ "params": match.group("params"), ++ "runs": int(match.group("runs")), ++ "us": float(match.group("us")), ++ }) ++ pending = None ++ return rows ++ ++rows = parse(sparse_log) + parse(dense_log) ++fallback_rows = parse(compact_fallback_log) ++component_rows = parse(compact_components_log) ++ ++def find(name, needle): ++ for row in rows: ++ if row["name"] == name and needle in row["params"]: ++ return row ++ return None ++ ++selected_1024 = find("GLM_DSA_SELECTED_ROW_FLASH", "kv=1024,top_k=768") ++sparse_1024 = find("DSA_SPARSE_ATTN", "n_kv=1024") ++compact_fallback_1024 = next(( ++ row for row in fallback_rows ++ if row["name"] == "GLM_DSA_SELECTED_ROW_FLASH" and "kv=1024,top_k=768" in row["params"] ++), None) ++compact_k_gather = next(( ++ row for row in component_rows ++ if row["name"] == "GLM_DSA_COMPACT_K_GATHER" and "kv=1024,top_k=768" in row["params"] ++), None) ++compact_flash = next(( ++ row for row in component_rows ++ if row["name"] == "FLASH_ATTN_EXT" and "hsk=576,hsv=512" in row["params"] and "kv=768,nb=1" in row["params"] and "mask=0" in row["params"] ++), None) ++ ++short_pairs = [] ++for batch in (4, 8, 16): ++ sparse = find("DSA_SPARSE_ATTN", f"n_kv=257,n_batch={batch}") ++ dense = find("FLASH_ATTN_EXT", f"kv=257,nb={batch}") ++ if sparse and dense: ++ short_pairs.append((batch, dense, sparse)) ++ ++lines = [ ++ "# GLM-5.2 Phase D Policy Report", ++ "", ++ "This report ties the GLM-DSA phase graph gates to focused backend microbench rows.", ++ "", ++ "## Inputs", ++ "", ++ f"- Phase gates: `{phase_log.name}`", ++ f"- Sparse policy perf: `{sparse_log.name}`", ++ f"- Compact gather fallback perf: `{compact_fallback_log.name}`", ++ f"- Compact decode components perf: `{compact_components_log.name}`", ++ f"- Dense flash perf: `{dense_log.name}`", ++ "", ++ "## Phase Gate", ++ "", ++] ++phase_text = phase_log.read_text(errors="replace") ++lines.append("- `test-llama-archs --arch glm-dsa` passed." if "OK" in phase_text else "- Phase gate did not report OK; inspect the raw log.") ++lines.extend([ ++ "- Covered phases: decode, short prefill, long prefill, verification.", ++ "- Native GLM-5.2-shaped gates use `top_k=768` where applicable.", ++ "", ++ "## Decode / Long-KV Policy", ++ "", ++]) ++if selected_1024 and sparse_1024: ++ ratio = sparse_1024["us"] / selected_1024["us"] ++ lines.append("| Shape | Selected-row flash | Direct sparse | Ratio |") ++ lines.append("| --- | ---: | ---: | ---: |") ++ lines.append(f"| `kv=1024, top_k=768` | {selected_1024['us']:.2f} us | {sparse_1024['us']:.2f} us | {ratio:.2f}x |") ++else: ++ lines.append("- Missing `kv=1024,top_k=768` decode comparison rows.") ++ ++lines.extend([ ++ "", ++ "## Compact Decode Kernel Decision", ++ "", ++]) ++if selected_1024 and compact_fallback_1024: ++ ratio = selected_1024["us"] / compact_fallback_1024["us"] ++ lines.append("| Shape | Selected-row fused | Compact gather + flash | Selected/fallback |") ++ lines.append("| --- | ---: | ---: | ---: |") ++ lines.append(f"| `kv=1024, top_k=768` | {selected_1024['us']:.2f} us | {compact_fallback_1024['us']:.2f} us | {ratio:.2f}x |") ++ if selected_1024["us"] <= compact_fallback_1024["us"]: ++ lines.append("- Decision: selected-row fused flash is eligible to promote for this shape.") ++ else: ++ lines.append("- Decision: keep selected-row fused flash opt-in; compact gather + flash is faster for this shape.") ++else: ++ lines.append("- Missing compact gather + flash fallback comparison row.") ++ ++lines.extend([ ++ "", ++ "## Compact Gather + Flash Components", ++ "", ++]) ++if compact_k_gather and compact_flash: ++ combined = compact_k_gather["us"] + compact_flash["us"] ++ lines.append("| Component | Shape | Time |") ++ lines.append("| --- | --- | ---: |") ++ lines.append(f"| Compact K gather | `f16 [576,1024] -> 768 rows` | {compact_k_gather['us']:.2f} us |") ++ lines.append(f"| Maskless compact flash | `dk=576, dv=512, kv=768, nb=1` | {compact_flash['us']:.2f} us |") ++ lines.append(f"| Combined component estimate | `gather + flash` | {combined:.2f} us |") ++ if compact_fallback_1024: ++ ratio = compact_fallback_1024["us"] / combined ++ lines.append(f"| Whole fallback graph | `kv=1024, top_k=768` | {compact_fallback_1024['us']:.2f} us |") ++ lines.append(f"| Whole/component ratio | `whole / components` | {ratio:.2f}x |") ++else: ++ lines.append("- Missing compact gather or maskless flash component rows.") ++ ++lines.extend([ ++ "", ++ "## Short Prefill / Verification Policy", ++ "", ++]) ++if short_pairs: ++ lines.append("| Batch | Dense flash | Direct sparse | Dense advantage |") ++ lines.append("| ---: | ---: | ---: | ---: |") ++ for batch, dense, sparse in short_pairs: ++ ratio = sparse["us"] / dense["us"] ++ lines.append(f"| {batch} | {dense['us']:.2f} us | {sparse['us']:.2f} us | {ratio:.2f}x |") ++else: ++ lines.append("- Missing short-prefill dense-vs-sparse comparison rows.") ++ ++lines.extend([ ++ "", ++ "## Raw Artifacts", ++ "", ++ f"- `{phase_log.name}`", ++ f"- `{sparse_log.name}`", ++ f"- `{Path(str(sparse_log).replace('.console.log', '.csv')).name}`", ++ f"- `{compact_fallback_log.name}`", ++ f"- `{Path(str(compact_fallback_log).replace('.console.log', '.csv')).name}`", ++ f"- `{compact_components_log.name}`", ++ f"- `{Path(str(compact_components_log).replace('.console.log', '.csv')).name}`", ++ f"- `{dense_log.name}`", ++ f"- `{Path(str(dense_log).replace('.console.log', '.csv')).name}`", ++ "", ++]) ++ ++summary.write_text("\n".join(lines)) ++print(summary.read_text()) ++PY ++ ++echo "wrote ${summary}" +diff --git a/scripts/glm52-phase-e-moe-report.sh b/scripts/glm52-phase-e-moe-report.sh +new file mode 100755 +index 00000000..8a1e0f38 +--- /dev/null ++++ b/scripts/glm52-phase-e-moe-report.sh +@@ -0,0 +1,950 @@ ++#!/usr/bin/env bash ++set -euo pipefail ++ ++script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" ++llama_root="$(cd "${script_dir}/.." && pwd)" ++ ++find_binary() { ++ local name="$1" ++ shift ++ for candidate in "$@"; do ++ if [[ -x "${candidate}/${name}" ]]; then ++ printf '%s\n' "${candidate}/${name}" ++ return 0 ++ fi ++ done ++ return 1 ++} ++ ++build_dir="${LLAMA_BUILD_DIR:-}" ++binary_dirs=() ++if [[ -n "${build_dir}" ]]; then ++ binary_dirs+=("${build_dir}/bin" "${build_dir}") ++fi ++binary_dirs+=( ++ "${llama_root}/build/bin" ++ "${llama_root}/build" ++ "${llama_root}/../llama-build/build-stage-abi-static-metal-tests/bin" ++ "${PWD}/.deps/llama-build/build-stage-abi-static-metal-tests/bin" ++) ++ ++test_backend_ops="${TEST_BACKEND_OPS:-}" ++if [[ -z "${test_backend_ops}" ]]; then ++ if ! test_backend_ops="$(find_binary test-backend-ops "${binary_dirs[@]}")"; then ++ echo "error: could not find test-backend-ops; set LLAMA_BUILD_DIR or TEST_BACKEND_OPS" >&2 ++ exit 1 ++ fi ++fi ++ ++timestamp="$(date -u +%Y%m%dT%H%M%SZ)" ++out_dir="${GLM52_PHASE_E_REPORT_DIR:-${PWD}/glm52-phase-e-moe-report-${timestamp}}" ++mkdir -p "${out_dir}" ++ ++moe_log="${out_dir}/moe-bottleneck-perf.console.log" ++moe_csv="${out_dir}/moe-bottleneck-perf.csv" ++summary="${out_dir}/summary.md" ++ ++backend_args=() ++if [[ -n "${GLM52_PHASE_E_BACKEND:-}" ]]; then ++ backend_args=(-b "${GLM52_PHASE_E_BACKEND}") ++fi ++ ++echo "== GLM-DSA MoE bottleneck perf rows ==" ++moe_filter='n_embd=6144|m=2048|m=4096|m=6144|n_tokens=1|n_tokens=4|n_embd=256' ++"${test_backend_ops}" perf \ ++ "${backend_args[@]}" \ ++ -o TOPK_MOE_GLM_ROUTE,TOPK_MOE_GLM_CONSUMER,TOPK_MOE_GLM_ACTIVATION,GLM_SHEXP_CONSUMER,GLM_SHEXP_COMPONENT,GLM_FFN_CONSUMER,MOE_WEIGHTED_SUM,MUL_MAT_ID \ ++ -p "${moe_filter}" \ ++ --output console 2>&1 | tee "${moe_log}" ++if [[ "${GLM52_PHASE_E_SKIP_CSV:-0}" == "0" ]]; then ++ "${test_backend_ops}" perf \ ++ "${backend_args[@]}" \ ++ -o TOPK_MOE_GLM_ROUTE,TOPK_MOE_GLM_CONSUMER,TOPK_MOE_GLM_ACTIVATION,GLM_SHEXP_CONSUMER,GLM_SHEXP_COMPONENT,GLM_FFN_CONSUMER,MOE_WEIGHTED_SUM,MUL_MAT_ID \ ++ -p "${moe_filter}" \ ++ --output csv > "${moe_csv}" ++fi ++ ++python3 - "${moe_log}" "${summary}" <<'PY' ++import os ++import re ++import sys ++from pathlib import Path ++ ++moe_log, summary = map(Path, sys.argv[1:]) ++ ++line_start_re = re.compile(r"^\s+[A-Z0-9_]+\([^)]*\):") ++line_re = re.compile( ++ r"^\s+(?P[A-Z0-9_]+)\((?P[^)]*)\):.*?" ++ r"(?P\d+) runs -\s+(?P[0-9.]+) us/run", ++ re.DOTALL, ++) ++ ++def parse(path): ++ rows = [] ++ pending = None ++ for line in path.read_text(errors="replace").splitlines(): ++ if line_start_re.match(line): ++ pending = line ++ elif pending is not None: ++ pending = f"{pending}\n{line}" ++ else: ++ continue ++ ++ if "not supported" in pending: ++ pending = None ++ continue ++ ++ match = line_re.search(pending) ++ if not match: ++ continue ++ ++ rows.append({ ++ "name": match.group("name"), ++ "params": match.group("params"), ++ "runs": int(match.group("runs")), ++ "us": float(match.group("us")), ++ }) ++ pending = None ++ return rows ++ ++rows = parse(moe_log) ++ ++def find(name, *needles): ++ for row in rows: ++ if row["name"] == name and all(needle in row["params"] for needle in needles): ++ return row ++ return None ++ ++route = find("TOPK_MOE_GLM_ROUTE") ++weighted_1 = find("MOE_WEIGHTED_SUM", "n_tokens=1") ++weighted_4 = find("MOE_WEIGHTED_SUM", "n_tokens=4") ++consumer_q = find("TOPK_MOE_GLM_CONSUMER", "gate_type=q2_K", "weighted_down=0", "fused_glu=0") ++consumer_q_fused = find("TOPK_MOE_GLM_CONSUMER", "gate_type=q2_K", "weighted_down=0", "fused_glu=1") ++consumer_q_weighted = find("TOPK_MOE_GLM_CONSUMER", "weighted_down=1", "fused_glu=0") ++consumer_q_weighted_fused = find("TOPK_MOE_GLM_CONSUMER", "weighted_down=1", "fused_glu=1") ++moe_act_unfused = find("TOPK_MOE_GLM_ACTIVATION", "n_ff=2048", "fused_glu=0") ++moe_act_fused = find("TOPK_MOE_GLM_ACTIVATION", "n_ff=2048", "fused_glu=1") ++shexp_unfused = find("GLM_SHEXP_CONSUMER", "n_embd=6144", "fused_glu=0") ++shexp_fused = find("GLM_SHEXP_CONSUMER", "n_embd=6144", "fused_glu=1", "merged_gate_up=0") ++shexp_merged = find("GLM_SHEXP_CONSUMER", "n_embd=6144", "fused_glu=1", "merged_gate_up=1") ++shexp_fused_q3 = find("GLM_SHEXP_CONSUMER", "gate_type=q3_K", "up_type=q3_K", "down_type=q3_K", "n_embd=6144", "fused_glu=1", "merged_gate_up=0") ++shexp_fused_q2 = find("GLM_SHEXP_CONSUMER", "gate_type=q2_K", "up_type=q2_K", "down_type=q2_K", "n_embd=6144", "fused_glu=1", "merged_gate_up=0") ++shexp_small = find("GLM_SHEXP_CONSUMER", "n_embd=256", "fused_glu=0") ++shexp_small_fused = find("GLM_SHEXP_CONSUMER", "n_embd=256", "fused_glu=1", "merged_gate_up=0") ++shexp_small_merged = find("GLM_SHEXP_CONSUMER", "n_embd=256", "fused_glu=1", "merged_gate_up=1") ++ffn_q3 = find("GLM_FFN_CONSUMER", "routed_down_type=q3_K", "n_embd=6144", "merged_shexp_gate_up=0") ++ffn_q2 = find("GLM_FFN_CONSUMER", "routed_down_type=q2_K", "n_embd=6144", "weighted_down=0", "merged_shexp_gate_up=0") ++ffn_q2_weighted = find("GLM_FFN_CONSUMER", "routed_down_type=q2_K", "n_embd=6144", "weighted_down=1", "merged_shexp_gate_up=0") ++ffn_q3_merged_shexp = find("GLM_FFN_CONSUMER", "routed_down_type=q3_K", "n_embd=6144", "merged_shexp_gate_up=1") ++ffn_small = find("GLM_FFN_CONSUMER", "n_embd=256") ++shexp_gate = find("GLM_SHEXP_COMPONENT", "component=0", "n_embd=6144") ++shexp_up = find("GLM_SHEXP_COMPONENT", "component=1", "n_embd=6144") ++shexp_gate_up = find("GLM_SHEXP_COMPONENT", "component=2", "n_embd=6144") ++shexp_down = find("GLM_SHEXP_COMPONENT", "component=3", "n_embd=6144") ++shexp_add = find("GLM_SHEXP_COMPONENT", "component=4", "n_embd=6144") ++shexp_swiglu = find("GLM_SHEXP_COMPONENT", "component=5", "n_embd=6144") ++shexp_merged_gate_up = find("GLM_SHEXP_COMPONENT", "component=6", "n_embd=6144") ++gate_up = find("MUL_MAT_ID", "type_a=q2_K,type_b=f32,n_mats=256,n_used=8,b=0,m=2048,n=1,k=6144") ++merged_gate_up = find("MUL_MAT_ID", "type_a=q2_K,type_b=f32,n_mats=256,n_used=8,b=0,m=4096,n=1,k=6144") ++down_q2 = find("MUL_MAT_ID", "type_a=q2_K,type_b=f32,n_mats=256,n_used=8,b=0,m=6144,n=1,k=2048") ++down = find("MUL_MAT_ID", "type_a=q3_K,type_b=f32,n_mats=256,n_used=8,b=0,m=6144,n=1,k=2048") ++ ++estimated_decode = None ++if route and weighted_1 and gate_up and down and moe_act_fused: ++ estimated_decode = route["us"] + 2.0 * gate_up["us"] + moe_act_fused["us"] + down["us"] + weighted_1["us"] ++ ++require_gates = os.environ.get("GLM52_PHASE_E_REQUIRE_GATES", "0") != "0" ++q2_ffn_min_speedup = float(os.environ.get("GLM52_PHASE_E_Q2_FFN_MIN_SPEEDUP", "1.08")) ++q2_weighted_review_speedup = float(os.environ.get("GLM52_PHASE_E_Q2_WEIGHTED_REVIEW_SPEEDUP", "1.02")) ++merged_shexp_review_speedup = float(os.environ.get("GLM52_PHASE_E_MERGED_SHEXP_REVIEW_SPEEDUP", "1.02")) ++ ++def gate_fail(message): ++ print(f"error: {message}", file=sys.stderr) ++ sys.exit(2) ++ ++if require_gates: ++ required_rows = { ++ "route/top-k": route, ++ "weighted sum": weighted_1, ++ "routed q2_K gate/up": gate_up, ++ "routed q3_K down": down, ++ "routed q2_K down alternative": down_q2, ++ "routed fused SwiGLU": moe_act_fused, ++ "small routed fused whole graph": consumer_q_fused, ++ "shared fused GLU": shexp_fused, ++ "shared q3_K fused GLU alternative": shexp_fused_q3, ++ "shared q2_K fused GLU alternative": shexp_fused_q2, ++ "combined q3_K routed down FFN": ffn_q3, ++ "combined q2_K routed down FFN": ffn_q2, ++ "combined q2_K weighted routed down FFN": ffn_q2_weighted, ++ "combined q3_K merged shared gate/up FFN": ffn_q3_merged_shexp, ++ "shared fused SwiGLU component": shexp_swiglu, ++ } ++ missing = [name for name, row in required_rows.items() if row is None] ++ if missing: ++ gate_fail("missing Phase E MoE gate rows: " + ", ".join(missing)) ++ if estimated_decode is None: ++ gate_fail("missing estimated routed FFN decode cost") ++ if down_q2["us"] >= down["us"]: ++ gate_fail(f"expected q2_K down alternative to be faster than q3_K down, got {down_q2['us']:.2f} us vs {down['us']:.2f} us") ++ q2_ffn_speedup = ffn_q3["us"] / ffn_q2["us"] ++ if q2_ffn_speedup < q2_ffn_min_speedup: ++ gate_fail(f"expected q2_K routed-down combined FFN to be at least {q2_ffn_min_speedup:.2f}x faster than q3_K, got {q2_ffn_speedup:.2f}x") ++ q2_weighted_speedup = ffn_q2["us"] / ffn_q2_weighted["us"] ++ if q2_weighted_speedup > q2_weighted_review_speedup: ++ gate_fail(f"q2_K weighted routed-down slots are now {q2_weighted_speedup:.2f}x faster; revisit q2 unweighted policy before accepting this result") ++ merged_shexp_speedup = ffn_q3["us"] / ffn_q3_merged_shexp["us"] ++ if merged_shexp_speedup > merged_shexp_review_speedup: ++ gate_fail(f"merged shared gate/up is now {merged_shexp_speedup:.2f}x faster; revisit disabled merged shared-expert policy before accepting this result") ++ if moe_act_fused["us"] >= estimated_decode * 0.05: ++ gate_fail(f"expected routed fused SwiGLU to stay below 5% of routed FFN estimate, got {moe_act_fused['us']:.2f} us of {estimated_decode:.2f} us") ++ ++lines = [ ++ "# GLM-5.2 Phase E MoE Report", ++ "", ++ "This report isolates GLM-DSA routed/shared MoE decode costs after the sparse-attention path has phase gates.", ++ "", ++ "## Inputs", ++ "", ++ f"- MoE perf: `{moe_log.name}`", ++ f"- CSV perf: `{'skipped' if os.environ.get('GLM52_PHASE_E_SKIP_CSV', '0') != '0' else 'moe-bottleneck-perf.csv'}`", ++ f"- Backend filter: `{os.environ.get('GLM52_PHASE_E_BACKEND', 'all')}`", ++ f"- Phase E gates: `{'required' if require_gates else 'advisory'}`", ++ f"- q2 routed-down combined FFN minimum speedup: `{q2_ffn_min_speedup:.2f}x`", ++ f"- q2 weighted routed-down review threshold: `{q2_weighted_review_speedup:.2f}x`", ++ f"- merged shared gate/up review threshold: `{merged_shexp_review_speedup:.2f}x`", ++ "", ++ "## Decode Components", ++ "", ++ "| Component | Shape | Time |", ++ "| --- | --- | ---: |", ++] ++ ++def add_row(label, row): ++ if row: ++ lines.append(f"| {label} | `{row['params']}` | {row['us']:.2f} us |") ++ else: ++ lines.append(f"| {label} | missing | missing |") ++ ++add_row("Route/top-k + weights", route) ++add_row("Weighted sum, 1 token", weighted_1) ++add_row("Weighted sum, 4 tokens", weighted_4) ++add_row("Small quantized whole-graph consumer, unfused diagnostic", consumer_q) ++add_row("Small quantized whole-graph consumer, fused GLU", consumer_q_fused) ++add_row("Small quantized whole-graph consumer, weighted down unfused diagnostic", consumer_q_weighted) ++add_row("Small quantized whole-graph consumer, weighted down fused GLU", consumer_q_weighted_fused) ++add_row("Routed activation/mul whole graph, unfused diagnostic", moe_act_unfused) ++add_row("Routed fused SwiGLU split", moe_act_fused) ++add_row("Small shared expert + final add, unfused diagnostic", shexp_small) ++add_row("Small shared expert + final add, fused GLU", shexp_small_fused) ++add_row("Small shared expert + final add, merged gate/up fused GLU", shexp_small_merged) ++add_row("Small combined routed+shared FFN", ffn_small) ++add_row("Shared expert + final add, unfused diagnostic", shexp_unfused) ++add_row("Shared expert + final add, fused GLU", shexp_fused) ++add_row("Shared expert + final add, merged gate/up fused GLU", shexp_merged) ++add_row("Shared expert + final add, q3_K fused GLU alternative", shexp_fused_q3) ++add_row("Shared expert + final add, q2_K fused GLU alternative", shexp_fused_q2) ++add_row("Shared gate matmul", shexp_gate) ++add_row("Shared up matmul", shexp_up) ++add_row("Shared activation/mul whole graph, unfused diagnostic", shexp_gate_up) ++add_row("Shared SwiGLU split", shexp_swiglu) ++add_row("Shared merged gate/up matmul + SwiGLU", shexp_merged_gate_up) ++add_row("Shared down matmul", shexp_down) ++add_row("Shared final add", shexp_add) ++add_row("Routed gate/up matmul", gate_up) ++add_row("Merged routed gate+up matmul", merged_gate_up) ++add_row("Routed down matmul, q2_K alternative", down_q2) ++add_row("Routed down matmul", down) ++add_row("Combined routed+shared FFN, q3_K routed down", ffn_q3) ++add_row("Combined routed+shared FFN, q2_K routed down", ffn_q2) ++add_row("Combined routed+shared FFN, q2_K routed down weighted slots", ffn_q2_weighted) ++add_row("Combined routed+shared FFN, q3_K routed down + merged shared gate/up", ffn_q3_merged_shexp) ++ ++lines.extend([ ++ "", ++ "## Estimated Routed FFN Decode Cost", ++ "", ++]) ++if estimated_decode is not None: ++ matmul_cost = 2.0 * gate_up["us"] + down["us"] ++ activation_cost = moe_act_fused["us"] ++ overhead = route["us"] + weighted_1["us"] ++ lines.append(f"- Estimated per-layer routed FFN decode cost: `{estimated_decode:.2f} us`.") ++ lines.append(f"- Expert matmuls: `{matmul_cost:.2f} us` ({100.0 * matmul_cost / estimated_decode:.1f}% of estimate).") ++ lines.append(f"- Routed fused SwiGLU split: `{activation_cost:.2f} us` ({100.0 * activation_cost / estimated_decode:.1f}% of estimate).") ++ lines.append(f"- Route + weighted sum: `{overhead:.2f} us` ({100.0 * overhead / estimated_decode:.1f}% of estimate).") ++ if merged_gate_up: ++ merged_cost = merged_gate_up["us"] + activation_cost + down["us"] + overhead ++ lines.append(f"- Merged gate+up estimate: `{merged_cost:.2f} us` ({estimated_decode / merged_cost:.2f}x faster than current estimate).") ++ if down_q2: ++ q2_down_cost = 2.0 * gate_up["us"] + activation_cost + down_q2["us"] + overhead ++ lines.append(f"- q2_K down estimate: `{q2_down_cost:.2f} us` ({estimated_decode / q2_down_cost:.2f}x faster than current estimate, quality tradeoff not measured here).") ++ if consumer_q and consumer_q_fused: ++ lines.append(f"- Small routed fused GLU whole-graph fixture: `{consumer_q_fused['us']:.2f} us` ({consumer_q['us'] / consumer_q_fused['us']:.2f}x faster than unfused diagnostic).") ++ if consumer_q_fused and consumer_q_weighted_fused: ++ lines.append(f"- Small weighted-down fused GLU whole-graph fixture: `{consumer_q_weighted_fused['us']:.2f} us` ({consumer_q_fused['us'] / consumer_q_weighted_fused['us']:.2f}x faster than fused baseline).") ++ if moe_act_unfused and moe_act_fused: ++ lines.append(f"- Routed unfused activation/mul diagnostic: `{moe_act_unfused['us']:.2f} us`; fused SwiGLU split: `{moe_act_fused['us']:.2f} us` (`{moe_act_unfused['us'] / moe_act_fused['us']:.2f}x` faster).") ++ if shexp_fused: ++ full_ffn = estimated_decode + shexp_fused["us"] ++ lines.append(f"- Shared expert + final add, fused GLU: `{shexp_fused['us']:.2f} us`; routed+shared FFN estimate: `{full_ffn:.2f} us`.") ++ lines.append(f"- Shared expert share of routed+shared estimate: `{100.0 * shexp_fused['us'] / full_ffn:.1f}%`.") ++ if shexp_unfused and shexp_fused: ++ lines.append(f"- Unfused diagnostic shared expert: `{shexp_unfused['us']:.2f} us`; fused GLU is `{shexp_unfused['us'] / shexp_fused['us']:.2f}x` faster.") ++ if shexp_fused and shexp_merged: ++ lines.append(f"- Merged shared gate/up fused GLU estimate: `{shexp_merged['us']:.2f} us` ({shexp_fused['us'] / shexp_merged['us']:.2f}x faster than separate shared gate/up).") ++ if shexp_fused and shexp_fused_q3: ++ lines.append(f"- Shared q3_K fused GLU alternative: `{shexp_fused_q3['us']:.2f} us` ({shexp_fused['us'] / shexp_fused_q3['us']:.2f}x versus q4_K, quality tradeoff not measured here).") ++ if shexp_fused and shexp_fused_q2: ++ lines.append(f"- Shared q2_K fused GLU alternative: `{shexp_fused_q2['us']:.2f} us` ({shexp_fused['us'] / shexp_fused_q2['us']:.2f}x versus q4_K, quality tradeoff not measured here).") ++ if all(row is not None for row in [shexp_gate, shexp_up, shexp_swiglu, shexp_down, shexp_add]): ++ shexp_components = shexp_gate["us"] + shexp_up["us"] + shexp_swiglu["us"] + shexp_down["us"] + shexp_add["us"] ++ shexp_matmuls = shexp_gate["us"] + shexp_up["us"] + shexp_down["us"] ++ lines.append(f"- Shared expert fused component row sum: `{shexp_components:.2f} us`.") ++ lines.append(f"- Shared expert matmuls: `{shexp_matmuls:.2f} us` ({100.0 * shexp_matmuls / shexp_components:.1f}% of component row sum).") ++ lines.append(f"- Shared SwiGLU split + final add: `{shexp_swiglu['us'] + shexp_add['us']:.2f} us`.") ++ if shexp_gate_up and shexp_swiglu: ++ lines.append(f"- Unfused activation/mul diagnostic: `{shexp_gate_up['us']:.2f} us`; fused SwiGLU split: `{shexp_swiglu['us']:.2f} us` (`{shexp_gate_up['us'] / shexp_swiglu['us']:.2f}x` faster).") ++ if shexp_merged_gate_up: ++ lines.append(f"- Merged shared gate/up matmul + SwiGLU component row: `{shexp_merged_gate_up['us']:.2f} us`.") ++ if ffn_q3: ++ lines.append(f"- Combined routed+shared FFN, q3_K routed down: `{ffn_q3['us']:.2f} us`.") ++ if ffn_q3 and ffn_q2: ++ lines.append(f"- Combined FFN q2_K routed-down alternative: `{ffn_q2['us']:.2f} us` (`{ffn_q3['us'] / ffn_q2['us']:.2f}x` faster, quality tradeoff not measured here).") ++ if ffn_q2 and ffn_q2_weighted: ++ lines.append(f"- Combined FFN q2_K weighted routed-down slots: `{ffn_q2_weighted['us']:.2f} us` (`{ffn_q2['us'] / ffn_q2_weighted['us']:.2f}x` versus q2_K baseline).") ++ if ffn_q3 and ffn_q3_merged_shexp: ++ lines.append(f"- Combined FFN merged shared gate/up: `{ffn_q3_merged_shexp['us']:.2f} us` (`{ffn_q3['us'] / ffn_q3_merged_shexp['us']:.2f}x` faster than separate shared gate/up).") ++else: ++ lines.append("- Missing rows prevented the per-layer estimate.") ++ ++lines.extend([ ++ "", ++ "## Raw Artifacts", ++ "", ++ f"- `{moe_log.name}`", ++ "- `moe-bottleneck-perf.csv`", ++ "", ++]) ++ ++summary.write_text("\n".join(lines)) ++print(summary.read_text()) ++PY ++ ++if [[ "${GLM52_PHASE_E_ROUTED_WHOLE_GRAPH_SANITY:-1}" != "0" ]]; then ++ sanity_log="${out_dir}/routed-whole-graph-sanity.console.log" ++ "${test_backend_ops}" perf \ ++ -o TOPK_MOE_GLM_CONSUMER_SANITY \ ++ -p 'n_embd=6144,n_ff=2048,.*fused_glu=1' \ ++ --output console > "${sanity_log}" 2>&1 ++ ++ python3 - "${summary}" "${sanity_log}" <<'PY' ++import os ++import re ++import sys ++from pathlib import Path ++ ++summary, sanity_log = map(Path, sys.argv[1:]) ++text = sanity_log.read_text(errors="replace") ++summary_text = summary.read_text() ++ ++row_re = re.compile( ++ r"TOPK_MOE_GLM_CONSUMER_SANITY\((?P[^)]*)\):.*?" ++ r"(?P\d+) runs -\s+(?P[0-9.]+) us/run", ++ re.DOTALL, ++) ++estimate_re = re.compile(r"Estimated per-layer routed FFN decode cost: `(?P[0-9.]+) us`") ++shared_re = re.compile(r"Shared expert \+ final add, fused GLU: `(?P[0-9.]+) us`") ++combined_q2_re = re.compile(r"Combined FFN q2_K routed-down alternative: `(?P[0-9.]+) us`") ++combined_q3_re = re.compile(r"Combined routed\+shared FFN, q3_K routed down: `(?P[0-9.]+) us`") ++ ++rows = [ ++ { ++ "params": match.group("params"), ++ "runs": int(match.group("runs")), ++ "us": float(match.group("us")), ++ } ++ for match in row_re.finditer(text) ++] ++estimate_match = estimate_re.search(summary_text) ++estimate_us = float(estimate_match.group("us")) if estimate_match else None ++shared_match = shared_re.search(summary_text) ++shared_us = float(shared_match.group("us")) if shared_match else None ++combined_q2_match = combined_q2_re.search(summary_text) ++combined_q2_us = float(combined_q2_match.group("us")) if combined_q2_match else None ++combined_q3_match = combined_q3_re.search(summary_text) ++combined_q3_us = float(combined_q3_match.group("us")) if combined_q3_match else None ++require_gates = os.environ.get("GLM52_PHASE_E_REQUIRE_GATES", "0") != "0" ++ ++q3 = next((row for row in rows if "down_type=q3_K" in row["params"] and "weighted_down=0" in row["params"]), None) ++q3_weighted = next((row for row in rows if "down_type=q3_K" in row["params"] and "weighted_down=1" in row["params"]), None) ++q2 = next((row for row in rows if "down_type=q2_K" in row["params"] and "weighted_down=0" in row["params"]), None) ++q2_weighted = next((row for row in rows if "down_type=q2_K" in row["params"] and "weighted_down=1" in row["params"]), None) ++ ++lines = [ ++ "", ++ "## Routed Whole-Graph Sanity Probe", ++ "", ++ "This probe is enabled by default and can be disabled with `GLM52_PHASE_E_ROUTED_WHOLE_GRAPH_SANITY=0`.", ++ "It times production-shaped routed MoE consumer graphs, but the row is accepted only if it is physically plausible against the isolated routed FFN estimate.", ++ "", ++ "| Case | Shape | Time | Verdict |", ++ "| --- | --- | ---: | --- |", ++] ++ ++def verdict(row): ++ if row is None: ++ return "missing" ++ if estimate_us is not None and row["us"] < estimate_us * 0.25: ++ return "rejected: below 25% of isolated routed FFN estimate" ++ return "plausible" ++ ++for label, row in (("q3_K down", q3), ("q3_K weighted-down", q3_weighted), ("q2_K down", q2), ("q2_K weighted-down", q2_weighted)): ++ if row is None: ++ lines.append(f"| {label} | missing | missing | missing |") ++ else: ++ lines.append(f"| {label} | `{row['params']}` | {row['us']:.2f} us | {verdict(row)} |") ++ ++if q3 and q2: ++ lines.append("") ++ q2_speedup = q3["us"] / q2["us"] ++ lines.append(f"- q2_K whole-graph speedup versus q3_K from the sanity probe: `{q2_speedup:.2f}x`.") ++if q3 and q3_weighted: ++ weighted_speedup = q3["us"] / q3_weighted["us"] ++ lines.append(f"- q3_K weighted-down whole-graph speedup versus q3_K baseline: `{weighted_speedup:.2f}x`.") ++if q2 and q2_weighted: ++ q2_weighted_speedup = q2["us"] / q2_weighted["us"] ++ lines.append(f"- q2_K weighted-down whole-graph speedup versus q2_K baseline: `{q2_weighted_speedup:.2f}x`.") ++if estimate_us is not None: ++ lines.append(f"- Isolated routed FFN estimate used for plausibility: `{estimate_us:.2f} us`.") ++if q2 and estimate_us is not None: ++ lines.append(f"- q2_K full routed graph is `{q2['us'] / estimate_us:.2f}x` the isolated routed FFN estimate.") ++if q2 and shared_us is not None and combined_q2_us is not None: ++ parts_sum = q2["us"] + shared_us ++ parts_max = max(q2["us"], shared_us) ++ overlap_saved = parts_sum - combined_q2_us ++ max_overhead = combined_q2_us - parts_max ++ lines.append(f"- q2_K full routed graph: `{q2['us']:.2f} us`; shared expert graph: `{shared_us:.2f} us`; combined q2_K FFN: `{combined_q2_us:.2f} us`.") ++ lines.append(f"- Naive routed+shared sum: `{parts_sum:.2f} us`; combined graph saves `{overlap_saved:.2f} us` versus serial execution.") ++ lines.append(f"- Combined graph overhead above the slower branch: `{max_overhead:.2f} us` (`{combined_q2_us / parts_max:.2f}x` slower than perfect overlap).") ++if q3 and shared_us is not None and combined_q3_us is not None: ++ parts_sum = q3["us"] + shared_us ++ parts_max = max(q3["us"], shared_us) ++ overlap_saved = parts_sum - combined_q3_us ++ max_overhead = combined_q3_us - parts_max ++ lines.append(f"- q3_K full routed graph: `{q3['us']:.2f} us`; shared expert graph: `{shared_us:.2f} us`; combined q3_K FFN: `{combined_q3_us:.2f} us`.") ++ lines.append(f"- q3_K combined graph saves `{overlap_saved:.2f} us` versus serial routed+shared execution and is `{combined_q3_us / parts_max:.2f}x` the slower branch.") ++lines.append(f"- Raw artifact: `{sanity_log.name}`.") ++ ++decision_lines = [] ++if q3 and q2: ++ if q2_speedup >= 1.05: ++ decision_lines.append( ++ f"- q2_K routed down is the only measured full-shape speed candidate here (`{q2_speedup:.2f}x`), " ++ "but it remains a quant-quality experiment rather than a default runtime change.") ++ else: ++ decision_lines.append( ++ f"- q2_K routed down is not a strong speed candidate in this run (`{q2_speedup:.2f}x`).") ++if q3 and q3_weighted: ++ if weighted_speedup > 1.02: ++ decision_lines.append( ++ f"- Weighted-down q3_K remains a possible runtime experiment (`{weighted_speedup:.2f}x`).") ++ else: ++ decision_lines.append( ++ f"- Weighted-down q3_K is not a speed win in this run (`{weighted_speedup:.2f}x`), so keep it off by default.") ++if q2 and q2_weighted: ++ if q2_weighted_speedup > 1.02: ++ decision_lines.append( ++ f"- Weighted-down q2_K should be tested in real decode (`{q2_weighted_speedup:.2f}x` in the whole-graph sanity probe).") ++ else: ++ decision_lines.append( ++ f"- Weighted-down q2_K is not a speed win in this run (`{q2_weighted_speedup:.2f}x`), so keep it off by default.") ++if estimate_us is not None and q3: ++ lines.append(f"- Full-shape q3_K sanity is `{q3['us'] / estimate_us:.2f}x` the isolated routed FFN estimate.") ++ decision_lines.append( ++ "- The remaining routed MoE work should focus on full routed-graph overhead and selected-expert matmul scheduling; " ++ "isolated route/top-k, weighted-sum, and component rows are not enough to predict decode cost.") ++if q2 and shared_us is not None and combined_q2_us is not None: ++ decision_lines.append( ++ "- Combined q2_K FFN already overlaps routed and shared work substantially, so the next high-impact target is " ++ "the routed full graph rather than shared-expert/combined scheduling alone.") ++decision_lines.append( ++ "- Merged shared gate/up rows are synthetic package-format probes. Treat them as evidence for a future " ++ "quant/package experiment only; they are not a runtime default unless the model package contains a real " ++ "merged shared-expert tensor.") ++if decision_lines: ++ lines.extend(["", "### Phase E Decision Hints", ""]) ++ lines.extend(decision_lines) ++ ++section = "\n".join(lines) + "\n" ++summary.write_text(summary_text + section) ++print(section) ++ ++if require_gates: ++ failures = [] ++ for label, row in (("q3_K down", q3), ("q3_K weighted-down", q3_weighted), ("q2_K down", q2), ("q2_K weighted-down", q2_weighted)): ++ row_verdict = verdict(row) ++ if row_verdict != "plausible": ++ failures.append(f"{label}: {row_verdict}") ++ if failures: ++ print("error: routed whole-graph sanity gate failed: " + ", ".join(failures), file=sys.stderr) ++ sys.exit(2) ++PY ++fi ++ ++if [[ "${GLM52_PHASE_E_ROUTE_COMPONENT_LADDER:-0}" != "0" ]]; then ++ route_ladder_log="${out_dir}/route-component-ladder.console.log" ++ "${test_backend_ops}" perf \ ++ "${backend_args[@]}" \ ++ -o TOPK_MOE_GLM_ROUTE_COMPONENT \ ++ --output console > "${route_ladder_log}" 2>&1 ++ ++ python3 - "${summary}" "${route_ladder_log}" <<'PY' ++import re ++import sys ++from pathlib import Path ++ ++summary, route_ladder_log = map(Path, sys.argv[1:]) ++text = route_ladder_log.read_text(errors="replace") ++ ++row_re = re.compile( ++ r"TOPK_MOE_GLM_ROUTE_COMPONENT\((?P[^)]*)\):.*?" ++ r"(?P\d+) runs -\s+(?P[0-9.]+) us/run", ++ re.DOTALL, ++) ++component_re = re.compile(r"component=(?P\d+)") ++ ++labels = { ++ 0: "sigmoid probs", ++ 1: "bias add", ++ 2: "top-k indices", ++ 3: "top-k weights get_rows", ++ 4: "weights sum", ++ 5: "sum clamp", ++ 6: "weights normalize", ++ 7: "final scale", ++} ++ ++rows = [] ++for match in row_re.finditer(text): ++ component_match = component_re.search(match.group("params")) ++ if not component_match: ++ continue ++ component = int(component_match.group("component")) ++ rows.append({ ++ "component": component, ++ "params": match.group("params"), ++ "runs": int(match.group("runs")), ++ "us": float(match.group("us")), ++ }) ++ ++rows.sort(key=lambda row: row["component"]) ++by_component = {row["component"]: row for row in rows} ++ ++lines = [ ++ "", ++ "## Optional Route Component Ladder", ++ "", ++ "This ladder is enabled by `GLM52_PHASE_E_ROUTE_COMPONENT_LADDER=1`.", ++ "It times cumulative GLM route/top-k subgraphs so decode routing overhead can be separated from selected expert matmuls.", ++ "", ++ "| Component | Case | Shape | Time | Delta from previous |", ++ "| ---: | --- | --- | ---: | ---: |", ++] ++ ++previous = None ++for component in range(8): ++ row = by_component.get(component) ++ if row is None: ++ lines.append(f"| {component} | {labels[component]} | missing | missing | missing |") ++ continue ++ delta = "n/a" if previous is None else f"{row['us'] - previous:.2f} us" ++ lines.append(f"| {component} | {labels[component]} | `{row['params']}` | {row['us']:.2f} us | {delta} |") ++ previous = row["us"] ++ ++if 0 in by_component and 7 in by_component: ++ lines.extend([ ++ "", ++ f"- Full route ladder is `{by_component[7]['us']:.2f} us`, `{by_component[7]['us'] / by_component[0]['us']:.2f}x` the sigmoid-only graph.", ++ ]) ++if 1 in by_component and 2 in by_component: ++ lines.append(f"- Top-k/index production adds `{by_component[2]['us'] - by_component[1]['us']:.2f} us` over sigmoid+bias.") ++if 2 in by_component and 3 in by_component: ++ lines.append(f"- Weight `get_rows` adds `{by_component[3]['us'] - by_component[2]['us']:.2f} us` after top-k.") ++if 3 in by_component and 7 in by_component: ++ lines.append(f"- Normalization and final scale add `{by_component[7]['us'] - by_component[3]['us']:.2f} us` after raw selected weights.") ++lines.append(f"- Raw artifact: `{route_ladder_log.name}`.") ++ ++section = "\n".join(lines) + "\n" ++summary.write_text(summary.read_text() + section) ++print(section) ++PY ++fi ++ ++if [[ "${GLM52_PHASE_E_ROUTED_COMPONENT_LADDER:-0}" != "0" ]]; then ++ ladder_log="${out_dir}/routed-component-ladder.console.log" ++ "${test_backend_ops}" perf \ ++ "${backend_args[@]}" \ ++ -o TOPK_MOE_GLM_COMPONENT \ ++ -p 'down_type=q2_K' \ ++ --output console > "${ladder_log}" 2>&1 ++ ++ python3 - "${summary}" "${ladder_log}" <<'PY' ++import re ++import sys ++from pathlib import Path ++ ++summary, ladder_log = map(Path, sys.argv[1:]) ++text = ladder_log.read_text(errors="replace") ++ ++row_re = re.compile( ++ r"TOPK_MOE_GLM_COMPONENT\((?P[^)]*)\):.*?" ++ r"(?P\d+) runs -\s+(?P[0-9.]+) us/run", ++ re.DOTALL, ++) ++component_re = re.compile(r"component=(?P\d+)") ++ ++labels = { ++ 0: "route weights", ++ 1: "route-anchored input", ++ 2: "route + gate matmul", ++ 3: "route + gate/up matmuls", ++ 4: "route + gate/up + SwiGLU", ++ 5: "route + gate/up + SwiGLU + down", ++ 6: "routed output", ++ 7: "routed output checksum", ++} ++ ++rows = [] ++for match in row_re.finditer(text): ++ component_match = component_re.search(match.group("params")) ++ if not component_match: ++ continue ++ component = int(component_match.group("component")) ++ rows.append({ ++ "component": component, ++ "params": match.group("params"), ++ "runs": int(match.group("runs")), ++ "us": float(match.group("us")), ++ }) ++ ++rows.sort(key=lambda row: row["component"]) ++by_component = {row["component"]: row for row in rows} ++ ++lines = [ ++ "", ++ "## Optional Routed Component Ladder", ++ "", ++ "This ladder is enabled by `GLM52_PHASE_E_ROUTED_COMPONENT_LADDER=1`.", ++ "It times cumulative production-shaped routed MoE subgraphs so isolated hot-cache matmul rows are not mistaken for full decode cost.", ++ "", ++ "| Component | Case | Shape | Time | Delta from previous |", ++ "| ---: | --- | --- | ---: | ---: |", ++] ++ ++previous = None ++for component in range(8): ++ row = by_component.get(component) ++ if row is None: ++ lines.append(f"| {component} | {labels[component]} | missing | missing | missing |") ++ continue ++ delta = "n/a" if previous is None else f"{row['us'] - previous:.2f} us" ++ lines.append(f"| {component} | {labels[component]} | `{row['params']}` | {row['us']:.2f} us | {delta} |") ++ previous = row["us"] ++ ++if 0 in by_component and 7 in by_component: ++ lines.extend([ ++ "", ++ f"- Routed checksum graph is `{by_component[7]['us']:.2f} us`, `{by_component[7]['us'] / by_component[0]['us']:.2f}x` the route-only ladder row.", ++ ]) ++if 1 in by_component and 2 in by_component: ++ lines.append(f"- First routed gate matmul adds `{by_component[2]['us'] - by_component[1]['us']:.2f} us` in this graph context.") ++if 2 in by_component and 3 in by_component: ++ lines.append(f"- Second routed up matmul plus output merge adds `{by_component[3]['us'] - by_component[2]['us']:.2f} us`.") ++if 4 in by_component and 5 in by_component: ++ lines.append(f"- q2_K routed down adds `{by_component[5]['us'] - by_component[4]['us']:.2f} us` in this graph context.") ++if 5 in by_component and 6 in by_component: ++ lines.append(f"- Weighted sum adds `{by_component[6]['us'] - by_component[5]['us']:.2f} us` after the down projection.") ++if 0 in by_component: ++ lines.append( ++ "- If the route-only ladder row is much slower than `TOPK_MOE_GLM_ROUTE`, investigate graph context, output materialization, and tensor residency before tuning standalone kernels." ++ ) ++lines.append(f"- Raw artifact: `{ladder_log.name}`.") ++ ++section = "\n".join(lines) + "\n" ++summary.write_text(summary.read_text() + section) ++print(section) ++PY ++fi ++ ++if [[ "${GLM52_PHASE_E_KERNEL_SWEEP:-0}" != "0" ]]; then ++ sweep_dir="${out_dir}/kernel-sweep" ++ mkdir -p "${sweep_dir}" ++ ++ forced_mm_log="${sweep_dir}/mul-mm-id-min-1.console.log" ++ SKIPPY_GLM_DSA_MUL_MM_ID_MIN_TOKENS=1 \ ++ "${test_backend_ops}" perf \ ++ -o MUL_MAT_ID \ ++ -p 'm=2048,n=1,k=6144|m=4096,n=1,k=6144|m=6144,n=1,k=2048' \ ++ --output console > "${forced_mm_log}" 2>&1 ++ ++ for nsg in 1 2 4; do ++ LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_GLM_DOWN=0 \ ++ LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_NSG="${nsg}" \ ++ "${test_backend_ops}" perf \ ++ -o MUL_MAT_ID \ ++ -p 'type_a=q3_K,type_b=f32,n_mats=256,n_used=8,b=0,m=6144,n=1,k=2048' \ ++ --output console > "${sweep_dir}/q3-k-mul-mv-id-nsg-${nsg}.console.log" 2>&1 ++ done ++ ++ for nr0 in 1 2 4 8; do ++ LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_GLM_DOWN=0 \ ++ LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_NR0="${nr0}" \ ++ "${test_backend_ops}" perf \ ++ -o MUL_MAT_ID \ ++ -p 'type_a=q3_K,type_b=f32,n_mats=256,n_used=8,b=0,m=6144,n=1,k=2048' \ ++ --output console > "${sweep_dir}/q3-k-mul-mv-id-nr0-${nr0}.console.log" 2>&1 ++ done ++ ++ for nsg in 2 4; do ++ LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_GLM_DOWN=1 \ ++ LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_NSG="${nsg}" \ ++ "${test_backend_ops}" perf \ ++ -o MUL_MAT_ID \ ++ -p 'type_a=q3_K,type_b=f32,n_mats=256,n_used=8,b=0,m=6144,n=1,k=2048' \ ++ --output console > "${sweep_dir}/q3-k-mul-mv-id-glm-down-nsg-${nsg}.console.log" 2>&1 ++ done ++ ++ python3 - "${summary}" "${sweep_dir}" <<'PY' ++import os ++import re ++import sys ++from pathlib import Path ++ ++summary, sweep_dir = Path(sys.argv[1]), Path(sys.argv[2]) ++ ++line_re = re.compile( ++ r"MUL_MAT_ID\((?P[^)]*)\):.*?" ++ r"(?P\d+) runs -\s+(?P[0-9.]+) us/run", ++ re.DOTALL, ++) ++ ++def rows(path): ++ text = path.read_text(errors="replace") ++ return [ ++ { ++ "params": match.group("params"), ++ "runs": int(match.group("runs")), ++ "us": float(match.group("us")), ++ } ++ for match in line_re.finditer(text) ++ ] ++ ++forced_rows = rows(sweep_dir / "mul-mm-id-min-1.console.log") ++q3_rows = [] ++for nsg in (1, 2, 4): ++ parsed = rows(sweep_dir / f"q3-k-mul-mv-id-nsg-{nsg}.console.log") ++ row = next((item for item in parsed if "type_a=q3_K" in item["params"]), None) ++ q3_rows.append((nsg, row)) ++ ++q3_nr0_rows = [] ++for nr0 in (1, 2, 4, 8): ++ parsed = rows(sweep_dir / f"q3-k-mul-mv-id-nr0-{nr0}.console.log") ++ row = next((item for item in parsed if "type_a=q3_K" in item["params"]), None) ++ q3_nr0_rows.append((nr0, row)) ++ ++q3_glm_down_rows = [] ++for nsg in (2, 4): ++ parsed = rows(sweep_dir / f"q3-k-mul-mv-id-glm-down-nsg-{nsg}.console.log") ++ row = next((item for item in parsed if "type_a=q3_K" in item["params"]), None) ++ q3_glm_down_rows.append((nsg, row)) ++ ++lines = [ ++ "", ++ "## Optional Kernel Sweep", ++ "", ++ "This sweep is enabled by `GLM52_PHASE_E_KERNEL_SWEEP=1` and records kernel-policy evidence for tempting GLM-DSA MoE alternatives.", ++ "The ordinary q3_K baseline rows explicitly set `LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_GLM_DOWN=0` so the fixed-block default is not compared to itself.", ++ "", ++ "| Case | Shape | Time |", ++ "| --- | --- | ---: |", ++] ++ ++for row in forced_rows: ++ label = "Forced matrix-matrix cutoff" ++ lines.append(f"| {label} | `{row['params']}` | {row['us']:.2f} us |") ++ ++for nsg, row in q3_rows: ++ if row: ++ lines.append(f"| q3_K `mul_mv_id` nsg={nsg} | `{row['params']}` | {row['us']:.2f} us |") ++ else: ++ lines.append(f"| q3_K `mul_mv_id` nsg={nsg} | missing | missing |") ++ ++for nr0, row in q3_nr0_rows: ++ if row: ++ lines.append(f"| q3_K `mul_mv_id` nr0={nr0} | `{row['params']}` | {row['us']:.2f} us |") ++ else: ++ lines.append(f"| q3_K `mul_mv_id` nr0={nr0} | missing | missing |") ++ ++for nsg, row in q3_glm_down_rows: ++ if row: ++ lines.append(f"| q3_K fixed-block GLM down nsg={nsg} | `{row['params']}` | {row['us']:.2f} us |") ++ else: ++ lines.append(f"| q3_K fixed-block GLM down nsg={nsg} | missing | missing |") ++ ++forced_q3 = next((row for row in forced_rows if "type_a=q3_K" in row["params"]), None) ++default_q3 = next((row for nsg, row in q3_rows if nsg == 2 and row), None) ++best_q3 = min((row for _, row in q3_rows if row), key=lambda row: row["us"], default=None) ++default_nr0_q3 = next((row for nr0, row in q3_nr0_rows if nr0 == 4 and row), None) ++best_nr0_q3 = min((row for _, row in q3_nr0_rows if row), key=lambda row: row["us"], default=None) ++default_glm_down_q3 = next((row for nsg, row in q3_glm_down_rows if nsg == 2 and row), None) ++best_glm_down_q3 = min((row for _, row in q3_glm_down_rows if row), key=lambda row: row["us"], default=None) ++ ++require_gates = os.environ.get("GLM52_PHASE_E_REQUIRE_GATES", "0") != "0" ++ ++def gate_fail(message): ++ print(f"error: {message}", file=sys.stderr) ++ sys.exit(2) ++ ++if require_gates: ++ missing = [] ++ if forced_q3 is None: ++ missing.append("forced q3_K mul_mm_id row") ++ missing.extend( ++ f"q3_K mul_mv_id nsg={nsg}" ++ for nsg, row in q3_rows ++ if row is None ++ ) ++ missing.extend( ++ f"q3_K mul_mv_id nr0={nr0}" ++ for nr0, row in q3_nr0_rows ++ if row is None ++ ) ++ missing.extend( ++ f"q3_K fixed-block GLM down nsg={nsg}" ++ for nsg, row in q3_glm_down_rows ++ if row is None ++ ) ++ if missing: ++ gate_fail("missing Phase E kernel sweep rows: " + ", ".join(missing)) ++ if forced_q3["us"] <= default_q3["us"] * 2.0: ++ gate_fail( ++ "expected forced q3_K mul_mm_id to be at least 2x slower than " ++ f"default mul_mv_id, got {forced_q3['us']:.2f} us vs " ++ f"{default_q3['us']:.2f} us" ++ ) ++ ++lines.extend(["", "### Kernel Sweep Notes", ""]) ++if forced_q3 and default_q3: ++ lines.append(f"- Forcing `mul_mm_id` at one token measured `{forced_q3['us']:.2f} us` for q3_K down versus `{default_q3['us']:.2f} us` on the default `mul_mv_id` path; keep one-token decode on matrix-vector kernels.") ++if default_q3 and best_q3: ++ lines.append(f"- q3_K `mul_mv_id` simdgroup tuning is noise-level here: default nsg=2 is `{default_q3['us']:.2f} us`, best measured row is `{best_q3['us']:.2f} us`.") ++if default_nr0_q3 and best_nr0_q3: ++ lines.append(f"- q3_K `mul_mv_id` row-height tuning: default nr0=4 is `{default_nr0_q3['us']:.2f} us`, best measured row is `{best_nr0_q3['us']:.2f} us`.") ++if default_q3 and default_glm_down_q3: ++ lines.append(f"- q3_K fixed-block GLM down specialization: default nsg=2 is `{default_glm_down_q3['us']:.2f} us` versus ordinary q3_K `mul_mv_id` at `{default_q3['us']:.2f} us`.") ++if best_glm_down_q3 and default_q3: ++ lines.append(f"- Best fixed-block GLM down row is `{best_glm_down_q3['us']:.2f} us`, `{default_q3['us'] / best_glm_down_q3['us']:.2f}x` versus ordinary default q3_K `mul_mv_id`.") ++ ++section = "\n".join(lines) + "\n" ++summary.write_text(summary.read_text() + section) ++print(section) ++PY ++fi ++ ++if [[ "${GLM52_PHASE_E_FIXED_BLOCK_DISPATCH:-${GLM52_PHASE_E_KERNEL_SWEEP:-0}}" != "0" ]]; then ++ dispatch_dir="${out_dir}/fixed-block-dispatch" ++ mkdir -p "${dispatch_dir}" ++ ++ default_dispatch_log="${dispatch_dir}/q3-k-glm-down-default.console.log" ++ disabled_dispatch_log="${dispatch_dir}/q3-k-glm-down-disabled.console.log" ++ ++ SKIPPY_GLM_DSA_LOG_METAL_DISPATCH=1 \ ++ "${test_backend_ops}" test \ ++ -o MUL_MAT_ID \ ++ -p 'type_a=q3_K,type_b=f32,n_mats=256,n_used=8,b=0,m=6144,n=1,k=2048' \ ++ -j 1 > "${default_dispatch_log}" 2>&1 ++ ++ LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_GLM_DOWN=0 \ ++ SKIPPY_GLM_DSA_LOG_METAL_DISPATCH=1 \ ++ "${test_backend_ops}" test \ ++ -o MUL_MAT_ID \ ++ -p 'type_a=q3_K,type_b=f32,n_mats=256,n_used=8,b=0,m=6144,n=1,k=2048' \ ++ -j 1 > "${disabled_dispatch_log}" 2>&1 ++ ++ python3 - "${summary}" "${default_dispatch_log}" "${disabled_dispatch_log}" <<'PY' ++import os ++import re ++import sys ++from pathlib import Path ++ ++summary, default_log, disabled_log = map(Path, sys.argv[1:]) ++default_text = default_log.read_text(errors="replace") ++disabled_text = disabled_log.read_text(errors="replace") ++require_gates = os.environ.get("GLM52_PHASE_E_REQUIRE_GATES", "0") != "0" ++ ++pipeline_re = re.compile(r"kernel_mul_mv_id_q3_K_f32(?:_glm_down)?_nsg=\d+") ++ ++def first_pipeline(text): ++ match = pipeline_re.search(text) ++ return match.group(0) if match else "missing" ++ ++default_pipeline = first_pipeline(default_text) ++disabled_pipeline = first_pipeline(disabled_text) ++default_passed = "1/1 tests passed" in default_text and "3/3 backends passed" in default_text ++disabled_passed = "1/1 tests passed" in disabled_text and "3/3 backends passed" in disabled_text ++default_fixed = "kernel_mul_mv_id_q3_K_f32_glm_down" in default_pipeline ++disabled_ordinary = ( ++ "kernel_mul_mv_id_q3_K_f32_nsg=" in disabled_pipeline and ++ "glm_down" not in disabled_pipeline ++) ++ ++lines = [ ++ "", ++ "## Fixed-Block Dispatch Gate", ++ "", ++ "This gate runs a focused `MUL_MAT_ID` correctness case with Metal dispatch logging.", ++ "It proves the default q3_K GLM down shape selects the fixed-block kernel and that the ordinary kernel remains available for A/B.", ++ "", ++ "| Mode | Pipeline | Test |", ++ "| --- | --- | --- |", ++ f"| default | `{default_pipeline}` | {'passed' if default_passed else 'failed'} |", ++ f"| disabled | `{disabled_pipeline}` | {'passed' if disabled_passed else 'failed'} |", ++ "", ++ f"- Default artifact: `{default_log.name}`.", ++ f"- Disabled artifact: `{disabled_log.name}`.", ++] ++ ++failures = [] ++if not default_passed: ++ failures.append("default dispatch correctness did not pass") ++if not disabled_passed: ++ failures.append("disabled dispatch correctness did not pass") ++if not default_fixed: ++ failures.append(f"default dispatch did not select fixed-block GLM down kernel: {default_pipeline}") ++if not disabled_ordinary: ++ failures.append(f"disabled dispatch did not select ordinary q3_K kernel: {disabled_pipeline}") ++ ++if failures: ++ lines.extend(["", "### Dispatch Gate Failures", ""]) ++ lines.extend(f"- {failure}" for failure in failures) ++else: ++ lines.extend([ ++ "", ++ "### Dispatch Gate Decision", ++ "", ++ "- The fixed-block q3_K GLM down kernel is the default dispatch path.", ++ "- The ordinary q3_K kernel remains available when `LLAMA_GLM_DSA_Q3_K_MUL_MV_ID_GLM_DOWN=0`.", ++ ]) ++ ++section = "\n".join(lines) + "\n" ++summary.write_text(summary.read_text() + section) ++print(section) ++ ++if require_gates and failures: ++ print("error: fixed-block dispatch gate failed: " + "; ".join(failures), file=sys.stderr) ++ sys.exit(2) ++PY ++fi ++ ++echo "wrote ${summary}" +diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp +index e3553c5e..18e06172 100644 +--- a/tests/test-backend-ops.cpp ++++ b/tests/test-backend-ops.cpp +@@ -51,6 +51,87 @@ + # define N_THREADS std::thread::hardware_concurrency() + #endif + ++static void set_test_env(const char * name, const char * value) { ++#if defined(_WIN32) ++ _putenv_s(name, value); ++#else ++ setenv(name, value, true); ++#endif ++} ++ ++static void unset_test_env(const char * name) { ++#if defined(_WIN32) ++ _putenv_s(name, ""); ++#else ++ unsetenv(name); ++#endif ++} ++ ++struct scoped_test_env { ++ std::string name; ++ bool had_value; ++ std::string old_value; ++ ++ scoped_test_env(const char * name, const char * value) ++ : name(name), ++ had_value(getenv(name) != nullptr), ++ old_value(had_value ? getenv(name) : "") { ++ set_test_env(name, value); ++ } ++ ++ ~scoped_test_env() { ++ if (had_value) { ++ set_test_env(name.c_str(), old_value.c_str()); ++ } else { ++ unset_test_env(name.c_str()); ++ } ++ } ++}; ++ ++static int64_t test_perf_min_us() { ++ const char * value = getenv("GGML_TEST_BACKEND_OPS_PERF_MIN_US"); ++ if (value == nullptr || value[0] == '\0') { ++ return 1000*1000; ++ } ++ ++ const int64_t parsed = strtoll(value, nullptr, 10); ++ return parsed > 0 ? parsed : 1000*1000; ++} ++ ++static int test_perf_samples() { ++ const char * value = getenv("GGML_TEST_BACKEND_OPS_PERF_SAMPLES"); ++ if (value == nullptr || value[0] == '\0') { ++ value = getenv("SKIPPY_GLM_DSA_PERF_SAMPLES"); ++ } ++ if (value == nullptr || value[0] == '\0') { ++ return 1; ++ } ++ ++ const int parsed = atoi(value); ++ return parsed > 0 ? parsed : 1; ++} ++ ++static int test_perf_warmup_samples() { ++ const char * value = getenv("GGML_TEST_BACKEND_OPS_PERF_WARMUP_SAMPLES"); ++ if (value == nullptr || value[0] == '\0') { ++ value = getenv("SKIPPY_GLM_DSA_PERF_WARMUP_SAMPLES"); ++ } ++ if (value == nullptr || value[0] == '\0') { ++ return 0; ++ } ++ ++ const int parsed = atoi(value); ++ return parsed > 0 ? parsed : 0; ++} ++ ++static bool test_perf_sync_backend() { ++ const char * value = getenv("GGML_TEST_BACKEND_OPS_PERF_SYNC_BACKEND"); ++ if (value == nullptr || value[0] == '\0') { ++ value = getenv("SKIPPY_GLM_DSA_PERF_SYNC_BACKEND"); ++ } ++ return value != nullptr && atoi(value) != 0; ++} ++ + static void init_tensor_uniform(ggml_tensor * tensor, float min = -1.0f, float max = 1.0f) { + size_t nels = ggml_nelements(tensor); + std::vector data(nels); +@@ -1215,7 +1296,13 @@ struct test_case { + return 0; + } + ++ virtual std::string perf_metadata(ggml_tensor * t) { ++ GGML_UNUSED(t); ++ return ""; ++ } ++ + virtual bool run_whole_graph() { return false; } ++ virtual bool perf_runs_whole_graph() { return false; } + virtual std::vector fusion_test_nodes() { return {}; } + virtual bool use_weight_context() { return false; } + +@@ -1313,14 +1400,14 @@ struct test_case { + } + } + +- test_status_t eval(ggml_backend_t backend1, +- ggml_backend_t backend2, +- const char * op_names_filter, +- printer * output_printer) { ++ virtual test_status_t eval(ggml_backend_t backend1, ++ ggml_backend_t backend2, ++ const char * op_names_filter, ++ printer * output_printer) { + mode = MODE_TEST; + + ggml_init_params params = { +- /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead(), ++ /* .mem_size = */ ggml_tensor_overhead()*512 + ggml_graph_overhead(), + /* .mem_base = */ NULL, + /* .no_alloc = */ true, + }; +@@ -1501,13 +1588,13 @@ struct test_case { + return test_passed ? test_status_t::OK : test_status_t::FAIL; + } + +- bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) { ++ virtual bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) { + mode = MODE_PERF; + + static const size_t graph_nodes = 8192; + + ggml_init_params params = { +- /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead_custom(graph_nodes, false), ++ /* .mem_size = */ ggml_tensor_overhead()*512 + ggml_graph_overhead_custom(graph_nodes, false), + /* .mem_base = */ NULL, + /* .no_alloc = */ true, + }; +@@ -1572,8 +1659,11 @@ struct test_case { + + // determine number of runs + int n_runs; ++ const bool perf_whole_graph = perf_runs_whole_graph(); + bool is_cpu = ggml_backend_dev_type(ggml_backend_get_device(backend)) == GGML_BACKEND_DEVICE_TYPE_CPU; +- if (op_flops(out) > 0) { ++ if (perf_whole_graph) { ++ n_runs = 1; ++ } else if (op_flops(out) > 0) { + // based on flops + const uint64_t GFLOP = 1000 * 1000 * 1000; + const uint64_t target_flops_cpu = 8ULL * GFLOP; +@@ -1590,8 +1680,10 @@ struct test_case { + } + + // duplicate the op +- for (int i = 1; i < n_runs; i++) { +- ggml_graph_add_node(gf, out); ++ if (!perf_whole_graph) { ++ for (int i = 1; i < n_runs; i++) { ++ ggml_graph_add_node(gf, out); ++ } + } + + // calculate memory +@@ -1617,19 +1709,87 @@ struct test_case { + int64_t total_time_us = 0; + int64_t total_mem = 0; + int total_runs = 0; +- do { +- int64_t start_time = ggml_time_us(); +- ggml_status status = ggml_backend_graph_compute(backend, gf); +- if (status != GGML_STATUS_SUCCESS) { +- fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); +- return false; ++ const int64_t perf_min_us = test_perf_min_us(); ++ const int perf_samples = test_perf_samples(); ++ const int warmup_samples = test_perf_warmup_samples(); ++ const int total_samples = warmup_samples + perf_samples; ++ std::vector sample_avg_us; ++ sample_avg_us.reserve(perf_samples); ++ for (int sample = 0; sample < total_samples; ++sample) { ++ int64_t sample_time_us = 0; ++ int64_t sample_mem = 0; ++ int sample_runs = 0; ++ do { ++ int64_t start_time = ggml_time_us(); ++ ggml_status status = ggml_backend_graph_compute(backend, gf); ++ if (status != GGML_STATUS_SUCCESS) { ++ fprintf(stderr, "%s: ggml_backend_graph_compute failed. status=%s \n", __func__, ggml_status_to_string(status)); ++ return false; ++ } ++ if (test_perf_sync_backend()) { ++ ggml_backend_synchronize(backend); ++ } ++ int64_t end_time = ggml_time_us(); ++ ++ sample_time_us += end_time - start_time; ++ sample_mem += mem; ++ sample_runs += n_runs; ++ } while (sample_time_us < perf_min_us); ++ ++ const double avg_us = (double) sample_time_us / sample_runs; ++ if (sample < warmup_samples) { ++ if (warmup_samples > 0) { ++ fprintf(stderr, ++ "ggml_test: perf_warmup op=%s params=%s sample=%d/%d runs=%d avg_us=%.2f\n", ++ current_op_name.c_str(), ++ vars().c_str(), ++ sample + 1, ++ warmup_samples, ++ sample_runs, ++ avg_us); ++ } ++ continue; + } +- int64_t end_time = ggml_time_us(); + +- total_time_us += end_time - start_time; +- total_mem += mem; +- total_runs += n_runs; +- } while (total_time_us < 1000*1000); // run for at least 1 second ++ total_time_us += sample_time_us; ++ total_mem += sample_mem; ++ total_runs += sample_runs; ++ sample_avg_us.push_back(avg_us); ++ if (perf_samples > 1 || warmup_samples > 0) { ++ const int measured_sample = sample - warmup_samples; ++ fprintf(stderr, ++ "ggml_test: perf_sample op=%s params=%s sample=%d/%d runs=%d avg_us=%.2f\n", ++ current_op_name.c_str(), ++ vars().c_str(), ++ measured_sample + 1, ++ perf_samples, ++ sample_runs, ++ avg_us); ++ } ++ } ++ if (perf_samples > 1 || warmup_samples > 0) { ++ std::vector sorted = sample_avg_us; ++ std::sort(sorted.begin(), sorted.end()); ++ const double min_us = sorted.front(); ++ const double med_us = sorted[sorted.size()/2]; ++ const double max_us = sorted.back(); ++ fprintf(stderr, ++ "ggml_test: perf_summary op=%s params=%s warmup_samples=%d samples=%d min_us=%.2f median_us=%.2f max_us=%.2f\n", ++ current_op_name.c_str(), ++ vars().c_str(), ++ warmup_samples, ++ perf_samples, ++ min_us, ++ med_us, ++ max_us); ++ } ++ const std::string metadata = perf_metadata(out); ++ if (!metadata.empty()) { ++ fprintf(stderr, "ggml_test: perf_metadata op=%s params=%s %s\n", ++ current_op_name.c_str(), ++ vars().c_str(), ++ metadata.c_str()); ++ } + + // Create test result + double avg_time_us = (double) total_time_us / total_runs; +@@ -2248,13 +2408,14 @@ struct test_get_rows : public test_case { + const int be1; // batch size + const int be2; // batch size + const bool v; // view (non-contiguous src1) ++ const bool typed; + + std::string vars() override { +- return VARS_TO_STR7(type, n, m, r, be1, be2, v); ++ return VARS_TO_STR8(type, n, m, r, be1, be2, v, typed); + } + +- test_get_rows(ggml_type type = GGML_TYPE_F32, int n = 10, int m = 5, int r = 3, int be1 = 1, int be2 = 1, bool v = false) +- : type(type), n(n), m(m), r(r), be1(be1), be2(be2), v(v) {} ++ test_get_rows(ggml_type type = GGML_TYPE_F32, int n = 10, int m = 5, int r = 3, int be1 = 1, int be2 = 1, bool v = false, bool typed = false) ++ : type(type), n(n), m(m), r(r), be1(be1), be2(be2), v(v), typed(typed) {} + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * in = ggml_new_tensor_4d(ctx, type, n, m, be1, be2); +@@ -2273,7 +2434,7 @@ struct test_get_rows : public test_case { + // rows is a constant input -> no gradients + } + +- ggml_tensor * out = ggml_get_rows(ctx, in, rows); ++ ggml_tensor * out = typed ? ggml_get_rows_typed(ctx, in, rows) : ggml_get_rows(ctx, in, rows); + ggml_set_name(out, "out"); + + return out; +@@ -4157,6 +4318,158 @@ struct test_gated_delta_net : public test_case { + } + }; + ++// GGML_OP_DSA_SPARSE_MASK ++struct test_dsa_sparse_mask : public test_case { ++ const ggml_type type; ++ ++ const int64_t n_kv; ++ const int64_t n_batch; ++ const int64_t n_stream; ++ const int64_t n_top_k; ++ const int64_t n_top_stream; ++ const bool runtime_top_k_layout; ++ ++ std::string vars() override { ++ return VARS_TO_STR7(type, n_kv, n_batch, n_stream, n_top_k, n_top_stream, runtime_top_k_layout); ++ } ++ ++ test_dsa_sparse_mask(ggml_type type = GGML_TYPE_F16, ++ int64_t n_kv = 33, int64_t n_batch = 2, int64_t n_stream = 1, ++ int64_t n_top_k = 4, int64_t n_top_stream = 1, bool runtime_top_k_layout = false) ++ : type(type), n_kv(n_kv), n_batch(n_batch), n_stream(n_stream), ++ n_top_k(n_top_k), n_top_stream(n_top_stream), runtime_top_k_layout(runtime_top_k_layout) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * kq_mask = ggml_new_tensor_4d(ctx, type, 1, n_kv, n_batch, n_stream); ++ ++ ggml_set_name(kq_mask, "kq_mask"); ++ ++ ggml_tensor * top_k = nullptr; ++ if (runtime_top_k_layout) { ++ ggml_tensor * top_k_runtime = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, n_batch, 1, n_top_stream); ++ ggml_set_name(top_k_runtime, "top_k"); ++ top_k = ggml_view_4d( ++ ctx, ++ top_k_runtime, ++ n_top_k, n_batch, n_top_stream, 1, ++ top_k_runtime->nb[1], top_k_runtime->nb[3], top_k_runtime->nb[2], ++ 0); ++ ggml_set_name(top_k, "top_k_sparse_view"); ++ } else { ++ top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, n_batch, n_top_stream, 1); ++ ggml_set_name(top_k, "top_k"); ++ } ++ ++ ggml_tensor * out = ggml_dsa_sparse_mask(ctx, kq_mask, top_k); ++ ggml_set_name(out, "out"); ++ return out; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "kq_mask") == 0) { ++ init_kq_mask(t); ++ } else if (strcmp(t->name, "top_k") == 0) { ++ init_top_k(t); ++ } else if (!ggml_is_view_op(t->op)) { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++ ++ double err(const float * a, const float * b, size_t n) override { ++ if (n != (size_t) (n_kv * n_batch * n_stream)) { ++ return 1.0; ++ } ++ ++ double err = 0.0; ++ for (int64_t i_stream = 0; i_stream < n_stream; ++i_stream) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ for (int64_t i_kv = 0; i_kv < n_kv; ++i_kv) { ++ const size_t idx = (i_stream*n_batch + i_batch)*n_kv + i_kv; ++ const float expected = expected_value(i_kv, i_batch, i_stream); ++ ++ err = std::max(err, value_err(a[idx], expected)); ++ err = std::max(err, value_err(b[idx], expected)); ++ } ++ } ++ } ++ return err; ++ } ++ ++ double max_err(ggml_backend_t backend) override { ++ (void) backend; ++ return type == GGML_TYPE_F16 ? 1e-3 : 0.0; ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } ++ ++private: ++ void init_kq_mask(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t i_stream = 0; i_stream < n_stream; ++i_stream) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ for (int64_t i_kv = 0; i_kv < n_kv; ++i_kv) { ++ const size_t idx = (i_stream*n_batch + i_batch)*n_kv + i_kv; ++ data[idx] = mask_value(i_kv, i_batch, i_stream); ++ } ++ } ++ } ++ ++ if (type == GGML_TYPE_F16) { ++ std::vector data_f16(data.size()); ++ ggml_fp32_to_fp16_row(data.data(), data_f16.data(), data.size()); ++ ggml_backend_tensor_set(t, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); ++ } else { ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float)); ++ } ++ } ++ ++ void init_top_k(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t i_stream = 0; i_stream < n_top_stream; ++i_stream) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ for (int64_t i_top = 0; i_top < n_top_k; ++i_top) { ++ const size_t idx = (i_stream*n_batch + i_batch)*n_top_k + i_top; ++ data[idx] = top_k_index(i_top, i_batch, i_stream); ++ } ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); ++ } ++ ++ int32_t top_k_index(int64_t i_top, int64_t i_batch, int64_t i_stream) const { ++ return (int32_t) ((i_top*7 + i_batch*3 + i_stream*5) % n_kv); ++ } ++ ++ bool is_selected(int64_t i_kv, int64_t i_batch, int64_t i_stream) const { ++ const int64_t i_top_stream = i_stream % n_top_stream; ++ for (int64_t i_top = 0; i_top < n_top_k; ++i_top) { ++ if (top_k_index(i_top, i_batch, i_top_stream) == i_kv) { ++ return true; ++ } ++ } ++ return false; ++ } ++ ++ float mask_value(int64_t i_kv, int64_t i_batch, int64_t i_stream) const { ++ return 0.01f*(float) (1 + i_kv + 10*i_batch + 100*i_stream); ++ } ++ ++ float expected_value(int64_t i_kv, int64_t i_batch, int64_t i_stream) const { ++ return is_selected(i_kv, i_batch, i_stream) ? mask_value(i_kv, i_batch, i_stream) : -INFINITY; ++ } ++ ++ static double value_err(float actual, float expected) { ++ if (std::isinf(expected)) { ++ return std::isinf(actual) && std::signbit(actual) == std::signbit(expected) ? 0.0 : 1.0; ++ } ++ return actual == expected ? 0.0 : std::abs((double) actual - (double) expected); ++ } ++}; ++ + // GGML_OP_GATED_LINEAR_ATTN + struct test_gla : public test_case { + const ggml_type type; +@@ -6060,120 +6373,4029 @@ struct test_topk_moe : public test_case { + } + }; + +-struct test_mul_mat_vec_fusion : public test_case { +- const ggml_type type; +- const ggml_glu_op glu_op; +- const int64_t m; +- const int64_t n; +- const int64_t k; +- const bool use_id; +- const int n_mats; +- const int n_used; +- const bool b; // broadcast b matrix (only for use_id) +- const bool with_bias; +- const bool with_gate; +- const bool with_lane_scale; +- std::array batch_dims; ++struct test_moe_weighted_sum : public test_case { ++ const int64_t n_embd; ++ const int64_t n_expert_used; ++ const int64_t n_tokens; + +- test_mul_mat_vec_fusion(ggml_type type, ggml_glu_op op, int64_t m, int64_t n, int64_t k, +- bool use_id = false, int n_mats = 1, int n_used = 1, bool b = false, bool with_bias = false, bool with_gate = true, +- bool with_lane_scale = false, std::array batch_dims = {4, 2}) +- : type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias), +- with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims) { +- if (use_id) { +- GGML_ASSERT(n_used <= n_mats); +- } ++ test_moe_weighted_sum(int64_t n_embd = 32, int64_t n_expert_used = 8, int64_t n_tokens = 1) ++ : n_embd(n_embd), n_expert_used(n_expert_used), n_tokens(n_tokens) {} ++ ++ std::string vars() override { ++ return VARS_TO_STR3(n_embd, n_expert_used, n_tokens); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "MOE_WEIGHTED_SUM"; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * experts = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, n_expert_used, n_tokens); ++ ggml_set_name(experts, "experts"); ++ ++ ggml_tensor * weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, n_expert_used, n_tokens); ++ ggml_set_name(weights, "weights"); ++ ++ ggml_tensor * out = ggml_moe_weighted_sum(ctx, experts, weights); ++ ggml_set_name(out, "out"); ++ ++ return out; ++ } ++}; ++ ++struct test_topk_mul_mat_id_minimal : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ const bool direct_top_k; ++ const ggml_type weight_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const int64_t n_tokens; ++ ++ test_topk_mul_mat_id_minimal( ++ bool direct_top_k, ++ ggml_type weight_type = GGML_TYPE_F32, ++ int64_t n_embd = 256, ++ int64_t n_ff = 256, ++ int64_t n_tokens = 1) ++ : direct_top_k(direct_top_k), ++ weight_type(weight_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ n_tokens(n_tokens) { ++ GGML_ASSERT(n_embd % ggml_blck_size(weight_type) == 0); ++ GGML_ASSERT(n_tokens > 0); + } + + std::string vars() override { +- return VARS_TO_STR13(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, batch_dims); ++ return VARS_TO_STR7(direct_top_k, weight_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff); + } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); +- return "MUL_MAT_VEC_FUSION"; ++ if (weight_type == GGML_TYPE_Q2_K && n_tokens > 1) { ++ return "TOPK_MUL_MAT_ID_MINIMAL_ARGSORT_Q2_MULTI"; ++ } ++ if (direct_top_k) { ++ return weight_type == GGML_TYPE_F32 ? ++ "TOPK_MUL_MAT_ID_MINIMAL_DIRECT_F32" : ++ "TOPK_MUL_MAT_ID_MINIMAL_DIRECT_Q2"; ++ } ++ return weight_type == GGML_TYPE_F32 ? ++ "TOPK_MUL_MAT_ID_MINIMAL_ARGSORT_F32" : ++ "TOPK_MUL_MAT_ID_MINIMAL_ARGSORT_Q2"; + } + + bool run_whole_graph() override { return true; } +- bool use_weight_context() override { return use_id && with_lane_scale; } ++ bool perf_runs_whole_graph() override { return true; } + +- ggml_tensor * build_gate(ggml_context * ctx, ggml_tensor * ffn_gate, ggml_tensor * ffn_up) { +- ggml_tensor * out = nullptr; +- if (with_gate) { +- if (glu_op == GGML_GLU_OP_SWIGLU_OAI) { +- constexpr float alpha = 1.702f; +- constexpr float limit = 7.0f; +- out = ggml_swiglu_oai(ctx, ffn_gate, ffn_up, alpha, limit); +- } else { +- out = ggml_glu_split(ctx, ffn_gate, ffn_up, glu_op); +- } ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * scores = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_set_name(scores, "minimal_route_scores"); ++ ++ ggml_tensor * selected = nullptr; ++ if (direct_top_k && getenv("GGML_TEST_TOPK_FULL_WIDTH_VIEW") != nullptr) { ++ selected = ggml_top_k(ctx, scores, n_expert); ++ selected = ggml_view_4d(ctx, selected, ++ n_expert_used, selected->ne[1], selected->ne[2], selected->ne[3], ++ selected->nb[1], selected->nb[2], selected->nb[3], ++ 0); ++ } else { ++ selected = direct_top_k ? ++ ggml_top_k(ctx, scores, n_expert_used) : ++ ggml_argsort_top_k(ctx, scores, n_expert_used); + } ++ ggml_set_name(selected, "minimal_topk_ids"); ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "minimal_cur"); ++ ++ ggml_tensor * experts = ggml_new_tensor_3d(ctx, weight_type, n_embd, n_ff, n_expert); ++ ggml_set_name(experts, n_tokens > 1 ? "ffn_gate_exps.weight" : "minimal_experts.weight"); ++ ++ ggml_tensor * out = ggml_mul_mat_id(ctx, experts, cur, selected); ++ ggml_set_name(out, "minimal_mul_mat_id_out"); + return out; + } + +- ggml_tensor * build_lane_scale_dense(ggml_context * ctx, ggml_tensor * out) { +- ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); +- return ggml_mul(ctx, out, scale); ++ double max_nmse_err() override { ++ return 5e-4; + } + +- ggml_tensor * build_lane_scale_id(ggml_context * ctx, ggml_context * ctx_weights, ggml_tensor * out, ggml_tensor * ids) { +- GGML_ASSERT(ctx_weights); +- ggml_tensor * scale = ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, n_mats); +- ggml_tensor * s = ggml_reshape_3d(ctx, scale, 1, n_mats, 1); +- s = ggml_repeat_4d(ctx, s, 1, n_mats, m, 1); +- s = ggml_get_rows(ctx, s, ids); +- return ggml_mul(ctx, out, s); ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "minimal_route_scores") == 0) { ++ std::vector data(n_expert * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[token*n_expert + i] = -1.0f + 0.0078125f*float(i) + 0.00001f*float(token); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "minimal_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strcmp(t->name, "minimal_experts.weight") == 0 || ++ strcmp(t->name, "ffn_gate_exps.weight") == 0) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } + } ++}; + +- ggml_tensor * build_graph(ggml_context * ctx) override { +- GGML_ASSERT(!use_weight_context()); +- return build_graph(ctx, nullptr); ++struct test_topk_moe_weighted_output_minimal : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ const bool direct_top_k; ++ const bool native_route_weights; ++ const ggml_type weight_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ ++ test_topk_moe_weighted_output_minimal( ++ bool direct_top_k, ++ bool native_route_weights, ++ ggml_type weight_type = GGML_TYPE_F32, ++ int64_t n_embd = 256, ++ int64_t n_ff = 256) ++ : direct_top_k(direct_top_k), ++ native_route_weights(native_route_weights), ++ weight_type(weight_type), ++ n_embd(n_embd), ++ n_ff(n_ff) { ++ GGML_ASSERT(n_embd % ggml_blck_size(weight_type) == 0); + } + +- ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) override { +- if (!use_id) { +- const int channels = batch_dims[0]; +- const int samples = batch_dims[1]; +- std::array ne = { k, m, channels, samples }; +- std::array ne0 = { k, n, channels, samples }; ++ std::string vars() override { ++ return VARS_TO_STR8(direct_top_k, native_route_weights, weight_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff); ++ } + +- ggml_tensor * cur = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); +- ggml_tensor * gate = with_gate ? ggml_new_tensor(ctx, type, 4, ne0.data()) : nullptr; +- ggml_tensor * up = ggml_new_tensor(ctx, type, 4, ne0.data()); ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (direct_top_k && native_route_weights) { ++ return weight_type == GGML_TYPE_F32 ? ++ "TOPK_MOE_WEIGHTED_OUT_MINIMAL_DIRECT_NATIVE_F32" : ++ "TOPK_MOE_WEIGHTED_OUT_MINIMAL_DIRECT_NATIVE_Q2"; ++ } ++ if (direct_top_k) { ++ return weight_type == GGML_TYPE_F32 ? ++ "TOPK_MOE_WEIGHTED_OUT_MINIMAL_DIRECT_F32" : ++ "TOPK_MOE_WEIGHTED_OUT_MINIMAL_DIRECT_Q2"; ++ } ++ return weight_type == GGML_TYPE_F32 ? ++ "TOPK_MOE_WEIGHTED_OUT_MINIMAL_ARGSORT_F32" : ++ "TOPK_MOE_WEIGHTED_OUT_MINIMAL_ARGSORT_Q2"; ++ } + +- auto build_lane_up = [&]() { +- ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur); +- if (with_lane_scale) { +- ffn_up = build_lane_scale_dense(ctx, ffn_up); +- } +- if (with_bias) { +- std::array bias_ne = { ffn_up->ne[0], 1, channels, samples }; +- ggml_tensor * up_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); +- ffn_up = ggml_add(ctx, ffn_up, up_bias); +- } +- return ffn_up; +- }; ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } + +- auto build_lane_gate = [&]() { +- ggml_tensor * ffn_gate = ggml_mul_mat(ctx, gate, cur); +- if (with_lane_scale) { +- ffn_gate = build_lane_scale_dense(ctx, ffn_gate); +- } +- if (with_bias) { +- std::array bias_ne = { ffn_gate->ne[0], 1, channels, samples }; +- ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); +- ffn_gate = ggml_add(ctx, ffn_gate, gate_bias); +- } +- return ffn_gate; +- }; ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * scores = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_set_name(scores, "minimal_moe_route_scores"); + +- ggml_tensor * ffn_up = build_lane_up(); +- ggml_tensor * ffn_gate = with_gate ? build_lane_gate() : nullptr; ++ ggml_tensor * probs = ggml_sigmoid(ctx, scores); ++ ggml_set_name(probs, "minimal_moe_probs"); + +- ggml_tensor * out = with_gate ? build_gate(ctx, ffn_gate, ffn_up) : ffn_up; ++ ggml_tensor * selected = direct_top_k ? ++ ggml_top_k(ctx, scores, n_expert_used) : ++ ggml_argsort_top_k(ctx, scores, n_expert_used); ++ ggml_set_name(selected, "minimal_moe_topk_ids"); + +- std::array bias2_ne = { out->ne[0], 1, channels, samples }; +- ggml_tensor * bias2 = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias2_ne.data()); ++ ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens); ++ ++ ggml_tensor * weights = nullptr; ++ if (native_route_weights) { ++ weights = ggml_moe_route_weights(ctx, probs_3d, selected, true, 6.103515625e-5, 1.0f); ++ ggml_set_name(weights, "minimal_moe_route_weights"); ++ } else { ++ weights = ggml_get_rows(ctx, probs_3d, selected); ++ ggml_set_name(weights, "minimal_moe_weights"); ++ weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); ++ ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); ++ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); ++ weights = ggml_div(ctx, weights, weights_sum); ++ weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); ++ ggml_set_name(weights, "minimal_moe_weights_norm"); ++ } ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "minimal_moe_cur"); ++ ++ ggml_tensor * experts = ggml_new_tensor_3d(ctx, weight_type, n_embd, n_ff, n_expert); ++ ggml_set_name(experts, "minimal_moe_experts.weight"); ++ ++ ggml_tensor * expert_out = ggml_mul_mat_id(ctx, experts, cur, selected); ++ ggml_set_name(expert_out, "minimal_moe_expert_out"); ++ ++ ggml_tensor * out = ggml_moe_weighted_sum(ctx, expert_out, weights); ++ ggml_set_name(out, "minimal_moe_weighted_out"); ++ return out; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "minimal_moe_route_scores") == 0) { ++ std::vector data(n_expert * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[token*n_expert + i] = -1.0f + 0.0078125f*float(i) + 0.00001f*float(token); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "minimal_moe_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strcmp(t->name, "minimal_moe_experts.weight") == 0) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_topk_moe_glm_route : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ const int64_t n_tokens; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * weights = nullptr; ++ ++ explicit test_topk_moe_glm_route(int64_t n_tokens = 1) : n_tokens(n_tokens) {} ++ ++ std::string vars() override { ++ return VARS_TO_STR3(n_expert, n_expert_used, n_tokens); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "TOPK_MOE_GLM_ROUTE"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_set_name(logits, "ffn_moe_logits"); ++ ++ ggml_tensor * probs = ggml_sigmoid(ctx, logits); ++ ggml_set_name(probs, "ffn_moe_probs"); ++ ++ ggml_tensor * exp_probs_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_expert); ++ ggml_set_name(exp_probs_b, "exp_probs_b"); ++ ++ ggml_tensor * selection_probs = ggml_add(ctx, probs, exp_probs_b); ++ ggml_set_name(selection_probs, "ffn_moe_probs_biased"); ++ ++ selected_experts = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ ++ weights = ggml_get_rows(ctx, ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), selected_experts); ++ ggml_set_name(weights, "ffn_moe_weights"); ++ ++ weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); ++ ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum"); ++ ++ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum_clamped"); ++ ++ weights = ggml_div(ctx, weights, weights_sum); ++ ggml_set_name(weights, "ffn_moe_weights_norm"); ++ ++ weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); ++ weights = ggml_scale(ctx, weights, 2.0f); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ ++ return weights; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ return { selected_experts, weights }; ++ } ++ ++ double max_nmse_err() override { ++ return 1e-6; ++ } ++ ++ double err(const float * a, const float * b, size_t n) override { ++ return nmse(a, b, n); ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_logits") == 0) { ++ std::vector data(n_expert * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[token*n_expert + i] = -3.0f + 0.03125f*float(i) + 0.00001f*float(token); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "exp_probs_b") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = 0.0001f*float((i*37) % 19); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_topk_moe_glm_route_component : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ const int component; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * weights = nullptr; ++ ++ test_topk_moe_glm_route_component(int component = 7) : component(component) { ++ GGML_ASSERT(component >= 0 && component <= 7); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR4(component, n_expert, n_expert_used, n_tokens); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "TOPK_MOE_GLM_ROUTE_COMPONENT"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_set_name(logits, "ffn_moe_logits"); ++ ++ ggml_tensor * probs = ggml_sigmoid(ctx, logits); ++ ggml_set_name(probs, "ffn_moe_probs"); ++ if (component == 0) { ++ return probs; ++ } ++ ++ ggml_tensor * exp_probs_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_expert); ++ ggml_set_name(exp_probs_b, "exp_probs_b"); ++ ++ ggml_tensor * selection_probs = ggml_add(ctx, probs, exp_probs_b); ++ ggml_set_name(selection_probs, "ffn_moe_probs_biased"); ++ if (component == 1) { ++ return selection_probs; ++ } ++ ++ selected_experts = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ if (component == 2) { ++ ggml_tensor * selected_experts_f32 = ggml_cast(ctx, selected_experts, GGML_TYPE_F32); ++ ggml_set_name(selected_experts_f32, "ffn_moe_topk_f32"); ++ return selected_experts_f32; ++ } ++ ++ weights = ggml_get_rows(ctx, ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), selected_experts); ++ ggml_set_name(weights, "ffn_moe_weights"); ++ if (component == 3) { ++ return weights; ++ } ++ ++ weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); ++ ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum"); ++ if (component == 4) { ++ return weights_sum; ++ } ++ ++ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum_clamped"); ++ if (component == 5) { ++ return weights_sum; ++ } ++ ++ weights = ggml_div(ctx, weights, weights_sum); ++ ggml_set_name(weights, "ffn_moe_weights_norm"); ++ if (component == 6) { ++ return weights; ++ } ++ ++ weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); ++ weights = ggml_scale(ctx, weights, 2.0f); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ ++ return weights; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ std::vector nodes; ++ if (selected_experts != nullptr) { ++ nodes.push_back(selected_experts); ++ } ++ if (weights != nullptr) { ++ nodes.push_back(weights); ++ } ++ return nodes; ++ } ++ ++ double max_nmse_err() override { ++ return 1e-6; ++ } ++ ++ double err(const float * a, const float * b, size_t n) override { ++ return nmse(a, b, n); ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_logits") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = -3.0f + 0.03125f*float(i); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "exp_probs_b") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = 0.0001f*float((i*37) % 19); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_topk_moe_glm_consumer : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ const ggml_type gate_type; ++ const ggml_type up_type; ++ const ggml_type down_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const bool weighted_down; ++ const bool fused_glu; ++ const bool sanity_probe; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * ffn_gate = nullptr; ++ ggml_tensor * ffn_up = nullptr; ++ ggml_tensor * ffn_down = nullptr; ++ ggml_tensor * out = nullptr; ++ ++ test_topk_moe_glm_consumer( ++ ggml_type gate_type = GGML_TYPE_F32, ++ ggml_type up_type = GGML_TYPE_F32, ++ ggml_type down_type = GGML_TYPE_F32, ++ int64_t n_embd = 64, ++ int64_t n_ff = 128, ++ bool weighted_down = false, ++ bool fused_glu = false, ++ bool sanity_probe = false) ++ : gate_type(gate_type), ++ up_type(up_type), ++ down_type(down_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ weighted_down(weighted_down), ++ fused_glu(fused_glu), ++ sanity_probe(sanity_probe) { ++ GGML_ASSERT(n_embd % ggml_blck_size(gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(down_type) == 0); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR10(gate_type, up_type, down_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff, weighted_down, fused_glu); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return sanity_probe ? "TOPK_MOE_GLM_CONSUMER_SANITY" : "TOPK_MOE_GLM_CONSUMER"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * build_glm_route(ggml_context * ctx) { ++ ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_set_name(logits, "ffn_moe_logits"); ++ ++ ggml_tensor * probs = ggml_sigmoid(ctx, logits); ++ ggml_set_name(probs, "ffn_moe_probs"); ++ ++ ggml_tensor * exp_probs_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_expert); ++ ggml_set_name(exp_probs_b, "exp_probs_b"); ++ ++ ggml_tensor * selection_probs = ggml_add(ctx, probs, exp_probs_b); ++ ggml_set_name(selection_probs, "ffn_moe_probs_biased"); ++ ++ selected_experts = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ ++ weights = ggml_get_rows(ctx, ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), selected_experts); ++ ggml_set_name(weights, "ffn_moe_weights"); ++ ++ weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); ++ ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum"); ++ ++ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum_clamped"); ++ ++ weights = ggml_div(ctx, weights, weights_sum); ++ ggml_set_name(weights, "ffn_moe_weights_norm"); ++ ++ weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); ++ weights = ggml_scale(ctx, weights, 2.0f); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ ++ return weights; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ build_glm_route(ctx); ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_moe_cur"); ++ ++ ggml_tensor * route_anchor = ggml_sum(ctx, weights); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_sum"); ++ ++ route_anchor = ggml_scale(ctx, route_anchor, 0.0f); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_zero"); ++ ++ route_anchor = ggml_repeat(ctx, route_anchor, cur); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_repeat"); ++ ++ cur = ggml_add(ctx, cur, route_anchor); ++ ggml_set_name(cur, "ffn_moe_cur_route_anchored"); ++ ++ ggml_tensor * gates = ggml_new_tensor_3d(ctx, gate_type, n_embd, n_ff, n_expert); ++ ggml_set_name(gates, "ffn_gate_exps.weight"); ++ ++ ggml_tensor * ups = ggml_new_tensor_3d(ctx, up_type, n_embd, n_ff, n_expert); ++ ggml_set_name(ups, "ffn_up_exps.weight"); ++ ++ ggml_tensor * downs = ggml_new_tensor_3d(ctx, down_type, n_ff, n_embd, n_expert); ++ ggml_set_name(downs, "ffn_down_exps.weight"); ++ ++ ffn_gate = ggml_mul_mat_id(ctx, gates, cur, selected_experts); ++ ggml_set_name(ffn_gate, "ffn_moe_gate"); ++ ++ ffn_up = ggml_mul_mat_id(ctx, ups, cur, selected_experts); ++ ggml_set_name(ffn_up, "ffn_moe_up"); ++ ++ ggml_tensor * ffn_act = nullptr; ++ if (fused_glu) { ++ ffn_act = ggml_swiglu_split(ctx, ffn_gate, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_swiglu"); ++ } else { ++ ffn_act = ggml_silu(ctx, ffn_gate); ++ ggml_set_name(ffn_act, "ffn_moe_gate_silu"); ++ ++ ffn_act = ggml_mul(ctx, ffn_act, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_gate_up"); ++ } ++ ++ if (weighted_down) { ++ ffn_act = ggml_mul(ctx, ffn_act, weights); ++ ggml_set_name(ffn_act, "ffn_moe_down_weighted_input"); ++ } ++ ++ ffn_down = ggml_mul_mat_id(ctx, downs, ffn_act, selected_experts); ++ ggml_set_name(ffn_down, "ffn_moe_down"); ++ ++ out = ggml_moe_weighted_sum(ctx, ffn_down, weights); ++ if (weighted_down) { ++ out->op_params[0] = 1; ++ } ++ ggml_set_name(out, "ffn_moe_out"); ++ ++ return out; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ return { selected_experts, weights, ffn_gate, ffn_up, ffn_down, out }; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_logits") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = -3.0f + 0.03125f*float(i); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "exp_probs_b") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = 0.0001f*float((i*37) % 19); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "ffn_moe_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_exps.weight") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++static bool glm_moe_q3_down_w0_default_direct_reduce(ggml_type down_type, int64_t n_ff, bool already_weighted) { ++ const char * disabled_w0 = getenv("GGML_METAL_DISABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8_W0"); ++ const char * disabled_w1 = getenv("GGML_METAL_DISABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8_W1"); ++ if (down_type != GGML_TYPE_Q3_K || n_ff != 2048) { ++ return false; ++ } ++ if (already_weighted) { ++ return !(disabled_w1 && atoi(disabled_w1) != 0); ++ } ++ return !(disabled_w0 && atoi(disabled_w0) != 0); ++} ++ ++static bool glm_moe_q2_down_default_direct_reduce(ggml_type down_type, int64_t n_ff, bool already_weighted) { ++ GGML_UNUSED(already_weighted); ++ const char * disabled = getenv("GGML_METAL_DISABLE_Q2_DOWN_SLOT_PARALLEL_REDUCE"); ++ return down_type == GGML_TYPE_Q2_K && n_ff == 2048 && ++ !(disabled && atoi(disabled) != 0); ++} ++ ++static bool glm_moe_q2_gate_up_pair_sg_default(ggml_type gate_type, ggml_type up_type, int64_t n_embd, int64_t n_ff) { ++ const char * disabled = getenv("GGML_METAL_DISABLE_Q2_GATE_UP_SWIGLU_PAIR_SG"); ++ return gate_type == GGML_TYPE_Q2_K && up_type == GGML_TYPE_Q2_K && ++ n_embd == 6144 && n_ff == 2048 && !(disabled && atoi(disabled) != 0); ++} ++ ++static std::string glm_moe_format_bytes(uint64_t bytes) { ++ char buf[64]; ++ snprintf(buf, sizeof(buf), "%.2fMiB", double(bytes) / (1024.0 * 1024.0)); ++ return buf; ++} ++ ++struct test_topk_moe_glm_component : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ const int component; ++ const ggml_type gate_type; ++ const ggml_type up_type; ++ const ggml_type down_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const bool fused_glu; ++ const bool weighted_down; ++ const bool route_anchored = true; ++ const bool native_route; ++ const bool direct_top_k; ++ const bool f16_down_input; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * ffn_gate = nullptr; ++ ggml_tensor * ffn_up = nullptr; ++ ggml_tensor * ffn_act = nullptr; ++ ggml_tensor * ffn_down = nullptr; ++ ggml_tensor * ffn_out = nullptr; ++ ++ test_topk_moe_glm_component( ++ int component, ++ ggml_type gate_type = GGML_TYPE_Q2_K, ++ ggml_type up_type = GGML_TYPE_Q2_K, ++ ggml_type down_type = GGML_TYPE_Q2_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048, ++ bool fused_glu = true, ++ bool weighted_down = false, ++ bool native_route = false, ++ bool direct_top_k = false, ++ bool f16_down_input = false) ++ : component(component), ++ gate_type(gate_type), ++ up_type(up_type), ++ down_type(down_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ fused_glu(fused_glu), ++ weighted_down(weighted_down), ++ native_route(native_route), ++ direct_top_k(direct_top_k), ++ f16_down_input(f16_down_input) { ++ GGML_ASSERT(component >= 0 && component <= 8); ++ GGML_ASSERT(n_embd % ggml_blck_size(gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(down_type) == 0); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR15(component, gate_type, up_type, down_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff, fused_glu, weighted_down, route_anchored, native_route, direct_top_k, f16_down_input); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "TOPK_MOE_GLM_COMPONENT"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ std::string perf_metadata(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ ++ const uint64_t gate_weight_bytes = ++ uint64_t(n_expert_used) * uint64_t(n_ff) * ggml_row_size(gate_type, n_embd); ++ const uint64_t up_weight_bytes = ++ uint64_t(n_expert_used) * uint64_t(n_ff) * ggml_row_size(up_type, n_embd); ++ const uint64_t down_weight_bytes = ++ uint64_t(n_expert_used) * uint64_t(n_embd) * ggml_row_size(down_type, n_ff); ++ const uint64_t gate_activation_bytes = ++ uint64_t(n_expert_used) * uint64_t(n_tokens) * uint64_t(n_ff) * sizeof(float); ++ const uint64_t up_activation_bytes = gate_activation_bytes; ++ const uint64_t swiglu_activation_bytes = gate_activation_bytes; ++ const uint64_t out_bytes = uint64_t(n_tokens) * uint64_t(n_embd) * sizeof(float); ++ const uint64_t route_bytes = ++ uint64_t(n_tokens) * uint64_t(n_expert) * sizeof(float) * 2 + ++ uint64_t(n_tokens) * uint64_t(n_expert_used) * (sizeof(float) + sizeof(int32_t)); ++ const uint64_t selected_weight_bytes = gate_weight_bytes + up_weight_bytes + down_weight_bytes; ++ const uint64_t unfused_intermediate_bytes = ++ gate_activation_bytes + up_activation_bytes + swiglu_activation_bytes + out_bytes; ++ const uint64_t current_fused_intermediate_bytes = swiglu_activation_bytes + out_bytes; ++ const uint64_t current_saved_intermediate_bytes = ++ unfused_intermediate_bytes > current_fused_intermediate_bytes ? ++ unfused_intermediate_bytes - current_fused_intermediate_bytes : 0; ++ ++ char buf[1024]; ++ snprintf( ++ buf, ++ sizeof(buf), ++ "glm_moe_shape=experts:%lld/topk:%lld/tokens:%lld/embd:%lld/ff:%lld " ++ "selected_weight_bytes=%s gate_weight=%s up_weight=%s down_weight=%s " ++ "route_bytes=%s unfused_intermediate=%s current_fused_intermediate=%s " ++ "current_saved_intermediate=%s activation_to_weight_ratio=%.5f", ++ (long long) n_expert, ++ (long long) n_expert_used, ++ (long long) n_tokens, ++ (long long) n_embd, ++ (long long) n_ff, ++ glm_moe_format_bytes(selected_weight_bytes).c_str(), ++ glm_moe_format_bytes(gate_weight_bytes).c_str(), ++ glm_moe_format_bytes(up_weight_bytes).c_str(), ++ glm_moe_format_bytes(down_weight_bytes).c_str(), ++ glm_moe_format_bytes(route_bytes).c_str(), ++ glm_moe_format_bytes(unfused_intermediate_bytes).c_str(), ++ glm_moe_format_bytes(current_fused_intermediate_bytes).c_str(), ++ glm_moe_format_bytes(current_saved_intermediate_bytes).c_str(), ++ selected_weight_bytes == 0 ? 0.0 : double(unfused_intermediate_bytes) / double(selected_weight_bytes)); ++ return buf; ++ } ++ ++ ggml_tensor * build_glm_route(ggml_context * ctx) { ++ ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_set_name(logits, "ffn_moe_logits"); ++ ++ ggml_tensor * probs = ggml_sigmoid(ctx, logits); ++ ggml_set_name(probs, "ffn_moe_probs"); ++ ++ ggml_tensor * exp_probs_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_expert); ++ ggml_set_name(exp_probs_b, "exp_probs_b"); ++ ++ ggml_tensor * selection_probs = ggml_add(ctx, probs, exp_probs_b); ++ ggml_set_name(selection_probs, "ffn_moe_probs_biased"); ++ ++ selected_experts = direct_top_k ? ++ ggml_top_k(ctx, selection_probs, n_expert_used) : ++ ggml_argsort_top_k(ctx, selection_probs, n_expert_used); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ if (direct_top_k && getenv("GGML_TEST_GLM_DSA_DUP_DIRECT_TOPK_IDS") != nullptr) { ++ selected_experts = ggml_dup(ctx, selected_experts); ++ ggml_set_name(selected_experts, "ffn_moe_topk_dup"); ++ } ++ ++ ggml_tensor * probs_3d = ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens); ++ if (native_route) { ++ weights = ggml_moe_route_weights(ctx, probs_3d, selected_experts, true, 6.103515625e-5, 2.0f); ++ ggml_set_name(weights, "ffn_moe_route_weights"); ++ return weights; ++ } ++ ++ weights = ggml_get_rows(ctx, probs_3d, selected_experts); ++ ggml_set_name(weights, "ffn_moe_weights"); ++ weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); ++ ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum"); ++ ++ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum_clamped"); ++ ++ weights = ggml_div(ctx, weights, weights_sum); ++ ggml_set_name(weights, "ffn_moe_weights_norm"); ++ ++ weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); ++ weights = ggml_scale(ctx, weights, 2.0f); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ ++ return weights; ++ } ++ ++ ggml_tensor * build_route_anchored_cur(ggml_context * ctx) { ++ build_glm_route(ctx); ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_moe_cur"); ++ ++ ggml_tensor * route_anchor = ggml_sum(ctx, weights); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_sum"); ++ ++ route_anchor = ggml_scale(ctx, route_anchor, 0.0f); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_zero"); ++ ++ route_anchor = ggml_repeat(ctx, route_anchor, cur); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_repeat"); ++ ++ cur = ggml_add(ctx, cur, route_anchor); ++ ggml_set_name(cur, "ffn_moe_cur_route_anchored"); ++ ++ return cur; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ if (component == 0) { ++ return build_glm_route(ctx); ++ } ++ ++ ggml_tensor * cur = build_route_anchored_cur(ctx); ++ if (component == 1) { ++ return cur; ++ } ++ ++ ggml_tensor * gates = ggml_new_tensor_3d(ctx, gate_type, n_embd, n_ff, n_expert); ++ ggml_set_name(gates, "ffn_gate_exps.weight"); ++ ++ ffn_gate = ggml_mul_mat_id(ctx, gates, cur, selected_experts); ++ ggml_set_name(ffn_gate, "ffn_moe_gate"); ++ if (component == 2) { ++ return ffn_gate; ++ } ++ ++ ggml_tensor * ups = ggml_new_tensor_3d(ctx, up_type, n_embd, n_ff, n_expert); ++ ggml_set_name(ups, "ffn_up_exps.weight"); ++ ++ ffn_up = ggml_mul_mat_id(ctx, ups, cur, selected_experts); ++ ggml_set_name(ffn_up, "ffn_moe_up"); ++ if (component == 3) { ++ ggml_tensor * gate_up_sum = ggml_add(ctx, ffn_gate, ffn_up); ++ ggml_set_name(gate_up_sum, "ffn_moe_gate_up_sum"); ++ return gate_up_sum; ++ } ++ ++ if (fused_glu) { ++ ffn_act = ggml_swiglu_split(ctx, ffn_gate, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_swiglu"); ++ } else { ++ ffn_act = ggml_silu(ctx, ffn_gate); ++ ggml_set_name(ffn_act, "ffn_moe_gate_silu"); ++ ++ ffn_act = ggml_mul(ctx, ffn_act, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_gate_up"); ++ } ++ if (component == 4) { ++ return ffn_act; ++ } ++ ++ if (f16_down_input && !weighted_down) { ++ ffn_act = ggml_cast(ctx, ffn_act, GGML_TYPE_F16); ++ ggml_set_name(ffn_act, "ffn_moe_swiglu_f16"); ++ } ++ ++ if (weighted_down) { ++ ffn_act = ggml_mul(ctx, ffn_act, weights); ++ ggml_set_name(ffn_act, "ffn_moe_down_weighted_input"); ++ if (component == 5) { ++ return ffn_act; ++ } ++ } ++ ++ ggml_tensor * downs = ggml_new_tensor_3d(ctx, down_type, n_ff, n_embd, n_expert); ++ ggml_set_name(downs, "ffn_down_exps.weight"); ++ ++ ffn_down = ggml_mul_mat_id(ctx, downs, ffn_act, selected_experts); ++ ggml_set_name(ffn_down, "ffn_moe_down"); ++ if ((!weighted_down && component == 5) || (weighted_down && component == 6)) { ++ return ffn_down; ++ } ++ ++ ggml_tensor * out = ggml_moe_weighted_sum(ctx, ffn_down, weights); ++ if (weighted_down) { ++ out->op_params[0] = 1; ++ } ++ ggml_set_name(out, "ffn_moe_out"); ++ ffn_out = out; ++ if ((!weighted_down && component == 6) || (weighted_down && component == 7)) { ++ return out; ++ } ++ ++ ggml_tensor * checksum = ggml_sum(ctx, out); ++ ggml_set_name(checksum, "ffn_moe_out_checksum"); ++ return checksum; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ const bool direct_reduce = ++ glm_moe_q2_down_default_direct_reduce(down_type, n_ff, weighted_down) || ++ glm_moe_q3_down_w0_default_direct_reduce(down_type, n_ff, weighted_down) || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_DOWN_WEIGHTED_REDUCE_DIRECT") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SLOT_PARALLEL_REDUCE_R16_W1") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_WEIGHTED_REDUCE_DIRECT") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8_W0") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R12_NB8_W0") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE_R16") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_SPLIT2_REDUCE") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_ATOMIC_ACCUM") != nullptr || ++ getenv("GGML_METAL_ENABLE_GLM_MOE_SWIGLU_Q3_DOWN_FUSION") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_MOE_SWIGLU_Q3_DOWN_FUSION") != nullptr; ++ const bool q2_gate_up_swiglu = ++ fused_glu && ++ gate_type == GGML_TYPE_Q2_K && ++ up_type == GGML_TYPE_Q2_K && ++ (glm_moe_q2_gate_up_pair_sg_default(gate_type, up_type, n_embd, n_ff) || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_GATE_UP_SWIGLU") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q2_GATE_UP_SWIGLU_FUSION") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_Q8_ACT") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_HALF_Y") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_GLM_MOE_ROUTE_GATE_UP_FUSION") != nullptr); ++ const bool swiglu_q3_down_fused = ++ getenv("GGML_METAL_ENABLE_GLM_MOE_SWIGLU_Q3_DOWN_FUSION") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_MOE_SWIGLU_Q3_DOWN_FUSION") != nullptr; ++ ++ std::vector nodes; ++ if (selected_experts != nullptr) { ++ nodes.push_back(selected_experts); ++ } ++ if (weights != nullptr) { ++ nodes.push_back(weights); ++ } ++ if (!q2_gate_up_swiglu) { ++ if (ffn_gate != nullptr) { ++ nodes.push_back(ffn_gate); ++ } ++ if (ffn_up != nullptr) { ++ nodes.push_back(ffn_up); ++ } ++ } ++ if (!swiglu_q3_down_fused && ffn_act != nullptr) { ++ nodes.push_back(ffn_act); ++ } ++ if (direct_reduce && ffn_out != nullptr) { ++ nodes.push_back(ffn_out); ++ return nodes; ++ } ++ if (ffn_down != nullptr) { ++ nodes.push_back(ffn_down); ++ } ++ return nodes; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_logits") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = -3.0f + 0.03125f*float(i); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "exp_probs_b") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = 0.0001f*float((i*37) % 19); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "ffn_moe_topk") == 0) { ++ std::vector data(n_expert_used); ++ for (int64_t i = 0; i < n_expert_used; ++i) { ++ data[i] = int32_t(n_expert - 1 - i); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "ffn_moe_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_exps.weight") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_glm_moe_decode_component : public test_topk_moe_glm_component { ++ explicit test_glm_moe_decode_component(int component, bool native_route = false, bool direct_top_k = false) ++ : test_topk_moe_glm_component( ++ component, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q3_K, ++ 6144, ++ 2048, ++ true, ++ true, ++ native_route, ++ direct_top_k) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (direct_top_k && native_route) { ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_CUR"; ++ case 2: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_GATE"; ++ case 3: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_WEIGHTED_ACT"; ++ case 6: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_DOWN"; ++ case 7: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_OUT"; ++ default: return "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_CHECKSUM"; ++ } ++ } ++ if (direct_top_k) { ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_DIRECT_TOPK_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_DIRECT_TOPK_CUR"; ++ case 2: return "GLM_MOE_DECODE_DIRECT_TOPK_GATE"; ++ case 3: return "GLM_MOE_DECODE_DIRECT_TOPK_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_DIRECT_TOPK_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_DIRECT_TOPK_WEIGHTED_ACT"; ++ case 6: return "GLM_MOE_DECODE_DIRECT_TOPK_DOWN"; ++ case 7: return "GLM_MOE_DECODE_DIRECT_TOPK_OUT"; ++ default: return "GLM_MOE_DECODE_DIRECT_TOPK_CHECKSUM"; ++ } ++ } ++ if (native_route) { ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_NATIVE_ROUTE_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_NATIVE_ROUTE_CUR"; ++ case 2: return "GLM_MOE_DECODE_NATIVE_ROUTE_GATE"; ++ case 3: return "GLM_MOE_DECODE_NATIVE_ROUTE_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_NATIVE_ROUTE_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_NATIVE_ROUTE_WEIGHTED_ACT"; ++ case 6: return "GLM_MOE_DECODE_NATIVE_ROUTE_DOWN"; ++ case 7: return "GLM_MOE_DECODE_NATIVE_ROUTE_OUT"; ++ default: return "GLM_MOE_DECODE_NATIVE_ROUTE_CHECKSUM"; ++ } ++ } ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_CUR"; ++ case 2: return "GLM_MOE_DECODE_GATE"; ++ case 3: return "GLM_MOE_DECODE_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_WEIGHTED_ACT"; ++ case 6: return "GLM_MOE_DECODE_DOWN"; ++ case 7: return "GLM_MOE_DECODE_OUT"; ++ default: return "GLM_MOE_DECODE_CHECKSUM"; ++ } ++ } ++}; ++ ++struct test_glm_moe_decode_post_weight_component : public test_topk_moe_glm_component { ++ const bool use_native_route; ++ const bool use_f16_down_input; ++ const bool use_weighted_down; ++ const ggml_type use_down_type; ++ ++ explicit test_glm_moe_decode_post_weight_component( ++ int component, ++ bool native_route = false, ++ bool f16_down_input = false, ++ ggml_type down_type = GGML_TYPE_Q3_K, ++ bool weighted_down = false) ++ : test_topk_moe_glm_component( ++ component, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q2_K, ++ down_type, ++ 6144, ++ 2048, ++ true, ++ weighted_down, ++ native_route, ++ false, ++ f16_down_input), ++ use_native_route(native_route), ++ use_f16_down_input(f16_down_input), ++ use_weighted_down(weighted_down), ++ use_down_type(down_type) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (use_weighted_down) { ++ if (use_native_route && use_down_type == GGML_TYPE_Q2_K) { ++ switch (component) { ++ case 6: return "GLM_MOE_DECODE_PREWEIGHT_Q2_NATIVE_ROUTE_DOWN"; ++ case 7: return "GLM_MOE_DECODE_PREWEIGHT_Q2_NATIVE_ROUTE_OUT"; ++ default: return "GLM_MOE_DECODE_PREWEIGHT_Q2_NATIVE_ROUTE_CHECKSUM"; ++ } ++ } ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_CUR"; ++ case 2: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_GATE"; ++ case 3: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_ACT"; ++ case 6: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_DOWN"; ++ case 7: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_OUT"; ++ default: return "GLM_MOE_DECODE_POST_WEIGHT_PREWEIGHT_CHECKSUM"; ++ } ++ } ++ if (use_f16_down_input) { ++ if (use_down_type == GGML_TYPE_Q2_K) { ++ if (use_native_route) { ++ switch (component) { ++ case 5: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_F16_ACT_DOWN"; ++ case 6: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_F16_ACT_OUT"; ++ default: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_F16_ACT_CHECKSUM"; ++ } ++ } ++ switch (component) { ++ case 5: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_F16_ACT_DOWN"; ++ case 6: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_F16_ACT_OUT"; ++ default: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_F16_ACT_CHECKSUM"; ++ } ++ } ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_POST_WEIGHT_F16_ACT_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_POST_WEIGHT_F16_ACT_CUR"; ++ case 2: return "GLM_MOE_DECODE_POST_WEIGHT_F16_ACT_GATE"; ++ case 3: return "GLM_MOE_DECODE_POST_WEIGHT_F16_ACT_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_POST_WEIGHT_F16_ACT_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_POST_WEIGHT_F16_ACT_DOWN"; ++ case 6: return "GLM_MOE_DECODE_POST_WEIGHT_F16_ACT_OUT"; ++ default: return "GLM_MOE_DECODE_POST_WEIGHT_F16_ACT_CHECKSUM"; ++ } ++ } ++ if (use_native_route) { ++ if (use_down_type == GGML_TYPE_Q2_K) { ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_CUR"; ++ case 2: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_GATE"; ++ case 3: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_DOWN"; ++ case 6: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_OUT"; ++ default: return "GLM_MOE_DECODE_POST_WEIGHT_Q2_NATIVE_ROUTE_CHECKSUM"; ++ } ++ } ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_POST_WEIGHT_NATIVE_ROUTE_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_POST_WEIGHT_NATIVE_ROUTE_CUR"; ++ case 2: return "GLM_MOE_DECODE_POST_WEIGHT_NATIVE_ROUTE_GATE"; ++ case 3: return "GLM_MOE_DECODE_POST_WEIGHT_NATIVE_ROUTE_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_POST_WEIGHT_NATIVE_ROUTE_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_POST_WEIGHT_NATIVE_ROUTE_DOWN"; ++ case 6: return "GLM_MOE_DECODE_POST_WEIGHT_NATIVE_ROUTE_OUT"; ++ default: return "GLM_MOE_DECODE_POST_WEIGHT_NATIVE_ROUTE_CHECKSUM"; ++ } ++ } ++ switch (component) { ++ case 0: return "GLM_MOE_DECODE_POST_WEIGHT_ROUTE"; ++ case 1: return "GLM_MOE_DECODE_POST_WEIGHT_CUR"; ++ case 2: return "GLM_MOE_DECODE_POST_WEIGHT_GATE"; ++ case 3: return "GLM_MOE_DECODE_POST_WEIGHT_GATE_UP"; ++ case 4: return "GLM_MOE_DECODE_POST_WEIGHT_SWIGLU"; ++ case 5: return "GLM_MOE_DECODE_POST_WEIGHT_DOWN"; ++ case 6: return "GLM_MOE_DECODE_POST_WEIGHT_OUT"; ++ default: return "GLM_MOE_DECODE_POST_WEIGHT_CHECKSUM"; ++ } ++ } ++ ++ double max_nmse_err() override { ++ return use_f16_down_input ? 2e-3 : test_topk_moe_glm_component::max_nmse_err(); ++ } ++}; ++ ++struct test_glm_moe_decode_direct_topk_final_output : public test_glm_moe_decode_component { ++ const bool use_native_route; ++ ++ explicit test_glm_moe_decode_direct_topk_final_output(bool use_native_route) ++ : test_glm_moe_decode_component(7, use_native_route, true), ++ use_native_route(use_native_route) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return use_native_route ? ++ "GLM_MOE_DECODE_DIRECT_TOPK_NATIVE_ROUTE_FINAL_OUT" : ++ "GLM_MOE_DECODE_DIRECT_TOPK_FINAL_OUT"; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ return {}; ++ } ++}; ++ ++struct test_glm_moe_routed_q2_final_output : public test_topk_moe_glm_component { ++ const bool use_native_route; ++ const bool use_direct_top_k; ++ ++ explicit test_glm_moe_routed_q2_final_output(bool native_route = false, bool direct_top_k = false) ++ : test_topk_moe_glm_component( ++ 6, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q2_K, ++ 6144, ++ 2048, ++ true, ++ false, ++ native_route, ++ direct_top_k), ++ use_native_route(native_route), ++ use_direct_top_k(direct_top_k) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (use_direct_top_k && use_native_route) { ++ return "GLM_MOE_ROUTED_MOTIF_Q2_DIRECT_TOPK_NATIVE_ROUTE_GLM"; ++ } ++ if (use_direct_top_k) { ++ return "GLM_MOE_ROUTED_MOTIF_Q2_DIRECT_TOPK_GLM"; ++ } ++ if (use_native_route) { ++ return "GLM_MOE_ROUTED_MOTIF_Q2_NATIVE_ROUTE_GLM"; ++ } ++ return "GLM_MOE_ROUTED_MOTIF_Q2_GLM"; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ return { selected_experts, weights, ffn_out }; ++ } ++}; ++ ++struct test_glm_moe_q2_weight_roofline : public test_glm_moe_routed_q2_final_output { ++ const bool bypass_route; ++ const char * variant; ++ ++ test_glm_moe_q2_weight_roofline(bool bypass_route, const char * variant) ++ : test_glm_moe_routed_q2_final_output(false, false), ++ bypass_route(bypass_route), ++ variant(variant) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return std::string("GLM_MOE_Q2_WEIGHT_ROOFLINE_") + variant; ++ } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env motif_reference_env("GGML_METAL_ENABLE_GLM_MOE_DECODE_MOTIF_REFERENCE", "1"); ++ scoped_test_env chunks_env("GGML_METAL_EXPERIMENTAL_GLM_MOE_Q2_WEIGHT_ROOFLINE_CHUNKS", "16"); ++ scoped_test_env bypass_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_MOE_Q2_WEIGHT_ROOFLINE_BYPASS_ROUTE", ++ bypass_route ? "1" : "0"); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++}; ++ ++struct test_glm_moe_selected_chain : public test_case { ++ static constexpr int64_t n_expert = 32; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ static constexpr int64_t n_embd = 6144; ++ static constexpr int64_t n_ff = 2048; ++ ++ const int n_repeats; ++ const bool byte_scan; ++ const int storage_block_bytes; ++ const char * variant; ++ ++ struct route_result { ++ ggml_tensor * ids; ++ ggml_tensor * weights; ++ }; ++ ++ test_glm_moe_selected_chain( ++ int n_repeats, ++ bool byte_scan, ++ const char * variant, ++ int storage_block_bytes = int(ggml_type_size(GGML_TYPE_Q2_K))) ++ : n_repeats(n_repeats), ++ byte_scan(byte_scan), ++ storage_block_bytes(storage_block_bytes), ++ variant(variant) { ++ GGML_ASSERT(n_repeats > 0); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR7(n_repeats, n_expert, n_expert_used, n_tokens, n_embd, n_ff, byte_scan) + ++ ",storage_block_bytes=" + std::to_string(storage_block_bytes) + ",variant=" + variant; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return std::string("GLM_MOE_SELECTED_CHAIN_") + variant; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ route_result build_route(ggml_context * ctx) { ++ ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_set_name(logits, "ffn_moe_logits"); ++ ++ ggml_tensor * probs = ggml_sigmoid(ctx, logits); ++ ggml_set_name(probs, "ffn_moe_probs"); ++ ++ ggml_tensor * bias = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_expert); ++ ggml_set_name(bias, "exp_probs_b"); ++ ++ ggml_tensor * selection_probs = ggml_add(ctx, probs, bias); ++ ggml_set_name(selection_probs, "ffn_moe_probs_biased"); ++ ++ ggml_tensor * ids = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); ++ ggml_set_name(ids, "ffn_moe_topk"); ++ ++ ggml_tensor * weights = ggml_get_rows( ++ ctx, ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), ids); ++ ggml_set_name(weights, "ffn_moe_weights"); ++ ++ weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); ++ ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum"); ++ ++ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum_clamped"); ++ ++ weights = ggml_div(ctx, weights, weights_sum); ++ ggml_set_name(weights, "ffn_moe_weights_norm"); ++ weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); ++ weights = ggml_scale(ctx, weights, 2.0f); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ return { ids, weights }; ++ } ++ ++ ggml_tensor * build_selected_ffn(ggml_context * ctx, ggml_tensor * cur) { ++ const route_result route = build_route(ctx); ++ ggml_tensor * ids = route.ids; ++ ggml_tensor * weights = route.weights; ++ ++ ggml_tensor * route_anchor = ggml_sum(ctx, weights); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_sum"); ++ route_anchor = ggml_scale(ctx, route_anchor, 0.0f); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_zero"); ++ route_anchor = ggml_repeat(ctx, route_anchor, cur); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_repeat"); ++ cur = ggml_add(ctx, cur, route_anchor); ++ ggml_set_name(cur, "ffn_moe_cur_route_anchored"); ++ ++ ggml_tensor * gates = ggml_new_tensor_3d(ctx, GGML_TYPE_Q2_K, n_embd, n_ff, n_expert); ++ ggml_set_name(gates, "ffn_gate_exps.weight"); ++ ggml_tensor * ups = ggml_new_tensor_3d(ctx, GGML_TYPE_Q2_K, n_embd, n_ff, n_expert); ++ ggml_set_name(ups, "ffn_up_exps.weight"); ++ ggml_tensor * downs = ggml_new_tensor_3d(ctx, GGML_TYPE_Q3_K, n_ff, n_embd, n_expert); ++ ggml_set_name(downs, "ffn_down_exps.weight"); ++ ++ ggml_tensor * gate = ggml_mul_mat_id(ctx, gates, cur, ids); ++ ggml_set_name(gate, "ffn_moe_gate"); ++ ggml_tensor * up = ggml_mul_mat_id(ctx, ups, cur, ids); ++ ggml_set_name(up, "ffn_moe_up"); ++ ggml_tensor * act = ggml_swiglu_split(ctx, gate, up); ++ ggml_set_name(act, "ffn_moe_swiglu"); ++ ggml_tensor * down = ggml_mul_mat_id(ctx, downs, act, ids); ++ ggml_set_name(down, "ffn_moe_down"); ++ ggml_tensor * out = ggml_moe_weighted_sum(ctx, down, weights); ++ ggml_set_name(out, "ffn_moe_out"); ++ return out; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_chain_cur"); ++ for (int i = 0; i < n_repeats; ++i) { ++ cur = build_selected_ffn(ctx, cur); ++ } ++ return cur; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_logits") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = -1.0f + 0.25f*float(i); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "exp_probs_b") == 0) { ++ std::vector data(n_expert, 0.0f); ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "ffn_chain_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_exps.weight") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++ ++ double max_nmse_err() override { return 5e-4; } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env motif_reference_env("GGML_METAL_ENABLE_GLM_MOE_DECODE_MOTIF_REFERENCE", "1"); ++ scoped_test_env chunks_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_MOE_Q2_WEIGHT_ROOFLINE_CHUNKS", ++ byte_scan ? "16" : "0"); ++ const std::string block_bytes = std::to_string(storage_block_bytes); ++ scoped_test_env block_bytes_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_MOE_Q2_WEIGHT_ROOFLINE_BLOCK_BYTES", ++ block_bytes.c_str()); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++}; ++ ++struct test_glm_attention_projection_chain : public test_case { ++ static constexpr int64_t n_tokens = 1; ++ static constexpr int64_t n_embd = 6144; ++ static constexpr int64_t n_q_lora = 2048; ++ static constexpr int64_t n_q = 16384; ++ static constexpr int64_t n_kv_lora = 576; ++ ++ const ggml_type weight_type; ++ const int n_repeats; ++ ++ test_glm_attention_projection_chain(ggml_type weight_type, int n_repeats) ++ : weight_type(weight_type), n_repeats(n_repeats) { ++ GGML_ASSERT(n_repeats > 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(weight_type) == 0); ++ GGML_ASSERT(n_q_lora % ggml_blck_size(weight_type) == 0); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR6(weight_type, n_repeats, n_tokens, n_embd, n_q_lora, n_q); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return std::string("GLM_ATTENTION_PROJECTION_CHAIN_") + ggml_type_name(weight_type); ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * cur = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_embd, n_tokens); ++ ggml_set_name(cur, "attention_chain_cur"); ++ ++ for (int repeat = 0; repeat < n_repeats; ++repeat) { ++ ggml_tensor * kv_a = ggml_new_tensor_2d(ctx, weight_type, n_embd, n_kv_lora); ++ ggml_format_name(kv_a, "attention_kv_a_%d", repeat); ++ ggml_tensor * q_a = ggml_new_tensor_2d(ctx, weight_type, n_embd, n_q_lora); ++ ggml_format_name(q_a, "attention_q_a_%d", repeat); ++ ggml_tensor * q_b = ggml_new_tensor_2d(ctx, weight_type, n_q_lora, n_q); ++ ggml_format_name(q_b, "attention_q_b_%d", repeat); ++ ++ ggml_tensor * kv_a_out = ggml_mul_mat(ctx, kv_a, cur); ++ ggml_format_name(kv_a_out, "attention_kv_a_out_%d", repeat); ++ ggml_tensor * q_a_out = ggml_mul_mat(ctx, q_a, cur); ++ ggml_format_name(q_a_out, "attention_q_a_out_%d", repeat); ++ ggml_tensor * q_b_out = ggml_mul_mat(ctx, q_b, q_a_out); ++ ggml_format_name(q_b_out, "attention_q_b_out_%d", repeat); ++ ++ ggml_tensor * anchor = ggml_add(ctx, ggml_sum(ctx, kv_a_out), ggml_sum(ctx, q_b_out)); ++ anchor = ggml_scale(ctx, anchor, 0.0f); ++ anchor = ggml_repeat(ctx, anchor, cur); ++ cur = ggml_add(ctx, cur, anchor); ++ ggml_format_name(cur, "attention_chain_cur_%d", repeat); ++ } ++ return cur; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "attention_chain_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "attention_") == t->name && strstr(t->name, "_out_") == nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++ ++ double max_nmse_err() override { return 5e-4; } ++}; ++ ++struct test_glm_decode_block_byte_ceiling : public test_case { ++ static constexpr int64_t row_width = 65536; ++ static constexpr int64_t rows = 27648; ++ ++ std::string vars() override { ++ return VARS_TO_STR2(row_width, rows); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_DECODE_BLOCK_BYTE_CEILING"; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * weights = ggml_new_tensor_2d(ctx, GGML_TYPE_Q8_0, row_width, rows); ++ ggml_set_name(weights, "glm_decode_block_byte_ceiling_weights"); ++ ggml_tensor * input = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, row_width); ++ ggml_set_name(input, "glm_decode_block_byte_ceiling_input"); ++ ggml_tensor * out = ggml_mul_mat(ctx, weights, input); ++ ggml_set_name(out, "glm_decode_block_byte_ceiling_output"); ++ return out; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ const uint8_t fill = strcmp(t->name, "glm_decode_block_byte_ceiling_weights") == 0 ? 0xa5 : 0; ++ std::vector data(std::min(ggml_nbytes(t), 4 * 1024 * 1024), fill); ++ for (size_t offset = 0; offset < ggml_nbytes(t); offset += data.size()) { ++ const size_t size = std::min(data.size(), ggml_nbytes(t) - offset); ++ ggml_backend_tensor_set(t, data.data(), offset, size); ++ } ++ } ++ } ++ ++ size_t op_size(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return 1908557824ull; ++ } ++}; ++ ++struct test_glm_projection_chain : public test_case { ++ const ggml_type weight_type; ++ const int64_t m; ++ const int64_t k; ++ const int n_repeats; ++ ++ test_glm_projection_chain(ggml_type weight_type, int64_t m, int64_t k, int n_repeats) ++ : weight_type(weight_type), m(m), k(k), n_repeats(n_repeats) { ++ GGML_ASSERT(m > 0 && k > 0 && n_repeats > 0); ++ GGML_ASSERT(k % ggml_blck_size(weight_type) == 0); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR4(weight_type, m, k, n_repeats); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return std::string("GLM_PROJECTION_CHAIN_") + ggml_type_name(weight_type) + ++ "_m" + std::to_string(m) + "_k" + std::to_string(k); ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * cur = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, k); ++ ggml_set_name(cur, "projection_chain_cur"); ++ for (int repeat = 0; repeat < n_repeats; ++repeat) { ++ ggml_tensor * weight = ggml_new_tensor_2d(ctx, weight_type, k, m); ++ ggml_format_name(weight, "projection_chain_weight_%d", repeat); ++ ggml_tensor * out = ggml_mul_mat(ctx, weight, cur); ++ ggml_format_name(out, "projection_chain_out_%d", repeat); ++ ++ ggml_tensor * anchor = ggml_scale(ctx, ggml_sum(ctx, out), 0.0f); ++ anchor = ggml_repeat(ctx, anchor, cur); ++ cur = ggml_add(ctx, cur, anchor); ++ ggml_format_name(cur, "projection_chain_cur_%d", repeat); ++ } ++ return cur; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "projection_chain_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "projection_chain_weight_") == t->name) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++ ++ double max_nmse_err() override { return 5e-4; } ++}; ++ ++struct test_glm_moe_routed_q2_rowtile_output : public test_glm_moe_routed_q2_final_output { ++ const bool row_tiled_layout; ++ ++ explicit test_glm_moe_routed_q2_rowtile_output(bool row_tiled_layout = true) ++ : test_glm_moe_routed_q2_final_output(false, false), ++ row_tiled_layout(row_tiled_layout) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return row_tiled_layout ? ++ "GLM_MOE_ROUTED_MOTIF_Q2_ROWTILE_GATE_UP_GLM" : ++ "GLM_MOE_ROUTED_MOTIF_Q2_SELECTED_GATE_UP_GLM"; ++ } ++ ++ test_status_t eval( ++ ggml_backend_t backend1, ++ ggml_backend_t backend2, ++ const char * op_names_filter, ++ printer * output_printer) override { ++ scoped_test_env motif_reference_env("GGML_METAL_ENABLE_GLM_MOE_DECODE_MOTIF_REFERENCE", "1"); ++ if (!row_tiled_layout) { ++ scoped_test_env rowtile_env("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_ROWTILE", "0"); ++ return test_case::eval(backend1, backend2, op_names_filter, output_printer); ++ } ++ ++ mode = MODE_TEST; ++ ++ ggml_init_params params = { ++ /* .mem_size = */ ggml_tensor_overhead()*256 + ggml_graph_overhead_custom(8192, false), ++ /* .mem_base = */ NULL, ++ /* .no_alloc = */ true, ++ }; ++ ggml_context_ptr ctx_metal(ggml_init(params)); ++ ggml_context_ptr ctx_ref(ggml_init(params)); ++ GGML_ASSERT(ctx_metal); ++ GGML_ASSERT(ctx_ref); ++ ++ ggml_cgraph * gf_metal = ggml_new_graph_custom(ctx_metal.get(), 8192, false); ++ ggml_cgraph * gf_ref = ggml_new_graph_custom(ctx_ref.get(), 8192, false); ++ ++ ggml_tensor * out_metal = build_graph(ctx_metal.get()); ++ current_op_name = op_desc(out_metal); ++ if (!matches_filter(out_metal, op_names_filter)) { ++ return test_status_t::SKIPPED; ++ } ++ scoped_test_env rowtile_env("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_ROWTILE", row_tiled_layout ? "1" : "0"); ++ ggml_tensor * out_ref = build_graph(ctx_ref.get()); ++ ++ if (!graph_supported(backend1, ctx_metal.get()) || !graph_supported(backend2, ctx_ref.get())) { ++ test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", ++ false, false, "not supported"); ++ print_test_result_locked(output_printer, result); ++ return test_status_t::NOT_SUPPORTED; ++ } ++ ++ ggml_backend_buffer_t buf_metal = ggml_backend_alloc_ctx_tensors(ctx_metal.get(), backend1); ++ ggml_backend_buffer_t buf_ref = ggml_backend_alloc_ctx_tensors(ctx_ref.get(), backend2); ++ if (buf_metal == nullptr || buf_ref == nullptr) { ++ if (buf_metal != nullptr) { ++ ggml_backend_buffer_free(buf_metal); ++ } ++ if (buf_ref != nullptr) { ++ ggml_backend_buffer_free(buf_ref); ++ } ++ test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", ++ true, false, "failed to allocate tensors"); ++ print_test_result_locked(output_printer, result); ++ return test_status_t::FAIL; ++ } ++ ++ ggml_build_forward_expand(gf_metal, out_metal); ++ ggml_build_forward_expand(gf_ref, out_ref); ++ ++ initialize_tensors_with_layout(ctx_metal.get(), true); ++ initialize_tensors_with_layout(ctx_ref.get(), false); ++ ++ const ggml_status metal_status = ggml_backend_graph_compute(backend1, gf_metal); ++ const ggml_status ref_status = ggml_backend_graph_compute(backend2, gf_ref); ++ ++ bool ok = metal_status == GGML_STATUS_SUCCESS && ref_status == GGML_STATUS_SUCCESS; ++ std::string error_msg; ++ if (!ok) { ++ error_msg = "graph compute failed"; ++ } else { ++ std::vector metal = tensor_to_float(out_metal); ++ std::vector ref = tensor_to_float(out_ref); ++ const double error = err(metal.data(), ref.data(), metal.size()); ++ if (error > max_err(backend1)) { ++ ok = false; ++ char buf[128]; ++ snprintf(buf, sizeof(buf), "ERR = %.9f > %.9f", error, max_err(backend1)); ++ error_msg = buf; ++ } ++ } ++ ++ ggml_backend_buffer_free(buf_metal); ++ ggml_backend_buffer_free(buf_ref); ++ ++ test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", ++ true, ok, error_msg); ++ print_test_result_locked(output_printer, result); ++ return ok ? test_status_t::OK : test_status_t::FAIL; ++ } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env motif_reference_env("GGML_METAL_ENABLE_GLM_MOE_DECODE_MOTIF_REFERENCE", "1"); ++ scoped_test_env rowtile_env("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_ROWTILE", row_tiled_layout ? "1" : "0"); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ initialize_tensors_with_layout(ctx, row_tiled_layout); ++ } ++ ++private: ++ static bool graph_supported(ggml_backend_t backend, ggml_context * ctx) { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (!ggml_backend_supports_op(backend, t)) { ++ return false; ++ } ++ } ++ return true; ++ } ++ ++ void initialize_tensors_with_layout(ggml_context * ctx, bool row_tiled) { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_logits") == 0) { ++ init_route_logits(t); ++ } else if (strcmp(t->name, "exp_probs_b") == 0) { ++ init_route_bias(t); ++ } else if (strcmp(t->name, "ffn_moe_cur") == 0) { ++ init_decode_vector(t); ++ } else if (strcmp(t->name, "ffn_gate_exps.weight") == 0) { ++ init_selected_q2_experts(t, row_tiled, 11); ++ } else if (strcmp(t->name, "ffn_up_exps.weight") == 0) { ++ init_selected_q2_experts(t, row_tiled, 29); ++ } else if (strcmp(t->name, "ffn_down_exps.weight") == 0) { ++ init_selected_q2_experts(t, false, 47); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++ ++ static void init_route_logits(ggml_tensor * t) { ++ std::vector data(ggml_nelements(t), -12.0f); ++ for (int64_t i = 0; i < n_expert_used; ++i) { ++ data[i] = 12.0f - 0.125f * float(i); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); ++ } ++ ++ static void init_route_bias(ggml_tensor * t) { ++ std::vector data(ggml_nelements(t), 0.0f); ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); ++ } ++ ++ void init_decode_vector(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t)); ++ for (size_t i = 0; i < data.size(); ++i) { ++ const int value = int((i*17 + 23) % 257) - 128; ++ data[i] = 0.003f * float(value); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); ++ } ++ ++ void init_selected_q2_experts(ggml_tensor * t, bool row_tiled, int salt) const { ++ GGML_ASSERT(t->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(t->ne[2] == n_expert); ++ ++ const int64_t input_dim = t->ne[0]; ++ const int64_t output_dim = t->ne[1]; ++ const int64_t nr0 = 8; ++ const size_t row_bytes = ggml_row_size(t->type, input_dim); ++ const size_t block_bytes = ggml_row_size(t->type, ggml_blck_size(t->type)); ++ const int64_t blocks_per_row = input_dim / ggml_blck_size(t->type); ++ const size_t expert_bytes = size_t(output_dim) * row_bytes; ++ ++ std::vector row(input_dim); ++ std::vector expert_row_major(expert_bytes); ++ std::vector expert_storage(expert_bytes); ++ ++ for (int64_t expert = 0; expert < n_expert_used; ++expert) { ++ for (int64_t out_row = 0; out_row < output_dim; ++out_row) { ++ fill_weight_row(row, expert, out_row, salt); ++ ggml_quantize_chunk(t->type, row.data(), ++ expert_row_major.data() + size_t(out_row) * row_bytes, ++ 0, 1, input_dim, nullptr); ++ } ++ ++ if (row_tiled) { ++ std::fill(expert_storage.begin(), expert_storage.end(), 0); ++ for (int64_t out_row = 0; out_row < output_dim; ++out_row) { ++ const int64_t tile = out_row / nr0; ++ const int64_t row_in_tile = out_row % nr0; ++ for (int64_t ib = 0; ib < blocks_per_row; ++ib) { ++ const size_t src = size_t(out_row) * row_bytes + size_t(ib) * block_bytes; ++ const size_t dst = size_t(tile*blocks_per_row*nr0 + ib*nr0 + row_in_tile) * block_bytes; ++ memcpy(expert_storage.data() + dst, expert_row_major.data() + src, block_bytes); ++ } ++ } ++ } else { ++ expert_storage = expert_row_major; ++ } ++ ++ ggml_backend_tensor_set(t, expert_storage.data(), size_t(expert) * t->nb[2], expert_storage.size()); ++ } ++ } ++ ++ void fill_weight_row(std::vector & row, int64_t expert, int64_t out_row, int salt) const { ++ for (int64_t col = 0; col < int64_t(row.size()); ++col) { ++ const int value = int((expert*37 + out_row*13 + col*7 + salt) % 401) - 200; ++ row[col] = 0.0002f * float(value); ++ } ++ } ++}; ++ ++struct test_glm_moe_decode_direct_topk_small_component : public test_topk_moe_glm_component { ++ const bool dense; ++ ++ explicit test_glm_moe_decode_direct_topk_small_component(int component, bool dense) ++ : test_topk_moe_glm_component( ++ component, ++ dense ? GGML_TYPE_F32 : GGML_TYPE_Q2_K, ++ dense ? GGML_TYPE_F32 : GGML_TYPE_Q2_K, ++ dense ? GGML_TYPE_F32 : GGML_TYPE_Q3_K, ++ 256, ++ 256, ++ true, ++ true, ++ false, ++ true), ++ dense(dense) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (dense) { ++ switch (component) { ++ case 2: return "GLM_MOE_DIRECT_TOPK_F32_SMALL_GATE"; ++ case 3: return "GLM_MOE_DIRECT_TOPK_F32_SMALL_GATE_UP"; ++ case 4: return "GLM_MOE_DIRECT_TOPK_F32_SMALL_SWIGLU"; ++ case 7: return "GLM_MOE_DIRECT_TOPK_F32_SMALL_OUT"; ++ default: return "GLM_MOE_DIRECT_TOPK_F32_SMALL"; ++ } ++ } ++ switch (component) { ++ case 2: return "GLM_MOE_DIRECT_TOPK_Q2_SMALL_GATE"; ++ case 3: return "GLM_MOE_DIRECT_TOPK_Q2_SMALL_GATE_UP"; ++ case 4: return "GLM_MOE_DIRECT_TOPK_Q2_SMALL_SWIGLU"; ++ case 7: return "GLM_MOE_DIRECT_TOPK_Q2_SMALL_OUT"; ++ default: return "GLM_MOE_DIRECT_TOPK_Q2_SMALL"; ++ } ++ } ++}; ++ ++struct test_glm_moe_motif_log_small : public test_topk_moe_glm_component { ++ explicit test_glm_moe_motif_log_small(bool weighted_down = false) ++ : test_topk_moe_glm_component( ++ weighted_down ? 7 : 6, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q3_K, ++ 2048, ++ 1024, ++ true, ++ weighted_down, ++ false, ++ false), ++ use_weighted_down(weighted_down) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return use_weighted_down ? "GLM_MOE_MOTIF_LOG_SMALL_PREWEIGHT" : "GLM_MOE_MOTIF_LOG_SMALL"; ++ } ++ ++private: ++ const bool use_weighted_down; ++}; ++ ++struct test_glm_moe_gate_up_pair : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ const int component; ++ const ggml_type gate_type; ++ const ggml_type up_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * ffn_gate = nullptr; ++ ggml_tensor * ffn_up = nullptr; ++ ggml_tensor * out = nullptr; ++ ++ test_glm_moe_gate_up_pair( ++ int component = 2, ++ ggml_type gate_type = GGML_TYPE_Q2_K, ++ ggml_type up_type = GGML_TYPE_Q2_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048) ++ : component(component), ++ gate_type(gate_type), ++ up_type(up_type), ++ n_embd(n_embd), ++ n_ff(n_ff) { ++ GGML_ASSERT(component >= 0 && component <= 2); ++ GGML_ASSERT(n_embd % ggml_blck_size(gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(up_type) == 0); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR8(component, gate_type, up_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_MOE_GATE_UP_PAIR"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ selected_experts = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, n_tokens); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_moe_cur"); ++ ++ ggml_tensor * gates = ggml_new_tensor_3d(ctx, gate_type, n_embd, n_ff, n_expert); ++ ggml_set_name(gates, "ffn_gate_exps.weight"); ++ ++ ffn_gate = ggml_mul_mat_id(ctx, gates, cur, selected_experts); ++ ggml_set_name(ffn_gate, "ffn_moe_gate"); ++ if (component == 0) { ++ out = ffn_gate; ++ return out; ++ } ++ ++ ggml_tensor * ups = ggml_new_tensor_3d(ctx, up_type, n_embd, n_ff, n_expert); ++ ggml_set_name(ups, "ffn_up_exps.weight"); ++ ++ ffn_up = ggml_mul_mat_id(ctx, ups, cur, selected_experts); ++ ggml_set_name(ffn_up, "ffn_moe_up"); ++ if (component == 1) { ++ out = ffn_up; ++ return out; ++ } ++ ++ out = ggml_add(ctx, ffn_gate, ffn_up); ++ ggml_set_name(out, "ffn_moe_gate_up_pair_sum"); ++ return out; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ return { selected_experts, ffn_gate, ffn_up, out }; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_topk") == 0) { ++ std::vector data(n_expert_used * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ data[token*n_expert_used + slot] = int32_t((slot * 17 + token * 29) % n_expert); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(int32_t)); ++ } else if (strcmp(t->name, "ffn_moe_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_exps.weight") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_glm_moe_gate_up_swiglu : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ const int64_t n_tokens; ++ const ggml_type gate_type; ++ const ggml_type up_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const bool slot_major_experts; ++ const bool all_expert_slot_major; ++ const bool replicate_decode_token; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * ffn_gate = nullptr; ++ ggml_tensor * ffn_up = nullptr; ++ ggml_tensor * ffn_act = nullptr; ++ ggml_tensor * ffn_q8 = nullptr; ++ ++ test_glm_moe_gate_up_swiglu( ++ ggml_type gate_type = GGML_TYPE_Q2_K, ++ ggml_type up_type = GGML_TYPE_Q2_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048, ++ bool slot_major_experts = false, ++ bool all_expert_slot_major = false, ++ int64_t n_tokens = 1, ++ bool replicate_decode_token = false) ++ : n_tokens(n_tokens), ++ gate_type(gate_type), ++ up_type(up_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ slot_major_experts(slot_major_experts), ++ all_expert_slot_major(all_expert_slot_major), ++ replicate_decode_token(replicate_decode_token) { ++ GGML_ASSERT(n_tokens > 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(up_type) == 0); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR9(gate_type, up_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff, slot_major_experts, all_expert_slot_major); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (replicate_decode_token) { ++ return "GLM_MOE_GATE_UP_SWIGLU_Q2_REPLICATED_R" + std::to_string(n_tokens) + "_GLM"; ++ } ++ if (all_expert_slot_major) { ++ return "GLM_MOE_GATE_UP_SWIGLU_ALL_EXPERT_SLOT_MAJOR"; ++ } ++ if (slot_major_experts) { ++ if (gate_type == up_type) { ++ switch (gate_type) { ++ case GGML_TYPE_F16: ++ return "GLM_MOE_GATE_UP_SWIGLU_SELECTED_F16_GLM"; ++ case GGML_TYPE_Q8_0: ++ return "GLM_MOE_GATE_UP_SWIGLU_SELECTED_Q8_0_GLM"; ++ case GGML_TYPE_Q4_0: ++ return "GLM_MOE_GATE_UP_SWIGLU_SELECTED_Q4_0_GLM"; ++ case GGML_TYPE_Q4_K: ++ return "GLM_MOE_GATE_UP_SWIGLU_SELECTED_Q4_K_GLM"; ++ default: ++ break; ++ } ++ } ++ return "GLM_MOE_GATE_UP_SWIGLU_SLOT_MAJOR"; ++ } ++ return "GLM_MOE_GATE_UP_SWIGLU"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * new_expert_tensor( ++ ggml_context * ctx, ++ ggml_type type, ++ int64_t input_dim, ++ int64_t output_dim, ++ const char * name) { ++ if (!slot_major_experts && !all_expert_slot_major) { ++ ggml_tensor * tensor = ggml_new_tensor_3d(ctx, type, input_dim, output_dim, n_expert); ++ ggml_set_name(tensor, name); ++ return tensor; ++ } ++ ++ const int64_t storage_experts = all_expert_slot_major ? n_expert : n_expert_used; ++ ggml_tensor * storage = ggml_new_tensor_3d(ctx, type, input_dim, storage_experts, output_dim); ++ ggml_format_name(storage, "%s.%s", name, all_expert_slot_major ? "all_expert_slot_major_storage" : "slot_major_storage"); ++ ++ ggml_tensor * view = ggml_view_3d(ctx, storage, ++ input_dim, output_dim, storage_experts, ++ storage->nb[2], storage->nb[1], 0); ++ ggml_set_name(view, name); ++ return view; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ selected_experts = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, n_tokens); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_moe_cur"); ++ ++ if (getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_PREQUANT_Q8") != nullptr) { ++ ffn_q8 = ggml_cast(ctx, cur, GGML_TYPE_Q8_0); ++ ggml_set_name(ffn_q8, "ffn_moe_gate_up_q8"); ++ } ++ ++ ggml_tensor * gates = new_expert_tensor(ctx, gate_type, n_embd, n_ff, "ffn_gate_exps.weight"); ++ ++ ggml_tensor * ups = new_expert_tensor(ctx, up_type, n_embd, n_ff, "ffn_up_exps.weight"); ++ ++ ffn_gate = ggml_mul_mat_id(ctx, gates, cur, selected_experts); ++ ffn_gate->src[3] = ffn_q8; ++ ggml_set_name(ffn_gate, "ffn_moe_gate"); ++ ++ ffn_up = ggml_mul_mat_id(ctx, ups, cur, selected_experts); ++ ffn_up->src[3] = ffn_q8; ++ ggml_set_name(ffn_up, "ffn_moe_up"); ++ ++ ffn_act = ggml_swiglu_split(ctx, ffn_gate, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_swiglu"); ++ return ffn_act; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ const bool q2_gate_up_swiglu = ++ gate_type == GGML_TYPE_Q2_K && ++ up_type == GGML_TYPE_Q2_K && ++ (glm_moe_q2_gate_up_pair_sg_default(gate_type, up_type, n_embd, n_ff) || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_GATE_UP_SWIGLU") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q2_GATE_UP_SWIGLU_FUSION") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_Q8_ACT") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_HALF_Y") != nullptr); ++ if (q2_gate_up_swiglu) { ++ return { selected_experts, ffn_act }; ++ } ++ return { selected_experts, ffn_gate, ffn_up, ffn_act }; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_topk") == 0) { ++ std::vector data(n_expert_used * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ data[token*n_expert_used + slot] = slot_major_experts ? ++ int32_t(slot) : ++ int32_t((slot * 17 + (replicate_decode_token ? 0 : token * 29)) % n_expert); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(int32_t)); ++ } else if (strcmp(t->name, "ffn_moe_cur") == 0) { ++ if (replicate_decode_token) { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t i = 0; i < n_embd; ++i) { ++ const int value = int((i*17 + 23) % 257) - 128; ++ data[token*n_embd + i] = 0.003f * float(value); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); ++ } else { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } ++ } else if (strstr(t->name, "_exps.weight") != nullptr || ++ strstr(t->name, "_exps.weight.slot_major_storage") != nullptr || ++ strstr(t->name, "_exps.weight.all_expert_slot_major_storage") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_glm_moe_gate_up_swiglu_rowtile_q2 : public test_glm_moe_gate_up_swiglu { ++ const bool row_tiled_layout; ++ const bool inblock_repacked; ++ ++ explicit test_glm_moe_gate_up_swiglu_rowtile_q2( ++ bool row_tiled_layout = true, ++ bool inblock_repacked = false) ++ : test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048), ++ row_tiled_layout(row_tiled_layout), ++ inblock_repacked(inblock_repacked) {} ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (inblock_repacked) { ++ return "GLM_MOE_GATE_UP_SWIGLU_INBLOCK_Q2_GLM"; ++ } ++ return row_tiled_layout ? "GLM_MOE_GATE_UP_SWIGLU_ROWTILE_Q2_GLM" : ++ "GLM_MOE_GATE_UP_SWIGLU_SELECTED_Q2_GLM"; ++ } ++ ++ test_status_t eval( ++ ggml_backend_t backend1, ++ ggml_backend_t backend2, ++ const char * op_names_filter, ++ printer * output_printer) override { ++ if (!row_tiled_layout && !inblock_repacked) { ++ scoped_test_env rowtile_env("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_ROWTILE", "0"); ++ return test_case::eval(backend1, backend2, op_names_filter, output_printer); ++ } ++ ++ mode = MODE_TEST; ++ ++ ggml_init_params params = { ++ /* .mem_size = */ ggml_tensor_overhead()*128 + ggml_graph_overhead(), ++ /* .mem_base = */ NULL, ++ /* .no_alloc = */ true, ++ }; ++ ggml_context_ptr ctx_metal(ggml_init(params)); ++ ggml_context_ptr ctx_ref(ggml_init(params)); ++ GGML_ASSERT(ctx_metal); ++ GGML_ASSERT(ctx_ref); ++ ++ ggml_cgraph * gf_metal = ggml_new_graph(ctx_metal.get()); ++ ggml_cgraph * gf_ref = ggml_new_graph(ctx_ref.get()); ++ ++ scoped_test_env rowtile_env("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_ROWTILE", row_tiled_layout ? "1" : "0"); ++ scoped_test_env inblock_env("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_INBLOCK_REPACK", inblock_repacked ? "1" : "0"); ++ ggml_tensor * out_metal = build_graph(ctx_metal.get()); ++ current_op_name = op_desc(out_metal); ++ if (!matches_filter(out_metal, op_names_filter)) { ++ return test_status_t::SKIPPED; ++ } ++ ggml_tensor * out_ref = build_graph(ctx_ref.get()); ++ ++ if (!graph_supported(backend1, ctx_metal.get()) || !graph_supported(backend2, ctx_ref.get())) { ++ test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", ++ false, false, "not supported"); ++ print_test_result_locked(output_printer, result); ++ return test_status_t::NOT_SUPPORTED; ++ } ++ ++ ggml_backend_buffer_t buf_metal = ggml_backend_alloc_ctx_tensors(ctx_metal.get(), backend1); ++ ggml_backend_buffer_t buf_ref = ggml_backend_alloc_ctx_tensors(ctx_ref.get(), backend2); ++ if (buf_metal == nullptr || buf_ref == nullptr) { ++ if (buf_metal != nullptr) { ++ ggml_backend_buffer_free(buf_metal); ++ } ++ if (buf_ref != nullptr) { ++ ggml_backend_buffer_free(buf_ref); ++ } ++ test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", ++ true, false, "failed to allocate tensors"); ++ print_test_result_locked(output_printer, result); ++ return test_status_t::FAIL; ++ } ++ ++ ggml_build_forward_expand(gf_metal, out_metal); ++ ggml_build_forward_expand(gf_ref, out_ref); ++ ++ initialize_tensors_with_layout(ctx_metal.get(), row_tiled_layout, inblock_repacked); ++ initialize_tensors_with_layout(ctx_ref.get(), false, false); ++ ++ const ggml_status metal_status = ggml_backend_graph_compute(backend1, gf_metal); ++ const ggml_status ref_status = ggml_backend_graph_compute(backend2, gf_ref); ++ ++ bool ok = metal_status == GGML_STATUS_SUCCESS && ref_status == GGML_STATUS_SUCCESS; ++ std::string error_msg; ++ if (!ok) { ++ error_msg = "graph compute failed"; ++ } else { ++ std::vector metal = tensor_to_float(out_metal); ++ std::vector ref = tensor_to_float(out_ref); ++ const double error = err(metal.data(), ref.data(), metal.size()); ++ if (error > max_err(backend1)) { ++ ok = false; ++ char buf[128]; ++ snprintf(buf, sizeof(buf), "ERR = %.9f > %.9f", error, max_err(backend1)); ++ error_msg = buf; ++ } ++ } ++ ++ ggml_backend_buffer_free(buf_metal); ++ ggml_backend_buffer_free(buf_ref); ++ ++ test_result result(ggml_backend_name(backend1), current_op_name, vars(), "test", ++ true, ok, error_msg); ++ print_test_result_locked(output_printer, result); ++ return ok ? test_status_t::OK : test_status_t::FAIL; ++ } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env rowtile_env("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG_ROWTILE", row_tiled_layout ? "1" : "0"); ++ scoped_test_env inblock_env("GGML_METAL_EXPERIMENTAL_Q2_GATE_UP_INBLOCK_REPACK", inblock_repacked ? "1" : "0"); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ initialize_tensors_with_layout(ctx, row_tiled_layout, inblock_repacked); ++ } ++ ++private: ++ static bool graph_supported(ggml_backend_t backend, ggml_context * ctx) { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (!ggml_backend_supports_op(backend, t)) { ++ return false; ++ } ++ } ++ return true; ++ } ++ ++ void initialize_tensors_with_layout(ggml_context * ctx, bool row_tiled, bool inblock) { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_topk") == 0) { ++ init_selected_experts(t); ++ } else if (strcmp(t->name, "ffn_moe_cur") == 0) { ++ init_decode_vector(t); ++ } else if (strcmp(t->name, "ffn_gate_exps.weight") == 0) { ++ init_selected_q2_experts(t, row_tiled, inblock, 11); ++ } else if (strcmp(t->name, "ffn_up_exps.weight") == 0) { ++ init_selected_q2_experts(t, row_tiled, inblock, 29); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++ ++ void init_selected_experts(ggml_tensor * t) const { ++ std::vector data(n_expert_used * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ data[token*n_expert_used + slot] = int32_t(slot); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(int32_t)); ++ } ++ ++ void init_decode_vector(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t)); ++ for (size_t i = 0; i < data.size(); ++i) { ++ const int value = int((i*17 + 23) % 257) - 128; ++ data[i] = 0.003f * float(value); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); ++ } ++ ++ void init_selected_q2_experts(ggml_tensor * t, bool row_tiled, bool inblock, int salt) const { ++ GGML_ASSERT(t->type == GGML_TYPE_Q2_K); ++ GGML_ASSERT(t->ne[0] == n_embd); ++ GGML_ASSERT(t->ne[1] == n_ff); ++ ++ const int64_t input_dim = t->ne[0]; ++ const int64_t output_dim = t->ne[1]; ++ const int64_t nr0 = 8; ++ const size_t row_bytes = ggml_row_size(t->type, input_dim); ++ const size_t block_bytes = ggml_row_size(t->type, ggml_blck_size(t->type)); ++ const int64_t blocks_per_row = input_dim / ggml_blck_size(t->type); ++ const size_t expert_bytes = size_t(output_dim) * row_bytes; ++ ++ std::vector row(input_dim); ++ std::vector expert_row_major(expert_bytes); ++ std::vector expert_storage(expert_bytes); ++ ++ for (int64_t expert = 0; expert < n_expert_used; ++expert) { ++ for (int64_t out_row = 0; out_row < output_dim; ++out_row) { ++ fill_weight_row(row, expert, out_row, salt); ++ ggml_quantize_chunk(t->type, row.data(), ++ expert_row_major.data() + size_t(out_row) * row_bytes, ++ 0, 1, input_dim, nullptr); ++ } ++ ++ if (row_tiled) { ++ std::fill(expert_storage.begin(), expert_storage.end(), 0); ++ for (int64_t out_row = 0; out_row < output_dim; ++out_row) { ++ const int64_t tile = out_row / nr0; ++ const int64_t row_in_tile = out_row % nr0; ++ for (int64_t ib = 0; ib < blocks_per_row; ++ib) { ++ const size_t src = size_t(out_row) * row_bytes + size_t(ib) * block_bytes; ++ const size_t dst = size_t(tile*blocks_per_row*nr0 + ib*nr0 + row_in_tile) * block_bytes; ++ memcpy(expert_storage.data() + dst, expert_row_major.data() + src, block_bytes); ++ } ++ } ++ } else { ++ expert_storage = expert_row_major; ++ } ++ ++ if (inblock) { ++ repack_q2_blocks_in_place(expert_storage, block_bytes); ++ } ++ ++ ggml_backend_tensor_set(t, expert_storage.data(), size_t(expert) * t->nb[2], expert_storage.size()); ++ } ++ } ++ ++ static void repack_q2_blocks_in_place(std::vector & bytes, size_t block_bytes) { ++ constexpr size_t q_offset = 16; ++ constexpr size_t q_bytes = 64; ++ GGML_ASSERT(block_bytes == q_offset + q_bytes + 2*sizeof(uint16_t)); ++ GGML_ASSERT(bytes.size() % block_bytes == 0); ++ ++ std::array original = {}; ++ for (size_t block = 0; block < bytes.size(); block += block_bytes) { ++ uint8_t * qs = bytes.data() + block + q_offset; ++ memcpy(original.data(), qs, q_bytes); ++ memset(qs, 0, q_bytes); ++ for (int logical = 0; logical < 256; ++logical) { ++ const int half = logical / 128; ++ const int within = logical % 128; ++ const int source_byte = half*32 + within%32; ++ const int source_shift = 2*(within/32); ++ const uint8_t q = (original[source_byte] >> source_shift) & 0x3; ++ const int subblock = logical/16; ++ const int in_subblock = logical%16; ++ qs[subblock*4 + in_subblock/4] |= q << (2*(in_subblock%4)); ++ } ++ } ++ } ++ ++ void fill_weight_row(std::vector & row, int64_t expert, int64_t out_row, int salt) const { ++ for (int64_t col = 0; col < int64_t(row.size()); ++col) { ++ const int value = int((expert*37 + out_row*13 + col*7 + salt) % 401) - 200; ++ row[col] = 0.0002f * float(value); ++ } ++ } ++}; ++ ++struct test_glm_moe_static_motif : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ const ggml_type gate_type; ++ const ggml_type up_type; ++ const ggml_type down_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const bool weighted_down; ++ const bool packed_selected_experts; ++ const bool slot_major_experts; ++ const int64_t active_experts; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * ffn_gate = nullptr; ++ ggml_tensor * ffn_up = nullptr; ++ ggml_tensor * ffn_act = nullptr; ++ ggml_tensor * ffn_down = nullptr; ++ ggml_tensor * ffn_out = nullptr; ++ ++ test_glm_moe_static_motif( ++ ggml_type gate_type = GGML_TYPE_Q2_K, ++ ggml_type up_type = GGML_TYPE_Q2_K, ++ ggml_type down_type = GGML_TYPE_Q3_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048, ++ bool weighted_down = true, ++ bool packed_selected_experts = false, ++ bool slot_major_experts = false, ++ int64_t active_experts = n_expert_used) ++ : gate_type(gate_type), ++ up_type(up_type), ++ down_type(down_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ weighted_down(weighted_down), ++ packed_selected_experts(packed_selected_experts), ++ slot_major_experts(slot_major_experts), ++ active_experts(active_experts) { ++ GGML_ASSERT(n_embd % ggml_blck_size(gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(down_type) == 0); ++ GGML_ASSERT(active_experts > 0 && active_experts <= n_expert_used); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR11(gate_type, up_type, down_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff, weighted_down, packed_selected_experts, slot_major_experts) + ++ "," + VAR_TO_STR(active_experts); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (gate_type == GGML_TYPE_Q2_K && up_type == GGML_TYPE_Q2_K && down_type == GGML_TYPE_Q2_K) { ++ if (n_embd == 6144 && n_ff == 2048) { ++ if (active_experts == 2) { ++ return "GLM_MOE_STATIC_MOTIF_Q2_ACTIVE2_GLM"; ++ } ++ if (active_experts == 4) { ++ return "GLM_MOE_STATIC_MOTIF_Q2_ACTIVE4_GLM"; ++ } ++ if (active_experts == 6) { ++ return "GLM_MOE_STATIC_MOTIF_Q2_ACTIVE6_GLM"; ++ } ++ if (slot_major_experts) { ++ return weighted_down ? "GLM_MOE_STATIC_MOTIF_Q2_WEIGHTED_SLOT_MAJOR_GLM" : "GLM_MOE_STATIC_MOTIF_Q2_SLOT_MAJOR_GLM"; ++ } ++ if (packed_selected_experts) { ++ return weighted_down ? "GLM_MOE_STATIC_MOTIF_Q2_WEIGHTED_PACKED_GLM" : "GLM_MOE_STATIC_MOTIF_Q2_PACKED_GLM"; ++ } ++ return weighted_down ? "GLM_MOE_STATIC_MOTIF_Q2_WEIGHTED_GLM" : "GLM_MOE_STATIC_MOTIF_Q2_GLM"; ++ } ++ return weighted_down ? "GLM_MOE_STATIC_MOTIF_Q2_WEIGHTED" : "GLM_MOE_STATIC_MOTIF_Q2"; ++ } ++ if (gate_type == GGML_TYPE_Q4_K && up_type == GGML_TYPE_Q4_K && down_type == GGML_TYPE_Q4_K) { ++ if (n_embd == 6144 && n_ff == 2048) { ++ return "GLM_MOE_STATIC_MOTIF_Q4_GLM"; ++ } ++ return "GLM_MOE_STATIC_MOTIF_Q4"; ++ } ++ return "GLM_MOE_STATIC_MOTIF"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * new_expert_tensor( ++ ggml_context * ctx, ++ ggml_type type, ++ int64_t input_dim, ++ int64_t output_dim, ++ const char * name) { ++ if (packed_selected_experts) { ++ ggml_tensor * tensor = ggml_new_tensor_3d(ctx, type, input_dim, output_dim, n_expert_used); ++ ggml_set_name(tensor, name); ++ return tensor; ++ } ++ ++ if (!slot_major_experts) { ++ ggml_tensor * tensor = ggml_new_tensor_3d(ctx, type, input_dim, output_dim, n_expert); ++ ggml_set_name(tensor, name); ++ return tensor; ++ } ++ ++ ggml_tensor * storage = ggml_new_tensor_3d(ctx, type, input_dim, n_expert_used, output_dim); ++ ggml_format_name(storage, "%s.slot_major_storage", name); ++ ++ ggml_tensor * view = ggml_view_3d(ctx, storage, ++ input_dim, output_dim, n_expert_used, ++ storage->nb[2], storage->nb[1], 0); ++ ggml_set_name(view, name); ++ return view; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ selected_experts = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, n_tokens); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ ++ weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, n_expert_used, n_tokens); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_moe_cur"); ++ ++ ggml_tensor * gates = new_expert_tensor(ctx, gate_type, n_embd, n_ff, "ffn_gate_exps.weight"); ++ ++ ggml_tensor * ups = new_expert_tensor(ctx, up_type, n_embd, n_ff, "ffn_up_exps.weight"); ++ ++ ggml_tensor * downs = new_expert_tensor(ctx, down_type, n_ff, n_embd, "ffn_down_exps.weight"); ++ ++ ffn_gate = ggml_mul_mat_id(ctx, gates, cur, selected_experts); ++ ggml_set_name(ffn_gate, "ffn_moe_gate"); ++ ++ ffn_up = ggml_mul_mat_id(ctx, ups, cur, selected_experts); ++ ggml_set_name(ffn_up, "ffn_moe_up"); ++ ++ ffn_act = ggml_swiglu_split(ctx, ffn_gate, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_swiglu"); ++ ++ if (weighted_down) { ++ ffn_act = ggml_mul(ctx, ffn_act, weights); ++ ggml_set_name(ffn_act, "ffn_moe_down_weighted_input"); ++ } ++ ++ ffn_down = ggml_mul_mat_id(ctx, downs, ffn_act, selected_experts); ++ ggml_set_name(ffn_down, "ffn_moe_down"); ++ ++ ffn_out = ggml_moe_weighted_sum(ctx, ffn_down, weights); ++ if (weighted_down) { ++ ffn_out->op_params[0] = 1; ++ } ++ ggml_set_name(ffn_out, "ffn_moe_out"); ++ ++ return ffn_out; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ const bool direct_reduce = ++ glm_moe_q2_down_default_direct_reduce(down_type, n_ff, weighted_down) || ++ glm_moe_q3_down_w0_default_direct_reduce(down_type, n_ff, weighted_down) || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_DOWN_WEIGHTED_REDUCE_DIRECT") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SLOT_PARALLEL_REDUCE_R16_W1") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_WEIGHTED_REDUCE_DIRECT") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8_W0") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R12_NB8_W0") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE_R16") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_SPLIT2_REDUCE") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_ATOMIC_ACCUM") != nullptr; ++ const bool gate_up_fused = glm_moe_q2_gate_up_pair_sg_default(gate_type, up_type, n_embd, n_ff) || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_GATE_UP_SWIGLU") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_GATE_UP_SWIGLU_PAIR_SG") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q2_GATE_UP_SWIGLU_FUSION") != nullptr; ++ if (direct_reduce && gate_up_fused) { ++ return { selected_experts, weights, ffn_out }; ++ } ++ if (direct_reduce) { ++ return { selected_experts, weights, ffn_gate, ffn_up, ffn_act, ffn_out }; ++ } ++ if (gate_up_fused) { ++ return { selected_experts, weights, ffn_act, ffn_down, ffn_out }; ++ } ++ return { selected_experts, weights, ffn_gate, ffn_up, ffn_act, ffn_down, ffn_out }; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_topk") == 0) { ++ std::vector data(n_expert_used * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ data[token*n_expert_used + slot] = ++ (packed_selected_experts || slot_major_experts) ? ++ int32_t(slot) : ++ int32_t((slot * 17 + token * 29) % n_expert); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(int32_t)); ++ } else if (strcmp(t->name, "ffn_moe_weights_scaled") == 0) { ++ std::vector data(n_expert_used * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ float sum = 0.0f; ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ const float value = slot < active_experts ? ++ 0.5f + 0.125f * float((slot * 7 + token * 3) % 11) : ++ 0.0f; ++ data[token*n_expert_used + slot] = value; ++ sum += value; ++ } ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ data[token*n_expert_used + slot] = 2.0f * data[token*n_expert_used + slot] / sum; ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); ++ } else if (strcmp(t->name, "ffn_moe_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_exps.weight") != nullptr || strstr(t->name, "_exps.weight.slot_major_storage") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_glm_moe_down_reduce : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ const ggml_type down_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const bool already_weighted; ++ const bool f16_act; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * ffn_act = nullptr; ++ ggml_tensor * ffn_down = nullptr; ++ ggml_tensor * ffn_out = nullptr; ++ ++ test_glm_moe_down_reduce( ++ ggml_type down_type = GGML_TYPE_Q3_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048, ++ bool already_weighted = false, ++ bool f16_act = false) ++ : down_type(down_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ already_weighted(already_weighted), ++ f16_act(f16_act) { ++ GGML_ASSERT(n_ff % ggml_blck_size(down_type) == 0); ++ GGML_ASSERT(!f16_act || !already_weighted); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR6(down_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff) + "," + ++ VARS_TO_STR2(already_weighted, f16_act); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ if (f16_act) { ++ return "GLM_MOE_DOWN_REDUCE_F16_ACT"; ++ } ++ if (down_type == GGML_TYPE_Q4_K) { ++ return "GLM_MOE_DOWN_REDUCE_Q4"; ++ } ++ return "GLM_MOE_DOWN_REDUCE"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ selected_experts = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_expert_used, n_tokens); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ ++ weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, n_expert_used, n_tokens); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ ++ ffn_act = ggml_new_tensor_3d(ctx, f16_act ? GGML_TYPE_F16 : GGML_TYPE_F32, n_ff, n_expert_used, n_tokens); ++ ggml_set_name(ffn_act, f16_act ? "ffn_moe_swiglu_f16" : "ffn_moe_swiglu"); ++ ++ if (already_weighted) { ++ ffn_act = ggml_mul(ctx, ffn_act, weights); ++ ggml_set_name(ffn_act, "ffn_moe_down_weighted_input"); ++ } ++ ggml_tensor * downs = ggml_new_tensor_3d(ctx, down_type, n_ff, n_embd, n_expert); ++ ggml_set_name(downs, "ffn_down_exps.weight"); ++ ++ ffn_down = ggml_mul_mat_id(ctx, downs, ffn_act, selected_experts); ++ ggml_set_name(ffn_down, "ffn_moe_down"); ++ ++ ffn_out = ggml_moe_weighted_sum(ctx, ffn_down, weights); ++ if (already_weighted) { ++ ffn_out->op_params[0] = 1; ++ } ++ ggml_set_name(ffn_out, "ffn_moe_out"); ++ ++ return ffn_out; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ const bool direct_reduce = ++ glm_moe_q2_down_default_direct_reduce(down_type, n_ff, already_weighted) || ++ glm_moe_q3_down_w0_default_direct_reduce(down_type, n_ff, already_weighted) || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q2_DOWN_WEIGHTED_REDUCE_DIRECT") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_F16_ACT") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SLOT_PARALLEL_REDUCE") != nullptr || ++ getenv("GGML_METAL_EXPERIMENTAL_Q2_DOWN_SLOT_PARALLEL_REDUCE_R16_W1") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_WEIGHTED_REDUCE_DIRECT") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R8_NB8_W0") != nullptr || ++ getenv("GGML_METAL_ENABLE_Q3_DOWN_SLOT_PARALLEL_REDUCE_R12_NB8_W0") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_PARALLEL_REDUCE_R16") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_SLOT_SPLIT2_REDUCE") != nullptr || ++ getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_Q3_DOWN_ATOMIC_ACCUM") != nullptr; ++ if (direct_reduce) { ++ return { selected_experts, weights, ffn_act, ffn_out }; ++ } ++ return { selected_experts, weights, ffn_act, ffn_down, ffn_out }; ++ } ++ ++ double max_nmse_err() override { ++ return f16_act ? 2e-3 : 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_topk") == 0) { ++ std::vector data(n_expert_used * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ data[token*n_expert_used + slot] = int32_t((slot * 17 + token * 29) % n_expert); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(int32_t)); ++ } else if (strcmp(t->name, "ffn_moe_weights_scaled") == 0) { ++ std::vector data(n_expert_used * n_tokens); ++ for (int64_t token = 0; token < n_tokens; ++token) { ++ float sum = 0.0f; ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ const float value = 0.5f + 0.125f * float((slot * 7 + token * 3) % 11); ++ data[token*n_expert_used + slot] = value; ++ sum += value; ++ } ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ data[token*n_expert_used + slot] = 2.0f * data[token*n_expert_used + slot] / sum; ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size() * sizeof(float)); ++ } else if (strcmp(t->name, "ffn_moe_swiglu") == 0) { ++ init_tensor_uniform(t, -0.25f, 0.25f); ++ } else if (strcmp(t->name, "ffn_down_exps.weight") == 0) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_topk_moe_glm_activation : public test_case { ++ const int64_t n_ff; ++ const int64_t n_expert_used; ++ const int64_t n_tokens; ++ const bool fused_glu; ++ const bool weighted; ++ ++ test_topk_moe_glm_activation( ++ int64_t n_ff = 2048, ++ int64_t n_expert_used = 8, ++ int64_t n_tokens = 1, ++ bool fused_glu = true, ++ bool weighted = false) ++ : n_ff(n_ff), ++ n_expert_used(n_expert_used), ++ n_tokens(n_tokens), ++ fused_glu(fused_glu), ++ weighted(weighted) { ++ GGML_ASSERT(!weighted || fused_glu); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR5(n_ff, n_expert_used, n_tokens, fused_glu, weighted); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "TOPK_MOE_GLM_ACTIVATION"; ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, n_expert_used, n_tokens); ++ ggml_set_name(gate, "ffn_moe_gate"); ++ ++ ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, n_expert_used, n_tokens); ++ ggml_set_name(up, "ffn_moe_up"); ++ ++ if (fused_glu) { ++ ggml_tensor * act = ggml_swiglu_split(ctx, gate, up); ++ ggml_set_name(act, "ffn_moe_swiglu"); ++ if (weighted) { ++ ggml_tensor * weights = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, 1, n_expert_used, n_tokens); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ ++ act = ggml_mul(ctx, act, weights); ++ ggml_set_name(act, "ffn_moe_down_weighted_input"); ++ } ++ return act; ++ } ++ ++ ggml_tensor * act = ggml_silu(ctx, gate); ++ ggml_set_name(act, "ffn_moe_gate_silu"); ++ ++ act = ggml_mul(ctx, act, up); ++ ggml_set_name(act, "ffn_moe_gate_up"); ++ return act; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } ++ } ++}; ++ ++struct test_glm_shexp_consumer : public test_case { ++ const ggml_type gate_type; ++ const ggml_type up_type; ++ const ggml_type down_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const int64_t n_tokens; ++ const bool fused_glu; ++ const bool merged_gate_up; ++ ++ test_glm_shexp_consumer( ++ ggml_type gate_type = GGML_TYPE_Q4_K, ++ ggml_type up_type = GGML_TYPE_Q4_K, ++ ggml_type down_type = GGML_TYPE_Q4_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048, ++ int64_t n_tokens = 1, ++ bool fused_glu = false, ++ bool merged_gate_up = false) ++ : gate_type(gate_type), ++ up_type(up_type), ++ down_type(down_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ n_tokens(n_tokens), ++ fused_glu(fused_glu), ++ merged_gate_up(merged_gate_up) { ++ GGML_ASSERT(n_embd % ggml_blck_size(gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(down_type) == 0); ++ GGML_ASSERT(!merged_gate_up || fused_glu); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR8(gate_type, up_type, down_type, n_embd, n_ff, n_tokens, fused_glu, merged_gate_up); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_SHEXP_CONSUMER"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_shexp_cur"); ++ ++ ggml_tensor * moe_out = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(moe_out, "ffn_moe_out"); ++ ++ ggml_tensor * downs = ggml_new_tensor_2d(ctx, down_type, n_ff, n_embd); ++ ggml_set_name(downs, "ffn_down_shexp.weight"); ++ ++ ggml_tensor * act = nullptr; ++ if (merged_gate_up) { ++ ggml_tensor * gate_ups = ggml_new_tensor_2d(ctx, gate_type, n_embd, 2*n_ff); ++ ggml_set_name(gate_ups, "ffn_gate_up_shexp.weight"); ++ ++ ggml_tensor * gate_up = ggml_mul_mat(ctx, gate_ups, cur); ++ ggml_set_name(gate_up, "ffn_shexp_gate_up_merged"); ++ ++ act = ggml_swiglu(ctx, gate_up); ++ ggml_set_name(act, "ffn_shexp_swiglu"); ++ } else { ++ ggml_tensor * gates = ggml_new_tensor_2d(ctx, gate_type, n_embd, n_ff); ++ ggml_set_name(gates, "ffn_gate_shexp.weight"); ++ ++ ggml_tensor * ups = ggml_new_tensor_2d(ctx, up_type, n_embd, n_ff); ++ ggml_set_name(ups, "ffn_up_shexp.weight"); ++ ++ ggml_tensor * gate = ggml_mul_mat(ctx, gates, cur); ++ ggml_set_name(gate, "ffn_shexp_gate"); ++ ++ ggml_tensor * up = ggml_mul_mat(ctx, ups, cur); ++ ggml_set_name(up, "ffn_shexp_up"); ++ ++ if (fused_glu) { ++ act = ggml_swiglu_split(ctx, gate, up); ++ ggml_set_name(act, "ffn_shexp_swiglu"); ++ } else { ++ act = ggml_silu(ctx, gate); ++ ggml_set_name(act, "ffn_shexp_gate_silu"); ++ ++ act = ggml_mul(ctx, act, up); ++ ggml_set_name(act, "ffn_shexp_gate_up"); ++ } ++ } ++ ++ ggml_tensor * shexp = ggml_mul_mat(ctx, downs, act); ++ ggml_set_name(shexp, "ffn_shexp"); ++ ++ ggml_tensor * out = ggml_add(ctx, moe_out, shexp); ++ ggml_set_name(out, "ffn_out"); ++ ++ return out; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_shexp_cur") == 0 || strcmp(t->name, "ffn_moe_out") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_shexp.weight") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_glm_shexp_component : public test_case { ++ const int component; ++ const ggml_type gate_type; ++ const ggml_type up_type; ++ const ggml_type down_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const int64_t n_tokens; ++ ++ test_glm_shexp_component( ++ int component, ++ ggml_type gate_type = GGML_TYPE_Q4_K, ++ ggml_type up_type = GGML_TYPE_Q4_K, ++ ggml_type down_type = GGML_TYPE_Q4_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048, ++ int64_t n_tokens = 1) ++ : component(component), ++ gate_type(gate_type), ++ up_type(up_type), ++ down_type(down_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ n_tokens(n_tokens) { ++ GGML_ASSERT(component >= 0 && component <= 6); ++ GGML_ASSERT(n_embd % ggml_blck_size(gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(down_type) == 0); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR7(component, gate_type, up_type, down_type, n_embd, n_ff, n_tokens); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_SHEXP_COMPONENT"; ++ } ++ ++ bool perf_runs_whole_graph() override { ++ return component == 2 || component == 6; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ switch (component) { ++ case 0: return build_gate_matmul(ctx); ++ case 1: return build_up_matmul(ctx); ++ case 2: return build_gate_up(ctx); ++ case 3: return build_down_matmul(ctx); ++ case 4: return build_output_add(ctx); ++ case 5: return build_swiglu(ctx); ++ case 6: return build_merged_gate_up_swiglu(ctx); ++ default: GGML_ABORT("invalid GLM shared expert component"); ++ } ++ } ++ ++ ggml_tensor * build_gate_matmul(ggml_context * ctx) { ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_shexp_cur"); ++ ++ ggml_tensor * gates = ggml_new_tensor_2d(ctx, gate_type, n_embd, n_ff); ++ ggml_set_name(gates, "ffn_gate_shexp.weight"); ++ ++ ggml_tensor * gate = ggml_mul_mat(ctx, gates, cur); ++ ggml_set_name(gate, "ffn_shexp_gate"); ++ return gate; ++ } ++ ++ ggml_tensor * build_up_matmul(ggml_context * ctx) { ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_shexp_cur"); ++ ++ ggml_tensor * ups = ggml_new_tensor_2d(ctx, up_type, n_embd, n_ff); ++ ggml_set_name(ups, "ffn_up_shexp.weight"); ++ ++ ggml_tensor * up = ggml_mul_mat(ctx, ups, cur); ++ ggml_set_name(up, "ffn_shexp_up"); ++ return up; ++ } ++ ++ ggml_tensor * build_gate_up(ggml_context * ctx) { ++ ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, 1, n_tokens); ++ ggml_set_name(gate, "ffn_shexp_gate"); ++ ++ ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, 1, n_tokens); ++ ggml_set_name(up, "ffn_shexp_up"); ++ ++ ggml_tensor * act = ggml_silu(ctx, gate); ++ ggml_set_name(act, "ffn_shexp_gate_silu"); ++ ++ act = ggml_mul(ctx, act, up); ++ ggml_set_name(act, "ffn_shexp_gate_up"); ++ return act; ++ } ++ ++ ggml_tensor * build_down_matmul(ggml_context * ctx) { ++ ggml_tensor * act = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, 1, n_tokens); ++ ggml_set_name(act, "ffn_shexp_gate_up"); ++ ++ ggml_tensor * downs = ggml_new_tensor_2d(ctx, down_type, n_ff, n_embd); ++ ggml_set_name(downs, "ffn_down_shexp.weight"); ++ ++ ggml_tensor * shexp = ggml_mul_mat(ctx, downs, act); ++ ggml_set_name(shexp, "ffn_shexp"); ++ return shexp; ++ } ++ ++ ggml_tensor * build_output_add(ggml_context * ctx) { ++ ggml_tensor * moe_out = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(moe_out, "ffn_moe_out"); ++ ++ ggml_tensor * shexp = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(shexp, "ffn_shexp"); ++ ++ ggml_tensor * out = ggml_add(ctx, moe_out, shexp); ++ ggml_set_name(out, "ffn_out"); ++ return out; ++ } ++ ++ ggml_tensor * build_swiglu(ggml_context * ctx) { ++ ggml_tensor * gate = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, 1, n_tokens); ++ ggml_set_name(gate, "ffn_shexp_gate"); ++ ++ ggml_tensor * up = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_ff, 1, n_tokens); ++ ggml_set_name(up, "ffn_shexp_up"); ++ ++ ggml_tensor * act = ggml_swiglu_split(ctx, gate, up); ++ ggml_set_name(act, "ffn_shexp_swiglu"); ++ return act; ++ } ++ ++ ggml_tensor * build_merged_gate_up_swiglu(ggml_context * ctx) { ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_shexp_cur"); ++ ++ ggml_tensor * gate_ups = ggml_new_tensor_2d(ctx, gate_type, n_embd, 2*n_ff); ++ ggml_set_name(gate_ups, "ffn_gate_up_shexp.weight"); ++ ++ ggml_tensor * gate_up = ggml_mul_mat(ctx, gate_ups, cur); ++ ggml_set_name(gate_up, "ffn_shexp_gate_up_merged"); ++ ++ ggml_tensor * act = ggml_swiglu(ctx, gate_up); ++ ggml_set_name(act, "ffn_shexp_swiglu"); ++ return act; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_shexp_cur") == 0 || ++ strcmp(t->name, "ffn_shexp_gate") == 0 || ++ strcmp(t->name, "ffn_shexp_up") == 0 || ++ strcmp(t->name, "ffn_shexp_gate_up") == 0 || ++ strcmp(t->name, "ffn_shexp_gate_up_merged") == 0 || ++ strcmp(t->name, "ffn_shexp_swiglu") == 0 || ++ strcmp(t->name, "ffn_shexp") == 0 || ++ strcmp(t->name, "ffn_moe_out") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_shexp.weight") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_glm_ffn_consumer : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ ++ const ggml_type routed_gate_type; ++ const ggml_type routed_up_type; ++ const ggml_type routed_down_type; ++ const ggml_type shexp_gate_type; ++ const ggml_type shexp_up_type; ++ const ggml_type shexp_down_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const bool fused_glu; ++ const bool weighted_down; ++ const bool merged_shexp_gate_up; ++ ++ struct route_result { ++ ggml_tensor * selected_experts; ++ ggml_tensor * weights; ++ }; ++ ++ ggml_tensor * selected_experts = nullptr; ++ ggml_tensor * weights = nullptr; ++ ggml_tensor * routed_out = nullptr; ++ ggml_tensor * out = nullptr; ++ ++ test_glm_ffn_consumer( ++ ggml_type routed_gate_type = GGML_TYPE_Q2_K, ++ ggml_type routed_up_type = GGML_TYPE_Q2_K, ++ ggml_type routed_down_type = GGML_TYPE_Q3_K, ++ ggml_type shexp_gate_type = GGML_TYPE_Q4_K, ++ ggml_type shexp_up_type = GGML_TYPE_Q4_K, ++ ggml_type shexp_down_type = GGML_TYPE_Q4_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048, ++ bool fused_glu = true, ++ bool weighted_down = false, ++ bool merged_shexp_gate_up = false) ++ : routed_gate_type(routed_gate_type), ++ routed_up_type(routed_up_type), ++ routed_down_type(routed_down_type), ++ shexp_gate_type(shexp_gate_type), ++ shexp_up_type(shexp_up_type), ++ shexp_down_type(shexp_down_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ fused_glu(fused_glu), ++ weighted_down(weighted_down), ++ merged_shexp_gate_up(merged_shexp_gate_up) { ++ GGML_ASSERT(n_embd % ggml_blck_size(routed_gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(routed_up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(routed_down_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(shexp_gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(shexp_up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(shexp_down_type) == 0); ++ GGML_ASSERT(!merged_shexp_gate_up || fused_glu); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR14(routed_gate_type, routed_up_type, routed_down_type, shexp_gate_type, shexp_up_type, ++ shexp_down_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff, fused_glu, weighted_down, ++ merged_shexp_gate_up); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_FFN_CONSUMER"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ route_result build_glm_route(ggml_context * ctx) { ++ ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_set_name(logits, "ffn_moe_logits"); ++ ++ ggml_tensor * probs = ggml_sigmoid(ctx, logits); ++ ggml_set_name(probs, "ffn_moe_probs"); ++ ++ ggml_tensor * exp_probs_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_expert); ++ ggml_set_name(exp_probs_b, "exp_probs_b"); ++ ++ ggml_tensor * selection_probs = ggml_add(ctx, probs, exp_probs_b); ++ ggml_set_name(selection_probs, "ffn_moe_probs_biased"); ++ ++ selected_experts = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); ++ ggml_set_name(selected_experts, "ffn_moe_topk"); ++ ++ weights = ggml_get_rows(ctx, ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), selected_experts); ++ ggml_set_name(weights, "ffn_moe_weights"); ++ ++ weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); ++ ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum"); ++ ++ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); ++ ggml_set_name(weights_sum, "ffn_moe_weights_sum_clamped"); ++ ++ weights = ggml_div(ctx, weights, weights_sum); ++ ggml_set_name(weights, "ffn_moe_weights_norm"); ++ ++ weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); ++ weights = ggml_scale(ctx, weights, 2.0f); ++ ggml_set_name(weights, "ffn_moe_weights_scaled"); ++ ++ return { selected_experts, weights }; ++ } ++ ++ ggml_tensor * build_routed_moe(ggml_context * ctx, ggml_tensor * cur) { ++ ggml_tensor * route_anchor = ggml_sum(ctx, weights); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_sum"); ++ ++ route_anchor = ggml_scale(ctx, route_anchor, 0.0f); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_zero"); ++ ++ route_anchor = ggml_repeat(ctx, route_anchor, cur); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_repeat"); ++ ++ cur = ggml_add(ctx, cur, route_anchor); ++ ggml_set_name(cur, "ffn_moe_cur_route_anchored"); ++ ++ ggml_tensor * gates = ggml_new_tensor_3d(ctx, routed_gate_type, n_embd, n_ff, n_expert); ++ ggml_set_name(gates, "ffn_gate_exps.weight"); ++ ++ ggml_tensor * ups = ggml_new_tensor_3d(ctx, routed_up_type, n_embd, n_ff, n_expert); ++ ggml_set_name(ups, "ffn_up_exps.weight"); ++ ++ ggml_tensor * downs = ggml_new_tensor_3d(ctx, routed_down_type, n_ff, n_embd, n_expert); ++ ggml_set_name(downs, "ffn_down_exps.weight"); ++ ++ ggml_tensor * ffn_gate = ggml_mul_mat_id(ctx, gates, cur, selected_experts); ++ ggml_set_name(ffn_gate, "ffn_moe_gate"); ++ ++ ggml_tensor * ffn_up = ggml_mul_mat_id(ctx, ups, cur, selected_experts); ++ ggml_set_name(ffn_up, "ffn_moe_up"); ++ ++ ggml_tensor * ffn_act = nullptr; ++ if (fused_glu) { ++ ffn_act = ggml_swiglu_split(ctx, ffn_gate, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_swiglu"); ++ } else { ++ ffn_act = ggml_silu(ctx, ffn_gate); ++ ggml_set_name(ffn_act, "ffn_moe_gate_silu"); ++ ++ ffn_act = ggml_mul(ctx, ffn_act, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_gate_up"); ++ } ++ ++ if (weighted_down) { ++ ffn_act = ggml_mul(ctx, ffn_act, weights); ++ ggml_set_name(ffn_act, "ffn_moe_down_weighted_input"); ++ } ++ ++ const char * native_down = getenv("GGML_TEST_GLM_FFN_NATIVE_DOWN"); ++ if (!weighted_down && native_down != nullptr && atoi(native_down) != 0) { ++ routed_out = ggml_moe_mul_mat_id(ctx, downs, ffn_act, selected_experts, weights); ++ ggml_set_name(routed_out, "ffn_moe_out"); ++ return routed_out; ++ } ++ ++ ggml_tensor * ffn_down = ggml_mul_mat_id(ctx, downs, ffn_act, selected_experts); ++ ggml_set_name(ffn_down, "ffn_moe_down"); ++ ++ routed_out = ggml_moe_weighted_sum(ctx, ffn_down, weights); ++ if (weighted_down) { ++ routed_out->op_params[0] = 1; ++ } ++ ggml_set_name(routed_out, "ffn_moe_out"); ++ ++ return routed_out; ++ } ++ ++ ggml_tensor * build_shexp(ggml_context * ctx, ggml_tensor * cur) { ++ ggml_tensor * downs = ggml_new_tensor_2d(ctx, shexp_down_type, n_ff, n_embd); ++ ggml_set_name(downs, "ffn_down_shexp.weight"); ++ ++ ggml_tensor * act = nullptr; ++ if (merged_shexp_gate_up) { ++ ggml_tensor * gate_ups = ggml_new_tensor_2d(ctx, shexp_gate_type, n_embd, 2*n_ff); ++ ggml_set_name(gate_ups, "ffn_gate_up_shexp.weight"); ++ ++ ggml_tensor * gate_up = ggml_mul_mat(ctx, gate_ups, cur); ++ ggml_set_name(gate_up, "ffn_shexp_gate_up_merged"); ++ ++ act = ggml_swiglu(ctx, gate_up); ++ ggml_set_name(act, "ffn_shexp_swiglu"); ++ } else { ++ ggml_tensor * gates = ggml_new_tensor_2d(ctx, shexp_gate_type, n_embd, n_ff); ++ ggml_set_name(gates, "ffn_gate_shexp.weight"); ++ ++ ggml_tensor * ups = ggml_new_tensor_2d(ctx, shexp_up_type, n_embd, n_ff); ++ ggml_set_name(ups, "ffn_up_shexp.weight"); ++ ++ ggml_tensor * gate = ggml_mul_mat(ctx, gates, cur); ++ ggml_set_name(gate, "ffn_shexp_gate"); ++ ++ ggml_tensor * up = ggml_mul_mat(ctx, ups, cur); ++ ggml_set_name(up, "ffn_shexp_up"); ++ ++ if (fused_glu) { ++ act = ggml_swiglu_split(ctx, gate, up); ++ ggml_set_name(act, "ffn_shexp_swiglu"); ++ } else { ++ act = ggml_silu(ctx, gate); ++ ggml_set_name(act, "ffn_shexp_gate_silu"); ++ ++ act = ggml_mul(ctx, act, up); ++ ggml_set_name(act, "ffn_shexp_gate_up"); ++ } ++ } ++ ++ ggml_tensor * shexp = ggml_mul_mat(ctx, downs, act); ++ ggml_set_name(shexp, "ffn_shexp"); ++ return shexp; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ build_glm_route(ctx); ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_cur"); ++ ++ ggml_tensor * moe_out = build_routed_moe(ctx, cur); ++ ggml_tensor * shexp = build_shexp(ctx, cur); ++ ++ out = ggml_add(ctx, moe_out, shexp); ++ ggml_set_name(out, "ffn_out"); ++ ++ return out; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ return { selected_experts, weights, routed_out, out }; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "ffn_moe_logits") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = -3.0f + 0.03125f*float(i); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "exp_probs_b") == 0) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = 0.0001f*float((i*37) % 19); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "ffn_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_exps.weight") != nullptr || strstr(t->name, "_shexp.weight") != nullptr) { ++ init_tensor_uniform(t, -0.05f, 0.05f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_glm_ffn_two_phase_consumer : public test_glm_ffn_consumer { ++ const bool two_phase; ++ const char * variant; ++ ++ test_glm_ffn_two_phase_consumer(bool two_phase, const char * variant) ++ : test_glm_ffn_consumer(), two_phase(two_phase), variant(variant) {} ++ ++ std::string vars() override { ++ return test_glm_ffn_consumer::vars() + ",two_phase=" + ++ std::to_string(two_phase ? 1 : 0) + ",variant=" + variant; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return std::string("GLM_FFN_TWO_PHASE_") + variant; ++ } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env motif_reference_env("GGML_METAL_ENABLE_GLM_MOE_DECODE_MOTIF_REFERENCE", "0"); ++ scoped_test_env two_phase_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_MOE_TWO_PHASE", two_phase ? "1" : "0"); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++}; ++ ++struct test_glm_ffn_dual_lane_consumer : public test_glm_ffn_consumer { ++ const bool dual_lane; ++ const char * variant; ++ ++ test_glm_ffn_dual_lane_consumer(bool dual_lane, const char * variant) ++ : test_glm_ffn_consumer(), dual_lane(dual_lane), variant(variant) {} ++ ++ std::string vars() override { ++ return test_glm_ffn_consumer::vars() + ",dual_lane=" + ++ std::to_string(dual_lane ? 1 : 0) + ",variant=" + variant; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return std::string("GLM_FFN_DUAL_LANE_") + variant; ++ } ++ ++ std::vector fusion_test_nodes() override { ++ return { out }; ++ } ++ ++ test_status_t eval( ++ ggml_backend_t backend1, ++ ggml_backend_t backend2, ++ const char * op_names_filter, ++ printer * output_printer) override { ++ scoped_test_env motif_reference_env("GGML_METAL_ENABLE_GLM_MOE_DECODE_MOTIF_REFERENCE", "0"); ++ scoped_test_env two_phase_env("GGML_METAL_EXPERIMENTAL_GLM_MOE_TWO_PHASE", "0"); ++ scoped_test_env dual_lane_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_MOE_DUAL_LANE", dual_lane ? "1" : "0"); ++ return test_case::eval(backend1, backend2, op_names_filter, output_printer); ++ } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env motif_reference_env("GGML_METAL_ENABLE_GLM_MOE_DECODE_MOTIF_REFERENCE", "0"); ++ scoped_test_env two_phase_env("GGML_METAL_EXPERIMENTAL_GLM_MOE_TWO_PHASE", "0"); ++ scoped_test_env dual_lane_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_MOE_DUAL_LANE", dual_lane ? "1" : "0"); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++}; ++ ++struct test_glm_ffn_native_down_consumer : public test_glm_ffn_consumer { ++ const bool native_down; ++ const char * variant; ++ ++ test_glm_ffn_native_down_consumer(bool native_down, const char * variant) ++ : test_glm_ffn_consumer(), native_down(native_down), variant(variant) {} ++ ++ std::string vars() override { ++ return test_glm_ffn_consumer::vars() + ",native_down=" + ++ std::to_string(native_down ? 1 : 0) + ",variant=" + variant; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_FFN_NATIVE_DOWN_AB"; ++ } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env native_down_env("GGML_TEST_GLM_FFN_NATIVE_DOWN", native_down ? "1" : "0"); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++}; ++ ++struct test_glm_ffn_chain : public test_case { ++ static constexpr int64_t n_expert = 256; ++ static constexpr int64_t n_expert_used = 8; ++ static constexpr int64_t n_tokens = 1; ++ ++ const int n_repeats; ++ const ggml_type routed_gate_type; ++ const ggml_type routed_up_type; ++ const ggml_type routed_down_type; ++ const ggml_type shexp_gate_type; ++ const ggml_type shexp_up_type; ++ const ggml_type shexp_down_type; ++ const int64_t n_embd; ++ const int64_t n_ff; ++ const bool fused_glu; ++ const bool weighted_down; ++ const bool merged_shexp_gate_up; ++ const bool distinct_routes; ++ ++ struct route_result { ++ ggml_tensor * selected_experts; ++ ggml_tensor * weights; ++ }; ++ ++ test_glm_ffn_chain( ++ int n_repeats = 4, ++ ggml_type routed_gate_type = GGML_TYPE_Q2_K, ++ ggml_type routed_up_type = GGML_TYPE_Q2_K, ++ ggml_type routed_down_type = GGML_TYPE_Q3_K, ++ ggml_type shexp_gate_type = GGML_TYPE_Q4_K, ++ ggml_type shexp_up_type = GGML_TYPE_Q4_K, ++ ggml_type shexp_down_type = GGML_TYPE_Q4_K, ++ int64_t n_embd = 6144, ++ int64_t n_ff = 2048, ++ bool fused_glu = true, ++ bool weighted_down = false, ++ bool merged_shexp_gate_up = false, ++ bool distinct_routes = false) ++ : n_repeats(n_repeats), ++ routed_gate_type(routed_gate_type), ++ routed_up_type(routed_up_type), ++ routed_down_type(routed_down_type), ++ shexp_gate_type(shexp_gate_type), ++ shexp_up_type(shexp_up_type), ++ shexp_down_type(shexp_down_type), ++ n_embd(n_embd), ++ n_ff(n_ff), ++ fused_glu(fused_glu), ++ weighted_down(weighted_down), ++ merged_shexp_gate_up(merged_shexp_gate_up), ++ distinct_routes(distinct_routes) { ++ GGML_ASSERT(n_repeats > 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(routed_gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(routed_up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(routed_down_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(shexp_gate_type) == 0); ++ GGML_ASSERT(n_embd % ggml_blck_size(shexp_up_type) == 0); ++ GGML_ASSERT(n_ff % ggml_blck_size(shexp_down_type) == 0); ++ GGML_ASSERT(!merged_shexp_gate_up || fused_glu); ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR15(n_repeats, routed_gate_type, routed_up_type, routed_down_type, shexp_gate_type, ++ shexp_up_type, shexp_down_type, n_expert, n_expert_used, n_tokens, n_embd, n_ff, fused_glu, ++ weighted_down, merged_shexp_gate_up) + ",distinct_routes=" + ++ std::to_string(distinct_routes ? 1 : 0); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_FFN_CHAIN"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool perf_runs_whole_graph() override { return true; } ++ ++ route_result build_glm_route(ggml_context * ctx, int repeat) { ++ ggml_tensor * logits = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_expert, n_tokens); ++ ggml_format_name(logits, "ffn_moe_logits_%d", repeat); ++ ++ ggml_tensor * probs = ggml_sigmoid(ctx, logits); ++ ggml_format_name(probs, "ffn_moe_probs_%d", repeat); ++ ++ ggml_tensor * exp_probs_b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_expert); ++ ggml_format_name(exp_probs_b, "exp_probs_b_%d", repeat); ++ ++ ggml_tensor * selection_probs = ggml_add(ctx, probs, exp_probs_b); ++ ggml_format_name(selection_probs, "ffn_moe_probs_biased_%d", repeat); ++ ++ ggml_tensor * selected_experts = ggml_argsort_top_k(ctx, selection_probs, n_expert_used); ++ ggml_format_name(selected_experts, "ffn_moe_topk_%d", repeat); ++ ++ ggml_tensor * weights = ggml_get_rows(ctx, ggml_reshape_3d(ctx, probs, 1, n_expert, n_tokens), selected_experts); ++ ggml_format_name(weights, "ffn_moe_weights_%d", repeat); ++ ++ weights = ggml_reshape_2d(ctx, weights, n_expert_used, n_tokens); ++ ggml_tensor * weights_sum = ggml_sum_rows(ctx, weights); ++ ggml_format_name(weights_sum, "ffn_moe_weights_sum_%d", repeat); ++ ++ weights_sum = ggml_clamp(ctx, weights_sum, 6.103515625e-5, INFINITY); ++ ggml_format_name(weights_sum, "ffn_moe_weights_sum_clamped_%d", repeat); ++ ++ weights = ggml_div(ctx, weights, weights_sum); ++ ggml_format_name(weights, "ffn_moe_weights_norm_%d", repeat); ++ ++ weights = ggml_reshape_3d(ctx, weights, 1, n_expert_used, n_tokens); ++ weights = ggml_scale(ctx, weights, 2.0f); ++ ggml_format_name(weights, "ffn_moe_weights_scaled_%d", repeat); ++ ++ return { selected_experts, weights }; ++ } ++ ++ ggml_tensor * build_routed_moe( ++ ggml_context * ctx, ++ ggml_tensor * cur, ++ ggml_tensor * routed_gates, ++ ggml_tensor * routed_ups, ++ ggml_tensor * routed_downs, ++ ggml_tensor * selected_experts, ++ ggml_tensor * weights) { ++ ggml_tensor * route_anchor = ggml_sum(ctx, weights); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_sum"); ++ ++ route_anchor = ggml_scale(ctx, route_anchor, 0.0f); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_zero"); ++ ++ route_anchor = ggml_repeat(ctx, route_anchor, cur); ++ ggml_set_name(route_anchor, "ffn_moe_route_anchor_repeat"); ++ ++ cur = ggml_add(ctx, cur, route_anchor); ++ ggml_set_name(cur, "ffn_moe_cur_route_anchored"); ++ ++ ggml_tensor * ffn_gate = ggml_mul_mat_id(ctx, routed_gates, cur, selected_experts); ++ ggml_set_name(ffn_gate, "ffn_moe_gate"); ++ ++ ggml_tensor * ffn_up = ggml_mul_mat_id(ctx, routed_ups, cur, selected_experts); ++ ggml_set_name(ffn_up, "ffn_moe_up"); ++ ++ ggml_tensor * ffn_act = nullptr; ++ if (fused_glu) { ++ ffn_act = ggml_swiglu_split(ctx, ffn_gate, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_swiglu"); ++ } else { ++ ffn_act = ggml_silu(ctx, ffn_gate); ++ ggml_set_name(ffn_act, "ffn_moe_gate_silu"); ++ ++ ffn_act = ggml_mul(ctx, ffn_act, ffn_up); ++ ggml_set_name(ffn_act, "ffn_moe_gate_up"); ++ } ++ ++ if (weighted_down) { ++ ffn_act = ggml_mul(ctx, ffn_act, weights); ++ ggml_set_name(ffn_act, "ffn_moe_down_weighted_input"); ++ } ++ ++ const char * native_down = getenv("GGML_TEST_GLM_FFN_CHAIN_NATIVE_DOWN"); ++ if (!weighted_down && native_down != nullptr && atoi(native_down) != 0) { ++ ggml_tensor * routed_out = ++ ggml_moe_mul_mat_id(ctx, routed_downs, ffn_act, selected_experts, weights); ++ ggml_set_name(routed_out, "ffn_moe_out"); ++ return routed_out; ++ } ++ ++ ggml_tensor * ffn_down = ggml_mul_mat_id(ctx, routed_downs, ffn_act, selected_experts); ++ ggml_set_name(ffn_down, "ffn_moe_down"); ++ ++ ggml_tensor * routed_out = ggml_moe_weighted_sum(ctx, ffn_down, weights); ++ if (weighted_down) { ++ routed_out->op_params[0] = 1; ++ } ++ ggml_set_name(routed_out, "ffn_moe_out"); ++ ++ return routed_out; ++ } ++ ++ ggml_tensor * build_shexp( ++ ggml_context * ctx, ++ ggml_tensor * cur, ++ ggml_tensor * shexp_gates, ++ ggml_tensor * shexp_ups, ++ ggml_tensor * shexp_downs, ++ ggml_tensor * shexp_gate_ups) { ++ ggml_tensor * act = nullptr; ++ if (merged_shexp_gate_up) { ++ ggml_tensor * gate_up = ggml_mul_mat(ctx, shexp_gate_ups, cur); ++ ggml_set_name(gate_up, "ffn_shexp_gate_up_merged"); ++ ++ act = ggml_swiglu(ctx, gate_up); ++ ggml_set_name(act, "ffn_shexp_swiglu"); ++ } else { ++ ggml_tensor * gate = ggml_mul_mat(ctx, shexp_gates, cur); ++ ggml_set_name(gate, "ffn_shexp_gate"); ++ ++ ggml_tensor * up = ggml_mul_mat(ctx, shexp_ups, cur); ++ ggml_set_name(up, "ffn_shexp_up"); ++ ++ if (fused_glu) { ++ act = ggml_swiglu_split(ctx, gate, up); ++ ggml_set_name(act, "ffn_shexp_swiglu"); ++ } else { ++ act = ggml_silu(ctx, gate); ++ ggml_set_name(act, "ffn_shexp_gate_silu"); ++ ++ act = ggml_mul(ctx, act, up); ++ ggml_set_name(act, "ffn_shexp_gate_up"); ++ } ++ } ++ ++ ggml_tensor * shexp = ggml_mul_mat(ctx, shexp_downs, act); ++ ggml_set_name(shexp, "ffn_shexp"); ++ return shexp; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * routed_gates = ggml_new_tensor_3d(ctx, routed_gate_type, n_embd, n_ff, n_expert); ++ ggml_set_name(routed_gates, "ffn_gate_exps.weight"); ++ ++ ggml_tensor * routed_ups = ggml_new_tensor_3d(ctx, routed_up_type, n_embd, n_ff, n_expert); ++ ggml_set_name(routed_ups, "ffn_up_exps.weight"); ++ ++ ggml_tensor * routed_downs = ggml_new_tensor_3d(ctx, routed_down_type, n_ff, n_embd, n_expert); ++ ggml_set_name(routed_downs, "ffn_down_exps.weight"); ++ ++ ggml_tensor * cur = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd, 1, n_tokens); ++ ggml_set_name(cur, "ffn_chain_cur"); ++ ++ for (int i = 0; i < n_repeats; ++i) { ++ ggml_tensor * shexp_downs = ggml_new_tensor_2d(ctx, shexp_down_type, n_ff, n_embd); ++ ggml_format_name(shexp_downs, "ffn_down_shexp.weight.%d", i); ++ ++ ggml_tensor * shexp_gates = nullptr; ++ ggml_tensor * shexp_ups = nullptr; ++ ggml_tensor * shexp_gate_ups = nullptr; ++ if (merged_shexp_gate_up) { ++ shexp_gate_ups = ggml_new_tensor_2d(ctx, shexp_gate_type, n_embd, 2*n_ff); ++ ggml_format_name(shexp_gate_ups, "ffn_gate_up_shexp.weight.%d", i); ++ } else { ++ shexp_gates = ggml_new_tensor_2d(ctx, shexp_gate_type, n_embd, n_ff); ++ ggml_format_name(shexp_gates, "ffn_gate_shexp.weight.%d", i); ++ ++ shexp_ups = ggml_new_tensor_2d(ctx, shexp_up_type, n_embd, n_ff); ++ ggml_format_name(shexp_ups, "ffn_up_shexp.weight.%d", i); ++ } ++ ++ route_result route = build_glm_route(ctx, i); ++ ++ ggml_tensor * moe_out = build_routed_moe(ctx, cur, routed_gates, routed_ups, routed_downs, route.selected_experts, route.weights); ++ ggml_tensor * shexp = build_shexp(ctx, cur, shexp_gates, shexp_ups, shexp_downs, shexp_gate_ups); ++ ++ cur = ggml_add(ctx, moe_out, shexp); ++ ggml_set_name(cur, "ffn_chain_out"); ++ } ++ ++ return cur; ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strstr(t->name, "ffn_moe_logits_") == t->name) { ++ std::vector data(n_expert, -8.0f); ++ if (distinct_routes) { ++ const int repeat = atoi(t->name + strlen("ffn_moe_logits_")); ++ const int first_expert = (repeat*32) % n_expert; ++ for (int64_t slot = 0; slot < n_expert_used; ++slot) { ++ data[first_expert + slot] = 2.0f + 0.1f*float(slot); ++ } ++ } else { ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = -3.0f + 0.03125f*float(i); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strstr(t->name, "exp_probs_b_") == t->name) { ++ std::vector data(n_expert); ++ for (int64_t i = 0; i < n_expert; ++i) { ++ data[i] = 0.0001f*float((i*37) % 19); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, ggml_nbytes(t)); ++ } else if (strcmp(t->name, "ffn_chain_cur") == 0) { ++ init_tensor_uniform(t, -0.5f, 0.5f); ++ } else if (strstr(t->name, "_exps.weight") != nullptr || strstr(t->name, "_shexp.weight") != nullptr) { ++ init_tensor_uniform(t, -0.005f, 0.005f); ++ } else { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++}; ++ ++struct test_glm_ffn_chain_native_down : public test_glm_ffn_chain { ++ const bool native_down; ++ const char * variant; ++ ++ test_glm_ffn_chain_native_down(bool native_down, const char * variant) ++ : test_glm_ffn_chain(8), native_down(native_down), variant(variant) {} ++ ++ std::string vars() override { ++ return test_glm_ffn_chain::vars() + ",native_down=" + ++ std::to_string(native_down ? 1 : 0) + ",variant=" + variant; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_FFN_CHAIN_NATIVE_DOWN_AB"; ++ } ++ ++ test_status_t eval( ++ ggml_backend_t backend1, ++ ggml_backend_t backend2, ++ const char * op_names_filter, ++ printer * output_printer) override { ++ scoped_test_env native_down_env( ++ "GGML_TEST_GLM_FFN_CHAIN_NATIVE_DOWN", native_down ? "1" : "0"); ++ return test_case::eval(backend1, backend2, op_names_filter, output_printer); ++ } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env native_down_env( ++ "GGML_TEST_GLM_FFN_CHAIN_NATIVE_DOWN", native_down ? "1" : "0"); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++}; ++ ++struct test_glm_ffn_chain_two_phase : public test_glm_ffn_chain { ++ const bool two_phase; ++ const char * variant; ++ ++ test_glm_ffn_chain_two_phase(bool two_phase, const char * variant) ++ : test_glm_ffn_chain( ++ 8, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q2_K, ++ GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, ++ GGML_TYPE_Q4_K, ++ GGML_TYPE_Q4_K, ++ 6144, ++ 2048, ++ true, ++ false, ++ false, ++ true), ++ two_phase(two_phase), ++ variant(variant) {} ++ ++ std::string vars() override { ++ return test_glm_ffn_chain::vars() + ",two_phase=" + ++ std::to_string(two_phase ? 1 : 0) + ",variant=" + variant; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_FFN_CHAIN_TWO_PHASE"; ++ } ++ ++ test_status_t eval( ++ ggml_backend_t backend1, ++ ggml_backend_t backend2, ++ const char * op_names_filter, ++ printer * output_printer) override { ++ scoped_test_env two_phase_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_MOE_TWO_PHASE", two_phase ? "1" : "0"); ++ return test_case::eval(backend1, backend2, op_names_filter, output_printer); ++ } ++ ++ bool eval_perf(ggml_backend_t backend, const char * op_names_filter, printer * output_printer) override { ++ scoped_test_env two_phase_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_MOE_TWO_PHASE", two_phase ? "1" : "0"); ++ return test_case::eval_perf(backend, op_names_filter, output_printer); ++ } ++}; ++ ++struct test_mul_mat_vec_fusion : public test_case { ++ const ggml_type type; ++ const ggml_glu_op glu_op; ++ const int64_t m; ++ const int64_t n; ++ const int64_t k; ++ const bool use_id; ++ const int n_mats; ++ const int n_used; ++ const bool b; // broadcast b matrix (only for use_id) ++ const bool with_bias; ++ const bool with_gate; ++ const bool with_lane_scale; ++ std::array batch_dims; ++ ++ test_mul_mat_vec_fusion(ggml_type type, ggml_glu_op op, int64_t m, int64_t n, int64_t k, ++ bool use_id = false, int n_mats = 1, int n_used = 1, bool b = false, bool with_bias = false, bool with_gate = true, ++ bool with_lane_scale = false, std::array batch_dims = {4, 2}) ++ : type(type), glu_op(op), m(m), n(n), k(k), use_id(use_id), n_mats(n_mats), n_used(n_used), b(b), with_bias(with_bias), ++ with_gate(with_gate), with_lane_scale(with_lane_scale), batch_dims(batch_dims) { ++ if (use_id) { ++ GGML_ASSERT(n_used <= n_mats); ++ } ++ } ++ ++ std::string vars() override { ++ return VARS_TO_STR13(type, glu_op, m, n, k, use_id, n_mats, n_used, b, with_bias, with_gate, with_lane_scale, batch_dims); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "MUL_MAT_VEC_FUSION"; ++ } ++ ++ bool run_whole_graph() override { return true; } ++ bool use_weight_context() override { return use_id && with_lane_scale; } ++ ++ ggml_tensor * build_gate(ggml_context * ctx, ggml_tensor * ffn_gate, ggml_tensor * ffn_up) { ++ ggml_tensor * out = nullptr; ++ if (with_gate) { ++ if (glu_op == GGML_GLU_OP_SWIGLU_OAI) { ++ constexpr float alpha = 1.702f; ++ constexpr float limit = 7.0f; ++ out = ggml_swiglu_oai(ctx, ffn_gate, ffn_up, alpha, limit); ++ } else { ++ out = ggml_glu_split(ctx, ffn_gate, ffn_up, glu_op); ++ } ++ } ++ return out; ++ } ++ ++ ggml_tensor * build_lane_scale_dense(ggml_context * ctx, ggml_tensor * out) { ++ ggml_tensor * scale = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); ++ return ggml_mul(ctx, out, scale); ++ } ++ ++ ggml_tensor * build_lane_scale_id(ggml_context * ctx, ggml_context * ctx_weights, ggml_tensor * out, ggml_tensor * ids) { ++ GGML_ASSERT(ctx_weights); ++ ggml_tensor * scale = ggml_new_tensor_1d(ctx_weights, GGML_TYPE_F32, n_mats); ++ ggml_tensor * s = ggml_reshape_3d(ctx, scale, 1, n_mats, 1); ++ s = ggml_repeat_4d(ctx, s, 1, n_mats, m, 1); ++ s = ggml_get_rows(ctx, s, ids); ++ return ggml_mul(ctx, out, s); ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ GGML_ASSERT(!use_weight_context()); ++ return build_graph(ctx, nullptr); ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx, ggml_context * ctx_weights) override { ++ if (!use_id) { ++ const int channels = batch_dims[0]; ++ const int samples = batch_dims[1]; ++ std::array ne = { k, m, channels, samples }; ++ std::array ne0 = { k, n, channels, samples }; ++ ++ ggml_tensor * cur = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne.data()); ++ ggml_tensor * gate = with_gate ? ggml_new_tensor(ctx, type, 4, ne0.data()) : nullptr; ++ ggml_tensor * up = ggml_new_tensor(ctx, type, 4, ne0.data()); ++ ++ auto build_lane_up = [&]() { ++ ggml_tensor * ffn_up = ggml_mul_mat(ctx, up, cur); ++ if (with_lane_scale) { ++ ffn_up = build_lane_scale_dense(ctx, ffn_up); ++ } ++ if (with_bias) { ++ std::array bias_ne = { ffn_up->ne[0], 1, channels, samples }; ++ ggml_tensor * up_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); ++ ffn_up = ggml_add(ctx, ffn_up, up_bias); ++ } ++ return ffn_up; ++ }; ++ ++ auto build_lane_gate = [&]() { ++ ggml_tensor * ffn_gate = ggml_mul_mat(ctx, gate, cur); ++ if (with_lane_scale) { ++ ffn_gate = build_lane_scale_dense(ctx, ffn_gate); ++ } ++ if (with_bias) { ++ std::array bias_ne = { ffn_gate->ne[0], 1, channels, samples }; ++ ggml_tensor * gate_bias = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias_ne.data()); ++ ffn_gate = ggml_add(ctx, ffn_gate, gate_bias); ++ } ++ return ffn_gate; ++ }; ++ ++ ggml_tensor * ffn_up = build_lane_up(); ++ ggml_tensor * ffn_gate = with_gate ? build_lane_gate() : nullptr; ++ ++ ggml_tensor * out = with_gate ? build_gate(ctx, ffn_gate, ffn_up) : ffn_up; ++ ++ std::array bias2_ne = { out->ne[0], 1, channels, samples }; ++ ggml_tensor * bias2 = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, bias2_ne.data()); + out = ggml_add(ctx, out, bias2); + + ggml_set_name(out, "out"); +@@ -6747,20 +10969,618 @@ struct test_timestep_embedding : public test_case { + return VARS_TO_STR4(type, ne_a, dim, max_period); + } + +- test_timestep_embedding(ggml_type type = GGML_TYPE_F32, +- std::array ne_a = {2, 1, 1, 1}, +- int dim = 320, int max_period=10000) +- : type(type), ne_a(ne_a), dim(dim), max_period(max_period) {} ++ test_timestep_embedding(ggml_type type = GGML_TYPE_F32, ++ std::array ne_a = {2, 1, 1, 1}, ++ int dim = 320, int max_period=10000) ++ : type(type), ne_a(ne_a), dim(dim), max_period(max_period) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); ++ ggml_set_name(a, "a"); ++ ++ ggml_tensor * out = ggml_timestep_embedding(ctx, a, dim, max_period); ++ ggml_set_name(out, "out"); ++ ++ return out; ++ } ++}; ++ ++// GGML_OP_DSA_SPARSE_ATTN ++struct test_dsa_sparse_attn : public test_case { ++ const ggml_type type_k; ++ const ggml_type type_v; ++ const ggml_type type_mask; ++ ++ const int64_t dk; ++ const int64_t dv; ++ const int64_t n_kv; ++ const int64_t n_batch; ++ const int64_t n_head; ++ const int64_t n_stream; ++ const int64_t n_top_k; ++ const int64_t n_top_stream; ++ const bool dense_equiv; ++ const bool graph_equiv; ++ const bool graph_expected; ++ const bool compact_equiv; ++ ++ std::string vars() override { ++ std::string base = VARS_TO_STR13(type_k, type_v, type_mask, dk, dv, n_kv, n_batch, n_head, n_stream, n_top_k, n_top_stream, dense_equiv, graph_equiv); ++ if (graph_expected) { ++ base += "," + VAR_TO_STR(graph_expected); ++ } ++ if (compact_equiv) { ++ base += "," + VAR_TO_STR(compact_equiv); ++ } ++ return base; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ if (((dense_equiv || graph_equiv || compact_equiv) && t->op == GGML_OP_SUB) || graph_expected) { ++ return ggml_op_name(GGML_OP_DSA_SPARSE_ATTN); ++ } ++ return ggml_op_desc(t); ++ } ++ ++ test_dsa_sparse_attn(ggml_type type_k = GGML_TYPE_F16, ggml_type type_v = GGML_TYPE_F16, ggml_type type_mask = GGML_TYPE_F16, ++ int64_t dk = 16, int64_t dv = 12, int64_t n_kv = 33, int64_t n_batch = 2, ++ int64_t n_head = 4, int64_t n_stream = 1, int64_t n_top_k = 4, int64_t n_top_stream = 1, bool dense_equiv = false, ++ bool graph_equiv = false, bool graph_expected = false, bool compact_equiv = false) ++ : type_k(type_k), type_v(type_v), type_mask(type_mask), dk(dk), dv(dv), n_kv(n_kv), ++ n_batch(n_batch), n_head(n_head), n_stream(n_stream), n_top_k(n_top_k), n_top_stream(n_top_stream), ++ dense_equiv(dense_equiv), graph_equiv(graph_equiv), graph_expected(graph_expected), compact_equiv(compact_equiv) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ if (compact_equiv) { ++ return build_compact_equiv_graph(ctx); ++ } ++ if (graph_equiv) { ++ return build_graph_equiv_graph(ctx); ++ } ++ if (graph_expected) { ++ return build_graph_sparse_expected(ctx); ++ } ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dk, n_batch, n_head, n_stream); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, type_k, dk, n_kv, 1, n_stream); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx, type_v, dv, n_kv, 1, n_stream); ++ ggml_tensor * kq_mask = ggml_new_tensor_4d(ctx, type_mask, 1, n_kv, n_batch, n_stream); ++ ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, n_batch, n_top_stream, 1); ++ ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "k"); ++ ggml_set_name(v, "v"); ++ ggml_set_name(kq_mask, "kq_mask"); ++ ggml_set_name(top_k, "top_k"); ++ ++ ggml_tensor * out = dense_equiv ++ ? build_dense_equiv_graph(ctx, q, k, v, kq_mask, top_k) ++ : ggml_dsa_sparse_attn(ctx, q, k, v, kq_mask, top_k, scale()); ++ ggml_set_name(out, "out"); ++ return out; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "q") == 0) { ++ if (graph_equiv || graph_expected) { ++ init_q_graph(t); ++ } else { ++ init_q(t); ++ } ++ } else if (strcmp(t->name, "k") == 0) { ++ init_k(t); ++ } else if (strcmp(t->name, "v") == 0) { ++ init_v(t); ++ } else if (strcmp(t->name, "kq_mask") == 0) { ++ init_kq_mask(t); ++ } else if (strcmp(t->name, "top_k") == 0) { ++ init_top_k(t); ++ } else if (strcmp(t->name, "v_mla") == 0) { ++ init_v_mla_identity(t); ++ } else if (!ggml_is_view_op(t->op)) { ++ init_tensor_uniform(t); ++ } ++ } ++ } ++ ++ double err(const float * a, const float * b, size_t n) override { ++ if (n != (size_t) (dv * n_batch * n_head * n_stream)) { ++ return 1.0; ++ } ++ ++ if (dense_equiv || graph_equiv || compact_equiv) { ++ double err = 0.0; ++ for (size_t i = 0; i < n; ++i) { ++ err = std::max(err, (double) std::abs(a[i])); ++ err = std::max(err, (double) std::abs(b[i])); ++ } ++ return err; ++ } ++ ++ if (graph_expected) { ++ double err = 0.0; ++ for (int64_t i_stream = 0; i_stream < n_stream; ++i_stream) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ for (int64_t i_head = 0; i_head < n_head; ++i_head) { ++ std::vector expected = expected_row(i_batch, i_head, i_stream); ++ for (int64_t i_dv = 0; i_dv < dv; ++i_dv) { ++ const size_t idx = ((i_stream*n_batch + i_batch)*(n_head*dv)) + i_head*dv + i_dv; ++ err = std::max(err, (double) std::abs(a[idx] - expected[i_dv])); ++ err = std::max(err, (double) std::abs(b[idx] - expected[i_dv])); ++ } ++ } ++ } ++ } ++ return err; ++ } ++ ++ double err = 0.0; ++ for (int64_t i_stream = 0; i_stream < n_stream; ++i_stream) { ++ for (int64_t i_head = 0; i_head < n_head; ++i_head) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ std::vector expected = expected_row(i_batch, i_head, i_stream); ++ for (int64_t i_dv = 0; i_dv < dv; ++i_dv) { ++ const size_t idx = ((i_stream*n_head + i_head)*n_batch + i_batch)*dv + i_dv; ++ err = std::max(err, (double) std::abs(a[idx] - expected[i_dv])); ++ err = std::max(err, (double) std::abs(b[idx] - expected[i_dv])); ++ } ++ } ++ } ++ } ++ return err; ++ } ++ ++ double max_err(ggml_backend_t backend) override { ++ (void) backend; ++ if (graph_expected) { ++ return type_k == GGML_TYPE_F16 || type_v == GGML_TYPE_F16 || type_mask == GGML_TYPE_F16 ? 2e-3 : 1e-5; ++ } ++ if (graph_equiv && type_k == GGML_TYPE_F32 && type_v == GGML_TYPE_F32 && type_mask == GGML_TYPE_F32 && dk >= 512) { ++ return 2e-4; ++ } ++ if (graph_equiv && (type_k == GGML_TYPE_F16 || type_v == GGML_TYPE_F16 || type_mask == GGML_TYPE_F16)) { ++ return 5e-3; ++ } ++ if (compact_equiv && (type_k == GGML_TYPE_F16 || type_v == GGML_TYPE_F16 || type_mask == GGML_TYPE_F16)) { ++ return 5e-3; ++ } ++ if (compact_equiv) { ++ return dk >= 512 ? 2e-4 : 1e-5; ++ } ++ return type_k == GGML_TYPE_F16 || type_v == GGML_TYPE_F16 || type_mask == GGML_TYPE_F16 ? 2e-3 : 1e-5; ++ } ++ ++ double max_nmse_err() override { ++ return type_k == GGML_TYPE_F16 || type_v == GGML_TYPE_F16 || type_mask == GGML_TYPE_F16 ? 2e-6 : 1e-7; ++ } ++ ++ uint64_t op_flops(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return 2ULL * (uint64_t) n_head * (uint64_t) n_stream * (uint64_t) n_batch * (uint64_t) n_top_k * (uint64_t) (dk + dv); ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } ++ ++private: ++ float scale() const { ++ return 1.0f / sqrtf((float) dk); ++ } ++ ++ ggml_tensor * build_graph_sparse_expected(ggml_context * ctx) { ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dk, n_head, n_batch*n_stream, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, type_k, dk, 1, n_kv, n_stream); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx, type_v, dv, 1, n_kv, n_stream); ++ ggml_tensor * kq_mask = ggml_new_tensor_4d(ctx, type_mask, 1, n_kv, n_batch, n_stream); ++ ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, n_batch, n_top_stream, 1); ++ ggml_tensor * v_mla = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dv, dv); ++ ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "k"); ++ ggml_set_name(v, "v"); ++ ggml_set_name(kq_mask, "kq_mask"); ++ ggml_set_name(top_k, "top_k"); ++ ggml_set_name(v_mla, "v_mla"); ++ ++ ggml_tensor * q_graph = ggml_view_4d(ctx, q, q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, ++ q->nb[1], q->nb[2], q->nb[3]/n_stream, 0); ++ q_graph = ggml_permute(ctx, q_graph, 0, 2, 1, 3); ++ ggml_tensor * k_graph = ggml_permute(ctx, k, 0, 2, 1, 3); ++ ggml_tensor * v_graph = ggml_permute(ctx, v, 0, 2, 1, 3); ++ ++ ggml_tensor * sparse = ggml_dsa_sparse_attn(ctx, q_graph, k_graph, v_graph, kq_mask, top_k, scale()); ++ sparse = ggml_mul_mat(ctx, v_mla, sparse); ++ sparse = ggml_permute(ctx, sparse, 0, 2, 1, 3); ++ sparse = ggml_cont_2d(ctx, sparse, sparse->ne[0]*sparse->ne[1], sparse->ne[2]*sparse->ne[3]); ++ ggml_set_name(sparse, "sparse_graph"); ++ ++ return sparse; ++ } ++ ++ ggml_tensor * build_graph_equiv_graph(ggml_context * ctx) { ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dk, n_head, n_batch*n_stream, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, type_k, dk, 1, n_kv, n_stream); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx, type_v, dv, 1, n_kv, n_stream); ++ ggml_tensor * kq_mask = ggml_new_tensor_4d(ctx, type_mask, 1, n_kv, n_batch, n_stream); ++ ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, n_batch, n_top_stream, 1); ++ ggml_tensor * v_mla = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dv, dv); ++ ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "k"); ++ ggml_set_name(v, "v"); ++ ggml_set_name(kq_mask, "kq_mask"); ++ ggml_set_name(top_k, "top_k"); ++ ggml_set_name(v_mla, "v_mla"); ++ ++ const bool v_trans = v->nb[1] > v->nb[2]; ++ ++ ggml_tensor * q_graph = ggml_view_4d(ctx, q, q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, ++ q->nb[1], q->nb[2], q->nb[3]/n_stream, 0); ++ q_graph = ggml_permute(ctx, q_graph, 0, 2, 1, 3); ++ ggml_tensor * k_graph = ggml_permute(ctx, k, 0, 2, 1, 3); ++ ggml_tensor * v_graph = ggml_permute(ctx, v, 0, 2, 1, 3); ++ ++ ggml_tensor * sparse = ggml_dsa_sparse_attn(ctx, q_graph, k_graph, v_graph, kq_mask, top_k, scale()); ++ sparse = ggml_mul_mat(ctx, v_mla, sparse); ++ sparse = ggml_permute(ctx, sparse, 0, 2, 1, 3); ++ sparse = ggml_cont_2d(ctx, sparse, sparse->ne[0]*sparse->ne[1], sparse->ne[2]*sparse->ne[3]); ++ ggml_set_name(sparse, "sparse_graph"); ++ ++ ggml_tensor * dense_mask = ggml_dsa_sparse_mask(ctx, kq_mask, top_k); ++ dense_mask = ggml_view_4d(ctx, dense_mask, ++ dense_mask->ne[1], dense_mask->ne[2], 1, dense_mask->ne[3], ++ dense_mask->nb[2], dense_mask->nb[3], dense_mask->nb[3], 0); ++ ++ ggml_tensor * kq = ggml_mul_mat(ctx, k_graph, q_graph); ++ ggml_mul_mat_set_prec(kq, GGML_PREC_F32); ++ kq = ggml_soft_max_ext(ctx, kq, dense_mask, scale(), 0.0f); ++ ++ ggml_tensor * v_dense = v_graph; ++ if (!v_trans) { ++ v_dense = ggml_cont(ctx, ggml_transpose(ctx, v_graph)); ++ } ++ ++ ggml_tensor * dense = ggml_mul_mat(ctx, v_dense, kq); ++ dense = ggml_mul_mat(ctx, v_mla, dense); ++ dense = ggml_permute(ctx, dense, 0, 2, 1, 3); ++ dense = ggml_cont_2d(ctx, dense, dense->ne[0]*dense->ne[1], dense->ne[2]*dense->ne[3]); ++ ggml_set_name(dense, "dense_graph"); ++ ++ return ggml_sub(ctx, sparse, dense); ++ } ++ ++ ggml_tensor * build_compact_equiv_graph(ggml_context * ctx) { ++ GGML_ASSERT(n_batch == 1); ++ GGML_ASSERT(n_head == 1); ++ GGML_ASSERT(n_stream == 1); ++ GGML_ASSERT(n_top_stream == 1); ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dk, 1, 1, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, type_k, dk, n_kv, 1, 1); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx, type_v, dv, n_kv, 1, 1); ++ ggml_tensor * kq_mask = ggml_new_tensor_4d(ctx, type_mask, 1, n_kv, 1, 1); ++ ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, 1, 1, 1); ++ ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "k"); ++ ggml_set_name(v, "v"); ++ ggml_set_name(kq_mask, "kq_mask"); ++ ggml_set_name(top_k, "top_k"); ++ ++ ggml_tensor * compact_k = ggml_get_rows_typed(ctx, k, top_k); ++ ggml_set_name(compact_k, "compact_k_topk_rows"); ++ ggml_tensor * compact_v = ggml_get_rows_typed(ctx, v, top_k); ++ ggml_set_name(compact_v, "compact_v_topk_rows"); ++ ++ ggml_tensor * compact = ggml_flash_attn_ext(ctx, q, compact_k, compact_v, nullptr, scale(), 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(compact, GGML_PREC_F32); ++ ggml_set_name(compact, "compact_no_mask_flash"); ++ ++ ggml_tensor * dense_mask = ggml_dsa_sparse_mask(ctx, kq_mask, top_k); ++ dense_mask = ggml_view_4d(ctx, dense_mask, ++ dense_mask->ne[1], dense_mask->ne[2], 1, dense_mask->ne[3], ++ dense_mask->nb[2], dense_mask->nb[3], dense_mask->nb[3], 0); ++ ggml_set_name(dense_mask, "dense_mask"); ++ ++ ggml_tensor * dense = ggml_flash_attn_ext(ctx, q, k, v, dense_mask, scale(), 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(dense, GGML_PREC_F32); ++ ggml_set_name(dense, "dense_mask_flash"); ++ ++ return ggml_sub(ctx, compact, dense); ++ } ++ ++ ggml_tensor * build_dense_equiv_graph( ++ ggml_context * ctx, ++ ggml_tensor * q, ++ ggml_tensor * k, ++ ggml_tensor * v, ++ ggml_tensor * kq_mask, ++ ggml_tensor * top_k) const { ++ ggml_tensor * sparse = ggml_dsa_sparse_attn(ctx, q, k, v, kq_mask, top_k, scale()); ++ ggml_set_name(sparse, "sparse"); ++ ++ ggml_tensor * dense_mask = ggml_dsa_sparse_mask(ctx, kq_mask, top_k); ++ dense_mask = ggml_view_4d(ctx, dense_mask, ++ dense_mask->ne[1], dense_mask->ne[2], 1, dense_mask->ne[3], ++ dense_mask->nb[2], dense_mask->nb[3], dense_mask->nb[3], 0); ++ ggml_set_name(dense_mask, "dense_mask"); ++ ++ ggml_tensor * kq = ggml_mul_mat(ctx, k, q); ++ ggml_mul_mat_set_prec(kq, GGML_PREC_F32); ++ ggml_set_name(kq, "dense_kq"); ++ ++ kq = ggml_soft_max_ext(ctx, kq, dense_mask, scale(), 0.0f); ++ ggml_set_name(kq, "dense_softmax"); ++ ++ ggml_tensor * v_dense = v; ++ if (!(v->nb[1] > v->nb[2])) { ++ v_dense = ggml_cont(ctx, ggml_transpose(ctx, v)); ++ ggml_set_name(v_dense, "dense_v_cont"); ++ } ++ ++ ggml_tensor * dense = ggml_mul_mat(ctx, v_dense, kq); ++ ggml_set_name(dense, "dense"); ++ ++ return ggml_sub(ctx, sparse, dense); ++ } ++ ++ float q_value(int64_t i_dk, int64_t i_batch, int64_t i_head, int64_t i_stream) const { ++ return 0.01f * (float) (1 + i_dk + 3*i_batch + 5*i_head + 7*i_stream); ++ } ++ ++ float k_value(int64_t i_dk, int64_t i_kv, int64_t i_stream) const { ++ return 0.02f * (float) (1 + i_dk + 2*i_kv + 11*i_stream); ++ } ++ ++ float v_value(int64_t i_dv, int64_t i_kv, int64_t i_stream) const { ++ return 0.03f * (float) (1 + i_dv + 3*i_kv + 13*i_stream); ++ } ++ ++ int32_t top_k_index(int64_t i_top, int64_t i_batch, int64_t i_stream) const { ++ return (int32_t) ((i_top + 5*i_batch + 9*i_stream) % n_kv); ++ } ++ ++ float mask_value(int64_t i_kv, int64_t i_batch, int64_t i_stream) const { ++ const int64_t valid_limit = n_kv - 1 - ((i_batch + i_stream) % std::max(1, n_top_k)); ++ return i_kv <= valid_limit ? 0.0f : -INFINITY; ++ } ++ ++ void init_q(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t i_stream = 0; i_stream < n_stream; ++i_stream) { ++ for (int64_t i_head = 0; i_head < n_head; ++i_head) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ for (int64_t i_dk = 0; i_dk < dk; ++i_dk) { ++ const size_t idx = ((i_stream*n_head + i_head)*n_batch + i_batch)*dk + i_dk; ++ data[idx] = q_value(i_dk, i_batch, i_head, i_stream); ++ } ++ } ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float)); ++ } ++ ++ void init_q_graph(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t i_stream = 0; i_stream < n_stream; ++i_stream) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ for (int64_t i_head = 0; i_head < n_head; ++i_head) { ++ for (int64_t i_dk = 0; i_dk < dk; ++i_dk) { ++ const size_t idx = ((i_stream*n_batch + i_batch)*n_head + i_head)*dk + i_dk; ++ data[idx] = q_value(i_dk, i_batch, i_head, i_stream); ++ } ++ } ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float)); ++ } ++ ++ template ++ void init_typed_row_tensor(ggml_tensor * t, ggml_type type, int64_t width, ValueFn value) const { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t i_stream = 0; i_stream < n_stream; ++i_stream) { ++ for (int64_t i_kv = 0; i_kv < n_kv; ++i_kv) { ++ for (int64_t i = 0; i < width; ++i) { ++ const size_t idx = (i_stream*n_kv + i_kv)*width + i; ++ data[idx] = value(i, i_kv, i_stream); ++ } ++ } ++ } ++ ++ if (type == GGML_TYPE_F16) { ++ std::vector data_f16(data.size()); ++ ggml_fp32_to_fp16_row(data.data(), data_f16.data(), data.size()); ++ ggml_backend_tensor_set(t, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); ++ } else { ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float)); ++ } ++ } ++ ++ void init_k(ggml_tensor * t) const { ++ init_typed_row_tensor(t, type_k, dk, [&](int64_t i, int64_t i_kv, int64_t i_stream) { ++ return k_value(i, i_kv, i_stream); ++ }); ++ } ++ ++ void init_v(ggml_tensor * t) const { ++ init_typed_row_tensor(t, type_v, dv, [&](int64_t i, int64_t i_kv, int64_t i_stream) { ++ return v_value(i, i_kv, i_stream); ++ }); ++ } ++ ++ void init_kq_mask(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t i_stream = 0; i_stream < n_stream; ++i_stream) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ for (int64_t i_kv = 0; i_kv < n_kv; ++i_kv) { ++ const size_t idx = (i_stream*n_batch + i_batch)*n_kv + i_kv; ++ data[idx] = mask_value(i_kv, i_batch, i_stream); ++ } ++ } ++ } ++ ++ if (type_mask == GGML_TYPE_F16) { ++ std::vector data_f16(data.size()); ++ ggml_fp32_to_fp16_row(data.data(), data_f16.data(), data.size()); ++ ggml_backend_tensor_set(t, data_f16.data(), 0, data_f16.size()*sizeof(ggml_fp16_t)); ++ } else { ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float)); ++ } ++ } ++ ++ void init_top_k(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t i_stream = 0; i_stream < n_top_stream; ++i_stream) { ++ for (int64_t i_batch = 0; i_batch < n_batch; ++i_batch) { ++ for (int64_t i_top = 0; i_top < n_top_k; ++i_top) { ++ const size_t idx = (i_stream*n_batch + i_batch)*n_top_k + i_top; ++ data[idx] = top_k_index(i_top, i_batch, i_stream); ++ } ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); ++ } ++ ++ void init_v_mla_identity(ggml_tensor * t) const { ++ std::vector data(ggml_nelements(t), 0.0f); ++ for (int64_t i = 0; i < dv; ++i) { ++ data[i*dv + i] = 1.0f; ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(float)); ++ } ++ ++ std::vector expected_row(int64_t i_batch, int64_t i_head, int64_t i_stream) const { ++ std::vector scores(n_top_k); ++ float max_score = -INFINITY; ++ const int64_t i_top_stream = i_stream % n_top_stream; ++ ++ for (int64_t i_top = 0; i_top < n_top_k; ++i_top) { ++ const int64_t i_kv = top_k_index(i_top, i_batch, i_top_stream); ++ float qk = 0.0f; ++ for (int64_t i_dk = 0; i_dk < dk; ++i_dk) { ++ qk += q_value(i_dk, i_batch, i_head, i_stream) * k_value(i_dk, i_kv, i_stream); ++ } ++ scores[i_top] = qk * scale() + mask_value(i_kv, i_batch, i_stream); ++ max_score = std::max(max_score, scores[i_top]); ++ } ++ ++ std::vector out(dv, 0.0f); ++ if (!std::isfinite(max_score)) { ++ return out; ++ } ++ ++ float sum = 0.0f; ++ for (float & score : scores) { ++ score = expf(score - max_score); ++ sum += score; ++ } ++ ++ for (int64_t i_top = 0; i_top < n_top_k; ++i_top) { ++ const int64_t i_kv = top_k_index(i_top, i_batch, i_top_stream); ++ const float p = scores[i_top] / sum; ++ for (int64_t i_dv = 0; i_dv < dv; ++i_dv) { ++ out[i_dv] += p * v_value(i_dv, i_kv, i_stream); ++ } ++ } ++ return out; ++ } ++}; ++ ++struct test_glm_dsa_selected_row_tiled_parity : public test_case { ++ static constexpr int64_t dk = 576; ++ static constexpr int64_t dv = 512; ++ static constexpr int64_t n_kv = 2304; ++ static constexpr int64_t n_batch = 16; ++ static constexpr int64_t n_head = 64; ++ static constexpr int64_t n_top_k = 2048; ++ ++ std::string vars() override { ++ return VARS_TO_STR6(dk, dv, n_kv, n_batch, n_head, n_top_k); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_DSA_SELECTED_ROW_TILED_PARITY"; ++ } + + ggml_tensor * build_graph(ggml_context * ctx) override { +- ggml_tensor * a = ggml_new_tensor(ctx, type, 4, ne_a.data()); +- ggml_set_name(a, "a"); +- +- ggml_tensor * out = ggml_timestep_embedding(ctx, a, dim, max_period); ++ set_test_env("SKIPPY_GLM_DSA_EXPERIMENTAL_SELECTED_ROW_FLASH", "1"); ++ set_test_env("LLAMA_GLM_DSA_EXPERIMENTAL_SELECTED_ROW_FLASH_TILED", "1"); ++ set_test_env("LLAMA_GLM_DSA_SELECTED_ROW_FLASH_TILED_TENSOR", "dsa_tiled"); ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, dk, n_batch, n_head, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, dk, n_kv, 1, 1); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, dv, n_kv, 1, 1); ++ ggml_tensor * mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 1, n_kv, n_batch, 1); ++ ggml_tensor * top_k = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, n_top_k, n_batch, 1, 1); ++ ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "k"); ++ ggml_set_name(v, "v"); ++ ggml_set_name(mask, "kq_mask"); ++ ggml_set_name(top_k, "top_k"); ++ ++ const float scale = 1.0f/sqrtf((float) dk); ++ ggml_tensor * vec = ggml_dsa_sparse_attn(ctx, q, k, v, mask, top_k, scale); ++ ggml_set_name(vec, "dsa_vec"); ++ ggml_tensor * tiled = ggml_dsa_sparse_attn(ctx, q, k, v, mask, top_k, scale); ++ ggml_set_name(tiled, "dsa_tiled"); ++ ++ ggml_tensor * out = ggml_sub(ctx, tiled, vec); + ggml_set_name(out, "out"); +- + return out; + } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "top_k") == 0) { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t batch = 0; batch < n_batch; ++batch) { ++ for (int64_t i = 0; i < n_top_k; ++i) { ++ data[batch*n_top_k + i] = (int32_t) ((37*i + 13*batch) % n_kv); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); ++ } else if (strcmp(t->name, "kq_mask") == 0) { ++ std::vector data(ggml_nelements(t)); ++ for (int64_t batch = 0; batch < n_batch; ++batch) { ++ const int64_t valid = n_kv - n_batch + batch + 1; ++ for (int64_t i = 0; i < n_kv; ++i) { ++ data[batch*n_kv + i] = ggml_fp32_to_fp16(i < valid ? 0.0f : -INFINITY); ++ } ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(ggml_fp16_t)); ++ } else if (!ggml_is_view_op(t->op)) { ++ init_tensor_uniform(t, -0.25f, 0.25f); ++ } ++ } ++ } ++ ++ double err(const float * a, const float * b, size_t n) override { ++ double error = 0.0; ++ for (size_t i = 0; i < n; ++i) { ++ error = std::max(error, (double) std::abs(a[i])); ++ error = std::max(error, (double) std::abs(b[i])); ++ } ++ if (getenv("SKIPPY_GLM_DSA_LOG_TILED_PARITY") != nullptr) { ++ fprintf(stderr, "GLM DSA selected-row tiled parity max_abs=%.9g\n", error); ++ } ++ return error; ++ } ++ ++ double max_err() override { ++ return 5e-3; ++ } ++ ++ double max_nmse_err() override { ++ return 1.0; ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } + }; + + // GGML_OP_LEAKY_RELU +@@ -6912,6 +11732,424 @@ struct test_flash_attn_ext : public test_case { + } + }; + ++struct test_glm_dsa_selected_row_flash : public test_case { ++ const int64_t kv; ++ const int64_t top_k; ++ const int64_t heads; ++ ++ std::string vars() override { ++ return VARS_TO_STR3(kv, top_k, heads); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_DSA_SELECTED_ROW_FLASH"; ++ } ++ ++ test_glm_dsa_selected_row_flash(int64_t kv = 257, int64_t top_k = 64, int64_t heads = 1) ++ : kv(kv), top_k(top_k), heads(heads) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ const char * disabled = getenv("SKIPPY_GLM_DSA_DISABLE_SELECTED_ROW_FLASH"); ++ const char * explicit_selected_row = getenv("SKIPPY_GLM_DSA_EXPERIMENTAL_SELECTED_ROW_FLASH"); ++ if (!(disabled && atoi(disabled) != 0) && explicit_selected_row == nullptr) { ++ set_test_env("SKIPPY_GLM_DSA_EXPERIMENTAL_SELECTED_ROW_FLASH", "1"); ++ } ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 576, 1, heads, 1); ++ ggml_tensor * packed_kv = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 576, kv, 1, 1); ++ ggml_tensor * top_k_ids = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, top_k, 1, 1, 1); ++ ++ ggml_set_name(q, "q"); ++ ggml_set_name(packed_kv, "packed_kv"); ++ ggml_set_name(top_k_ids, "top_k"); ++ ++ ggml_tensor * k_top_k = ggml_get_rows_typed(ctx, packed_kv, top_k_ids); ++ ggml_set_name(k_top_k, "dsa_compact_k_topk_rows"); ++ ++ ggml_tensor * v_top_k = ggml_view_4d( ++ ctx, ++ k_top_k, ++ 512, k_top_k->ne[1], k_top_k->ne[2], k_top_k->ne[3], ++ k_top_k->nb[1], k_top_k->nb[2], k_top_k->nb[3], ++ 0); ++ ggml_set_name(v_top_k, "dsa_compact_v_topk_view"); ++ ++ ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k_top_k, v_top_k, nullptr, 1.0f/sqrtf(576.0f), 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); ++ ggml_set_name(out, "out"); ++ ++ return out; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "top_k") == 0) { ++ std::vector data(top_k); ++ for (int64_t i = 0; i < top_k; ++i) { ++ data[i] = (int32_t) ((i*37) % kv); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); ++ } else { ++ init_tensor_uniform(t, -0.25f, 0.25f); ++ } ++ } ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ double max_err() override { ++ return 5e-3; ++ } ++ ++ uint64_t op_flops(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return 2ULL * (uint64_t) top_k * (576 + 512); ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } ++}; ++ ++struct test_glm_dsa_compact_multihead_flash : public test_case { ++ const int64_t kv; ++ const int64_t heads; ++ ++ std::string vars() override { ++ return VARS_TO_STR2(kv, heads); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_DSA_COMPACT_MULTIHEAD_FLASH"; ++ } ++ ++ test_glm_dsa_compact_multihead_flash(int64_t kv = 2048, int64_t heads = 64) ++ : kv(kv), heads(heads) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 576, 1, heads, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 576, kv, 1, 1); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 512, kv, 1, 1); ++ ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "compact_k"); ++ ggml_set_name(v, "compact_v"); ++ ++ ggml_tensor * out = ggml_flash_attn_ext(ctx, q, k, v, nullptr, 1.0f/sqrtf(576.0f), 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(out, GGML_PREC_F32); ++ ggml_set_name(out, "glm_compact_multihead_flash"); ++ ++ return out; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ init_tensor_uniform(t, -0.25f, 0.25f); ++ } ++ } ++ ++ double max_nmse_err() override { ++ return 5e-4; ++ } ++ ++ double max_err() override { ++ return 5e-3; ++ } ++ ++ uint64_t op_flops(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return 2ULL*(uint64_t) heads*(uint64_t) kv*(576 + 512); ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } ++}; ++ ++struct test_glm_dsa_compact_multihead_parity : public test_case { ++ std::string vars() override { ++ return "kv=2048,heads=64,stock_vs_multihead"; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_DSA_COMPACT_MULTIHEAD_PARITY"; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ set_test_env("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_MULTIHEAD_FLASH", "1"); ++ set_test_env( ++ "GGML_METAL_EXPERIMENTAL_GLM_COMPACT_MULTIHEAD_FLASH_TENSOR", ++ "compact_multihead_candidate"); ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 576, 1, 64, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 576, 2048, 1, 1); ++ ggml_tensor * v = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 512, 2048, 1, 1); ++ ggml_set_name(q, "real_q"); ++ ggml_set_name(k, "real_compact_k"); ++ ggml_set_name(v, "real_compact_v"); ++ ++ const float scale = 0.0625f; ++ ggml_tensor * control = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(control, GGML_PREC_F32); ++ ggml_set_name(control, "compact_multihead_control"); ++ ++ ggml_tensor * candidate = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(candidate, GGML_PREC_F32); ++ ggml_set_name(candidate, "compact_multihead_candidate"); ++ ++ ggml_tensor * delta = ggml_sub(ctx, candidate, control); ++ ggml_set_name(delta, "compact_multihead_delta"); ++ return delta; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ const char * fixture_dir = getenv("GGML_GLM_DSA_REAL_COMPACT_FIXTURE_DIR"); ++ if (fixture_dir != nullptr && fixture_dir[0] != '\0') { ++ auto read_fixture = [fixture_dir](const char * filename) { ++ std::ifstream input(std::string(fixture_dir) + "/" + filename, std::ios::binary); ++ GGML_ASSERT(input.good()); ++ return std::vector( ++ (std::istreambuf_iterator(input)), std::istreambuf_iterator()); ++ }; ++ ++ const std::vector compact_k = ++ read_fixture("dsa_compact_k_topk_rows-30_576_2048_1_1_f16.bin"); ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "real_q") == 0) { ++ const std::vector q = ++ read_fixture("Qcur-30__view___permuted__576_1_64_1_f32.bin"); ++ GGML_ASSERT(q.size() == ggml_nbytes(t)); ++ ggml_backend_tensor_set(t, q.data(), 0, q.size()); ++ } else if (strcmp(t->name, "real_compact_k") == 0) { ++ GGML_ASSERT(compact_k.size() == ggml_nbytes(t)); ++ ggml_backend_tensor_set(t, compact_k.data(), 0, compact_k.size()); ++ } else if (strcmp(t->name, "real_compact_v") == 0) { ++ std::vector compact_v(ggml_nbytes(t)); ++ constexpr size_t k_row_bytes = 576*sizeof(ggml_fp16_t); ++ constexpr size_t v_row_bytes = 512*sizeof(ggml_fp16_t); ++ GGML_ASSERT(compact_k.size() == 2048*k_row_bytes); ++ for (size_t row = 0; row < 2048; ++row) { ++ memcpy(compact_v.data() + row*v_row_bytes, compact_k.data() + row*k_row_bytes, v_row_bytes); ++ } ++ ggml_backend_tensor_set(t, compact_v.data(), 0, compact_v.size()); ++ } ++ } ++ return; ++ } ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ init_tensor_uniform(t, -0.25f, 0.25f); ++ } ++ } ++ ++ double max_nmse_err() override { ++ return 1.0; ++ } ++ ++ double max_err() override { ++ return 0.0; ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } ++}; ++ ++struct test_glm_dsa_compact_nwg_parity : public test_case { ++ std::string vars() override { ++ return "kv=2048,heads=64,nwg=4_vs_8"; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_DSA_COMPACT_NWG_PARITY"; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ set_test_env("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_NWG8_TENSOR", "compact_nwg8_candidate"); ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 576, 1, 64, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 576, 2048, 1, 1); ++ ggml_tensor * v = ggml_view_4d( ++ ctx, k, 512, 2048, 1, 1, k->nb[1], k->nb[2], k->nb[3], 0); ++ ggml_set_name(q, "real_q"); ++ ggml_set_name(k, "real_compact_k"); ++ ggml_set_name(v, "real_compact_v"); ++ ++ ggml_tensor * q_candidate = ggml_dup(ctx, q); ++ ggml_tensor * k_candidate = ggml_dup(ctx, k); ++ ggml_tensor * v_candidate = ggml_view_4d( ++ ctx, ++ k_candidate, ++ 512, ++ 2048, ++ 1, ++ 1, ++ k_candidate->nb[1], ++ k_candidate->nb[2], ++ k_candidate->nb[3], ++ 0); ++ ggml_set_name(q_candidate, "real_q_candidate"); ++ ggml_set_name(k_candidate, "real_compact_k_candidate"); ++ ggml_set_name(v_candidate, "real_compact_v_candidate"); ++ ++ ggml_set_name(q, "q"); ++ ggml_set_name(k, "compact_k"); ++ ggml_set_name(v, "compact_v"); ++ ++ const float scale = 0.0625f; ++ ggml_tensor * control = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(control, GGML_PREC_F32); ++ ggml_set_name(control, "compact_nwg4_control"); ++ ++ ggml_tensor * candidate = ++ ggml_flash_attn_ext(ctx, q_candidate, k_candidate, v_candidate, nullptr, scale, 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(candidate, GGML_PREC_F32); ++ ggml_set_name(candidate, "compact_nwg8_candidate"); ++ ++ ggml_tensor * delta = ggml_sub(ctx, candidate, control); ++ ggml_set_name(delta, "compact_nwg_delta"); ++ return delta; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ init_tensor_uniform(t, -0.25f, 0.25f); ++ } ++ } ++ ++ double max_nmse_err() override { ++ return 1.0; ++ } ++ ++ double max_err() override { ++ return 0.0; ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } ++}; ++ ++struct test_glm_dsa_compact_split_parity : public test_case { ++ std::string vars() override { ++ return "kv=2048,heads=64,split=control_vs_exact"; ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_DSA_COMPACT_SPLIT_PARITY"; ++ } ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ set_test_env("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_SPLIT_EXACT", "1"); ++ set_test_env("GGML_METAL_EXPERIMENTAL_GLM_COMPACT_SPLIT_EXACT_TENSOR", "compact_split_candidate"); ++ ++ ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, 576, 1, 64, 1); ++ ggml_tensor * k = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 576, 2048, 1, 1); ++ ggml_tensor * v = ggml_view_4d( ++ ctx, k, 512, 2048, 1, 1, k->nb[1], k->nb[2], k->nb[3], 0); ++ ++ const float scale = 0.0625f; ++ ggml_tensor * control = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(control, GGML_PREC_F32); ++ ggml_set_name(control, "compact_split_control"); ++ ++ ggml_tensor * candidate = ggml_flash_attn_ext(ctx, q, k, v, nullptr, scale, 0.0f, 0.0f); ++ ggml_flash_attn_ext_set_prec(candidate, GGML_PREC_F32); ++ ggml_set_name(candidate, "compact_split_candidate"); ++ ++ ggml_tensor * delta = ggml_sub(ctx, candidate, control); ++ ggml_set_name(delta, "compact_split_delta"); ++ return delta; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ const char * fixture_dir = getenv("GGML_GLM_DSA_REAL_COMPACT_FIXTURE_DIR"); ++ if (fixture_dir != nullptr && fixture_dir[0] != '\0') { ++ auto load_tensor = [fixture_dir](ggml_tensor * tensor, const char * filename) { ++ std::ifstream input(std::string(fixture_dir) + "/" + filename, std::ios::binary); ++ GGML_ASSERT(input.good()); ++ std::vector bytes( ++ (std::istreambuf_iterator(input)), std::istreambuf_iterator()); ++ GGML_ASSERT(bytes.size() == ggml_nbytes(tensor)); ++ ggml_backend_tensor_set(tensor, bytes.data(), 0, bytes.size()); ++ }; ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "real_q") == 0) { ++ load_tensor(t, "Qcur-30__view___permuted__576_1_64_1_f32.bin"); ++ } else if (strcmp(t->name, "real_compact_k") == 0) { ++ load_tensor(t, "dsa_compact_k_topk_rows-30_576_2048_1_1_f16.bin"); ++ } ++ } ++ return; ++ } ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ init_tensor_uniform(t, -2.0f, 2.0f); ++ } ++ } ++ ++ double max_nmse_err() override { ++ return 1.0; ++ } ++ ++ double max_err() override { ++ return 0.0; ++ } ++ ++ bool run_whole_graph() override { ++ return true; ++ } ++}; ++ ++struct test_glm_dsa_compact_k_gather : public test_case { ++ const int64_t kv; ++ const int64_t top_k; ++ ++ std::string vars() override { ++ return VARS_TO_STR2(kv, top_k); ++ } ++ ++ std::string op_desc(ggml_tensor * t) override { ++ GGML_UNUSED(t); ++ return "GLM_DSA_COMPACT_K_GATHER"; ++ } ++ ++ test_glm_dsa_compact_k_gather(int64_t kv = 1024, int64_t top_k = 768) ++ : kv(kv), top_k(top_k) {} ++ ++ ggml_tensor * build_graph(ggml_context * ctx) override { ++ ggml_tensor * packed_kv = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, 576, kv, 1, 1); ++ ggml_tensor * top_k_ids = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, top_k, 1, 1, 1); ++ ++ ggml_set_name(packed_kv, "packed_kv"); ++ ggml_set_name(top_k_ids, "top_k"); ++ ++ ggml_tensor * out = ggml_get_rows_typed(ctx, packed_kv, top_k_ids); ++ ggml_set_name(out, "dsa_compact_k_topk_rows"); ++ ++ return out; ++ } ++ ++ void initialize_tensors(ggml_context * ctx) override { ++ for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) { ++ if (strcmp(t->name, "top_k") == 0) { ++ std::vector data(top_k); ++ for (int64_t i = 0; i < top_k; ++i) { ++ data[i] = (int32_t) ((i*37) % kv); ++ } ++ ggml_backend_tensor_set(t, data.data(), 0, data.size()*sizeof(int32_t)); ++ } else { ++ init_tensor_uniform(t, -0.25f, 0.25f); ++ } ++ } ++ } ++}; ++ + // GGML_OP_CROSS_ENTROPY_LOSS + struct test_cross_entropy_loss : public test_case { + const ggml_type type; +@@ -8097,6 +13335,13 @@ static std::vector> make_test_cases_eval() { + } + + test_cases.emplace_back(new test_get_rows(GGML_TYPE_F32, 1, 8, 2, 1, 1, false)); ++ test_cases.emplace_back(new test_get_rows(GGML_TYPE_F16, 576, 2048, 128, 1, 1, false, true)); ++ test_cases.emplace_back(new test_get_rows(GGML_TYPE_F16, 512, 2048, 128, 1, 1, false, true)); ++ test_cases.emplace_back(new test_get_rows(GGML_TYPE_F16, 576, 1024, 768, 1, 1, false, true)); ++ test_cases.emplace_back(new test_get_rows(GGML_TYPE_F32, 256, 5, 4, 1, 1, false, true)); ++ test_cases.emplace_back(new test_get_rows(GGML_TYPE_I32, 256, 5, 4, 1, 1, false, true)); ++ test_cases.emplace_back(new test_get_rows(GGML_TYPE_BF16, 256, 5, 4, 1, 1, false, true)); ++ test_cases.emplace_back(new test_get_rows(GGML_TYPE_Q4_0, 256, 5, 4, 1, 1, false, true)); + for (ggml_type type : all_types) { + for (int b : {1, 7}) { + for (bool v : {false, true}) { +@@ -8805,6 +14050,18 @@ static std::vector> make_test_cases_eval() { + test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 8192, 512, 5120, {128, 1}, {1, 1})); + #endif + ++ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q2_K, GGML_TYPE_F32, 256, 8, false, 2048, 1, 6144)); // GLM-DSA routed gate/up decode ++ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q3_K, GGML_TYPE_F32, 256, 8, false, 6144, 1, 2048)); // GLM-DSA routed down decode ++ for (ggml_type type : {GGML_TYPE_Q8_0, GGML_TYPE_Q6_K, GGML_TYPE_Q5_K, GGML_TYPE_Q4_K}) { ++ test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 576, 1, 6144, {1, 1}, {1, 1})); // GLM-DSA kv_a decode ++ test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 2048, 1, 6144, {1, 1}, {1, 1})); // GLM-DSA q_a decode ++ test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 16384, 1, 2048, {1, 1}, {1, 1})); // GLM-DSA q_b decode ++ } ++ for (int n = 1; n <= 8; ++n) { ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q2_K, GGML_TYPE_F32, 2048, n, 6144, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q3_K, GGML_TYPE_F32, 6144, n, 2048, {1, 1}, {1, 1})); ++ } ++ + for (ggml_type type_a : all_types) { + for (int i = 1; i < 10; ++i) { + test_cases.emplace_back(new test_mul_mat(type_a, GGML_TYPE_F32, 16, i, 256, { 1, 1}, {1, 1})); +@@ -9601,6 +14858,218 @@ static std::vector> make_test_cases_eval() { + } + } + } ++ test_cases.emplace_back(new test_moe_weighted_sum(31, 5, 3)); ++ test_cases.emplace_back(new test_moe_weighted_sum(6144, 8, 1)); ++ test_cases.emplace_back(new test_moe_weighted_sum(6144, 8, 4)); ++ for (bool direct_top_k : {false, true}) { ++ test_cases.emplace_back(new test_topk_mul_mat_id_minimal(direct_top_k, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_topk_mul_mat_id_minimal(direct_top_k, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_topk_moe_weighted_output_minimal(direct_top_k, false, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_topk_moe_weighted_output_minimal(direct_top_k, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_topk_moe_weighted_output_minimal(direct_top_k, true, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_topk_moe_weighted_output_minimal(direct_top_k, true, GGML_TYPE_Q2_K)); ++ } ++ test_cases.emplace_back(new test_topk_mul_mat_id_minimal(false, GGML_TYPE_Q2_K, 6144, 2048, 16)); ++ test_cases.emplace_back(new test_topk_moe_glm_route()); ++ for (int64_t n_tokens : {2, 4, 8, 16}) { ++ test_cases.emplace_back(new test_topk_moe_glm_route(n_tokens)); ++ } ++ for (int component = 0; component <= 7; ++component) { ++ test_cases.emplace_back(new test_topk_moe_glm_route_component(component)); ++ } ++ test_cases.emplace_back(new test_topk_moe_glm_consumer()); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256, false, true)); ++ for (int component = 0; component <= 7; ++component) { ++ test_cases.emplace_back(new test_topk_moe_glm_component(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256)); ++ } ++ for (int component = 4; component <= 8; ++component) { ++ test_cases.emplace_back(new test_topk_moe_glm_component(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256, true, true)); ++ } ++ for (int component = 0; component <= 7; ++component) { ++ test_cases.emplace_back(new test_glm_moe_decode_component(component)); ++ test_cases.emplace_back(new test_glm_moe_decode_component(component, true)); ++ } ++ for (int component : { 2, 3, 4, 7 }) { ++ test_cases.emplace_back(new test_glm_moe_decode_component(component, false, true)); ++ test_cases.emplace_back(new test_glm_moe_decode_component(component, true, true)); ++ } ++ test_cases.emplace_back(new test_glm_moe_decode_direct_topk_final_output(false)); ++ test_cases.emplace_back(new test_glm_moe_decode_direct_topk_final_output(true)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_final_output(false, false)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_final_output(true, false)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_final_output(false, true)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_final_output(true, true)); ++ test_cases.emplace_back(new test_glm_moe_q2_weight_roofline(false, "ROUTE_A")); ++ test_cases.emplace_back(new test_glm_moe_q2_weight_roofline(true, "FIXED_IDS")); ++ test_cases.emplace_back(new test_glm_moe_q2_weight_roofline(false, "ROUTE_B")); ++ test_cases.emplace_back(new test_glm_moe_selected_chain(1, false, "CORRECTNESS")); ++ for (ggml_type type : {GGML_TYPE_Q8_0, GGML_TYPE_Q6_K, GGML_TYPE_Q5_K, GGML_TYPE_Q4_K}) { ++ test_cases.emplace_back(new test_glm_attention_projection_chain(type, 2)); ++ } ++ test_cases.emplace_back(new test_glm_projection_chain(GGML_TYPE_Q3_K, 6144, 16384, 2)); ++ test_cases.emplace_back(new test_glm_projection_chain(GGML_TYPE_Q4_K, 2048, 6144, 2)); ++ test_cases.emplace_back(new test_glm_projection_chain(GGML_TYPE_Q4_K, 6144, 2048, 2)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_rowtile_output(false)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_rowtile_output(true)); ++ for (int component : { 2, 3, 4, 7 }) { ++ test_cases.emplace_back(new test_glm_moe_decode_direct_topk_small_component(component, false)); ++ test_cases.emplace_back(new test_glm_moe_decode_direct_topk_small_component(component, true)); ++ } ++ for (int component = 0; component <= 6; ++component) { ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(component)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(component, true)); ++ if (component != 5) { ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(component, false, true)); ++ } ++ } ++ for (int component = 0; component <= 4; ++component) { ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(component, true, false, GGML_TYPE_Q2_K)); ++ } ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(5, false, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, false, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(5, true, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, true, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(5, false, true, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, false, true, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(5, true, true, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, true, true, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, false, false, GGML_TYPE_Q2_K, true)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(7, false, false, GGML_TYPE_Q2_K, true)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, true, false, GGML_TYPE_Q2_K, true)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(7, true, false, GGML_TYPE_Q2_K, true)); ++ test_cases.emplace_back(new test_glm_moe_motif_log_small(false)); ++ test_cases.emplace_back(new test_glm_moe_motif_log_small(true)); ++ for (int component = 0; component <= 2; ++component) { ++ test_cases.emplace_back(new test_glm_moe_gate_up_pair(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 256, 256)); ++ } ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 256, 256)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, 256, 256)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_F16, GGML_TYPE_F16, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q4_0, GGML_TYPE_Q4_0, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 1, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 4, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 8, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 16, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 32, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu_rowtile_q2(false)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu_rowtile_q2()); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu_rowtile_q2(false, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true, false, 2)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true, false, 4)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true, false, 6)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, false, 2)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, false, 6)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, false, 4)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true, true, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, 256, 256, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q1_0, GGML_TYPE_Q1_0, GGML_TYPE_Q1_0, 256, 256, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, false)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q3_K, 256, 256)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q3_K, 256, 256, true)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q3_K, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q3_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q2_K, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q2_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q2_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_TQ2_0, 256, 256)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q4_K, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256, 1, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256, 1, true, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q3_K, GGML_TYPE_Q3_K, GGML_TYPE_Q3_K, 6144, 2048, 1, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, 1, true)); ++ ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 2048, 1024)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 2048, 1024, true, true)); ++ test_cases.emplace_back(new test_glm_ffn_consumer()); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, true)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256, true, false)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q1_0, GGML_TYPE_Q1_0, GGML_TYPE_Q1_0, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256, true, false)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_activation(256, 8, 1, false)); ++ test_cases.emplace_back(new test_topk_moe_glm_activation(256, 8, 1, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_activation(256, 8, 1, true, true)); ++ for (int component = 0; component <= 6; ++component) { ++ test_cases.emplace_back(new test_glm_shexp_component(component, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256)); ++ } ++ ++ test_cases.emplace_back(new test_lightning_indexer(16, 4, 33, 2, 1, 1, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_lightning_indexer(16, 4, 33, 2, 1, 1, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_lightning_indexer(32, 8, 65, 4, 2, 2, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_lightning_indexer(128, 32, 65, 2, 1, 1, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_lightning_indexer(32, 4, 33, 2, 1, 1, GGML_TYPE_Q4_0)); ++ test_cases.emplace_back(new test_lightning_indexer(32, 4, 33, 2, 1, 1, GGML_TYPE_Q8_0)); ++ test_cases.emplace_back(new test_lightning_indexer(128, 32, 65, 2, 1, 1, GGML_TYPE_Q8_0)); ++ test_cases.emplace_back(new test_lightning_indexer(256, 4, 33, 2, 1, 1, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_lightning_indexer(256, 4, 33, 2, 1, 1, GGML_TYPE_Q3_K)); ++ test_cases.emplace_back(new test_lightning_indexer(256, 4, 33, 2, 1, 1, GGML_TYPE_Q4_K)); ++ test_cases.emplace_back(new test_lightning_indexer(256, 4, 33, 2, 1, 1, GGML_TYPE_Q5_K)); ++ test_cases.emplace_back(new test_lightning_indexer(256, 4, 33, 2, 1, 1, GGML_TYPE_Q6_K)); ++ test_cases.emplace_back(new test_dsa_sparse_mask(GGML_TYPE_F32, 33, 2, 1, 4, 1)); ++ test_cases.emplace_back(new test_dsa_sparse_mask(GGML_TYPE_F16, 33, 2, 1, 4, 1)); ++ test_cases.emplace_back(new test_dsa_sparse_mask(GGML_TYPE_F16, 65, 4, 2, 8, 2)); ++ test_cases.emplace_back(new test_dsa_sparse_mask(GGML_TYPE_F16, 65, 4, 2, 8, 2, true)); ++ test_cases.emplace_back(new test_dsa_sparse_mask(GGML_TYPE_F16, 65, 4, 4, 8, 2, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 16, 12, 33, 2, 4, 1, 4, 1)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, 16, 12, 33, 2, 4, 1, 4, 1)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 16, 12, 65, 33, 4, 1, 8, 1)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, 16, 12, 65, 33, 4, 1, 8, 1)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, 32, 24, 65, 4, 8, 2, 8, 2)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 16, 12, 33, 2, 4, 1, 4, 1, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, 32, 24, 65, 4, 8, 2, 8, 2, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 16, 12, 33, 2, 4, 1, 4, 1, false, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, 32, 24, 65, 4, 8, 2, 8, 2, false, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 16, 12, 65, 33, 4, 1, 8, 1, false, false, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, 16, 12, 65, 33, 4, 1, 8, 1, false, false, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, 257, 1, 4, 1, 64, 1, false, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, 513, 1, 8, 1, 513, 1, false, true)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, 257, 64, 8, 1, 64, 1)); ++ for (int64_t batch : {2, 4, 8, 16}) { ++ test_cases.emplace_back(new test_dsa_sparse_attn( ++ GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, ++ 576, 512, 257, batch, 8, 1, 64, 1)); ++ } ++ test_cases.emplace_back(new test_dsa_sparse_attn( ++ GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, ++ 576, 512, 2304, 16, 64, 1, 2048, 1)); ++ test_cases.emplace_back(new test_glm_dsa_selected_row_tiled_parity()); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F16, 576, 512, 1024, 1, 1, 1, 768, 1, false, false, false, true)); ++ // GLM-5.2 native decode shape after the 2048-token DSA threshold. ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F16, 576, 512, 2304, 1, 1, 1, 2048, 1, false, false, false, true)); ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(257, 64)); ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(257, 64, 8)); ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(1024, 768)); ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(1024, 768, 8)); ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(2304, 2048)); ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(2304, 2048, 64)); ++ test_cases.emplace_back(new test_glm_dsa_compact_multihead_flash()); ++ test_cases.emplace_back(new test_glm_dsa_compact_multihead_parity()); ++ test_cases.emplace_back(new test_glm_dsa_compact_nwg_parity()); ++ test_cases.emplace_back(new test_glm_dsa_compact_split_parity()); + + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 1, 1)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 16, 1, 1)); +@@ -9668,6 +15137,10 @@ static std::vector> make_test_cases_eval() { + } + } + ++ test_cases.emplace_back(new test_glm_ffn_chain_native_down(true, "CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(true, "CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_dual_lane_consumer(true, "CANDIDATE")); ++ + return test_cases; + } + #ifdef _MSC_VER +@@ -9824,6 +15297,21 @@ static std::vector> make_test_cases_perf() { + } + } + ++ for (int n = 1; n <= 8; ++n) { ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q2_K, GGML_TYPE_F32, 2048, n, 6144, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q3_K, GGML_TYPE_F32, 6144, n, 2048, {1, 1}, {1, 1})); ++ } ++ for (ggml_type type : {GGML_TYPE_Q8_0, GGML_TYPE_Q6_K, GGML_TYPE_Q5_K, GGML_TYPE_Q4_K}) { ++ test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 576, 1, 6144, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 2048, 1, 6144, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 16384, 1, 2048, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_glm_attention_projection_chain(type, 8)); ++ } ++ test_cases.emplace_back(new test_glm_decode_block_byte_ceiling()); ++ test_cases.emplace_back(new test_glm_projection_chain(GGML_TYPE_Q3_K, 6144, 16384, 8)); ++ test_cases.emplace_back(new test_glm_projection_chain(GGML_TYPE_Q4_K, 2048, 6144, 16)); ++ test_cases.emplace_back(new test_glm_projection_chain(GGML_TYPE_Q4_K, 6144, 2048, 16)); ++ + // qwen3-30b-a3b + for (int bs : {1, 4, 8, 32, 64, 128, 256, 512}) { + for (ggml_type type_a : {GGML_TYPE_F32, GGML_TYPE_F16, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_IQ2_XS}) { +@@ -9843,6 +15331,266 @@ static std::vector> make_test_cases_perf() { + } + } + ++ // GLM-DSA routed expert decode shapes. ++ test_cases.emplace_back(new test_topk_moe_glm_route()); ++ for (int64_t n_tokens : {2, 4, 8, 16}) { ++ test_cases.emplace_back(new test_topk_moe_glm_route(n_tokens)); ++ } ++ for (int component = 0; component <= 7; ++component) { ++ test_cases.emplace_back(new test_topk_moe_glm_route_component(component)); ++ } ++ test_cases.emplace_back(new test_moe_weighted_sum(6144, 8, 1)); ++ test_cases.emplace_back(new test_moe_weighted_sum(6144, 8, 4)); ++ for (bool direct_top_k : {false, true}) { ++ test_cases.emplace_back(new test_topk_mul_mat_id_minimal(direct_top_k, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_topk_mul_mat_id_minimal(direct_top_k, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_topk_moe_weighted_output_minimal(direct_top_k, false, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_topk_moe_weighted_output_minimal(direct_top_k, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_topk_moe_weighted_output_minimal(direct_top_k, true, GGML_TYPE_F32)); ++ test_cases.emplace_back(new test_topk_moe_weighted_output_minimal(direct_top_k, true, GGML_TYPE_Q2_K)); ++ } ++ test_cases.emplace_back(new test_topk_mul_mat_id_minimal(false, GGML_TYPE_Q2_K, 6144, 2048, 16)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer()); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256, false, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 256, 256, true, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, false, true, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true, true, true)); ++ for (int component = 0; component <= 7; ++component) { ++ test_cases.emplace_back(new test_topk_moe_glm_component(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 2048, 1024)); ++ test_cases.emplace_back(new test_topk_moe_glm_component(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K)); ++ test_cases.emplace_back(new test_topk_moe_glm_component(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K)); ++ } ++ for (int component = 4; component <= 8; ++component) { ++ test_cases.emplace_back(new test_topk_moe_glm_component(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 2048, 1024, true, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_component(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true)); ++ } ++ for (int component = 0; component <= 7; ++component) { ++ test_cases.emplace_back(new test_glm_moe_decode_component(component)); ++ test_cases.emplace_back(new test_glm_moe_decode_component(component, true)); ++ } ++ for (int component : { 2, 3, 4, 7 }) { ++ test_cases.emplace_back(new test_glm_moe_decode_component(component, false, true)); ++ test_cases.emplace_back(new test_glm_moe_decode_component(component, true, true)); ++ } ++ test_cases.emplace_back(new test_glm_moe_decode_direct_topk_final_output(false)); ++ test_cases.emplace_back(new test_glm_moe_decode_direct_topk_final_output(true)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_final_output(false, false)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_final_output(true, false)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_final_output(false, true)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_final_output(true, true)); ++ test_cases.emplace_back(new test_glm_moe_q2_weight_roofline(false, "ROUTE_A")); ++ test_cases.emplace_back(new test_glm_moe_q2_weight_roofline(true, "FIXED_IDS")); ++ test_cases.emplace_back(new test_glm_moe_q2_weight_roofline(false, "ROUTE_B")); ++ test_cases.emplace_back(new test_glm_moe_selected_chain(3, false, "USEFUL_A")); ++ test_cases.emplace_back(new test_glm_moe_selected_chain(3, true, "BYTE84")); ++ test_cases.emplace_back(new test_glm_moe_selected_chain(3, true, "BYTE128", 128)); ++ test_cases.emplace_back(new test_glm_moe_selected_chain(3, false, "USEFUL_B")); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_rowtile_output(false)); ++ test_cases.emplace_back(new test_glm_moe_routed_q2_rowtile_output(true)); ++ for (int component : { 2, 3, 4, 7 }) { ++ test_cases.emplace_back(new test_glm_moe_decode_direct_topk_small_component(component, false)); ++ test_cases.emplace_back(new test_glm_moe_decode_direct_topk_small_component(component, true)); ++ } ++ for (int component = 0; component <= 6; ++component) { ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(component)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(component, true)); ++ if (component != 5) { ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(component, false, true)); ++ } ++ } ++ for (int component = 0; component <= 4; ++component) { ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(component, true, false, GGML_TYPE_Q2_K)); ++ } ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(5, false, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, false, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(5, true, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, true, false, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(5, false, true, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, false, true, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(5, true, true, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, true, true, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, false, false, GGML_TYPE_Q2_K, true)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(7, false, false, GGML_TYPE_Q2_K, true)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(6, true, false, GGML_TYPE_Q2_K, true)); ++ test_cases.emplace_back(new test_glm_moe_decode_post_weight_component(7, true, false, GGML_TYPE_Q2_K, true)); ++ test_cases.emplace_back(new test_glm_moe_motif_log_small(false)); ++ test_cases.emplace_back(new test_glm_moe_motif_log_small(true)); ++ for (int component = 0; component <= 2; ++component) { ++ test_cases.emplace_back(new test_glm_moe_gate_up_pair(component, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048)); ++ } ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_F16, GGML_TYPE_F16, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q8_0, GGML_TYPE_Q8_0, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q4_0, GGML_TYPE_Q4_0, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 1, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 4, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 8, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 16, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, false, 32, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu_rowtile_q2(false)); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu_rowtile_q2()); ++ test_cases.emplace_back(new test_glm_moe_gate_up_swiglu_rowtile_q2(false, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true, false, 2)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true, false, 4)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true, false, 6)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, 6144, 2048, true, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_IQ2_XXS, GGML_TYPE_IQ2_XXS, GGML_TYPE_Q3_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_IQ2_XS, GGML_TYPE_IQ2_XS, GGML_TYPE_Q3_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_IQ2_S, GGML_TYPE_IQ2_S, GGML_TYPE_Q3_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, false, 2)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, false, 6)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, false, 4)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, false, true, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, true, true, true)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, 6144, 2048, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q1_0, GGML_TYPE_Q1_0, GGML_TYPE_Q1_0, 6144, 2048, false)); ++ test_cases.emplace_back(new test_glm_moe_static_motif(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, false)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q3_K, 2048, 1024)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q3_K, 2048, 1024, true)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q3_K, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q3_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_TQ2_0, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q2_K, 6144, 2048)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q2_K, 6144, 2048, true)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q2_K, 6144, 2048, false, true)); ++ test_cases.emplace_back(new test_glm_moe_down_reduce(GGML_TYPE_Q4_K, 6144, 2048)); ++ test_cases.emplace_back(new test_topk_moe_glm_activation(2048, 8, 1, false)); ++ test_cases.emplace_back(new test_topk_moe_glm_activation(2048, 8, 1, true)); ++ test_cases.emplace_back(new test_topk_moe_glm_activation(2048, 8, 1, true, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256, 1, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256, 1, true, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer()); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, 1, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, 1, true, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q3_K, GGML_TYPE_Q3_K, GGML_TYPE_Q3_K, 6144, 2048, 1, true)); ++ test_cases.emplace_back(new test_glm_shexp_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, 6144, 2048, 1, true)); ++ ++ // GLM-5.2 Shared-layer absorbed-MLA projection shapes. Keep these as ++ // independent matvecs so command-buffer GPU timestamps can attribute the ++ // real layer without forcing per-op synchronization. ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2048, 1, 6144, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 16384, 1, 2048, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 576, 1, 6144, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q2_K, GGML_TYPE_F32, 6144, 1, 16384, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q3_K, GGML_TYPE_F32, 6144, 1, 16384, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_K, GGML_TYPE_F32, 6144, 1, 16384, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q5_K, GGML_TYPE_F32, 6144, 1, 16384, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q6_K, GGML_TYPE_F32, 6144, 1, 16384, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 6144, 1, 16384, {1, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q2_K, GGML_TYPE_F32, 256, 1, 512, {64, 1}, {1, 1})); ++ test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q2_K, GGML_TYPE_F32, 6144, 1, 512, {64, 1}, {1, 1})); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 2048, 1024)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 2048, 1024, true, true)); ++ test_cases.emplace_back(new test_glm_ffn_consumer()); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, true)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, true)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, false, true)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, GGML_TYPE_TQ2_0, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, false)); ++ test_cases.emplace_back(new test_glm_ffn_consumer(GGML_TYPE_Q1_0, GGML_TYPE_Q1_0, GGML_TYPE_Q1_0, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, false)); ++ test_cases.emplace_back(new test_glm_ffn_chain(4)); ++ test_cases.emplace_back(new test_glm_ffn_chain(4, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, false, true)); ++ test_cases.emplace_back(new test_glm_ffn_chain(4, GGML_TYPE_Q2_K, GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, ++ GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 6144, 2048, true, true, true)); ++ test_cases.emplace_back(new test_glm_ffn_chain_native_down(false, "CONTROL_A")); ++ test_cases.emplace_back(new test_glm_ffn_chain_native_down(true, "CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_chain_native_down(false, "CONTROL_B")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(false, "RAMP_CONTROL_A")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(true, "RAMP_CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(false, "RAMP_CONTROL_B")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(false, "CONTROL_A")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(true, "CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(false, "CONTROL_B")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(true, "REVERSE_CANDIDATE_A")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(false, "REVERSE_CONTROL")); ++ test_cases.emplace_back(new test_glm_ffn_chain_two_phase(true, "REVERSE_CANDIDATE_B")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(false, "RAMP_CONTROL_A")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(true, "RAMP_CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(false, "RAMP_CONTROL_B")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(false, "CONTROL_A")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(true, "CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(false, "CONTROL_B")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(true, "REVERSE_CANDIDATE_A")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(false, "REVERSE_CONTROL")); ++ test_cases.emplace_back(new test_glm_ffn_two_phase_consumer(true, "REVERSE_CANDIDATE_B")); ++ test_cases.emplace_back(new test_glm_ffn_dual_lane_consumer(false, "RAMP_CONTROL_A")); ++ test_cases.emplace_back(new test_glm_ffn_dual_lane_consumer(true, "RAMP_CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_dual_lane_consumer(false, "RAMP_CONTROL_B")); ++ test_cases.emplace_back(new test_glm_ffn_dual_lane_consumer(false, "CONTROL_A")); ++ test_cases.emplace_back(new test_glm_ffn_dual_lane_consumer(true, "CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_dual_lane_consumer(false, "CONTROL_B")); ++ test_cases.emplace_back(new test_glm_ffn_native_down_consumer(false, "CONTROL_A")); ++ test_cases.emplace_back(new test_glm_ffn_native_down_consumer(true, "CANDIDATE")); ++ test_cases.emplace_back(new test_glm_ffn_native_down_consumer(false, "CONTROL_B")); ++ for (int component = 0; component <= 6; ++component) { ++ test_cases.emplace_back(new test_glm_shexp_component(component, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, GGML_TYPE_Q4_K, 256, 256)); ++ test_cases.emplace_back(new test_glm_shexp_component(component)); ++ } ++ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q2_K, GGML_TYPE_F32, 256, 8, false, 2048, 1, 6144)); ++ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q2_K, GGML_TYPE_F32, 256, 8, false, 4096, 1, 6144)); ++ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q2_K, GGML_TYPE_F32, 256, 8, false, 6144, 1, 2048)); ++ test_cases.emplace_back(new test_mul_mat_id(GGML_TYPE_Q3_K, GGML_TYPE_F32, 256, 8, false, 6144, 1, 2048)); ++ ++ // GLM-DSA sparse-attention decode threshold shapes. ++ for (int64_t kv : {128, 256, 257, 513}) { ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(kv, 64)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, kv, 1, 1, 1, 64, 1)); ++ test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {1, 1}, kv, 1, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ } ++ for (int64_t kv : {128, 256, 257, 513}) { ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(kv, kv)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, kv, 1, 1, 1, kv, 1)); ++ } ++ // Smaller post-threshold shapes keep the performance sweep cheap. ++ for (int64_t kv : {1024, 2048}) { ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(kv, 768)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, kv, 1, 1, 1, 768, 1)); ++ } ++ // GLM-5.2 native IndexShare width after visible KV exceeds index_topk=2048. ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(2304, 2048)); ++ test_cases.emplace_back(new test_glm_dsa_selected_row_flash(2304, 2048, 64)); ++ test_cases.emplace_back(new test_glm_dsa_compact_multihead_flash()); ++ test_cases.emplace_back(new test_glm_dsa_compact_multihead_parity()); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, 2304, 1, 1, 1, 2048, 1)); ++ test_cases.emplace_back(new test_glm_dsa_compact_k_gather(1024, 768)); ++ test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {1, 1}, 768, 1, false, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ ++ // GLM-DSA sparse-attention short phase shapes. ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, 257, 4, 1, 1, 64, 1)); ++ test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {1, 1}, 257, 4, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, 257, 8, 1, 1, 64, 1)); ++ test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {1, 1}, 257, 8, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_dsa_sparse_attn(GGML_TYPE_F32, GGML_TYPE_F32, GGML_TYPE_F32, 576, 512, 257, 16, 1, 1, 64, 1)); ++ test_cases.emplace_back(new test_flash_attn_ext(576, 512, 1, {1, 1}, 257, 16, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_F16, GGML_TYPE_F16)); ++ test_cases.emplace_back(new test_dsa_sparse_attn( ++ GGML_TYPE_F16, GGML_TYPE_F16, GGML_TYPE_F16, ++ 576, 512, 2304, 16, 64, 1, 2048, 1)); + + // gpt-oss-20b + for (int bs : {1, 4, 8, 512}) { +diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp +index 57d33a62..fe0accc8 100644 +--- a/tests/test-llama-archs.cpp ++++ b/tests/test-llama-archs.cpp +@@ -14,6 +14,7 @@ + #include + #include + #include ++#include + #include + #include + #include +@@ -79,10 +80,24 @@ static std::vector get_tokens(const uint32_t n_tokens, const uint32 + return ret; + } + +-static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ++enum class glm_dsa_indexshare_fixture { ++ DEFAULT, ++ CONFLICTING_METADATA, ++ SHARED_FIRST, ++ FULL_SHARED, ++ FULL_SHARED_SHARED, ++ PERIODIC_FULL_SHARED, ++}; ++ ++static gguf_context_ptr get_gguf_ctx( ++ const llm_arch arch, ++ const bool moe, ++ const glm_dsa_indexshare_fixture glm_dsa_indexshare = glm_dsa_indexshare_fixture::DEFAULT, ++ const uint32_t context_length = 128, ++ const uint32_t glm_dsa_indexer_top_k = 8) { + gguf_context_ptr ret(gguf_init_empty()); + llama_model_saver ms(arch, ret.get()); +- const uint32_t n_ctx = 128; ++ const uint32_t n_ctx = context_length; + + uint32_t n_vocab = 128; + uint32_t n_embd = 256; +@@ -109,6 +124,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { + n_embd = 128; + n_head = 1; + n_ff = 192; ++ if (arch == LLM_ARCH_GLM_DSA && glm_dsa_indexshare == glm_dsa_indexshare_fixture::FULL_SHARED_SHARED) { ++ n_layer = 3; ++ } else if (arch == LLM_ARCH_GLM_DSA && glm_dsa_indexshare == glm_dsa_indexshare_fixture::PERIODIC_FULL_SHARED) { ++ n_layer = 4; ++ } + } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { + n_layer = 3; + } else if (arch == LLM_ARCH_CHAMELEON) { +@@ -199,8 +219,33 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { + } + + ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, uint32_t(1)); +- ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64)); +- ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, arch == LLM_ARCH_GLM_DSA ? uint32_t(128) : uint32_t(64)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, arch == LLM_ARCH_GLM_DSA ? glm_dsa_indexer_top_k : uint32_t(8)); ++ if (arch == LLM_ARCH_GLM_DSA) { ++ if (glm_dsa_indexshare == glm_dsa_indexshare_fixture::CONFLICTING_METADATA) { ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K_FREQUENCY, uint32_t(1)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_SKIP_TOP_K_OFFSET, uint32_t(0)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, std::vector({"full", "shared"})); ++ } else if (glm_dsa_indexshare == glm_dsa_indexshare_fixture::SHARED_FIRST) { ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, std::vector({"shared", "full"})); ++ } else if (glm_dsa_indexshare == glm_dsa_indexshare_fixture::FULL_SHARED) { ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K_FREQUENCY, uint32_t(2)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_SKIP_TOP_K_OFFSET, uint32_t(1)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, std::vector({"full", "shared"})); ++ } else if (glm_dsa_indexshare == glm_dsa_indexshare_fixture::FULL_SHARED_SHARED) { ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K_FREQUENCY, uint32_t(3)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_SKIP_TOP_K_OFFSET, uint32_t(1)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, std::vector({"full", "shared", "shared"})); ++ } else if (glm_dsa_indexshare == glm_dsa_indexshare_fixture::PERIODIC_FULL_SHARED) { ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K_FREQUENCY, uint32_t(3)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_SKIP_TOP_K_OFFSET, uint32_t(1)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, std::vector({"full", "shared", "shared", "full"})); ++ } else { ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K_FREQUENCY, uint32_t(1)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_SKIP_TOP_K_OFFSET, uint32_t(0)); ++ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, std::vector({"full", "full"})); ++ } ++ } + ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4})); + ms.add_kv(LLM_KV_TOKENIZER_MODEL, "no_vocab"); + // ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT, n_embd); +@@ -254,9 +299,242 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) + return true; + } + ++static void set_test_env(const char * name, const char * value) { ++#if defined(_WIN32) ++ _putenv_s(name, value); ++#else ++ setenv(name, value, true); ++#endif ++} ++ ++static void unset_test_env(const char * name) { ++#if defined(_WIN32) ++ _putenv_s(name, ""); ++#else ++ unsetenv(name); ++#endif ++} ++ ++static bool test_env_enabled(const char * name) { ++ const char * value = getenv(name); ++ return value != nullptr && strcmp(value, "0") != 0 && strcmp(value, "false") != 0 && strcmp(value, "FALSE") != 0; ++} ++ ++class scoped_test_env { ++public: ++ scoped_test_env(const char * name, const char * value) : name(name) { ++ const char * old_value = getenv(name); ++ if (old_value != nullptr) { ++ had_old_value = true; ++ this->old_value = old_value; ++ } ++ set_test_env(name, value); ++ } ++ ++ ~scoped_test_env() { ++ if (had_old_value) { ++ set_test_env(name, old_value.c_str()); ++ } else { ++ unset_test_env(name); ++ } ++ } ++ ++private: ++ const char * name; ++ bool had_old_value = false; ++ std::string old_value; ++}; ++ ++static bool test_glm_dsa_rejects_conflicting_indexshare_metadata(const size_t seed) { ++ static constexpr const char * expected_error = ++ "attention.indexer.types conflicts with attention.indexer.top_k_frequency"; ++ ++ try { ++ gguf_context_ptr gguf_ctx = get_gguf_ctx( ++ LLM_ARCH_GLM_DSA, ++ true, ++ glm_dsa_indexshare_fixture::CONFLICTING_METADATA); ++ ++ llama_model_params model_params = llama_model_default_params(); ++ model_params.progress_callback = silent_model_load_progress; ++ size_t tmp = seed; ++ llama_model_ptr model(llama_model_init_from_user(gguf_ctx.get(), set_tensor_data, &tmp, model_params)); ++ if (!model) { ++ return true; ++ } ++ GGML_UNUSED(model); ++ } catch (const std::exception & err) { ++ const std::string message = err.what(); ++ return message.find(expected_error) != std::string::npos; ++ } ++ ++ return false; ++} ++ ++static bool test_glm_dsa_rejects_shared_first_indexshare_metadata(const size_t seed) { ++ static constexpr const char * expected_error = ++ "GLM_DSA IndexShare Shared layer 0 has no preceding Full layer"; ++ ++ try { ++ gguf_context_ptr gguf_ctx = get_gguf_ctx( ++ LLM_ARCH_GLM_DSA, ++ true, ++ glm_dsa_indexshare_fixture::SHARED_FIRST); ++ ++ llama_model_params model_params = llama_model_default_params(); ++ model_params.progress_callback = silent_model_load_progress; ++ size_t tmp = seed; ++ llama_model_ptr model(llama_model_init_from_user(gguf_ctx.get(), set_tensor_data, &tmp, model_params)); ++ if (!model) { ++ return true; ++ } ++ GGML_UNUSED(model); ++ } catch (const std::exception & err) { ++ const std::string message = err.what(); ++ return message.find(expected_error) != std::string::npos; ++ } ++ ++ return false; ++} ++ ++struct glm_dsa_graph_op_counts { ++ int total = 0; ++ int lightning_indexer = 0; ++ int top_k = 0; ++ int top_k_consumers = 0; ++ int compact_get_rows = 0; ++ int compact_get_rows_ops = 0; ++ int compact_k_get_rows = 0; ++ int compact_v_get_rows = 0; ++ int compact_v_views = 0; ++ int compact_flash_attn = 0; ++ int compact_flash_no_mask = 0; ++ int compact_flash_v_view = 0; ++ int flash_attn = 0; ++ int sparse_mask = 0; ++ int sparse_mask_nodes = 0; ++ int dsa_sparse_attn = 0; ++ int64_t max_top_k_width = 0; ++ int64_t max_compact_rows = 0; ++}; ++ ++static bool glm_dsa_tensor_name_starts_with(const ggml_tensor * tensor, const char * prefix) { ++ return tensor != nullptr && ++ tensor->name[0] != '\0' && ++ strncmp(tensor->name, prefix, strlen(prefix)) == 0; ++} ++ ++static bool glm_dsa_is_indexshare_top_k_tensor(const ggml_tensor * tensor) { ++ if (tensor == nullptr) { ++ return false; ++ } ++ if (strncmp(tensor->name, "ffn_moe_topk", strlen("ffn_moe_topk")) == 0) { ++ return false; ++ } ++ return tensor->op == GGML_OP_TOP_K || ++ strncmp(tensor->name, "top_k", strlen("top_k")) == 0; ++} ++ ++static bool glm_dsa_tensor_wraps_indexshare_top_k(const ggml_tensor * tensor, int depth = 0) { ++ if (tensor == nullptr || depth > 4) { ++ return false; ++ } ++ if (glm_dsa_is_indexshare_top_k_tensor(tensor)) { ++ return true; ++ } ++ if (tensor->op == GGML_OP_VIEW || ++ tensor->op == GGML_OP_RESHAPE || ++ tensor->op == GGML_OP_CONT || ++ tensor->op == GGML_OP_PERMUTE) { ++ return glm_dsa_tensor_wraps_indexshare_top_k(tensor->src[0], depth + 1); ++ } ++ return false; ++} ++ ++static bool glm_dsa_tensor_consumes_indexshare_top_k(const ggml_tensor * tensor) { ++ if (tensor == nullptr) { ++ return false; ++ } ++ if (tensor->op == GGML_OP_DSA_SPARSE_MASK) { ++ return glm_dsa_tensor_wraps_indexshare_top_k(tensor->src[1]); ++ } ++ if (tensor->op == GGML_OP_DSA_SPARSE_ATTN) { ++ return glm_dsa_tensor_wraps_indexshare_top_k(tensor->src[4]); ++ } ++ if (tensor->op == GGML_OP_DSA_TOP1_ATTN) { ++ return glm_dsa_tensor_wraps_indexshare_top_k(tensor->src[2]); ++ } ++ if (tensor->op == GGML_OP_GET_ROWS && ++ (strncmp(tensor->name, "dsa_compact_", strlen("dsa_compact_")) == 0 || ++ strncmp(tensor->name, "dsa_sparse_mask_topk", strlen("dsa_sparse_mask_topk")) == 0)) { ++ return glm_dsa_tensor_wraps_indexshare_top_k(tensor->src[1]); ++ } ++ return false; ++} ++ ++static bool glm_dsa_graph_op_counter_callback(ggml_tensor * tensor, bool ask, void * user_data) { ++ glm_dsa_graph_op_counts * counts = static_cast(user_data); ++ ++counts->total; ++ if (!ask) { ++ return true; ++ } ++ ++ if (tensor->op == GGML_OP_LIGHTNING_INDEXER) { ++ ++counts->lightning_indexer; ++ } else if (tensor->op == GGML_OP_TOP_K) { ++ if (strncmp(tensor->name, "ffn_moe_topk", strlen("ffn_moe_topk")) != 0) { ++ ++counts->top_k; ++ counts->max_top_k_width = std::max(counts->max_top_k_width, tensor->ne[0]); ++ } ++ } else if (tensor->op == GGML_OP_FLASH_ATTN_EXT) { ++ ++counts->flash_attn; ++ if (glm_dsa_tensor_name_starts_with(tensor->src[1], "dsa_compact_k") && ++ glm_dsa_tensor_name_starts_with(tensor->src[2], "dsa_compact_v")) { ++ ++counts->compact_flash_attn; ++ if (tensor->src[3] == nullptr) { ++ ++counts->compact_flash_no_mask; ++ } ++ if (tensor->src[2]->op == GGML_OP_VIEW && tensor->src[2]->src[0] == tensor->src[1]) { ++ ++counts->compact_flash_v_view; ++ } ++ } ++ } else if (tensor->op == GGML_OP_DSA_SPARSE_MASK) { ++ ++counts->sparse_mask; ++ } else if (tensor->op == GGML_OP_DSA_SPARSE_ATTN) { ++ ++counts->dsa_sparse_attn; ++ } ++ ++ if (glm_dsa_tensor_consumes_indexshare_top_k(tensor)) { ++ ++counts->top_k_consumers; ++ } ++ if (glm_dsa_tensor_name_starts_with(tensor, "dsa_compact_")) { ++ ++counts->compact_get_rows; ++ counts->max_compact_rows = std::max(counts->max_compact_rows, tensor->ne[1]); ++ if (tensor->op == GGML_OP_GET_ROWS) { ++ ++counts->compact_get_rows_ops; ++ if (glm_dsa_tensor_name_starts_with(tensor, "dsa_compact_k")) { ++ ++counts->compact_k_get_rows; ++ } else if (glm_dsa_tensor_name_starts_with(tensor, "dsa_compact_v")) { ++ ++counts->compact_v_get_rows; ++ } ++ } else if (tensor->op == GGML_OP_VIEW && glm_dsa_tensor_name_starts_with(tensor, "dsa_compact_v")) { ++ ++counts->compact_v_views; ++ } ++ } ++ if (strncmp(tensor->name, "dsa_sparse_mask", strlen("dsa_sparse_mask")) == 0) { ++ ++counts->sparse_mask_nodes; ++ } ++ return true; ++} ++ ++static std::vector get_logits( ++ llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode); ++ + static std::pair get_model_and_ctx( + struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, +- const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { ++ const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, ++ ggml_backend_sched_eval_callback cb_eval = nullptr, void * cb_eval_user_data = nullptr, ++ const llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO) { + GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); + llama_model_params model_params = llama_model_default_params(); + model_params.progress_callback = silent_model_load_progress; +@@ -272,6 +550,9 @@ static std::pair get_model_and_ctx( + if (!encode) { + ctx_params.n_ubatch = 64; + } ++ ctx_params.flash_attn_type = flash_attn_type; ++ ctx_params.cb_eval = cb_eval; ++ ctx_params.cb_eval_user_data = cb_eval_user_data; + + size_t tmp = seed; + llama_model_ptr model(gguf_ctx != nullptr ? +@@ -287,6 +568,656 @@ static std::pair get_model_and_ctx( + return std::make_pair(std::move(model), std::move(lctx)); + } + ++static bool test_glm_dsa_indexshare_graph( ++ const size_t seed, ++ const std::vector & devs, ++ const glm_dsa_indexshare_fixture fixture, ++ const char * label, ++ const int expected_producers = 1, ++ const int expected_consumers = 2) { ++ try { ++ gguf_context_ptr gguf_ctx = get_gguf_ctx( ++ LLM_ARCH_GLM_DSA, ++ true, ++ fixture); ++ glm_dsa_graph_op_counts counts; ++ auto model_ctx = get_model_and_ctx( ++ gguf_ctx.get(), ++ nullptr, ++ seed, ++ devs, ++ LLAMA_SPLIT_MODE_LAYER, ++ false, ++ glm_dsa_graph_op_counter_callback, ++ &counts); ++ ++ const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_ctx.first.get())); ++ const std::vector tokens = get_tokens(8, n_vocab, seed); ++ get_logits(model_ctx.first.get(), model_ctx.second.get(), tokens, false); ++ ++ if (counts.top_k != expected_producers || ++ counts.lightning_indexer != expected_producers || ++ counts.top_k_consumers != expected_consumers) { ++ fprintf(stderr, "GLM_DSA %s graph counted top_k=%d lightning_indexer=%d top_k_consumers=%d\n", ++ label, counts.top_k, counts.lightning_indexer, counts.top_k_consumers); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA %s graph test failed: %s\n", label, err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_full_shared_indexshare_graph( ++ const size_t seed, const std::vector & devs) { ++ return test_glm_dsa_indexshare_graph( ++ seed, ++ devs, ++ glm_dsa_indexshare_fixture::FULL_SHARED, ++ "Full->Shared", ++ 1, ++ 2); ++} ++ ++static bool test_glm_dsa_full_shared_shared_indexshare_graph( ++ const size_t seed, const std::vector & devs) { ++ return test_glm_dsa_indexshare_graph( ++ seed, ++ devs, ++ glm_dsa_indexshare_fixture::FULL_SHARED_SHARED, ++ "Full->Shared->Shared", ++ 1, ++ 3); ++} ++ ++static bool test_glm_dsa_periodic_full_shared_indexshare_graph( ++ const size_t seed, const std::vector & devs) { ++ return test_glm_dsa_indexshare_graph( ++ seed, ++ devs, ++ glm_dsa_indexshare_fixture::PERIODIC_FULL_SHARED, ++ "Full->Shared->Shared->Full", ++ 2, ++ 4); ++} ++ ++static void decode_tokens(llama_context * lctx, const std::vector & tokens, llama_pos start_pos, bool output_all = true) { ++ llama_batch batch = llama_batch_init(tokens.size(), 0, 1); ++ for (size_t i = 0; i < tokens.size(); ++i) { ++ const bool output = output_all || i + 1 == tokens.size(); ++ common_batch_add(batch, tokens[i], start_pos + (llama_pos) i, {0}, output); ++ } ++ batch.n_tokens = tokens.size(); ++ if (llama_decode(lctx, batch)) { ++ llama_batch_free(batch); ++ throw std::runtime_error("failed to decode batch"); ++ } ++ llama_batch_free(batch); ++} ++ ++static glm_dsa_graph_op_counts run_glm_dsa_decode_graph_after_prefill( ++ const size_t seed, ++ const std::vector & devs, ++ uint32_t n_prefill = 16, ++ uint32_t context_length = 128, ++ uint32_t indexer_top_k = 8) { ++ gguf_context_ptr gguf_ctx = get_gguf_ctx( ++ LLM_ARCH_GLM_DSA, ++ true, ++ glm_dsa_indexshare_fixture::FULL_SHARED, ++ context_length, ++ indexer_top_k); ++ glm_dsa_graph_op_counts counts; ++ auto model_ctx = get_model_and_ctx( ++ gguf_ctx.get(), ++ nullptr, ++ seed, ++ devs, ++ LLAMA_SPLIT_MODE_LAYER, ++ false, ++ glm_dsa_graph_op_counter_callback, ++ &counts, ++ LLAMA_FLASH_ATTN_TYPE_ENABLED); ++ ++ const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_ctx.first.get())); ++ const std::vector tokens = get_tokens(n_prefill + 1, n_vocab, seed); ++ decode_tokens(model_ctx.second.get(), std::vector(tokens.begin(), tokens.begin() + n_prefill), 0); ++ ++ counts = {}; ++ decode_tokens(model_ctx.second.get(), std::vector{tokens[n_prefill]}, n_prefill); ++ ++ return counts; ++} ++ ++static std::vector glm_dsa_native_sparse_devs() { ++ const size_t device_count = ggml_backend_dev_count(); ++ for (size_t i = 0; i < device_count; ++i) { ++ ggml_backend_dev_t dev = ggml_backend_dev_get(i); ++ switch (ggml_backend_dev_type(dev)) { ++ case GGML_BACKEND_DEVICE_TYPE_GPU: ++ case GGML_BACKEND_DEVICE_TYPE_IGPU: ++ case GGML_BACKEND_DEVICE_TYPE_ACCEL: ++ return {dev}; ++ case GGML_BACKEND_DEVICE_TYPE_CPU: ++ case GGML_BACKEND_DEVICE_TYPE_META: ++ break; ++ } ++ } ++ return {}; ++} ++ ++static glm_dsa_graph_op_counts run_glm_dsa_verify_graph_after_prefill( ++ const size_t seed, ++ const std::vector & devs, ++ uint32_t n_prefill = 16, ++ uint32_t n_verify = 4, ++ uint32_t context_length = 128, ++ uint32_t indexer_top_k = 8) { ++ gguf_context_ptr gguf_ctx = get_gguf_ctx( ++ LLM_ARCH_GLM_DSA, ++ true, ++ glm_dsa_indexshare_fixture::FULL_SHARED, ++ context_length, ++ indexer_top_k); ++ glm_dsa_graph_op_counts counts; ++ auto model_ctx = get_model_and_ctx( ++ gguf_ctx.get(), ++ nullptr, ++ seed, ++ devs, ++ LLAMA_SPLIT_MODE_LAYER, ++ false, ++ glm_dsa_graph_op_counter_callback, ++ &counts, ++ LLAMA_FLASH_ATTN_TYPE_ENABLED); ++ ++ const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_ctx.first.get())); ++ const std::vector tokens = get_tokens(n_prefill + n_verify, n_vocab, seed); ++ decode_tokens(model_ctx.second.get(), std::vector(tokens.begin(), tokens.begin() + n_prefill), 0, false); ++ ++ counts = {}; ++ decode_tokens(model_ctx.second.get(), std::vector(tokens.begin() + n_prefill, tokens.end()), n_prefill, true); ++ ++ return counts; ++} ++ ++static glm_dsa_graph_op_counts run_glm_dsa_prefill_graph( ++ const size_t seed, ++ const std::vector & devs, ++ uint32_t n_tokens, ++ uint32_t context_length = 128, ++ uint32_t indexer_top_k = 8, ++ bool output_all = true) { ++ gguf_context_ptr gguf_ctx = get_gguf_ctx( ++ LLM_ARCH_GLM_DSA, ++ true, ++ glm_dsa_indexshare_fixture::FULL_SHARED, ++ context_length, ++ indexer_top_k); ++ glm_dsa_graph_op_counts counts; ++ auto model_ctx = get_model_and_ctx( ++ gguf_ctx.get(), ++ nullptr, ++ seed, ++ devs, ++ LLAMA_SPLIT_MODE_LAYER, ++ false, ++ glm_dsa_graph_op_counter_callback, ++ &counts, ++ LLAMA_FLASH_ATTN_TYPE_ENABLED); ++ ++ const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_ctx.first.get())); ++ const std::vector tokens = get_tokens(n_tokens, n_vocab, seed); ++ decode_tokens(model_ctx.second.get(), tokens, 0, output_all); ++ ++ return counts; ++} ++ ++static glm_dsa_graph_op_counts run_glm_dsa_prefill_graph_after_prefill( ++ const size_t seed, ++ const std::vector & devs, ++ uint32_t n_prefill, ++ uint32_t n_tokens, ++ uint32_t context_length, ++ uint32_t indexer_top_k) { ++ gguf_context_ptr gguf_ctx = get_gguf_ctx( ++ LLM_ARCH_GLM_DSA, ++ true, ++ glm_dsa_indexshare_fixture::FULL_SHARED, ++ context_length, ++ indexer_top_k); ++ glm_dsa_graph_op_counts counts; ++ auto model_ctx = get_model_and_ctx( ++ gguf_ctx.get(), ++ nullptr, ++ seed, ++ devs, ++ LLAMA_SPLIT_MODE_LAYER, ++ false, ++ glm_dsa_graph_op_counter_callback, ++ &counts, ++ LLAMA_FLASH_ATTN_TYPE_ENABLED); ++ ++ const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model_ctx.first.get())); ++ const std::vector tokens = get_tokens(n_prefill + n_tokens, n_vocab, seed); ++ decode_tokens(model_ctx.second.get(), std::vector(tokens.begin(), tokens.begin() + n_prefill), 0, false); ++ ++ counts = {}; ++ decode_tokens(model_ctx.second.get(), std::vector(tokens.begin() + n_prefill, tokens.end()), n_prefill, false); ++ ++ return counts; ++} ++ ++static bool validate_glm_dsa_compact_decode_graph_counts( ++ const char * label, ++ const glm_dsa_graph_op_counts & counts, ++ int64_t expected_top_k_width = 0) { ++ const bool top_k_width_ok = expected_top_k_width <= 0 || ++ (counts.top_k > 0 && ++ counts.max_top_k_width == expected_top_k_width && ++ counts.max_compact_rows == expected_top_k_width); ++ const bool compact_kv_ok = ++ counts.compact_k_get_rows > 0 && ++ (counts.compact_v_views > 0 || counts.compact_v_get_rows > 0) && ++ counts.top_k_consumers > 0; ++ const bool flash_ok = ++ counts.compact_flash_attn > 0 && ++ counts.compact_flash_no_mask > 0 && ++ counts.flash_attn > 0; ++ const bool dense_sparse_absent = ++ counts.sparse_mask == 0 && ++ counts.sparse_mask_nodes == 0 && ++ counts.dsa_sparse_attn == 0; ++ ++ if (!top_k_width_ok || !compact_kv_ok || !flash_ok || !dense_sparse_absent) { ++ fprintf(stderr, ++ "GLM_DSA %s compact decode graph counted total=%d top_k=%d max_top_k_width=%lld top_k_consumers=%d compact_get_rows=%d compact_get_rows_ops=%d compact_k_get_rows=%d compact_v_get_rows=%d compact_v_views=%d max_compact_rows=%lld flash_attn=%d compact_flash_attn=%d compact_flash_no_mask=%d compact_flash_v_view=%d sparse_mask=%d sparse_mask_nodes=%d dsa_sparse_attn=%d expected_top_k_width=%lld\n", ++ label, ++ counts.total, ++ counts.top_k, ++ (long long) counts.max_top_k_width, ++ counts.top_k_consumers, ++ counts.compact_get_rows, ++ counts.compact_get_rows_ops, ++ counts.compact_k_get_rows, ++ counts.compact_v_get_rows, ++ counts.compact_v_views, ++ (long long) counts.max_compact_rows, ++ counts.flash_attn, ++ counts.compact_flash_attn, ++ counts.compact_flash_no_mask, ++ counts.compact_flash_v_view, ++ counts.sparse_mask, ++ counts.sparse_mask_nodes, ++ counts.dsa_sparse_attn, ++ (long long) expected_top_k_width); ++ return false; ++ } ++ ++ if (test_env_enabled("LLAMA_GLM_DSA_PRINT_COMPACT_DECODE_COUNTS")) { ++ fprintf(stderr, ++ "GLM_DSA compact-decode graph validated label=%s total=%d top_k=%d max_top_k_width=%lld top_k_consumers=%d compact_get_rows=%d compact_get_rows_ops=%d compact_k_get_rows=%d compact_v_get_rows=%d compact_v_views=%d max_compact_rows=%lld flash_attn=%d compact_flash_attn=%d compact_flash_no_mask=%d compact_flash_v_view=%d sparse_mask=%d sparse_mask_nodes=%d dsa_sparse_attn=%d expected_top_k_width=%lld\n", ++ label, ++ counts.total, ++ counts.top_k, ++ (long long) counts.max_top_k_width, ++ counts.top_k_consumers, ++ counts.compact_get_rows, ++ counts.compact_get_rows_ops, ++ counts.compact_k_get_rows, ++ counts.compact_v_get_rows, ++ counts.compact_v_views, ++ (long long) counts.max_compact_rows, ++ counts.flash_attn, ++ counts.compact_flash_attn, ++ counts.compact_flash_no_mask, ++ counts.compact_flash_v_view, ++ counts.sparse_mask, ++ counts.sparse_mask_nodes, ++ counts.dsa_sparse_attn, ++ (long long) expected_top_k_width); ++ } ++ ++ return true; ++} ++ ++static bool test_glm_dsa_direct_sparse_decode_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ if (devs.empty()) { ++ fprintf(stderr, "GLM_DSA direct sparse decode graph test skipped: no native sparse backend device available\n"); ++ return true; ++ } ++ scoped_test_env direct_decode_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_MAX_TOP_K", "1024"); ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, devs); ++ ++ if (counts.dsa_sparse_attn == 0 || counts.compact_get_rows != 0 || counts.sparse_mask != 0 || counts.sparse_mask_nodes != 0) { ++ fprintf(stderr, ++ "GLM_DSA direct sparse decode graph counted total=%d dsa_sparse_attn=%d compact_get_rows=%d sparse_mask=%d sparse_mask_nodes=%d flash_attn=%d\n", ++ counts.total, counts.dsa_sparse_attn, counts.compact_get_rows, counts.sparse_mask, counts.sparse_mask_nodes, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA direct sparse decode graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_compact_decode_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env compact_decode_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_MAX_TOP_K", "4"); ++ scoped_test_env compact_min_kv("SKIPPY_GLM_DSA_COMPACT_FLASH_MIN_KV", "1"); ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, devs); ++ ++ if (!validate_glm_dsa_compact_decode_graph_counts("compact", counts)) { ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA compact decode graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static std::vector glm_dsa_cpu_devs() { ++ ggml_backend_dev_t cpu_dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); ++ if (cpu_dev == nullptr) { ++ throw std::runtime_error("CPU backend device is unavailable"); ++ } ++ return {cpu_dev}; ++} ++ ++static bool test_glm_dsa_cpu_compact_decode_graph(const size_t seed) { ++ try { ++ scoped_test_env compact_decode_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_MAX_TOP_K", "4"); ++ scoped_test_env compact_min_kv("SKIPPY_GLM_DSA_COMPACT_FLASH_MIN_KV", "1"); ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, glm_dsa_cpu_devs()); ++ ++ if (!validate_glm_dsa_compact_decode_graph_counts("CPU compact", counts)) { ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA CPU compact decode graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_cpu_direct_sparse_decode_falls_back_graph(const size_t seed) { ++ try { ++ scoped_test_env direct_decode_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_MAX_TOP_K", "1024"); ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, glm_dsa_cpu_devs()); ++ ++ if (counts.sparse_mask_nodes == 0 || counts.dsa_sparse_attn != 0 || counts.compact_get_rows != 0) { ++ fprintf(stderr, ++ "GLM_DSA CPU direct-sparse fallback graph counted total=%d sparse_mask=%d sparse_mask_nodes=%d dsa_sparse_attn=%d compact_get_rows=%d flash_attn=%d\n", ++ counts.total, counts.sparse_mask, counts.sparse_mask_nodes, counts.dsa_sparse_attn, counts.compact_get_rows, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA CPU direct-sparse fallback graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_compact_min_kv_decode_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env compact_decode_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_MAX_TOP_K", "4"); ++ scoped_test_env compact_min_kv("SKIPPY_GLM_DSA_COMPACT_FLASH_MIN_KV", "1024"); ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, devs); ++ ++ if (counts.sparse_mask_nodes == 0 || counts.compact_get_rows != 0 || counts.dsa_sparse_attn != 0) { ++ fprintf(stderr, ++ "GLM_DSA compact min-kv decode graph counted total=%d sparse_mask=%d sparse_mask_nodes=%d compact_get_rows=%d dsa_sparse_attn=%d flash_attn=%d\n", ++ counts.total, counts.sparse_mask, counts.sparse_mask_nodes, counts.compact_get_rows, counts.dsa_sparse_attn, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA compact min-kv decode graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_default_compact_decode_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ // GLM-5.2's native top_k=768 exceeds the default direct sparse decode ++ // cap, so default decode policy should select compact selected-KV flash. ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, devs, 768, 1024, 768); ++ ++ if (!validate_glm_dsa_compact_decode_graph_counts("default", counts, 768)) { ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA default compact decode graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_short_history_compact_decode_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env compact_decode_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_MAX_TOP_K", "4"); ++ scoped_test_env compact_min_kv("SKIPPY_GLM_DSA_COMPACT_FLASH_MIN_KV", "1"); ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ // During one-token decode, IndexShare top-k is selected from the previous ++ // KV state while attention sees previous+current KV. Even at short ++ // history, the forced compact policy should exercise selected-KV flash. ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, devs, 4); ++ ++ if (!validate_glm_dsa_compact_decode_graph_counts("short-history", counts)) { ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA short-history compact decode graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_native_top_k_compact_decode_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ // GLM-5.2 uses top_k=768. After a 768-token prefix, the next token sees ++ // 769 visible KV entries while IndexShare produces 768 selected rows. ++ // This is the model-native compact selected-KV decode shape. ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, devs, 768, 1024, 768); ++ ++ if (!validate_glm_dsa_compact_decode_graph_counts("native top-k", counts, 768)) { ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA native top-k compact decode graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_compact_disabled_decode_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env compact_decode_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_DECODE_MAX_TOP_K", "4"); ++ scoped_test_env compact_disabled("SKIPPY_GLM_DSA_DISABLE_COMPACT_FLASH_ATTN", "1"); ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_decode_graph_after_prefill(seed, devs); ++ ++ if (counts.sparse_mask_nodes == 0 || counts.compact_get_rows != 0 || counts.dsa_sparse_attn != 0) { ++ fprintf(stderr, ++ "GLM_DSA compact-disabled decode graph counted total=%d sparse_mask=%d sparse_mask_nodes=%d compact_get_rows=%d dsa_sparse_attn=%d flash_attn=%d\n", ++ counts.total, counts.sparse_mask, counts.sparse_mask_nodes, counts.compact_get_rows, counts.dsa_sparse_attn, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA compact-disabled decode graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_verify_dense_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ scoped_test_env prefill_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_PREFILL_MAX_TOKENS", "8"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_verify_graph_after_prefill(seed, devs); ++ ++ if (counts.sparse_mask_nodes == 0 || counts.compact_get_rows != 0 || counts.dsa_sparse_attn != 0) { ++ fprintf(stderr, ++ "GLM_DSA verification graph counted total=%d dsa_sparse_attn=%d compact_get_rows=%d sparse_mask=%d sparse_mask_nodes=%d flash_attn=%d\n", ++ counts.total, counts.dsa_sparse_attn, counts.compact_get_rows, counts.sparse_mask, counts.sparse_mask_nodes, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA verification dense graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_native_top_k_verify_dense_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ scoped_test_env prefill_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_PREFILL_MAX_TOKENS", "8"); ++ // Verification requests logits for the whole candidate span. Even with ++ // GLM-5.2's native top_k=768 shape, this phase is intentionally kept on ++ // the dense/auto path until verifier-specific sparse parity is proven. ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_verify_graph_after_prefill(seed, devs, 768, 4, 1024, 768); ++ ++ if (counts.top_k == 0 || counts.max_top_k_width != 768 || ++ counts.sparse_mask_nodes == 0 || counts.compact_get_rows != 0 || counts.dsa_sparse_attn != 0) { ++ fprintf(stderr, ++ "GLM_DSA native top-k verification graph counted total=%d top_k=%d max_top_k_width=%lld dsa_sparse_attn=%d compact_get_rows=%d sparse_mask=%d sparse_mask_nodes=%d flash_attn=%d\n", ++ counts.total, counts.top_k, (long long) counts.max_top_k_width, counts.dsa_sparse_attn, counts.compact_get_rows, ++ counts.sparse_mask, counts.sparse_mask_nodes, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA native top-k verification dense graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_short_prefill_dense_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ scoped_test_env prefill_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_PREFILL_MAX_TOKENS", "8"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_prefill_graph(seed, devs, 8); ++ ++ if (counts.sparse_mask_nodes == 0 || counts.compact_get_rows != 0 || counts.dsa_sparse_attn != 0) { ++ fprintf(stderr, ++ "GLM_DSA short prefill graph counted total=%d dsa_sparse_attn=%d compact_get_rows=%d sparse_mask=%d sparse_mask_nodes=%d flash_attn=%d\n", ++ counts.total, counts.dsa_sparse_attn, counts.compact_get_rows, counts.sparse_mask, counts.sparse_mask_nodes, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA short prefill dense graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_native_top_k_short_prefill_dense_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ scoped_test_env prefill_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_PREFILL_MAX_TOKENS", "2048"); ++ // GLM-5.2's package policy keeps short prefill dense below the ++ // threshold. Use output_all=false so this remains prefill rather than ++ // speculative verification, while still exercising native top_k=768. ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_prefill_graph(seed, devs, 768, 1024, 768, false); ++ ++ if (counts.top_k == 0 || counts.max_top_k_width != 768 || ++ counts.sparse_mask_nodes == 0 || counts.compact_get_rows != 0 || counts.dsa_sparse_attn != 0) { ++ fprintf(stderr, ++ "GLM_DSA native top-k short prefill graph counted total=%d top_k=%d max_top_k_width=%lld dsa_sparse_attn=%d compact_get_rows=%d sparse_mask=%d sparse_mask_nodes=%d flash_attn=%d\n", ++ counts.total, counts.top_k, (long long) counts.max_top_k_width, counts.dsa_sparse_attn, counts.compact_get_rows, ++ counts.sparse_mask, counts.sparse_mask_nodes, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA native top-k short prefill dense graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_long_prefill_dense_mask_guard_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ if (devs.empty()) { ++ fprintf(stderr, "GLM_DSA long prefill dense-mask guard graph test skipped: no native sparse backend device available\n"); ++ return true; ++ } ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ scoped_test_env direct_sparse_prefill("SKIPPY_GLM_DSA_ENABLE_DIRECT_SPARSE_PREFILL", "1"); ++ scoped_test_env prefill_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_PREFILL_MAX_TOKENS", "8"); ++ scoped_test_env dense_mask_limit("SKIPPY_GLM_DSA_DENSE_SPARSE_MASK_MAX_BYTES", "1"); ++ scoped_test_env large_prefill_enabled("SKIPPY_GLM_DSA_ENABLE_UNPROVEN_LARGE_DIRECT_SPARSE_PREFILL", "1"); ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_prefill_graph(seed, devs, 16); ++ ++ if (counts.dsa_sparse_attn == 0 || counts.compact_get_rows != 0 || counts.sparse_mask != 0 || counts.sparse_mask_nodes != 0) { ++ fprintf(stderr, ++ "GLM_DSA long prefill dense-mask guard graph counted total=%d dsa_sparse_attn=%d compact_get_rows=%d sparse_mask=%d sparse_mask_nodes=%d flash_attn=%d\n", ++ counts.total, counts.dsa_sparse_attn, counts.compact_get_rows, counts.sparse_mask, counts.sparse_mask_nodes, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA long prefill dense-mask guard graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ ++static bool test_glm_dsa_native_top_k_long_prefill_dense_mask_guard_graph( ++ const size_t seed, const std::vector & devs) { ++ try { ++ if (devs.empty()) { ++ fprintf(stderr, "GLM_DSA native top-k long prefill dense-mask guard graph test skipped: no native sparse backend device available\n"); ++ return true; ++ } ++ scoped_test_env graph_reuse_disabled("LLAMA_GRAPH_REUSE_DISABLE", "1"); ++ scoped_test_env direct_sparse_prefill("SKIPPY_GLM_DSA_ENABLE_DIRECT_SPARSE_PREFILL", "1"); ++ scoped_test_env prefill_cap("SKIPPY_GLM_DSA_DIRECT_SPARSE_PREFILL_MAX_TOKENS", "8"); ++ scoped_test_env dense_mask_limit("SKIPPY_GLM_DSA_DENSE_SPARSE_MASK_MAX_BYTES", "1"); ++ scoped_test_env large_prefill_enabled("SKIPPY_GLM_DSA_ENABLE_UNPROVEN_LARGE_DIRECT_SPARSE_PREFILL", "1"); ++ // GLM-5.2's configured top_k=768 would otherwise require a large dense ++ // sparse mask during long-context prefill. Warm the KV first, then ++ // assert the measured continuation selects native direct sparse ++ // attention and avoids materializing that mask. ++ const glm_dsa_graph_op_counts counts = run_glm_dsa_prefill_graph_after_prefill(seed, devs, 768, 64, 1024, 768); ++ ++ if (counts.top_k == 0 || counts.max_top_k_width != 768 || counts.dsa_sparse_attn == 0 || ++ counts.compact_get_rows != 0 || counts.sparse_mask != 0 || counts.sparse_mask_nodes != 0) { ++ fprintf(stderr, ++ "GLM_DSA native top-k long prefill dense-mask guard graph counted total=%d top_k=%d max_top_k_width=%lld dsa_sparse_attn=%d compact_get_rows=%d sparse_mask=%d sparse_mask_nodes=%d flash_attn=%d\n", ++ counts.total, counts.top_k, (long long) counts.max_top_k_width, counts.dsa_sparse_attn, counts.compact_get_rows, ++ counts.sparse_mask, counts.sparse_mask_nodes, counts.flash_attn); ++ return false; ++ } ++ return true; ++ } catch (const std::exception & err) { ++ fprintf(stderr, "GLM_DSA native top-k long prefill dense-mask guard graph test failed: %s\n", err.what()); ++ return false; ++ } ++} ++ + static std::vector get_logits( + llama_model * model, llama_context * lctx, const std::vector & tokens, bool encode = false) { + const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); +@@ -429,6 +1360,20 @@ static bool arch_supported(const llm_arch arch) { + return true; + } + ++static bool glm_dsa_backend_config_supported(const std::vector & devs) { ++ for (ggml_backend_dev_t dev : devs) { ++ switch (ggml_backend_dev_type(dev)) { ++ case GGML_BACKEND_DEVICE_TYPE_CPU: ++ case GGML_BACKEND_DEVICE_TYPE_GPU: ++ case GGML_BACKEND_DEVICE_TYPE_IGPU: ++ break; ++ default: ++ return false; ++ } ++ } ++ return true; ++} ++ + static int save_models(const llm_arch target_arch, const size_t seed, const ggml_log_level log_level, const std::string & dir) { + struct user_data_t { + struct { +@@ -542,6 +1487,33 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg + const std::string template_row_res = "%15s %10s|%20s|\n"; + + bool all_ok = true; ++ if (target_arch == LLM_ARCH_UNKNOWN || target_arch == LLM_ARCH_GLM_DSA) { ++ const bool glm_dsa_contract_ok = ++ test_glm_dsa_rejects_conflicting_indexshare_metadata(seed) && ++ test_glm_dsa_rejects_shared_first_indexshare_metadata(seed) && ++ test_glm_dsa_full_shared_indexshare_graph(seed, {}) && ++ test_glm_dsa_full_shared_shared_indexshare_graph(seed, {}) && ++ test_glm_dsa_periodic_full_shared_indexshare_graph(seed, {}) && ++ test_glm_dsa_direct_sparse_decode_graph(seed, glm_dsa_native_sparse_devs()) && ++ test_glm_dsa_compact_decode_graph(seed, glm_dsa_cpu_devs()) && ++ test_glm_dsa_cpu_compact_decode_graph(seed) && ++ test_glm_dsa_cpu_direct_sparse_decode_falls_back_graph(seed) && ++ test_glm_dsa_compact_min_kv_decode_graph(seed, glm_dsa_cpu_devs()) && ++ test_glm_dsa_default_compact_decode_graph(seed, glm_dsa_cpu_devs()) && ++ test_glm_dsa_short_history_compact_decode_graph(seed, glm_dsa_cpu_devs()) && ++ test_glm_dsa_native_top_k_compact_decode_graph(seed, glm_dsa_cpu_devs()) && ++ test_glm_dsa_compact_disabled_decode_graph(seed, glm_dsa_cpu_devs()) && ++ test_glm_dsa_verify_dense_graph(seed, {}) && ++ test_glm_dsa_native_top_k_verify_dense_graph(seed, {}) && ++ test_glm_dsa_short_prefill_dense_graph(seed, {}) && ++ test_glm_dsa_native_top_k_short_prefill_dense_graph(seed, {}) && ++ test_glm_dsa_long_prefill_dense_mask_guard_graph(seed, glm_dsa_native_sparse_devs()) && ++ test_glm_dsa_native_top_k_long_prefill_dense_mask_guard_graph(seed, glm_dsa_native_sparse_devs()); ++ if (!glm_dsa_contract_ok) { ++ fprintf(stderr, "GLM_DSA IndexShare metadata contract test failed\n"); ++ all_ok = false; ++ } ++ } + common_log_flush(common_log_main()); + printf(template_header.c_str(), "Model arch.", "Device", "Config", "NMSE vs. CPU", "Roundtrip"); + printf("|"); +@@ -591,6 +1563,9 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg + std::string status_roundtrip = "\033[1;33mSKIP\033[0m"; + char nmse_str[12] = {0}; + bool skip = !arch_supported(arch) || (dc.split_mode == LLAMA_SPLIT_MODE_TENSOR && dc.devs.empty()); ++ if (arch == LLM_ARCH_GLM_DSA && !glm_dsa_backend_config_supported(dc.devs)) { ++ skip = true; ++ } + #if defined(GGML_USE_WEBGPU) + skip = true; // FIXME + #endif // GGML_USE_WEBGPU +-- +2.54.0 (Apple Git-157) diff --git a/third_party/llama.cpp/patches/0025-Support-GLM-DSA-fused-KV_B-tensors.patch b/third_party/llama.cpp/patches/0025-Support-GLM-DSA-fused-KV_B-tensors.patch new file mode 100644 index 000000000..71600508d --- /dev/null +++ b/third_party/llama.cpp/patches/0025-Support-GLM-DSA-fused-KV_B-tensors.patch @@ -0,0 +1,89 @@ +From 320112e2be2f38c92559b4eaeb482816c6037874 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 24 Jun 2026 15:35:23 +1000 +Subject: [PATCH 22/24] Support GLM-DSA fused KV_B tensors + +--- + src/models/deepseek2.cpp | 23 +++++++++++++++++------ + src/models/glm-dsa.cpp | 10 +++++++--- + 2 files changed, 24 insertions(+), 9 deletions(-) + +diff --git a/src/models/deepseek2.cpp b/src/models/deepseek2.cpp +index 180fb082..feb05570 100644 +--- a/src/models/deepseek2.cpp ++++ b/src/models/deepseek2.cpp +@@ -210,8 +210,18 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p + // inp_pos - contains the positions + ggml_tensor * inp_pos = build_inp_pos(); + +- auto * inp_attn_kv = !is_mla ? build_attn_inp_kv() : nullptr; +- auto * inp_attn_k = is_mla ? build_attn_inp_k() : nullptr; ++ bool use_absorbed_mla = is_mla; ++ if (use_absorbed_mla) { ++ for (int il = il_start; il < il_end; ++il) { ++ if (!model.layers[il].wk_b || !model.layers[il].wv_b) { ++ use_absorbed_mla = false; ++ break; ++ } ++ } ++ } ++ ++ auto * inp_attn_kv = !use_absorbed_mla ? build_attn_inp_kv() : nullptr; ++ auto * inp_attn_k = use_absorbed_mla ? build_attn_inp_k() : nullptr; + + ggml_tensor * inp_out_ids = (!stage_filtered || stage_filter.include_output) ? build_inp_out_ids() : nullptr; + +@@ -311,7 +321,7 @@ llama_model_deepseek2::graph::graph(const llama_model & model, const llm_graph_p + kv_cmpr = build_norm(kv_cmpr, model.layers[il].attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + cb(kv_cmpr, "kv_cmpr", il); + +- if (is_mla) { ++ if (use_absorbed_mla) { + // {n_embd_head_qk_nope, n_tokens, n_head} + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + cb(q_nope, "q_nope_perm", il); +@@ -531,8 +541,9 @@ llama_model_deepseek2::graph_mtp::graph_mtp(const llama_model & model, const llm + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); +- auto * inp_attn_kv = !is_mla ? build_attn_inp_kv() : nullptr; +- auto * inp_attn_k = is_mla ? build_attn_inp_k() : nullptr; ++ const bool use_absorbed_mla = is_mla && layer.wk_b && layer.wv_b; ++ auto * inp_attn_kv = !use_absorbed_mla ? build_attn_inp_kv() : nullptr; ++ auto * inp_attn_k = use_absorbed_mla ? build_attn_inp_k() : nullptr; + + ggml_tensor * h_norm = build_norm(h_input, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); +@@ -603,7 +614,7 @@ llama_model_deepseek2::graph_mtp::graph_mtp(const llama_model & model, const llm + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + cb(kv_cmpr, "mtp_kv_cmpr_norm", il); + +- if (is_mla) { ++ if (use_absorbed_mla) { + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); + cb(q_nope, "mtp_q_nope_perm", il); + +diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp +index 10bf235f..a975da0d 100644 +--- a/src/models/glm-dsa.cpp ++++ b/src/models/glm-dsa.cpp +@@ -87,9 +87,13 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, flags); + +- // note: only old legacy GGUF files will have the unsplit wkv_b tensor in +- layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, flags); +- layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, flags); ++ layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, flags | TENSOR_NOT_REQUIRED); ++ layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, flags | TENSOR_NOT_REQUIRED); ++ if (!layer.wk_b || !layer.wv_b) { ++ layer.wk_b = nullptr; ++ layer.wv_b = nullptr; ++ layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), {kv_lora_rank, n_head * (n_embd_head_qk_nope + n_embd_head_v_mla)}, flags); ++ } + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, flags); + +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0026-Run-GLM-DSA-through-IndexShare-graph.patch b/third_party/llama.cpp/patches/0026-Run-GLM-DSA-through-IndexShare-graph.patch new file mode 100644 index 000000000..79eb6fbb0 --- /dev/null +++ b/third_party/llama.cpp/patches/0026-Run-GLM-DSA-through-IndexShare-graph.patch @@ -0,0 +1,370 @@ +From e3ec7639933e124f86f056e61b5a6b375f108e70 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Thu, 25 Jun 2026 15:51:20 +1000 +Subject: [PATCH 23/24] Run GLM-DSA through IndexShare graph + +--- + src/models/glm-dsa.cpp | 322 +++++++++++++++++++++++++++++++++++++++++ + src/models/models.h | 5 +- + 2 files changed, 326 insertions(+), 1 deletion(-) + +diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp +index a975da0d..adde84fd 100644 +--- a/src/models/glm-dsa.cpp ++++ b/src/models/glm-dsa.cpp +@@ -1,5 +1,7 @@ + #include "models.h" + ++#include "llama-kv-cache-dsa.h" ++ + void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); +@@ -145,6 +147,326 @@ void llama_model_glm_dsa::load_arch_tensors(llama_model_loader &) { + } + } + ++llama_model_glm_dsa::graph::graph(const llama_model & model, const llm_graph_params & params) : ++ llm_graph_context(params) { ++ const bool is_mla = hparams.is_mla(); ++ GGML_ASSERT(is_mla); ++ ++ const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); ++ const int64_t n_embd_head_v = hparams.n_embd_head_v_mla(); ++ GGML_UNUSED(n_embd_head_v); ++ ++ const int64_t n_embd_head_qk_rope = hparams.n_rot(); ++ const int64_t n_embd_head_qk_nope = n_embd_head_k - n_embd_head_qk_rope; ++ ++ const int64_t n_indexer_head = hparams.indexer_n_head; ++ const int64_t n_embd_indexer_head = hparams.indexer_head_size; ++ const int64_t n_embd_indexer_head_rope = hparams.n_rot(); ++ const int64_t n_embd_indexer_head_nope = n_embd_indexer_head - n_embd_indexer_head_rope; ++ const uint32_t n_indexer_top_k = hparams.indexer_top_k; ++ ++ const uint32_t kv_lora_rank = hparams.n_lora_kv; ++ ++ GGML_ASSERT(ext_factor >= 0.0f); ++ const float attn_factor_org = attn_factor * (1.0f + 0.1f * logf(1.0f / freq_scale)); ++ const float mscale = attn_factor_org * (1.0f + 0.1f * hparams.rope_yarn_log_mul * logf(1.0f / freq_scale)); ++ const float kq_scale = 1.0f * mscale * mscale / sqrtf(float(n_embd_head_k)); ++ ++ ggml_tensor * cur; ++ ggml_tensor * inpL; ++ ++ const int effective_n_layers = n_layer; ++ const skippy_graph_filter & stage_filter = skippy_graph_get_filter(); ++ const bool stage_filtered = stage_filter.enabled; ++ const int il_start = stage_filtered ? stage_filter.layer_start : 0; ++ const int il_end = stage_filtered ? ++ std::min(stage_filter.layer_end, effective_n_layers) : ++ effective_n_layers; ++ ++ for (int il = il_start; il < il_end; ++il) { ++ if (!model.layers[il].wk_b || !model.layers[il].wv_b) { ++ throw std::runtime_error("GLM_DSA sparse attention requires split K_B and V_B tensors"); ++ } ++ } ++ ++ inpL = build_inp_embd(stage_filtered && il_start > 0 ? nullptr : model.tok_embd); ++ ++ ggml_tensor * inp_pos = build_inp_pos(); ++ llm_graph_input_attn_k_dsa * inp_attn_dsa = build_attn_inp_k_dsa(); ++ ggml_tensor * inp_out_ids = (!stage_filtered || stage_filter.include_output) ? build_inp_out_ids() : nullptr; ++ ++ for (int il = il_start; il < il_end; ++il) { ++ ggml_tensor * inpSA = inpL; ++ ++ cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); ++ cb(cur, "attn_norm", il); ++ ++ { ++ ggml_tensor * qr = ggml_mul_mat(ctx0, model.layers[il].wq_a, cur); ++ cb(qr, "qr", il); ++ ++ qr = build_norm(qr, model.layers[il].attn_q_a_norm, nullptr, LLM_NORM_RMS, il); ++ cb(qr, "qr", il); ++ ++ ggml_tensor * top_k = nullptr; ++ ++ { ++ ggml_tensor * indexer_q = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_q_b, qr); ++ cb(indexer_q, "indexer_q", il); ++ ++ ggml_tensor * indexer_q_pe = ++ ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_rope, n_indexer_head, n_tokens, ++ ggml_row_size(indexer_q->type, n_embd_indexer_head), ++ ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, 0); ++ cb(indexer_q_pe, "indexer_q_pe", il); ++ ++ ggml_tensor * indexer_q_nope = ++ ggml_view_3d(ctx0, indexer_q, n_embd_indexer_head_nope, n_indexer_head, n_tokens, ++ ggml_row_size(indexer_q->type, n_embd_indexer_head), ++ ggml_row_size(indexer_q->type, n_embd_indexer_head) * n_indexer_head, ++ ggml_row_size(indexer_q->type, n_embd_indexer_head_nope)); ++ cb(indexer_q_nope, "indexer_q_nope", il); ++ ++ indexer_q_pe = ggml_rope_ext(ctx0, indexer_q_pe, inp_pos, nullptr, n_rot, ++ LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale, ++ ext_factor, attn_factor, beta_fast, beta_slow); ++ cb(indexer_q_pe, "indexer_q_pe", il); ++ ++ indexer_q = ggml_concat(ctx0, indexer_q_pe, indexer_q_nope, 0); ++ cb(indexer_q, "indexer_q", il); ++ ++ ggml_tensor * indexer_k = ggml_mul_mat(ctx0, model.layers[il].indexer_attn_k, cur); ++ cb(indexer_k, "indexer_k", il); ++ ++ indexer_k = build_norm(indexer_k, model.layers[il].indexer_k_norm, model.layers[il].indexer_k_norm_b, LLM_NORM, il); ++ cb(indexer_k, "indexer_k", il); ++ ++ ggml_tensor * indexer_k_pe = ++ ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_rope, 1, n_tokens, ++ ggml_row_size(indexer_k->type, n_embd_indexer_head), ++ ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, 0); ++ cb(indexer_k_pe, "indexer_k_pe", il); ++ ++ ggml_tensor * indexer_k_nope = ++ ggml_view_3d(ctx0, indexer_k, n_embd_indexer_head_nope, 1, n_tokens, ++ ggml_row_size(indexer_k->type, n_embd_indexer_head), ++ ggml_row_size(indexer_k->type, n_embd_indexer_head) * 1, ++ ggml_row_size(indexer_k->type, n_embd_indexer_head_nope)); ++ cb(indexer_k_nope, "indexer_k_nope", il); ++ ++ indexer_k_pe = ggml_rope_ext(ctx0, indexer_k_pe, inp_pos, nullptr, n_rot, ++ LLAMA_ROPE_TYPE_NEOX, n_ctx_orig, freq_base, freq_scale, ++ ext_factor, attn_factor, beta_fast, beta_slow); ++ cb(indexer_k_pe, "indexer_k_pe", il); ++ ++ indexer_k = ggml_concat(ctx0, indexer_k_pe, indexer_k_nope, 0); ++ cb(indexer_k, "indexer_k", il); ++ ++ indexer_q = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_q); ++ cb(indexer_q, "indexer_q", il); ++ indexer_k = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_k); ++ cb(indexer_k, "indexer_k", il); ++ ++ const auto * mctx_lid = inp_attn_dsa->mctx->get_lid(); ++ const auto & k_idxs_lid = inp_attn_dsa->get_k_idxs_lid(); ++ ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, indexer_k, k_idxs_lid, il)); ++ ++ ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, model.layers[il].indexer_proj, cur); ++ cb(indexer_weights, "indexer_weights", il); ++ ++ indexer_k = mctx_lid->get_k(ctx0, il); ++ ++ const auto n_stream = indexer_k->ne[3]; ++ indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0); ++ indexer_weights = ggml_view_4d(ctx0, indexer_weights, indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0); ++ ++ indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3); ++ cb(indexer_q, "indexer_q", il); ++ indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3); ++ cb(indexer_k, "indexer_k", il); ++ ++ ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q); ++ cb(indexer_kq, "indexer_kq", il); ++ ++ indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3)); ++ cb(indexer_kq, "indexer_kq", il); ++ ++ ggml_tensor * indexer_score = ggml_relu(ctx0, indexer_kq); ++ cb(indexer_score, "indexer_score", il); ++ ++ indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head))); ++ cb(indexer_weights, "indexer_weights", il); ++ ++ indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights); ++ cb(indexer_score, "indexer_score", il); ++ ++ indexer_score = ggml_sum_rows(ctx0, indexer_score); ++ cb(indexer_score, "indexer_score", il); ++ ++ indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3)); ++ cb(indexer_score, "indexer_score", il); ++ ++ ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid(); ++ indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask); ++ cb(indexer_score, "indexer_score", il); ++ ++ uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k; ++ top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k)); ++ cb(top_k, "top_k", il); ++ } ++ ++ ggml_tensor * q = ggml_mul_mat(ctx0, model.layers[il].wq_b, qr); ++ cb(q, "q", il); ++ ++ ggml_tensor * q_nope = ++ ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k), ++ ggml_row_size(q->type, n_embd_head_k) * n_head, 0); ++ cb(q_nope, "q_nope", il); ++ ++ ggml_tensor * q_pe = ggml_view_3d( ++ ctx0, q, n_embd_head_qk_rope, n_head, n_tokens, ggml_row_size(q->type, n_embd_head_k), ++ ggml_row_size(q->type, n_embd_head_k) * n_head, ggml_row_size(q->type, n_embd_head_qk_nope)); ++ cb(q_pe, "q_pe", il); ++ ++ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, model.layers[il].wkv_a_mqa, cur); ++ cb(kv_cmpr_pe, "kv_cmpr_pe", il); ++ ++ ggml_tensor * kv_cmpr = ++ ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens, ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), 0); ++ cb(kv_cmpr, "kv_cmpr", il); ++ ++ ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_embd_head_qk_rope, 1, n_tokens, ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_embd_head_qk_rope), ++ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank)); ++ cb(k_pe, "k_pe", il); ++ ++ q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ++ ext_factor, attn_factor, beta_fast, beta_slow); ++ cb(q_pe, "q_pe", il); ++ ++ k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ++ ext_factor, attn_factor, beta_fast, beta_slow); ++ cb(k_pe, "k_pe", il); ++ ++ kv_cmpr = build_norm(kv_cmpr, model.layers[il].attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); ++ cb(kv_cmpr, "kv_cmpr", il); ++ ++ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); ++ cb(q_nope, "q_nope_perm", il); ++ ++ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, model.layers[il].wk_b, q_nope); ++ cb(q_nope_absorbed, "q_nope_absorbed", il); ++ ++ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3); ++ cb(q_nope_absorbed, "q_nope_absorbed_perm", il); ++ ++ ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0); ++ cb(Qcur, "Qcur", il); ++ ++ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); ++ cb(kv_cmpr, "kv_cmpr_reshape", il); ++ ++ ggml_tensor * Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0); ++ cb(Kcur, "Kcur", il); ++ ++ ggml_tensor * Vcur = kv_cmpr; ++ cb(Vcur, "Vcur", il); ++ ++ cur = build_attn(inp_attn_dsa, ++ model.layers[il].wo, NULL, model.layers[il].wo_s, ++ Qcur, Kcur, Vcur, nullptr, nullptr, model.layers[il].wv_b, top_k, kq_scale, il); ++ } ++ ++ if (il == il_end - 1 && inp_out_ids && cparams.embeddings_nextn_masked) { ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); ++ } ++ ++ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); ++ cb(ffn_inp, "ffn_inp", il); ++ ++ cur = build_norm(ffn_inp, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); ++ cb(cur, "ffn_norm", il); ++ ++ if ((uint32_t) il < hparams.n_layer_dense_lead) { ++ cur = build_ffn(cur, ++ model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s, ++ model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s, ++ model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s, ++ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); ++ cb(cur, "ffn_out", il); ++ } else { ++ ggml_tensor * moe_out = build_moe_ffn(cur, ++ model.layers[il].ffn_gate_inp, ++ model.layers[il].ffn_up_exps, ++ model.layers[il].ffn_gate_exps, ++ model.layers[il].ffn_down_exps, ++ model.layers[il].ffn_exp_probs_b, ++ n_expert, n_expert_used, ++ LLM_FFN_SILU, hparams.expert_weights_norm, ++ hparams.expert_weights_scale, ++ (llama_expert_gating_func_type) hparams.expert_gating_func, ++ il, ++ nullptr, ++ model.layers[il].ffn_gate_up_exps, ++ model.layers[il].ffn_up_exps_s, ++ model.layers[il].ffn_gate_exps_s, ++ model.layers[il].ffn_down_exps_s); ++ cb(moe_out, "ffn_moe_out", il); ++ ++ ggml_tensor * ffn_shexp = ++ build_ffn(cur, ++ model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s, ++ model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s, ++ model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s, ++ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il); ++ cb(ffn_shexp, "ffn_shexp", il); ++ ++ cur = ggml_add(ctx0, moe_out, ffn_shexp); ++ cb(cur, "ffn_out", il); ++ } ++ cur = ggml_add(ctx0, cur, ffn_inp); ++ ++ cur = build_cvec(cur, il); ++ cb(cur, "l_out", il); ++ ++ inpL = cur; ++ } ++ cur = inpL; ++ ++ if (stage_filtered && !stage_filter.include_output) { ++ if (!cparams.embeddings_nextn_masked && inp_out_ids) { ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ } ++ ++ cb(cur, "stage_boundary", il_end - 1); ++ res->t_embd = cur; ++ ggml_build_forward_expand(gf, cur); ++ return; ++ } ++ ++ cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); ++ ++ cb(cur, "h_nextn", -1); ++ res->t_h_nextn = cur; ++ ++ if (!cparams.embeddings_nextn_masked && inp_out_ids) { ++ cur = ggml_get_rows(ctx0, cur, inp_out_ids); ++ } ++ ++ cb(cur, "result_norm", -1); ++ res->t_embd = cur; ++ ++ cur = ggml_mul_mat(ctx0, model.output, cur); ++ ++ cb(cur, "result_output", -1); ++ res->t_logits = cur; ++ ++ ggml_build_forward_expand(gf, cur); ++} ++ + std::unique_ptr llama_model_glm_dsa::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); +diff --git a/src/models/models.h b/src/models/models.h +index e9a7c0c2..d8120d9d 100644 +--- a/src/models/models.h ++++ b/src/models/models.h +@@ -1221,9 +1221,12 @@ struct llama_model_glm_dsa : public llama_model_base { + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + +- using graph = llama_model_deepseek2::graph; + using graph_mtp = llama_model_deepseek2::graph_mtp; + ++ struct graph : public llm_graph_context { ++ graph(const llama_model & model, const llm_graph_params & params); ++ }; ++ + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; + }; + +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0027-Use-DSA-KV-cache-for-GLM-DSA.patch b/third_party/llama.cpp/patches/0027-Use-DSA-KV-cache-for-GLM-DSA.patch new file mode 100644 index 000000000..3004688e4 --- /dev/null +++ b/third_party/llama.cpp/patches/0027-Use-DSA-KV-cache-for-GLM-DSA.patch @@ -0,0 +1,41 @@ +From 847934296d90daed13b31e8b76b0a5f44b801b05 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Thu, 25 Jun 2026 15:58:34 +1000 +Subject: [PATCH 24/24] Use DSA KV cache for GLM-DSA + +--- + src/llama-kv-cache.cpp | 5 +++-- + src/llama-model.cpp | 1 + + 2 files changed, 4 insertions(+), 2 deletions(-) + +diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp +index 43e9492a..9df4984b 100644 +--- a/src/llama-kv-cache.cpp ++++ b/src/llama-kv-cache.cpp +@@ -321,8 +321,9 @@ llama_kv_cache::llama_kv_cache( + ggml_is_quantized(type_k) && + hparams.n_embd_head_k() % 64 == 0; + +- // always create Hadamard rotation tensors for DeepSeek V3.2 DSA lightning indexer +- if (model.arch == LLM_ARCH_DEEPSEEK32 && hparams.n_embd_head_k_full == hparams.indexer_head_size) { ++ // always create Hadamard rotation tensors for the DSA lightning indexer ++ if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_GLM_DSA) && ++ hparams.n_embd_head_k_full == hparams.indexer_head_size) { + attn_rot_k = true; + } + +diff --git a/src/llama-model.cpp b/src/llama-model.cpp +index 5b83c0c0..a325a4ed 100644 +--- a/src/llama-model.cpp ++++ b/src/llama-model.cpp +@@ -2061,6 +2061,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, + res = nullptr; + } break; + case LLM_ARCH_DEEPSEEK32: ++ case LLM_ARCH_GLM_DSA: + { + res = new llama_kv_cache_dsa( + *this, +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0028-Bump-Skippy-ABI-for-GLM-DSA-runtime-config.patch b/third_party/llama.cpp/patches/0028-Bump-Skippy-ABI-for-GLM-DSA-runtime-config.patch new file mode 100644 index 000000000..b5eb4614f --- /dev/null +++ b/third_party/llama.cpp/patches/0028-Bump-Skippy-ABI-for-GLM-DSA-runtime-config.patch @@ -0,0 +1,25 @@ +From 0fb232f0f632084ab4fce7ed82ad12e4afe6a56f Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Tue, 21 Jul 2026 15:31:56 +1000 +Subject: [PATCH] Bump Skippy ABI for GLM DSA runtime config + +--- + include/skippy/common.h | 2 +- + 1 file changed, 1 insertion(+), 1 deletion(-) + +diff --git a/include/skippy/common.h b/include/skippy/common.h +index 2d73bf63..3346b6a5 100644 +--- a/include/skippy/common.h ++++ b/include/skippy/common.h +@@ -26,7 +26,7 @@ extern "C" { + + #define SKIPPY_ABI_VERSION_MAJOR 0 + #define SKIPPY_ABI_VERSION_MINOR 1 +-#define SKIPPY_ABI_VERSION_PATCH 31 ++#define SKIPPY_ABI_VERSION_PATCH 32 + + enum skippy_feature { + SKIPPY_FEATURE_RUNTIME_SLICE = 1 << 0, +-- +2.54.0 (Apple Git-157) + diff --git a/third_party/llama.cpp/patches/0029-Fix-GLM-DSA-Metal-get_rows-placement.patch b/third_party/llama.cpp/patches/0029-Fix-GLM-DSA-Metal-get_rows-placement.patch new file mode 100644 index 000000000..4f8b563e4 --- /dev/null +++ b/third_party/llama.cpp/patches/0029-Fix-GLM-DSA-Metal-get_rows-placement.patch @@ -0,0 +1,45 @@ +From c982f6c4823c99f194690c87f5498bdb03c90a65 Mon Sep 17 00:00:00 2001 +From: Mesh-LLM CI +Date: Wed, 22 Jul 2026 15:13:16 +1000 +Subject: [PATCH] Fix GLM DSA Metal get_rows placement + +--- + ggml/src/ggml-metal/ggml-metal.metal | 22 ---------------------- + 1 file changed, 22 deletions(-) + +diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal +index b7b7e065..7a6df158 100644 +--- a/ggml/src/ggml-metal/ggml-metal.metal ++++ b/ggml/src/ggml-metal/ggml-metal.metal +@@ -16815,28 +16815,6 @@ kernel void kernel_dsa_sparse_attn_decode_grouped_impl( + reduce[reduce_offset + tid] = local_active_top_end; + threadgroup_barrier(mem_flags::mem_threadgroup); + +-template [[host_name("kernel_get_rows_q1_0")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q2_0")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q4_0")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q4_1")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q5_0")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q5_1")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q8_0")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_mxfp4")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q2_K")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q3_K")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q4_K")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q5_K")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_q6_K")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq2_xxs")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq2_xs")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq3_xxs")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq3_s")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq2_s")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq1_s")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq1_m")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq4_nl")]] kernel get_rows_q_t kernel_get_rows_q; +-template [[host_name("kernel_get_rows_iq4_xs")]] kernel get_rows_q_t kernel_get_rows_q; + for (int32_t stride = nth/2; stride > 0; stride >>= 1) { + if (tid < stride) { + reduce[reduce_offset + tid] = max( +-- +2.54.0 (Apple Git-157) +