From 79de1f5750ebdaaba4b26e85fac4b29f9ab243dc Mon Sep 17 00:00:00 2001
From: AlliDev <285906080+AIalliAI@users.noreply.github.com>
Date: Fri, 11 Sep 2026 01:22:57 +0000
Subject: [PATCH 1/2] model : fix k-pool indexer softmax gridDim.y overflow

Reshape the k-pool gate logits before softmax so n_new does not map to
gridDim.y, which is capped at 65535 on CUDA. Fixes prompt processing
abort at n_kv >= 262144 reported in #27773.

Assisted-by: Cursor
---
 src/models/glm5-next.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index b86d914e4c5f..9ed3095a464d 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -617,7 +617,9 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
 
         ggml_tensor * logits = ggml_add(ctx0, pg, layer.indexer_kpool_ape);
         logits = ggml_cont(ctx0, ggml_permute(ctx0, logits, 1, 0, 2, 3)); // [kpool, head_dim, n_new]
+        logits = ggml_reshape_2d(ctx0, logits, kpool, n_embd_indexer * n_new);
         ggml_tensor * probs = ggml_soft_max(ctx0, logits);
+        probs  = ggml_reshape_3d(ctx0, probs, kpool, n_embd_indexer, n_new);
 
         pk = ggml_cont(ctx0, ggml_permute(ctx0, pk, 1, 0, 2, 3));
         pooled_new = ggml_sum_rows(ctx0, ggml_mul(ctx0, probs, pk)); // [1, head_dim, n_new]

From 798d6cecb5c0baa824640776299f8751feb9b7fe Mon Sep 17 00:00:00 2001
From: AlliDev <285906080+AIalliAI@users.noreply.github.com>
Date: Fri, 11 Sep 2026 01:43:23 +0000
Subject: [PATCH 2/2] model : comment k-pool softmax gridDim.y reshape

Assisted-by: Cursor
---
 src/models/glm5-next.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp
index 9ed3095a464d..63d38643e69b 100644
--- a/src/models/glm5-next.cpp
+++ b/src/models/glm5-next.cpp
@@ -617,6 +617,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
 
         ggml_tensor * logits = ggml_add(ctx0, pg, layer.indexer_kpool_ape);
         logits = ggml_cont(ctx0, ggml_permute(ctx0, logits, 1, 0, 2, 3)); // [kpool, head_dim, n_new]
+        // soft_max launches gridDim.y = ne2, capped at 65535, and 262144/4 = 65536
         logits = ggml_reshape_2d(ctx0, logits, kpool, n_embd_indexer * n_new);
         ggml_tensor * probs = ggml_soft_max(ctx0, logits);
         probs  = ggml_reshape_3d(ctx0, probs, kpool, n_embd_indexer, n_new);