Commit 8baabd53b for llama.cpp
commit 8baabd53b77d1c48c6ef8d3b620436292f3f6231
Author: Georgi Gerganov <ggerganov@gmail.com>
Date: Wed Oct 21 09:01:00 2026 +0300
llama : scale moe down projection per expert token (#30861)
* llama : scale moe down projection per expert token
Replace the global amax rescaling for the MoE down projection with a per-expert-token factor computed by a POOL_1D max reduction along the intermediate dimension. Undo the scale on the output with the same exact power-of-two factor. Handle POOL_1D in the Meta split-state logic so a split along axis 0 becomes a replicated per-device partial result, enabling tensor parallelism.
Assisted-by: pi:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)
* cont : simplify comments
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
* llama : refactor amax src1 scaling result
Return the rescaled activation and the optional scale tensor through a small struct instead of an output pointer.
Assisted-by: pi:llama.cpp/ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF (Mac Studio, M2 Ultra, 192GB)
* cont : add TODO + clean-up
diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp
index 679d5216d..edbe4da41 100644
--- a/ggml/src/ggml-backend-meta.cpp
+++ b/ggml/src/ggml-backend-meta.cpp
@@ -774,6 +774,13 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
return src_ss[0];
};
+ auto handle_pool_1d = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
+ if (src_ss[0].axis == GGML_BACKEND_SPLIT_AXIS_0) {
+ return {GGML_BACKEND_SPLIT_AXIS_MIRRORED, {0}, {1}, 1};
+ }
+ return src_ss[0];
+ };
+
auto handle_flash_attn_ext = [&](const std::vector<ggml_backend_meta_split_state> & src_ss) -> ggml_backend_meta_split_state {
GGML_ASSERT(tensor->src[3] == nullptr || src_ss[3].axis == GGML_BACKEND_SPLIT_AXIS_MIRRORED);
@@ -990,12 +997,14 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(
case GGML_OP_CONV_3D:
case GGML_OP_CONV_2D_DW:
case GGML_OP_CONV_TRANSPOSE_2D:
- case GGML_OP_POOL_1D:
case GGML_OP_POOL_2D:
case GGML_OP_POOL_2D_BACK:
case GGML_OP_UPSCALE: {
split_state = handle_generic(src_ss, /*scalar_only =*/ true);
} break;
+ case GGML_OP_POOL_1D: {
+ split_state = handle_pool_1d(src_ss);
+ } break;
case GGML_OP_PAD: {
split_state = handle_pad(src_ss);
} break;
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 2bc3a291a..59caffe87 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -1666,6 +1666,40 @@ ggml_tensor * llm_graph_context::build_lora_mm(
return res;
}
+llm_graph_amax_src1 llm_graph_context::build_amax_src1_f16(
+ ggml_tensor * cur) const {
+ llm_graph_amax_src1 result = {
+ /*.res =*/ cur,
+ /*.scale_s =*/ nullptr,
+ };
+
+ const int64_t n_embd = cur->ne[0];
+
+ // amax = max|cur| along the rows for each [expert, token] -> [1, n_expert, n_tokens]
+ ggml_tensor * ab = ggml_abs(ctx0, cur);
+ ggml_tensor * amax = ggml_pool_1d(ctx0, ab, GGML_OP_POOL_MAX, n_embd, n_embd, 0);
+
+ // ref: https://github.com/ggml-org/llama.cpp/pull/26223
+ constexpr float k_log2_e = 1.4426950408889634f; // 1 / ln(2)
+ constexpr float k_ln_2 = 0.6931471805599453f;
+
+ ggml_tensor * e;
+ e = ggml_scale(ctx0, amax, 1.0f/32768.0f);
+ e = ggml_log (ctx0, e);
+ e = ggml_scale(ctx0, e, k_log2_e);
+ e = ggml_ceil (ctx0, e);
+ e = ggml_clamp(ctx0, e, 0.0f, 127.0f);
+ e = ggml_scale(ctx0, e, k_ln_2);
+ e = ggml_exp (ctx0, e);
+
+ result = {
+ /*.res =*/ ggml_div(ctx0, cur, e),
+ /*.scale_s =*/ e,
+ };
+
+ return result;
+}
+
ggml_tensor * llm_graph_context::build_lora_mm_id(
ggml_tensor * w, // ggml_tensor * as
ggml_tensor * cur, // ggml_tensor * b
@@ -2465,12 +2499,25 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
GGML_ABORT("fatal error");
}
- experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
- if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM) {
- // src1 can exceed F16 range
- ggml_prec_set_src(experts, GGML_PREC_F32, 1);
+ {
+ // TODO: this graph-level rescaling should only be applied when the corresponding backend does not support
+ // ggml_prec_set_src(experts, GGML_PREC_BF16, 1);
+ // ref: https://github.com/ggml-org/llama.cpp/pull/30656#issuecomment-6269597922
+ llm_graph_amax_src1 amax_src1 = { nullptr, nullptr };
+
+ if (arch == LLM_ARCH_MISTRAL4 || arch == LLM_ARCH_MELLUM) {
+ amax_src1 = build_amax_src1_f16(cur);
+ cur = amax_src1.res;
+ }
+
+ experts = build_lora_mm_id(down_exps, cur, selected_experts, down_exps_s, slots); // [n_embd, n_expert_used, n_tokens]
+ //ggml_prec_set_src(experts, GGML_PREC_BF16, 1); // (see TODO above)
+ cb(experts, "ffn_moe_down", il);
+
+ if (amax_src1.scale_s) {
+ experts = ggml_mul(ctx0, experts, amax_src1.scale_s);
+ }
}
- cb(experts, "ffn_moe_down", il);
if (down_exps_s) {
cb(experts, "ffn_moe_down_scaled", il);
diff --git a/src/llama-graph.h b/src/llama-graph.h
index f66bedb32..abe5133c9 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -1012,6 +1012,11 @@ struct llm_graph_qkv {
ggml_tensor * v; // [n_embd_head, n_head_kv, n_tokens]
};
+struct llm_graph_amax_src1 {
+ ggml_tensor * res; // cur rescaled
+ ggml_tensor * scale_s; // scale to undo on the matmul result, or nullptr
+};
+
struct llm_graph_context {
const llm_arch arch;
@@ -1115,6 +1120,11 @@ struct llm_graph_context {
ggml_tensor * w_s = nullptr,
ggml_tensor * slots = nullptr) const;
+ // rescale `cur` (the src1 of a matmul) by a power of two, computed from its maximum absolute
+ // when the values already fit in F16, the factors are exactly 1.0f
+ llm_graph_amax_src1 build_amax_src1_f16(
+ ggml_tensor * cur) const;
+
ggml_tensor * build_norm(
ggml_tensor * cur,
ggml_tensor * mw,