llama : scale moe down projection per expert token by ggerganov · Pull Request #30861 · ggml-org/llama.cpp

ggerganov · 2026-10-20T07:46:55Z

Overview

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. The same exact power-of-two factor is undone on the output, so the roundtrip preserves the matmul result. Handle POOL_1D in the Meta split-state logic.

Additional information

The scaling factor is exact (power of two), so the input is divided by the same scale used on the output.

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES
- Harness: pi:llama.cpp
- Hardware: Mac Studio, M2 Ultra, 192GB
- Model: ggml-org/DeepSeek-V4-Flash-Vision-Exp-GGUF
- Stats: ◷ 3.4h · ⚙ 10.6m · ↑ 497k · ↓ 259k · ≡ 915 · ⌨ 28

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)

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

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)

ggerganov · 2026-10-20T16:53:18Z

Tested on CUDA - all good. Should be ready to merge.

Trivaxy · 2026-10-20T18:58:51Z

Tested on Vulkan with Mellum2.1, tps and prefill are much better now

ggerganov · 2026-10-21T10:44:56Z

There is one remaining issue - some quantization implementations use the Q8_1 to quantize the activations. The block_q8_1 computes F16 block-wise sums. These sums can still overflow, even if the individual values do not overflow.

@jeffbolznv Let me know if you have any thoughts. For now I think I'll adjust the scaling to be more conservative.

Resolve the conflict with 8baabd5 (llama : scale moe down projection per expert token, ggml-org#30861): master replaced the MISTRAL4/MELLUM F32 precision override on the MoE down projection with an amax rescale of src1. Move that rescale into the PR's build_experts lambda, the single path every expert computation goes through (plain, cache, hybrid hot/cold, mailbox), in place of the old ggml_prec_set_src there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

jeffbolznv · 2026-10-21T13:49:35Z

I think there was some proposal to use bf16 for the q8_1 scale factor, maybe we should do that?

ggerganov · 2026-10-21T14:16:33Z

I think there was some proposal to use bf16 for the q8_1 scale factor, maybe we should do that?

Ok, I'll look into that.

297 upstream commits. Conflicts and how they were resolved: - ggml ops: upstream added GGML_OP_TOP_K_MASK next to the fork's GGML_OP_TURBO_WHT; both are kept, upstream's first, GGML_OP_COUNT 103. - MoE down projection: upstream replaced the MISTRAL4 F32 src1 precision hint (its own, ggml-org#29029) with a per-row amax rescale (ggml-org#30861) for MISTRAL4 and MELLUM. It now sits in build_moe_ffn_exps, so every tiered call (hot, cold, each arena window) rescales its own rows; rows skipped by a partial mul_mat_id are zero, so their scale is 1. - MoE post_exp_norm (new upstream): applied in build_moe_ffn after the hot/cold merge and before the router weights; the two halves fill disjoint expert slots, so the norm sees whole rows. - build_lora_mm_id: upstream's Hadamard input transform (prism) after the fork's partial/slots assert. Expert tensors split by LLAMA_MOE_HOT are not in the Hadamard map, so prism and tiering do not stack on one layer. - kv-cache attention rotation: the fork's default (off unless LLAMA_ATTN_ROT_DISABLE=0) is kept; upstream's llama_model_set_attn_rot can still turn it off per model. - model loader: hugetlb per-tensor populate kept before upstream's multi-buffer lookup of the mmap buffer. - llama-model: MoE arena setup and upstream's load_tensors_hadamard both run at the end of load_tensors. - qwen4exp: turbo Q pre-rotation kept; the selection mask is now `sel`. - server: upstream added draft-token logprobs (ggml-org#27196) read back from the batch logits, without post_sampling_probs. The fork's on_sample path, which records them during verification and covers post_sampling_probs, is kept; upstream's test_draft_token_probs passes with it. Upstream's local spec_i_batch is taken. - mmq: fork's IQ1_M MMQ gains upstream's prec_src1 parameter. - fattn-tile: the fork's D=640 case and HIP guard kept, upstream's default-case style. - argsort.cu: upstream's CCCL guard; the fork's hipcub arm was dead since HIP does not define GGML_CUDA_USE_CUB. - tests/CMakeLists.txt: fork tests kept, upstream's LINK and consolidated test-llama-context taken. Tested on gfx1100: ctest 78/80 (the two failures are environmental: the ggml-vocabs clone has git-lfs pointers, and system python has no jinja2). All fork tests pass (turbo-kv x8, moe-tier, fattn-turbo4, fattn-long-kv, mmvq-src1-cache, fusions, mul-mat-f16-range, top-k-graph-capture, op-offload-devices). test-llama-archs 390/390. Server test_speculative.py 10/10, completion/probs pytests 14/14. Needle at 20.5k tokens, Qwen3.8-27B IQ4_NL, turbo4 K/V: correct, prefill 925 tok/s vs 909 before the merge (single run).

…#30861) Cherry-pick of upstream 8baabd5, minus the MoE-cache slots argument mx does not have. The F32 src1 precision on the mistral4 down projection (from ggml-org#29029) made the CUDA backend reject the MUL_MAT_ID, so every layer's down projection ran on CPU. With layer split that staged each layer's expert weights in a ~101 GiB host buffer, and copying out of the repacked buffer (no get_tensor) segfaulted. The graph now rescales src1 per expert token instead and the op stays on the GPU. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

JohannesGaessler · 2026-10-23T16:43:46Z

+    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];
+    };
+


