llama : fix unexpected graph reallocation in the k-pool models by ggerganov · Pull Request #29958 · ggml-org/llama.cpp

ggerganov · 2026-10-04T17:26:35Z

Overview

The k-pool models built their graph shape from state the full-context reserve cannot know, so the first decode with a different shape forced a scheduler re-reserve that dropped the worst-case sizing; the next state growth then needed more room at an unchanged graph size and aborted under GGML_SCHED_DEBUG_REALLOC=1.

Two branches caused it:

qwen4exp branched on inp->cache_safe, which turns false as soon as llama_memory_seq_cp shares cells (e.g. llama-batched-bench -pps): the 12 QSA layers swapped scatter+gather for fill+concat and dropped the new_pool_rep leaf, so the decode graph had 12 fewer nodes and one fewer leaf than the reserved one
glm5-next branched on gather = n_tokens <= 16 && n_kv > n_sel, so a TG decode built the gather shape (7564 nodes) while the last reserve, the PP one, had the dense shape (7762 nodes)

The fix makes the shape constant per context:

qwen4exp and glm5-next always create new_pool_rep and always scatter the fresh pooled keys before gathering all n_pool of them by cell; without sharing every pool is re-pooled anyway (n_new == n_pool_real, layout order), so new_pool_rep[ip] == pool_cells[ip] and the gather returns exactly what the removed concat produced for every visible pool, while padded pools stay -inf masked as before
glm5-next picks gather from context constants: cparams.n_ubatch bounds every ubatch, top_k + kpool - 1 bounds n_sel, so the shape follows neither n_tokens nor n_kv
set_input_kpool now requires new_pool_rep, and the dead get_kpool_cache_safe() API plus the unread cache_safe copy in kpool_build_sizes() are gone

Additional information

Repro, aborts before this PR:

GGML_SCHED_DEBUG_REALLOC=1 ./bin/llama-batched-bench -hf ggml-org/GLM-5.3-Flash-GGUF:Q2_K -npp 2500 -ntg 32 -npl 1,2 -c 32768 -pps -kvu

ggml_backend_sched_alloc_splits: unexpected graph reallocation (graph size = 8116, nodes = 8116, leafs = 1520), debug_realloc = 1

Verified after the change, all exit 0: the command above, the same with -npp 512, with -ub 16 (exercises the gather shape, all graphs 7564 nodes), the qwen4exp equivalent, plus test-llama-archs (502) and test-recurrent-state-rollback / test-save-load-state / test-opt, all under GGML_SCHED_DEBUG_REALLOC=1. Reserves and every decode now share one node count per context, so no decode-time re-reserve happens any more.

The dense path also measured faster than the gather path on Apple M2 Ultra, so the constant choice costs nothing: at 16k context PP/TG are 259/23.9 t/s dense vs 110/21.6 t/s gather, at 2.5k context 284/24.4 vs 273/21.5.

Requirements

I have read and agree with the contributing guidelines
AI usage disclosure: YES. pi:llama.cpp/MiMo-V2.6-Flash-MOPD

Both k-pool models built a graph shape that depends on state the full-context reserve cannot know: - qwen4exp branched on inp->cache_safe, which turns false as soon as llama_memory_seq_cp shares cells (e.g. batched-bench -pps): the QSA layers swapped scatter+gather for fill+concat and dropped the new_pool_rep leaf, so the decode graph had 12 fewer nodes than the reserved one - glm5-next branched on gather = n_tokens <= 16 && n_kv > n_sel, so the TG decode built the gather shape (7564 nodes) while the last reserve, the PP one, had the dense shape (7762 nodes) Either mismatch forces a decode-time re-reserve that drops the worst-case sizing and bakes in the current state, so the next state growth (n_pool, n_kv, n_new) needs more room at an unchanged graph size and aborts under GGML_SCHED_DEBUG_REALLOC=1. Reproduce with, e.g.: GGML_SCHED_DEBUG_REALLOC=1 ./bin/llama-batched-bench \ -hf ggml-org/GLM-5.3-Flash-GGUF:Q2_K -npp 2500 -ntg 32 -npl 1,2 \ -c 32768 -pps -kvu Always scatter+gather the pooled keys, and pick gather from context constants only: n_ubatch bounds every ubatch, top_k + kpool - 1 bounds n_sel. Every graph of a context then shares one shape, which the reserve covers, and the dense path measured faster than the gather path at 2.5k and 16k context. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD

The k-pool graphs no longer branch on cache_safe, so nothing reads get_kpool_cache_safe() or the conditional new_pool_rep any more: both models always pass the scatter target, which set_input_kpool now requires instead of merely preferring. Also drop the cache_safe copy in kpool_build_sizes(), a sizes-only helper. The layout and state flag itself stays, it still decides which pools a layout with shared cells must re-pool. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD

Decode a prompt into seq 0, share its cells with seq 1 via llama_memory_seq_cp (what llama-batched-bench does for -pps), then keep decoding both sequences. For the k-pool models sharing clears cache_safe, which changes the graph topology while the pools keep growing, so a scheduler that re-reserves with the current state instead of the worst-case one aborts under GGML_SCHED_DEBUG_REALLOC=1. The test registration sets that flag, and the test aborts on both k-pool models before 2220411. kimi-linear and minimax-01 are skipped: they reserve the final pp graph with n_seqs = 1 (see [TAG_RESERVE_DIAG_DECAY] in llama-context.cpp), so every multi-seq graph has a different layout and re-reserves by design. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD

ggerganov · 2026-10-04T19:03:24Z

+        // both terms are context constants: n_ubatch bounds every ubatch and top_k + kpool - 1 bounds
+        // n_sel, so the graph shape follows neither n_tokens nor n_kv, which the reserve cannot predict
+        // TODO: remove "gather" logic and everything related. the backends now support sparse attension so this is obsolete
+        inp->gather = (int64_t) cparams.n_ubatch <= max_ub && (int64_t) cparams.n_ctx > hparams.indexer_top_k + kpool - 1;


@am17an This is the TODO to remove the gather logic in a follow-up

ggml_cuda_match_moe_weighted_reduction rejected tensors with zero rows. A ubatch without outputs shrinks the last layer to zero rows through inp_out_ids, so graph_optimize dropped its alloc dep there and the scheduler graph lost one node compared to the reserved one. The scheduler then re-reserved at the size of that ubatch, and the next ubatch with the same node count but larger tensors aborted under GGML_SCHED_DEBUG_REALLOC=1. The compute loop already skips empty nodes before trying any fusion, so the guard only made the alloc deps depend on the row count.

ServeurpersoCom · 2026-10-05T05:16:23Z

I pushed a one-line CUDA fix on top: ggml_cuda_match_moe_weighted_reduction rejected zero-row tensors, so an output-less ubatch dropped one alloc dep and the scheduler graph lost a node versus the reserve, which made the realloc repro abort on CUDA even with this PR. Empty nodes are already skipped at compute, so the guard only made the alloc deps depend on the row count.

Tested on an RTX PRO 6000, master -> PR:

Qwen3.8-Flash-Next UD-Q4_K_XL

	master -> PR
greedy 16k / 64k	identical
llama-parallel 4 slots, 16 seqs	identical
pp512	3970 -> 3980 t/s
tg128	115.4 -> 115.9 t/s
tg128 @ 64k	85.4 -> 85.7 t/s
realloc repro	abort -> pass

GLM-5.3-Flash UD-IQ1_S**

	master -> PR
greedy / parallel	near-tie word choices only (dense vs gather)
tg greedy 16k	76.1 -> 84.2 t/s
tg greedy 64k	72.5 -> 77.3 t/s
tg llama-parallel 4x16	118.1 -> 124.9 t/s
realloc repro	abort -> pass

