ClaudeBot/llama.cpp
0
0
Fork 1
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-17 20:31:47 +02:00
Code Issues Packages Projects Releases Wiki Activity

metal : use function constant for top-k MoE with_norm

Browse Source 
Replaces the runtime with_norm karg with a Metal function constant. The
top-k MoE kernel is compiled separately for the normalized and non-normalized
routing variants, removing the runtime branch.

Assisted-by: pi:llama.cpp/DeepSeek-V4-Flash-Vision-Exp
This commit is contained in:
Georgi Gerganov
2026-09-15 15:17:50 +03:00
parent 8dbc7c5b6a
commit 1de7aa5d12
5 changed files with 15 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
ggml/src/ggml-metal/ggml-metal-device.cpp
+8 -3
View File
@@ -1499,16 +1499,21 @@ 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(ggml_metal_library_t lib) {
ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe(ggml_metal_library_t lib, bool with_norm) {
char base[256];
char name[256];
snprintf(base, 256, "kernel_topk_moe_f32");
snprintf(name, 256, "%s", base);
snprintf(name, 256, "%s_with_norm=%d", base, with_norm ? 1 : 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);
ggml_metal_cv_t cv = ggml_metal_cv_init();
ggml_metal_cv_set_bool(cv, with_norm, FC_TOPK_MOE + 0);
res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
ggml_metal_cv_free(cv);
}
return res;
ggml/src/ggml-metal/ggml-metal-device.h
+1 -1
View File
@@ -147,7 +147,7 @@ struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht
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_radix (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 (ggml_metal_library_t lib);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_topk_moe (ggml_metal_library_t lib, bool with_norm);
struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_moe_weighted_reduction (ggml_metal_library_t lib);
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);
ggml/src/ggml-metal/ggml-metal-impl.h
+1 -1
View File
@@ -115,6 +115,7 @@
#define FC_UPSCALE 1500
#define FC_GATED_DELTA_NET 1600
#define FC_NORM 1700
#define FC_TOPK_MOE 1800
// op-specific constants
#define OP_FLASH_ATTN_EXT_NQPSG 8
@@ -1229,7 +1230,6 @@ typedef struct {
uint64_t nb01; // logits row stride
uint64_t nb1_ids; // ids row stride
int32_t top_k; // n_expert_used
int32_t with_norm; // 0/1
float clamp_val;
float scale_val;
} ggml_metal_kargs_topk_moe;
ggml/src/ggml-metal/ggml-metal-ops.cpp
+1 -2
View File
@@ -5420,12 +5420,11 @@ int ggml_metal_op_topk_moe(ggml_metal_op_t ctx, int idx) {
/*.nb01 =*/ logits->nb[1],
/*.nb1_ids =*/ ids->nb[1],
/*.top_k =*/ (int32_t) n_expert_used,
/*.with_norm =*/ with_norm ? 1 : 0,
/*.clamp_val =*/ clamp_val,
/*.scale_val =*/ scale_val,
};
auto pipeline = ggml_metal_library_get_pipeline_topk_moe(lib);
auto pipeline = ggml_metal_library_get_pipeline_topk_moe(lib, with_norm);
ggml_metal_encoder_set_pipeline(enc, pipeline);
ggml_metal_encoder_set_bytes (enc, &args, sizeof(args), 0);
ggml/src/ggml-metal/kernels/argsort.metal
+4 -2
View File
@@ -1,5 +1,7 @@
#include "common.h"
constant bool FC_topk_moe_with_norm [[function_constant(FC_TOPK_MOE + 0)]];
// bitonic sort implementation following the CUDA kernels as reference
typedef void (argsort_t)(
constant ggml_metal_kargs_argsort & args,
@@ -423,13 +425,13 @@ kernel void kernel_topk_moe_f32(
if ((best_expert & 31) == lane) {
ids_row[k] = best_expert;
if (args.with_norm) {
if (FC_topk_moe_with_norm) {
wt_sum += best_val;
}
}
}
if (args.with_norm) {
if (FC_topk_moe_with_norm) {
wt_sum = simd_sum(wt_sum);
wt_sum = max(wt_sum, args.clamp_val);
const float inv = 1.0f / wt_sum;