I don't think this is correct. Let's say we have data like [0, 1, 2, 3]. If you do a max pooling with 2 elements each you would get [1, 3]. If you have a split along axis 0 you would go from backend0=[0, 1], backend1=[2, 3] to backend0=[1], backend1=[3]. That is still a split along axis 0 and not mirrored because the data per backend is not the same.

Thanks for flagging. So my understanding of the meta split states was a bit incorrect and looking in more details, I now realize that the implemented max pool is technically wrong with -sm tensor. The activations before the down projection are indeed split along axis 0 and we want each device to compute the max along the entire rows - not just its own split.

I don't see an easy fix for that, so probably I'll disable -sm tensor for the arches that need this correction. The amax scaling is a temporary workaround either way, so eventually it won't be necessary.

github-actions Bot added the ggml changes relating to the ggml tensor library for machine learning label Oct 20, 2026

ggerganov commented Oct 20, 2026

View reviewed changes

Comment thread src/llama-graph.h Outdated

Comment thread src/llama-graph.cpp Outdated

Comment thread ggml/src/ggml-backend-meta.cpp Outdated

ggerganov force-pushed the gg/amax-pool1d-meta branch from 1f6a17e to 9e7f76c Compare October 20, 2026 07:57

This was referenced Oct 20, 2026

vulkan: ggml src1 precision support #30656

Merged

Revert "graph : add precision handling for MELLUM architecture (#30584)" #30733

Open

metal : optimize argmax and pool_1d #30862

Merged

ggerganov marked this pull request as ready for review October 20, 2026 09:35

ggerganov requested review from CISC and JohannesGaessler as code owners October 20, 2026 09:35

ggerganov added 2 commits October 20, 2026 16:47

cont : simplify comments

0c5b7cb

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>

ggerganov force-pushed the gg/amax-pool1d-meta branch from 9e7f76c to 0c5b7cb Compare October 20, 2026 13:47

ggerganov added 2 commits October 20, 2026 19:45

llama : refactor amax src1 scaling result

4019bf9

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

6bd4567

ggerganov added the merge ready A maintainer can use this label to indicate that they consider the changes final and ready to merge. label Oct 20, 2026

ggerganov merged commit 8baabd5 into master Oct 21, 2026
33 of 34 checks passed

ggerganov deleted the gg/amax-pool1d-meta branch October 21, 2026 06:01

leehack mentioned this pull request Oct 21, 2026

Qualify post-v0.6.0 upstream correctness fixes and consumer compatibility leehack/llamadart-native#112

Open

6 tasks

ggerganov mentioned this pull request Oct 21, 2026

metal : remove mul_mm_id amax rescaling #30924

Open

alexey-rsi mentioned this pull request Oct 21, 2026

models: Add Flint model (FlintForCausalLM) #30928

Merged

Siddhesh2377 mentioned this pull request Oct 21, 2026

models: Flint (FlintForCausalLM) on ra/main RunanywhereAI/llama.cpp#6

Merged

denizyuret mentioned this pull request Oct 21, 2026

Eval bug: CUDA gpt-oss-120b with IQ1_M (or TQ1_0/TQ2_0) experts produces NaN ("GGGG") at prompt batch > 8: fp16 accumulation and result in the dequantize + cuBLAS fallback overflow #30951

Open

JohannesGaessler reviewed Oct 23, 2026

View reviewed changes

0cc4m mentioned this pull request Oct 24, 2026

vulkan : keep F32-precision MUL_MAT_ID on the GPU for mat-vec, skip MMVQ #30544

Open

ggerganov mentioned this pull request Oct 24, 2026

llama : deny -sm tensor for arches with moe amax rescaling #31081

Open

Navigation Menu

Repository navigation

llama : scale moe down projection per expert token - #30861

llama : scale moe down projection per expert token#30861
ggerganov merged 4 commits into
masterfrom
gg/amax-pool1d-meta

ggerganov commented Oct 20, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ggerganov commented Oct 20, 2026

Uh oh!

Trivaxy commented Oct 20, 2026

Uh oh!

Uh oh!

ggerganov commented Oct 21, 2026

Uh oh!

jeffbolznv commented Oct 21, 2026

Uh oh!

ggerganov commented Oct 21, 2026

Uh oh!

JohannesGaessler Oct 23, 2026

Uh oh!

ggerganov Oct 24, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

ggerganov commented Oct 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Overview

Additional information

Requirements

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ggerganov commented Oct 20, 2026

Uh oh!

Trivaxy commented Oct 20, 2026

Uh oh!

Uh oh!

ggerganov commented Oct 21, 2026

Uh oh!

jeffbolznv commented Oct 21, 2026

Uh oh!

ggerganov commented Oct 21, 2026

Uh oh!

JohannesGaessler Oct 23, 2026

Choose a reason for hiding this comment

Uh oh!

ggerganov Oct 24, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

ggerganov commented Oct 20, 2026 •

edited

Loading