No quality regression on either model. The GLM file is an early glm5next conversion, so it was loaded with the compat commit from unsloth (rename).

ServeurpersoCom · 2026-10-05T05:21:05Z

Windows CI fails to link test-recurrent-state-rollback because the new skip calls llm_arch_from_string, which is not exported from the DLL: moving its llama_build into the NOT WIN32 OR NOT BUILD_SHARED_LIBS block, next to test-llama-archs, should fix it.

The webgpu-nvidia segfault and the openvino thread-safety failure are already red on master.

The shared-seq case calls llm_arch_from_string, which libllama does not export through LLAMA_API, so linking test-recurrent-state-rollback fails on Windows with shared libraries. Its build now sits in the NOT WIN32 OR NOT BUILD_SHARED_LIBS block, next to test-llama-archs and the test registration it already lives under.

ggerganov · 2026-10-05T05:27:16Z

    const int64_t n_tokens      = weighted->ne[2] * weighted->ne[3];
-    if (n_expert_used < 2 || n_expert_used > MOE_WEIGHTED_REDUCTION_MAX_EXPERTS || n_tokens <= 0) {
+    if (n_expert_used < 2 || n_expert_used > MOE_WEIGHTED_REDUCTION_MAX_EXPERTS) {


@am17an FYI - this keeps the graphs static when there are no outputs in the ubatch.

The skip of kimi-linear and minimax-01 went through llm_arch_from_string, which libllama does not export through LLAMA_API, so the test could not link on Windows with shared libraries. It now compares the general.architecture string directly, and the test builds on every platform again.

ggerganov · 2026-10-05T06:58:09Z

@ggml-org/ggml-vulkan I think that the gpu-vulkan-intel-linux job failing with the new test is likely a bug that is specific to Vulkan with Intel hardware. I can't reproduce with Metal, CUDA and Vulkan (RTX 5090).

ggerganov · 2026-10-05T08:57:51Z

@ggml-org/ggml-vulkan And one more fail that is probably worth looking into: https://github.com/ggml-org/llama.cpp/actions/runs/37283758139/job/111677663012#step:3:1232. This is from the test added in #28498

0cc4m · 2026-10-05T09:03:16Z

@rillomas I can't reproduce the Intel failure either.

I'll look at the Nvidia one.

0cc4m · 2026-10-05T09:29:51Z

I can reproduce the Nvidia issue, will look into it.

0cc4m · 2026-10-05T10:28:52Z

Found the issue, should be fixed by #29988

rillomas · 2026-10-07T03:37:03Z

I think that the gpu-vulkan-intel-linux job failing with the new test is likely a bug that is specific to Vulkan with Intel hardware.

I've updated the runner's image with latest mesa/packages and I don't see the error anymore. Should be fine now.

rillomas · 2026-10-09T03:19:38Z

I'm still seeing flaky test results on the LunarLake instance so made an issue.
#30199

…org#29958) * llama : fix unexpected graph reallocation in the k-pool models Both k-pool models built a graph shape that depends on state the full-context reserve cannot know: - qwen4exp branched on inp->cache_safe, which turns false as soon as llama_memory_seq_cp shares cells (e.g. batched-bench -pps): the QSA layers swapped scatter+gather for fill+concat and dropped the new_pool_rep leaf, so the decode graph had 12 fewer nodes than the reserved one - glm5-next branched on gather = n_tokens <= 16 && n_kv > n_sel, so the TG decode built the gather shape (7564 nodes) while the last reserve, the PP one, had the dense shape (7762 nodes) Either mismatch forces a decode-time re-reserve that drops the worst-case sizing and bakes in the current state, so the next state growth (n_pool, n_kv, n_new) needs more room at an unchanged graph size and aborts under GGML_SCHED_DEBUG_REALLOC=1. Reproduce with, e.g.: GGML_SCHED_DEBUG_REALLOC=1 ./bin/llama-batched-bench \ -hf ggml-org/GLM-5.3-Flash-GGUF:Q2_K -npp 2500 -ntg 32 -npl 1,2 \ -c 32768 -pps -kvu Always scatter+gather the pooled keys, and pick gather from context constants only: n_ubatch bounds every ubatch, top_k + kpool - 1 bounds n_sel. Every graph of a context then shares one shape, which the reserve covers, and the dense path measured faster than the gather path at 2.5k and 16k context. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * llama : drop the unused k-pool cache_safe graph API The k-pool graphs no longer branch on cache_safe, so nothing reads get_kpool_cache_safe() or the conditional new_pool_rep any more: both models always pass the scatter target, which set_input_kpool now requires instead of merely preferring. Also drop the cache_safe copy in kpool_build_sizes(), a sizes-only helper. The layout and state flag itself stays, it still decides which pools a layout with shared cells must re-pool. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * tests : add a shared-seq graph reserve regression test Decode a prompt into seq 0, share its cells with seq 1 via llama_memory_seq_cp (what llama-batched-bench does for -pps), then keep decoding both sequences. For the k-pool models sharing clears cache_safe, which changes the graph topology while the pools keep growing, so a scheduler that re-reserves with the current state instead of the worst-case one aborts under GGML_SCHED_DEBUG_REALLOC=1. The test registration sets that flag, and the test aborts on both k-pool models before 2220411. kimi-linear and minimax-01 are skipped: they reserve the final pp graph with n_seqs = 1 (see [TAG_RESERVE_DIAG_DECAY] in llama-context.cpp), so every multi-seq graph has a different layout and re-reserves by design. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * cont : add TODOs * cont : fix comment * cuda: match the moe weighted reduction on empty ubatches ggml_cuda_match_moe_weighted_reduction rejected tensors with zero rows. A ubatch without outputs shrinks the last layer to zero rows through inp_out_ids, so graph_optimize dropped its alloc dep there and the scheduler graph lost one node compared to the reserved one. The scheduler then re-reserved at the size of that ubatch, and the next ubatch with the same node count but larger tensors aborted under GGML_SCHED_DEBUG_REALLOC=1. The compute loop already skips empty nodes before trying any fusion, so the guard only made the alloc deps depend on the row count. * tests: build the rollback test only where internal symbols link The shared-seq case calls llm_arch_from_string, which libllama does not export through LLAMA_API, so linking test-recurrent-state-rollback fails on Windows with shared libraries. Its build now sits in the NOT WIN32 OR NOT BUILD_SHARED_LIBS block, next to test-llama-archs and the test registration it already lives under. * tests: skip archs by name in the shared-seq reserve test The skip of kimi-linear and minimax-01 went through llm_arch_from_string, which libllama does not export through LLAMA_API, so the test could not link on Windows with shared libraries. It now compares the general.architecture string directly, and the test builds on every platform again. --------- Co-authored-by: Pascal <admin@serveurperso.com>

…org#29958) * llama : fix unexpected graph reallocation in the k-pool models Both k-pool models built a graph shape that depends on state the full-context reserve cannot know: - qwen4exp branched on inp->cache_safe, which turns false as soon as llama_memory_seq_cp shares cells (e.g. batched-bench -pps): the QSA layers swapped scatter+gather for fill+concat and dropped the new_pool_rep leaf, so the decode graph had 12 fewer nodes than the reserved one - glm5-next branched on gather = n_tokens <= 16 && n_kv > n_sel, so the TG decode built the gather shape (7564 nodes) while the last reserve, the PP one, had the dense shape (7762 nodes) Either mismatch forces a decode-time re-reserve that drops the worst-case sizing and bakes in the current state, so the next state growth (n_pool, n_kv, n_new) needs more room at an unchanged graph size and aborts under GGML_SCHED_DEBUG_REALLOC=1. Reproduce with, e.g.: GGML_SCHED_DEBUG_REALLOC=1 ./bin/llama-batched-bench \ -hf ggml-org/GLM-5.3-Flash-GGUF:Q2_K -npp 2500 -ntg 32 -npl 1,2 \ -c 32768 -pps -kvu Always scatter+gather the pooled keys, and pick gather from context constants only: n_ubatch bounds every ubatch, top_k + kpool - 1 bounds n_sel. Every graph of a context then shares one shape, which the reserve covers, and the dense path measured faster than the gather path at 2.5k and 16k context. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * llama : drop the unused k-pool cache_safe graph API The k-pool graphs no longer branch on cache_safe, so nothing reads get_kpool_cache_safe() or the conditional new_pool_rep any more: both models always pass the scatter target, which set_input_kpool now requires instead of merely preferring. Also drop the cache_safe copy in kpool_build_sizes(), a sizes-only helper. The layout and state flag itself stays, it still decides which pools a layout with shared cells must re-pool. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * tests : add a shared-seq graph reserve regression test Decode a prompt into seq 0, share its cells with seq 1 via llama_memory_seq_cp (what llama-batched-bench does for -pps), then keep decoding both sequences. For the k-pool models sharing clears cache_safe, which changes the graph topology while the pools keep growing, so a scheduler that re-reserves with the current state instead of the worst-case one aborts under GGML_SCHED_DEBUG_REALLOC=1. The test registration sets that flag, and the test aborts on both k-pool models before 2220411. kimi-linear and minimax-01 are skipped: they reserve the final pp graph with n_seqs = 1 (see [TAG_RESERVE_DIAG_DECAY] in llama-context.cpp), so every multi-seq graph has a different layout and re-reserves by design. Assisted-by: pi:llama.cpp/MiMo-V2.6-Flash-MOPD * cont : add TODOs * cont : fix comment * cuda: match the moe weighted reduction on empty ubatches ggml_cuda_match_moe_weighted_reduction rejected tensors with zero rows. A ubatch without outputs shrinks the last layer to zero rows through inp_out_ids, so graph_optimize dropped its alloc dep there and the scheduler graph lost one node compared to the reserved one. The scheduler then re-reserved at the size of that ubatch, and the next ubatch with the same node count but larger tensors aborted under GGML_SCHED_DEBUG_REALLOC=1. The compute loop already skips empty nodes before trying any fusion, so the guard only made the alloc deps depend on the row count. * tests: build the rollback test only where internal symbols link The shared-seq case calls llm_arch_from_string, which libllama does not export through LLAMA_API, so linking test-recurrent-state-rollback fails on Windows with shared libraries. Its build now sits in the NOT WIN32 OR NOT BUILD_SHARED_LIBS block, next to test-llama-archs and the test registration it already lives under. * tests: skip archs by name in the shared-seq reserve test The skip of kimi-linear and minimax-01 went through llm_arch_from_string, which libllama does not export through LLAMA_API, so the test could not link on Windows with shared libraries. It now compares the general.architecture string directly, and the test builds on every platform again. --------- Co-authored-by: Pascal <admin@serveurperso.com> (cherry picked from commit c173a53) [ dropped the qwen4exp.cpp hunks, our qwen4exp does not use cache_safe. Ported test_shared_seq_reserve to llama_batch/llama_decode, since this base has no common_batch. ] Assisted-by: Claude Opus 5.5

ggml-cuda : keep the MoE weighted-reduction dependency on ubatches without outputs (upstream ggml-org#29958)

ggerganov added 2 commits October 4, 2026 20:12

github-actions Bot added the model Model specific label Oct 4, 2026

ggerganov added 2 commits October 4, 2026 21:56

cont : add TODOs

d69cf37

github-actions Bot added the testing Everything test related label Oct 4, 2026

cont : fix comment

30e9f97

ggerganov marked this pull request as ready for review October 4, 2026 19:02

ggerganov requested a review from CISC as a code owner October 4, 2026 19:02

ggerganov commented Oct 4, 2026

View reviewed changes

am17an approved these changes Oct 5, 2026

View reviewed changes

ServeurpersoCom requested a review from a team as a code owner October 5, 2026 05:12

github-actions Bot added ggml changes relating to the ggml tensor library for machine learning CUDA Related to the CUDA backend labels Oct 5, 2026

ServeurpersoCom approved these changes Oct 5, 2026

View reviewed changes

ggerganov commented Oct 5, 2026

View reviewed changes

ggerganov merged commit c173a53 into master Oct 5, 2026
29 of 36 checks passed

ggerganov deleted the gg/sparse-fix-realloc branch October 5, 2026 08:36

0cc4m mentioned this pull request Oct 5, 2026

vulkan: fix Flash Attention shmem write out of bounds #29988

Merged

ServeurpersoCom mentioned this pull request Oct 5, 2026

llama: fix k-pool scatter data race on shared sequences #29994

Merged

This was referenced Oct 5, 2026

ggml : bump version to 0.26.0 ggml-org/ggml#1652

Merged

llama.cpp : bump version to 0.6.0 #29997

Merged

BrewTestBot mentioned this pull request Oct 5, 2026

ggml 0.26.0 Homebrew/homebrew-core#315641

Merged

1 task

pwilkin mentioned this pull request Oct 5, 2026

llama : re-reserve the sched when the nextn extraction flags change #30020

Merged

gaetan-puleo mentioned this pull request Oct 14, 2026

Rebase onto upstream llama.cpp c114bd4c5 (2026-10-14) halo-box/strix-llama.cpp#178

Closed

3 of 5 tasks

pratiknarola-t mentioned this pull request Oct 15, 2026

ggml-cuda : keep the MoE weighted-reduction dependency on ubatches without outputs (upstream #29958) tetherto/qvac-fabric-llm.cpp#394

Merged

Navigation Menu

Repository navigation

llama : fix unexpected graph reallocation in the k-pool models - #29958

llama : fix unexpected graph reallocation in the k-pool models#29958
ggerganov merged 8 commits into
masterfrom
gg/sparse-fix-realloc

ggerganov commented Oct 4, 2026

Uh oh!

ggerganov Oct 4, 2026

Uh oh!

ServeurpersoCom commented Oct 5, 2026

Uh oh!

ServeurpersoCom commented Oct 5, 2026

Uh oh!

ggerganov Oct 5, 2026

Uh oh!

ggerganov commented Oct 5, 2026

Uh oh!

Uh oh!

ggerganov commented Oct 5, 2026

Uh oh!

0cc4m commented Oct 5, 2026

Uh oh!

0cc4m commented Oct 5, 2026

Uh oh!

0cc4m commented Oct 5, 2026

Uh oh!

rillomas commented Oct 7, 2026

Uh oh!

rillomas commented Oct 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

Conversation

ggerganov commented Oct 4, 2026

Overview

Additional information

Requirements

Uh oh!

ggerganov Oct 4, 2026

Choose a reason for hiding this comment

Uh oh!

ServeurpersoCom commented Oct 5, 2026

Qwen3.8-Flash-Next UD-Q4_K_XL

GLM-5.3-Flash UD-IQ1_S**

Uh oh!

ServeurpersoCom commented Oct 5, 2026

Uh oh!

ggerganov Oct 5, 2026

Choose a reason for hiding this comment

Uh oh!

ggerganov commented Oct 5, 2026

Uh oh!

Uh oh!

ggerganov commented Oct 5, 2026

Uh oh!

0cc4m commented Oct 5, 2026

Uh oh!

0cc4m commented Oct 5, 2026

Uh oh!

0cc4m commented Oct 5, 2026

Uh oh!

rillomas commented Oct 7, 2026

Uh oh!

rillomas commented Oct 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants