From 601cfe44334819a0a9ae4c1b8b32df8a81dc40a0 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 5 Sep 2026 19:49:31 +0000
Subject: [PATCH 01/33] cuda: add GGML_CUDA_BATCH_INVARIANT so a row does not
 depend on its batch

The number of tokens in a batch selects the matmul implementation, the flash
attention kernel, and inside several of them how the K loop or the KV cache is
divided between threads and blocks. All of those change the order in which the
partial products of one destination element are summed, so a request decoding
next to three others produces different bits than the same request decoding
alone, even at temperature 0.

GGML_CUDA_BATCH_INVARIANT=1 computes every destination column, and attends every
query row, with the configuration a batch of one would use. =2 does the same but
only where the batch-of-one configuration actually differs, which leaves the
quantized projections batched because MMVQ already uses the same nwarps for one
to four columns.

Flash attention additionally pins the vector kernel, pins the split over the KV
cache to one block per tile, and scans the mask for the sequence's own extent,
so neither the query count nor the length of a shared KV cache selects the
algorithm.

Measured on a B200 with Qwen3.5-4B-UD-Q4_K_XL: with the KV cache state held
equal, the 1831 node decode graph goes from 1190 nodes whose sequence-0 row
differs between a one token and a four token ubatch to 0, and 256 greedy decode
steps that diverged at step 125 become identical.
---
 ggml/src/ggml-cuda/common.cuh       |   3 +
 ggml/src/ggml-cuda/fattn-common.cuh |  14 ++-
 ggml/src/ggml-cuda/fattn.cu         |  37 ++++++
 ggml/src/ggml-cuda/ggml-cuda.cu     | 173 +++++++++++++++++++++++-----
 ggml/src/ggml-cuda/mmvq.cu          |  15 +++
 ggml/src/ggml-cuda/mmvq.cuh         |   5 +
 6 files changed, 217 insertions(+), 30 deletions(-)

diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
index 14dd1098c97..51ad1d6aa4d 100644
--- a/ggml/src/ggml-cuda/common.cuh
+++ b/ggml/src/ggml-cuda/common.cuh
@@ -49,6 +49,9 @@
 
 #define GGML_CUDA_CC_PASCAL          600
 #define GGML_CUDA_CC_DP4A            610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products
+// [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits
+int ggml_cuda_batch_invariant();
+
 #define GGML_CUDA_CC_VOLTA           700
 #define GGML_CUDA_CC_TURING          750
 #define GGML_CUDA_CC_AMPERE          800
diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh
index e67cc7fdf78..123e5071911 100644
--- a/ggml/src/ggml-cuda/fattn-common.cuh
+++ b/ggml/src/ggml-cuda/fattn-common.cuh
@@ -1091,7 +1091,10 @@ void launch_fattn(
     // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
     // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
     //     multiple sequences of possibly different lengths.
-    if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1)) {
+    // [TAG_BATCH_INVARIANT] Without this scan the KV loop runs to K->ne[1], which grows with the
+    // other sequences sharing the cache. Scanning the mask bounds it by the sequence's own extent.
+    const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0;
+    if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) {
         const int64_t s31 = mask->nb[1] / sizeof(half2);
         const int64_t s33 = mask->nb[3] / sizeof(half2);
 
@@ -1148,6 +1151,15 @@ void launch_fattn(
         if (ntiles_dst % blocks_num.x != 0) { // Fixup is only needed if the SMs work on fractional tiles.
             dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2)));
         }
+    } else if (ggml_cuda_batch_invariant()) {
+        // [TAG_BATCH_INVARIANT] How the KV cache is split between blocks, and therefore the order
+        // in which the partial attention results are combined, follows K->ne[1]. That length grows
+        // with the other sequences sharing the cache, so pin the split to a single block per tile.
+        parallel_blocks = 1;
+
+        blocks_num.x = ntiles_x;
+        blocks_num.y = parallel_blocks;
+        blocks_num.z = ntiles_z_gqa*K->ne[2]*Q->ne[3];
     } else {
         // parallel_blocks must not be larger than what the tensor size allows:
         parallel_blocks = std::min(parallel_blocks, ntiles_KV);
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index ab7a3b297c0..b2f6f0660ee 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -457,6 +457,13 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
     // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path.
     const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0;
 
+    // [TAG_BATCH_INVARIANT] Every choice below switches on Q->ne[1] or on K->ne[1], and both
+    // grow with the other sequences in the batch and in the shared KV cache. Pin the kernel a
+    // batch of one would use so a request is never moved onto a different algorithm by its neighbours.
+    if (ggml_cuda_batch_invariant() && can_use_vector_kernel && Q->ne[1] == 1) {
+        return BEST_FATTN_KERNEL_VEC;
+    }
+
     // If Turing tensor cores are available, use them:
     if (turing_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) {
         if (can_use_vector_kernel) {
@@ -569,6 +576,36 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d
 
 void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     ggml_cuda_set_device(ctx.device);
+
+    // [TAG_BATCH_INVARIANT] Attend one query row at a time, as a batch of one would.
+    if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1) {
+        const ggml_tensor * Q    = dst->src[0];
+        const ggml_tensor * mask = dst->src[3];
+
+        for (int64_t i = 0; i < Q->ne[1]; ++i) {
+            ggml_tensor Q_row = *Q;
+            Q_row.ne[1] = 1;
+            Q_row.data  = (char *) Q->data + i*Q->nb[1];
+
+            ggml_tensor mask_row;
+            ggml_tensor dst_row = *dst;
+            // ne[2] keeps running to the end of dst so that the scratch space for F16 copies of
+            // K and V, which is placed right behind dst, is still put in the same place.
+            dst_row.ne[2] = dst->ne[2] - i;
+            dst_row.data  = (char *) dst->data + i*dst->nb[2];
+            dst_row.src[0] = &Q_row;
+            if (mask) {
+                mask_row = *mask;
+                mask_row.ne[1] = 1;
+                mask_row.data  = (char *) mask->data + i*mask->nb[1];
+                dst_row.src[3] = &mask_row;
+            }
+
+            ggml_cuda_flash_attn_ext(ctx, &dst_row);
+        }
+        return;
+    }
+
     switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) {
         case BEST_FATTN_KERNEL_NONE:
             GGML_ABORT("fatal error");
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 2456f7dcc62..f9aa99ad003 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1758,6 +1758,12 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up,
 }
 
 static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) {
+    // [TAG_BATCH_INVARIANT] mul_mat+GLU is only fused for a single destination column, so
+    // leaving it on would give a solo request a different code path from a batched one.
+    if (ggml_cuda_batch_invariant()) {
+        return false;
+    }
+
     ggml_tensor *       src0 = tensor->src[0];
     ggml_tensor *       src1 = tensor->src[1];
     const ggml_tensor * dst  = tensor;
@@ -1785,6 +1791,12 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) {
 }
 
 static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
+    // [TAG_BATCH_INVARIANT] mul_mat+GLU is only fused for a single destination column, so
+    // leaving it on would give a solo request a different code path from a batched one.
+    if (ggml_cuda_batch_invariant()) {
+        return false;
+    }
+
     ggml_tensor *       src0 = tensor->src[0];
     ggml_tensor *       src1 = tensor->src[1];
     const ggml_tensor * dst  = tensor;
@@ -1813,60 +1825,163 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
     return use_mul_mat_vec_q;
 }
 
-static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
-    GGML_TENSOR_BINARY_OP_LOCALS
+// [TAG_BATCH_INVARIANT]
+// The number of tokens in a batch picks both the matmul implementation below and, inside
+// several of them, how the K loop is divided between threads. Both change the order in
+// which the partial products of one destination element are summed, so the same request
+// produces different bits depending on how many other requests decode alongside it.
+//
+// GGML_CUDA_BATCH_INVARIANT removes that dependency:
+//   1 - compute every destination column on its own, exactly as a batch of one would.
+//   2 - split off only the columns whose batch-of-one configuration differs from the
+//       batched one, leaving the already invariant matmuls batched.
+int ggml_cuda_batch_invariant() {
+    static const int mode = []() {
+        const char * val = getenv("GGML_CUDA_BATCH_INVARIANT");
+        return val ? atoi(val) : 0;
+    }();
+    return mode;
+}
 
-    const int32_t hint = ggml_get_op_params_i32(dst, 1);
-    if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) {
-        return;
-    }
+enum ggml_cuda_mm_path {
+    GGML_CUDA_MM_CUBLAS_UNSUPPORTED,
+    GGML_CUDA_MM_MMVF,
+    GGML_CUDA_MM_MMVF_TRANSPOSED,
+    GGML_CUDA_MM_MMF,
+    GGML_CUDA_MM_MMVQ,
+    GGML_CUDA_MM_MMQ,
+    GGML_CUDA_MM_CUBLAS,
+};
 
+// The implementation ggml_cuda_mul_mat would pick for a batch of ne11 columns.
+static ggml_cuda_mm_path ggml_cuda_mul_mat_path(
+        int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int64_t ne11) {
     // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q.
     // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data.
     // Therefore, in such cases use cuBLAS.
     const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE
         && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src;
     if (bad_padding_clear || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) {
-        ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);
-        return;
+        return GGML_CUDA_MM_CUBLAS_UNSUPPORTED;
     }
-
-    const int cc        = ggml_cuda_info().devices[ctx.device].cc;
-    const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size;
-
     if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) {
         // The custom F16 vector kernel can be used over batched cuBLAS GEMM.
         // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention)
-        ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst);
-        return;
+        return GGML_CUDA_MM_MMVF;
     }
     // A transposed vector can still use MMVQ (i.e. ne01 == 1)
-    if (ne01 == 1 && ne11 > MMVF_MAX_BATCH_SIZE && ne2 == 1 && ne3 == 1
+    if (src0->ne[1] == 1 && ne11 > MMVF_MAX_BATCH_SIZE && dst->ne[2] == 1 && dst->ne[3] == 1
             && src0->type == GGML_TYPE_F32
             && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst)
             && ggml_cuda_should_use_mmvf(src1->type, cc, src1->ne, src1->nb, /*ne11 =*/ 1)) {
-        ggml_tensor dst_vec = *dst;
-        dst_vec.ne[0] = ne11;
-        dst_vec.ne[1] = 1;
-        dst_vec.nb[1] = dst_vec.nb[0]*ne11;
-        dst_vec.nb[2] = dst_vec.nb[1];
-        dst_vec.nb[3] = dst_vec.nb[1];
-        ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec);
-        return;
+        return GGML_CUDA_MM_MMVF_TRANSPOSED;
     }
     if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) {
-        ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst);
-        return;
+        return GGML_CUDA_MM_MMF;
     }
     if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) {
-        ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst);
-        return;
+        return GGML_CUDA_MM_MMVQ;
     }
     if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) {
-        ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst);
+        return GGML_CUDA_MM_MMQ;
+    }
+    return GGML_CUDA_MM_CUBLAS;
+}
+
+static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
+
+// Recompute dst one column at a time so that each column sees the batch-of-one configuration.
+// Returns false when the batched launch already gives every column that same value.
+static bool ggml_cuda_mul_mat_split_columns(
+        ggml_backend_cuda_context & ctx, int cc, int warp_size,
+        const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    const int64_t ncols_dst = dst->ne[1];
+    if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) {
+        return false;
+    }
+    // Only the token dimension is split, batched matmuls (attention) keep their shape.
+    if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) {
+        return false;
+    }
+
+    if (ggml_cuda_batch_invariant() >= 2) {
+        const ggml_cuda_mm_path path_one     = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1);
+        const ggml_cuda_mm_path path_batched = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ncols_dst);
+        if (path_one == path_batched) {
+            // Same implementation, but it still has to sum each destination element in the same order.
+            if (path_batched == GGML_CUDA_MM_MMVF) {
+                return false; // the block size follows K alone
+            }
+            if (path_batched == GGML_CUDA_MM_MMVQ &&
+                    ggml_cuda_mmvq_matches_single_column(src0->type, cc, ncols_dst)) {
+                return false;
+            }
+        }
+    }
+
+    for (int64_t i = 0; i < ncols_dst; ++i) {
+        ggml_tensor src1_col = *src1;
+        ggml_tensor dst_col  = *dst;
+
+        src1_col.ne[1] = 1;
+        src1_col.nb[2] = src1_col.nb[1];
+        src1_col.nb[3] = src1_col.nb[1];
+        src1_col.data  = (char *) src1->data + i*src1->nb[1];
+
+        dst_col.ne[1] = 1;
+        dst_col.nb[2] = dst_col.nb[1];
+        dst_col.nb[3] = dst_col.nb[1];
+        dst_col.data  = (char *) dst->data + i*dst->nb[1];
+
+        ggml_cuda_mul_mat(ctx, src0, &src1_col, &dst_col);
+    }
+    return true;
+}
+
+static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    GGML_TENSOR_BINARY_OP_LOCALS
+
+    const int32_t hint = ggml_get_op_params_i32(dst, 1);
+    if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) {
+        return;
+    }
+
+    const int cc        = ggml_cuda_info().devices[ctx.device].cc;
+    const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size;
+
+    if (ggml_cuda_batch_invariant() && ggml_cuda_mul_mat_split_columns(ctx, cc, warp_size, src0, src1, dst)) {
         return;
     }
-    ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);
+
+    switch (ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ne11)) {
+        case GGML_CUDA_MM_CUBLAS_UNSUPPORTED:
+        case GGML_CUDA_MM_CUBLAS:
+            ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst);
+            return;
+        case GGML_CUDA_MM_MMVF:
+            ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst);
+            return;
+        case GGML_CUDA_MM_MMVF_TRANSPOSED: {
+            ggml_tensor dst_vec = *dst;
+            dst_vec.ne[0] = ne11;
+            dst_vec.ne[1] = 1;
+            dst_vec.nb[1] = dst_vec.nb[0]*ne11;
+            dst_vec.nb[2] = dst_vec.nb[1];
+            dst_vec.nb[3] = dst_vec.nb[1];
+            ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec);
+            return;
+        }
+        case GGML_CUDA_MM_MMF:
+            ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst);
+            return;
+        case GGML_CUDA_MM_MMVQ:
+            ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst);
+            return;
+        case GGML_CUDA_MM_MMQ:
+            ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst);
+            return;
+    }
+    GGML_ABORT("fatal error");
 }
 
 // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization
diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu
index 97053480980..b14ef9681c5 100644
--- a/ggml/src/ggml-cuda/mmvq.cu
+++ b/ggml/src/ggml-cuda/mmvq.cu
@@ -541,6 +541,21 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int
     return 1;
 }
 
+// [TAG_BATCH_INVARIANT]
+bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst) {
+    if (ncols_dst < 1 || ncols_dst > MMVQ_MAX_BATCH_SIZE) {
+        return false;
+    }
+    const mmvq_parameter_table_id table_id = get_device_table_id(cc);
+    if (table_id == MMVQ_PARAMETERS_GB10) {
+        // There nwarps also depends on the K loop trip count, which the caller does not pass in.
+        return ncols_dst == 1;
+    }
+    // blocks_per_iter, which is what assigns K blocks to threads, is proportional to nwarps.
+    // rows_per_cuda_block only changes which rows a block owns, not the order within a row.
+    return calc_nwarps(type, 1, table_id) == calc_nwarps(type, (int) ncols_dst, table_id);
+}
+
 template <ggml_type type, int ncols_dst, bool has_fusion, bool small_k = false, bool halve_iters = false>
 __launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1)
 static __global__ void mul_mat_vec_q(
diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh
index 5605bf7a4e6..61a88b851ec 100644
--- a/ggml/src/ggml-cuda/mmvq.cuh
+++ b/ggml/src/ggml-cuda/mmvq.cuh
@@ -4,6 +4,11 @@
 
 bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11);
 
+// [TAG_BATCH_INVARIANT]
+// True when an MMVQ launch of ncols_dst columns sums each destination element in the same
+// order as a launch of a single column, i.e. when the column count leaves nwarps unchanged.
+bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst);
+
 // Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID,
 // based on the quantization type and GPU architecture (compute capability).
 int get_mmvq_mmid_max_batch(ggml_type type, int cc);

From b5c10293842a3e51dafce13b95829ea53af7146f Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 5 Sep 2026 20:52:21 +0000
Subject: [PATCH 02/33] cuda: bound the batch-invariant split with
 GGML_CUDA_BATCH_INVARIANT_MAX_COLS

Splitting a prompt-sized batch costs far more than splitting a decode-sized one:
a 273 token prefill becomes 273 single column matmuls and 273 single row
attention launches per layer, which took prompt processing from 2731 to 225
tok/s on a B200 while four-chat decode only lost 7 percent.

GGML_CUDA_BATCH_INVARIANT_MAX_COLS caps the width the split applies to, 0 keeps
the previous unbounded behaviour. At 8 it covers every decode batch the server
can form and leaves prefill alone, which restores prompt processing to 2696
tok/s and four-chat wall throughput to 127.2 against 136.4 unpatched.

The bound gives up invariance for the prompt phase, so it is opt-in rather than
the default.
---
 ggml/src/ggml-cuda/common.cuh   |  3 +++
 ggml/src/ggml-cuda/fattn.cu     |  4 +++-
 ggml/src/ggml-cuda/ggml-cuda.cu | 12 ++++++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)

diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
index 51ad1d6aa4d..728aa08dcb5 100644
--- a/ggml/src/ggml-cuda/common.cuh
+++ b/ggml/src/ggml-cuda/common.cuh
@@ -51,6 +51,9 @@
 #define GGML_CUDA_CC_DP4A            610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products
 // [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits
 int ggml_cuda_batch_invariant();
+// Widest batch the split is applied to, 0 = no bound. Prompt-sized batches cost far more to
+// split than decode-sized ones, and only prompt-phase invariance is given up by bounding it.
+int ggml_cuda_batch_invariant_max_cols();
 
 #define GGML_CUDA_CC_VOLTA           700
 #define GGML_CUDA_CC_TURING          750
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index b2f6f0660ee..6d20a756016 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -578,7 +578,9 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
     ggml_cuda_set_device(ctx.device);
 
     // [TAG_BATCH_INVARIANT] Attend one query row at a time, as a batch of one would.
-    if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1) {
+    const int fattn_max_cols = ggml_cuda_batch_invariant_max_cols();
+    if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1 &&
+            (fattn_max_cols <= 0 || dst->src[0]->ne[1] <= fattn_max_cols)) {
         const ggml_tensor * Q    = dst->src[0];
         const ggml_tensor * mask = dst->src[3];
 
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index f9aa99ad003..a007bad9f44 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1843,6 +1843,14 @@ int ggml_cuda_batch_invariant() {
     return mode;
 }
 
+int ggml_cuda_batch_invariant_max_cols() {
+    static const int max_cols = []() {
+        const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
+        return val ? atoi(val) : 0;
+    }();
+    return max_cols;
+}
+
 enum ggml_cuda_mm_path {
     GGML_CUDA_MM_CUBLAS_UNSUPPORTED,
     GGML_CUDA_MM_MMVF,
@@ -1903,6 +1911,10 @@ static bool ggml_cuda_mul_mat_split_columns(
     if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) {
         return false;
     }
+    const int max_cols = ggml_cuda_batch_invariant_max_cols();
+    if (max_cols > 0 && ncols_dst > max_cols) {
+        return false;
+    }
 
     if (ggml_cuda_batch_invariant() >= 2) {
         const ggml_cuda_mm_path path_one     = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1);

From f3ce9725e6ef793aa6b4a4d0fda89ff1896c2678 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 5 Sep 2026 21:39:00 +0000
Subject: [PATCH 03/33] cuda: add exact concurrency with canonical paged
 attention

---
 ggml/src/ggml-cuda/fattn-common.cuh |   6 +-
 ggml/src/ggml-cuda/fattn-vec.cuh    |  18 +-
 ggml/src/ggml-cuda/fattn.cu         |  15 ++
 ggml/src/ggml-cuda/ggml-cuda.cu     |  30 +++
 scripts/batchinv/README.md          |  52 ++++
 scripts/batchinv/bench.py           |  53 ++++
 scripts/batchinv/divergence.py      | 164 +++++++++++++
 scripts/batchinv/probe.cpp          | 368 ++++++++++++++++++++++++++++
 scripts/batchinv/prompts.py         |  53 ++++
 src/llama-graph.cpp                 |  10 +-
 src/llama-graph.h                   |   4 +-
 src/llama-kv-cache.cpp              |  97 ++++++++
 src/llama-kv-cache.h                |   7 +
 tests/test-backend-ops.cpp          |  46 ++++
 14 files changed, 914 insertions(+), 9 deletions(-)
 create mode 100644 scripts/batchinv/README.md
 create mode 100644 scripts/batchinv/bench.py
 create mode 100644 scripts/batchinv/divergence.py
 create mode 100644 scripts/batchinv/probe.cpp
 create mode 100644 scripts/batchinv/prompts.py

diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh
index 123e5071911..f6aa5b03ec1 100644
--- a/ggml/src/ggml-cuda/fattn-common.cuh
+++ b/ggml/src/ggml-cuda/fattn-common.cuh
@@ -1094,7 +1094,7 @@ void launch_fattn(
     // [TAG_BATCH_INVARIANT] Without this scan the KV loop runs to K->ne[1], which grows with the
     // other sequences sharing the cache. Scanning the mask bounds it by the sequence's own extent.
     const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0;
-    if (mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) {
+    if (!dst->src[5] && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) {
         const int64_t s31 = mask->nb[1] / sizeof(half2);
         const int64_t s33 = mask->nb[3] / sizeof(half2);
 
@@ -1151,7 +1151,7 @@ void launch_fattn(
         if (ntiles_dst % blocks_num.x != 0) { // Fixup is only needed if the SMs work on fractional tiles.
             dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2)));
         }
-    } else if (ggml_cuda_batch_invariant()) {
+    } else if (dst->src[5] || ggml_cuda_batch_invariant()) {
         // [TAG_BATCH_INVARIANT] How the KV cache is split between blocks, and therefore the order
         // in which the partial attention results are combined, follows K->ne[1]. That length grows
         // with the other sequences sharing the cache, so pin the split to a single block per tile.
@@ -1226,7 +1226,7 @@ void launch_fattn(
         V_data,
         mask ? ((const char *) mask->data) : nullptr,
         sinks ? ((const char *) sinks->data) : nullptr,
-        KV_max.ptr,
+        dst->src[5] ? (const int *) dst->src[5]->data : KV_max.ptr,
         !stream_k && parallel_blocks > 1 ? dst_tmp.ptr : (float *) KQV->data, dst_tmp_meta.ptr,
         scale, max_bias, m0, m1, n_head_log2, logit_softcap,
         Q->ne[0], ne01,     Q->ne[2], Q->ne[3], Q->nb[1], Q->nb[2], Q->nb[3],
diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh
index 69dd9368624..f402795942c 100644
--- a/ggml/src/ggml-cuda/fattn-vec.cuh
+++ b/ggml/src/ggml-cuda/fattn-vec.cuh
@@ -16,7 +16,7 @@ static constexpr __device__ int ggml_cuda_fattn_vec_get_nthreads_device() {
 #pragma clang diagnostic push
 #pragma clang diagnostic ignored "-Wpass-failed"
 #endif // __clang__
-template<int D, int ncols, ggml_type type_K, ggml_type type_V, bool use_logit_softcap> // D == head size
+template<int D, int ncols, ggml_type type_K, ggml_type type_V, bool use_logit_softcap, bool paged = false> // D == head size
 __launch_bounds__(ggml_cuda_fattn_vec_get_nthreads_device(), 1)
 static __global__ void flash_attn_ext_vec(
         const char * Q_ptr,
@@ -247,13 +247,25 @@ static __global__ void flash_attn_ext_vec(
 #endif // V_DOT2_F32_F16_AVAILABLE
     }
 
-    const int k_VKQ_max = KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11;
+    // In the paged specialization KV_max carries [count, physical page IDs...] per query.
+    // The loop and each warp's recurrence follow logical positions, never physical addresses.
+    static_assert(!paged || ncols == 1, "paged attention has one query per block");
+    const int * pages = paged ? KV_max + (sequence*int(ne01.z) + ic0)*(1 + ne11/FATTN_KQ_STRIDE) : nullptr;
+    const int k_VKQ_max = paged ? pages[0]*FATTN_KQ_STRIDE : (KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11);
+    const char * K_base = K;
+    const char * V_base = V;
+    const half * mask_base = maskh;
     K     += blockIdx.y*nthreads * nb11;
     V     += blockIdx.y*nthreads * nb21;
     maskh += blockIdx.y*nthreads;
     for (int k_VKQ_0 = blockIdx.y*nthreads; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nthreads,
-             // Increment pointers after each loop:
              K += gridDim.y*nthreads*nb11, V += gridDim.y*nthreads*nb21, maskh += gridDim.y*nthreads) {
+        if constexpr (paged) {
+            const int physical = pages[1 + k_VKQ_0/FATTN_KQ_STRIDE]*FATTN_KQ_STRIDE + k_VKQ_0%FATTN_KQ_STRIDE;
+            K = K_base + int64_t(physical)*nb11;
+            V = V_base + int64_t(physical)*nb21;
+            maskh = mask_base + physical;
+        }
 
         // Calculate KQ tile and keep track of new maximum KQ values:
         float KQ_reg[ncols]; // KQ in registers.
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index 6d20a756016..eff18212272 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -577,6 +577,21 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d
 void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     ggml_cuda_set_device(ctx.device);
 
+    if (dst->src[5]) {
+        GGML_ASSERT(dst->src[0]->ne[0] == 256 && dst->src[2]->ne[0] == 256);
+        GGML_ASSERT(dst->src[1]->type == GGML_TYPE_F16 && dst->src[2]->type == GGML_TYPE_F16);
+        GGML_ASSERT(dst->src[3] && dst->src[0]->ne[3] == 1);
+        GGML_ASSERT(dst->src[5]->type == GGML_TYPE_I32 && ggml_is_contiguous(dst->src[5]));
+        GGML_ASSERT(dst->src[5]->ne[0] == 1 + dst->src[1]->ne[1]/FATTN_KQ_STRIDE);
+        GGML_ASSERT(dst->src[5]->ne[1] == dst->src[0]->ne[1]);
+        float softcap;
+        memcpy(&softcap, (const float *) dst->op_params + 2, sizeof(softcap));
+        GGML_ASSERT(softcap == 0.0f);
+        fattn_kernel_t kernel = flash_attn_ext_vec<256, 1, GGML_TYPE_F16, GGML_TYPE_F16, false, true>;
+        launch_fattn<256, 1, 1>(ctx, dst, kernel, 4, 0, 128, false, false, false);
+        return;
+    }
+
     // [TAG_BATCH_INVARIANT] Attend one query row at a time, as a batch of one would.
     const int fattn_max_cols = ggml_cuda_batch_invariant_max_cols();
     if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1 &&
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index a007bad9f44..251758c60ee 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1835,8 +1835,17 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
 //   1 - compute every destination column on its own, exactly as a batch of one would.
 //   2 - split off only the columns whose batch-of-one configuration differs from the
 //       batched one, leaving the already invariant matmuls batched.
+static bool ggml_cuda_exact_concurrency() {
+    static const bool exact = []() {
+        const char * value = getenv("LLAMA_EXACT_CONCURRENCY");
+        return value && atoi(value) != 0;
+    }();
+    return exact;
+}
+
 int ggml_cuda_batch_invariant() {
     static const int mode = []() {
+        if (ggml_cuda_exact_concurrency()) { return 2; }
         const char * val = getenv("GGML_CUDA_BATCH_INVARIANT");
         return val ? atoi(val) : 0;
     }();
@@ -1845,6 +1854,7 @@ int ggml_cuda_batch_invariant() {
 
 int ggml_cuda_batch_invariant_max_cols() {
     static const int max_cols = []() {
+        if (ggml_cuda_exact_concurrency()) { return 0; }
         const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
         return val ? atoi(val) : 0;
     }();
@@ -1903,6 +1913,26 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
 static bool ggml_cuda_mul_mat_split_columns(
         ggml_backend_cuda_context & ctx, int cc, int warp_size,
         const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
+    // Recurrent-model output projections broadcast one weight matrix over sequence
+    // planes. These are token projections too, even though ne[2] or ne[3] is > 1.
+    // Normalize each plane before applying the existing selective column policy.
+    if (ggml_cuda_exact_concurrency() && src0->ne[2] == 1 && src0->ne[3] == 1 &&
+            (dst->ne[2] > 1 || dst->ne[3] > 1) &&
+            src1->ne[2] == dst->ne[2] && src1->ne[3] == dst->ne[3]) {
+        for (int64_t i3 = 0; i3 < dst->ne[3]; ++i3) {
+            for (int64_t i2 = 0; i2 < dst->ne[2]; ++i2) {
+                ggml_tensor src_plane = *src1;
+                ggml_tensor dst_plane = *dst;
+                src_plane.ne[2] = src_plane.ne[3] = 1;
+                dst_plane.ne[2] = dst_plane.ne[3] = 1;
+                src_plane.data = (char *) src1->data + i2*src1->nb[2] + i3*src1->nb[3];
+                dst_plane.data = (char *) dst->data + i2*dst->nb[2] + i3*dst->nb[3];
+                ggml_cuda_mul_mat(ctx, src0, &src_plane, &dst_plane);
+            }
+        }
+        return true;
+    }
+
     const int64_t ncols_dst = dst->ne[1];
     if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) {
         return false;
diff --git a/scripts/batchinv/README.md b/scripts/batchinv/README.md
new file mode 100644
index 00000000000..62a51487047
--- /dev/null
+++ b/scripts/batchinv/README.md
@@ -0,0 +1,52 @@
+# Exact concurrency experiment p
+
+Opt in before loading the model with `LLAMA_EXACT_CONCURRENCY=1`. This also forces
+`GGML_CUDA_BATCH_INVARIANT=2` with no column limit, including during prefill.
+
+The experimental policy supports unified, offloaded F16 K/V, causal flash attention,
+256-dimensional K and V heads, no attention soft cap, and no sliding window.
+Shared-weight matmuls over multiple sequence planes are normalized to one plane
+before the inherited selective column dispatcher. Without this, the recurrent
+output projection bypasses batch invariance during concurrent prefill.
+It is measured on text prompts with Qwen3.5-4B on one B200. Context shifting,
+position division, cross-sequence prefix copies, shared-prefix input tokens, and
+whole-context state loading are unsupported. Per-sequence state save and restore
+is supported. Unsupported cache transformations assert instead of silently
+violating the page invariant.
+
+The allocator owns pages of 256 cells on behalf of one (sequence, position/256).
+Position modulo 256 fixes the cell offset. Empty pages remain in the unified pool
+and can be allocated by any sequence. The metadata is derived from live cells so
+allocation rollback, tail removal, and sequence removal do not need another
+transaction log. Restoring a sequence allocates free pages from this same pool.
+The cost is up to 255 reserved cells per active sequence tail, plus holes introduced
+by partial range removal.
+
+Attention receives an I32 page table in source 5: `[count, physical page IDs...]`
+for each query, sorted by logical position. The physical K/V view and mask span
+the pool, but the attention loop only visits the query's logical pages. The
+final page is padded to 256 cells using the existing causal mask. Wholly future
+pages in a prefill ubatch are excluded from the query's table.
+
+The vector attention specialization runs one query per block, four warps, and
+`parallel_blocks=1`. It reads K/V directly from the physical pages, with two
+128-cell softmax iterations per page, in logical page order. There is no K/V
+gather and no split-K combine. The default vector specialization has no page
+lookup. The ordinary path allocates no page metadata and launches no extra kernels.
+
+`FATTN_KQ_STRIDE=256` is a mask-scan stride, not the actual MMA rescaling tile.
+For K/V head size 256, the Ampere-or-newer MMA configurations use 64 KV rows at
+8 query/head columns and 32 KV rows at 16/32/64 columns. The retained vector path
+has a 128-cell iteration. Both divide the 256-cell placement page.
+
+The probe is adapted from the existing batch-invariant harness and rejects
+nonfinite logits and attention. `PROBE_B_REVERSE=1` fills neighbours before P0;
+`PROBE_RESTORE=1` parks P0, releases a neighbour, restores P0, then rebuilds the
+neighbour. Compute rows can be compared across this relocation; physical cache
+views and index tensors must not be mistaken for sequence-0 compute outputs.
+
+`divergence.py --reference FILE` compares against an existing unparked solo token
+reference. `bench.py --modes 0,1 --pairs 3` measures default off against exact mode
+on, with 256 predicted tokens. Set `UNSLOTH_WORKSPACE` to the model parent workspace
+and `LD_LIBRARY_PATH` to this build's bin directory. The harness uses GPU 3;
+select a port in 9601-9610 explicitly.
diff --git a/scripts/batchinv/bench.py b/scripts/batchinv/bench.py
new file mode 100644
index 00000000000..64d8b32244b
--- /dev/null
+++ b/scripts/batchinv/bench.py
@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+"""Cost of the knob: solo tok/s and four-chat aggregate tok/s, knob off and on, back to back."""
+import argparse, json, os, sys, time
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from divergence import Server, completion, run_concurrent
+from prompts import PROMPTS
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--binary", required=True)
+    ap.add_argument("--spec", default="none")
+    ap.add_argument("--n-predict", type=int, default=256)
+    ap.add_argument("--pairs", type=int, default=3)
+    ap.add_argument("--port", type=int, default=9602)
+    ap.add_argument("--modes", default="0,1")
+    ap.add_argument("--out", required=True)
+    a = ap.parse_args()
+
+    modes = a.modes.split(",")
+    rows = []
+    for pair in range(a.pairs):
+        for mode in modes:
+            env = {"LLAMA_EXACT_CONCURRENCY": mode, "GGML_CUDA_BATCH_INVARIANT": "0" if mode == "0" else "2"}
+            with Server(a.port, a.binary, [], env, a.out + ".server.log", a.spec) as s:
+                completion(a.port, PROMPTS["P0"], 32)  # warm
+                solo = completion(a.port, PROMPTS["P0"], a.n_predict)
+                outs, wall = run_concurrent(a.port, ["P0", "P1", "P2", "P3"], a.n_predict)
+                row = {
+                    "pair": pair, "mode": mode, "spec": a.spec,
+                    "solo_tok_per_s": solo["timings"]["predicted_per_second"],
+                    "solo_prompt_tok_per_s": solo["timings"]["prompt_per_second"],
+                    "four_aggregate_tok_per_s": sum(o["timings"]["predicted_per_second"] for o in outs.values()),
+                    "four_wall_s": wall,
+                    "four_total_tokens": sum(len(o["tokens"]) for o in outs.values()),
+                }
+                row["four_wall_tok_per_s"] = row["four_total_tokens"] / wall
+                rows.append(row)
+                print(json.dumps(row), flush=True)
+    with open(a.out, "w") as f:
+        json.dump(rows, f, indent=2)
+
+    print("\n=== summary ===", flush=True)
+    for mode in modes:
+        rs = [r for r in rows if r["mode"] == mode]
+        for k in ("solo_tok_per_s", "four_aggregate_tok_per_s", "four_wall_tok_per_s", "solo_prompt_tok_per_s"):
+            vals = sorted(r[k] for r in rs)
+            print(f"mode={mode} {k}: median {vals[len(vals)//2]:.1f}  values {[round(v,1) for v in vals]}", flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py
new file mode 100644
index 00000000000..e451afca0e1
--- /dev/null
+++ b/scripts/batchinv/divergence.py
@@ -0,0 +1,164 @@
+#!/usr/bin/env python3
+"""Baseline / patched divergence harness: solo P0 vs P0 sharing batches with P1..P3."""
+import argparse, json, os, signal, subprocess, sys, threading, time, urllib.request, urllib.error
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from prompts import PROMPTS
+
+WS = os.environ["UNSLOTH_WORKSPACE"]
+MODEL = f"{WS}/models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf"
+
+
+def post(port, path, payload, timeout=1800):
+    req = urllib.request.Request(f"http://127.0.0.1:{port}{path}",
+                                 data=json.dumps(payload).encode(),
+                                 headers={"Content-Type": "application/json"})
+    with urllib.request.urlopen(req, timeout=timeout) as r:
+        return json.loads(r.read().decode())
+
+
+def get(port, path, timeout=10):
+    with urllib.request.urlopen(f"http://127.0.0.1:{port}{path}", timeout=timeout) as r:
+        return json.loads(r.read().decode())
+
+
+def completion(port, prompt, n_predict):
+    return post(port, "/completion", {
+        "prompt": prompt, "n_predict": n_predict, "temperature": 0.0, "top_k": 1,
+        "top_p": 1.0, "min_p": 0.0, "typical_p": 1.0, "seed": 0,
+        "repeat_penalty": 1.0, "presence_penalty": 0.0, "frequency_penalty": 0.0,
+        "cache_prompt": False, "return_tokens": True, "samplers": ["top_k", "temperature"],
+    })
+
+
+class Server:
+    def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=True):
+        self.port, self.log_path = port, log_path
+        self.args = [binary, "-m", MODEL, "--port", str(port), "--host", "127.0.0.1",
+                     "--parallel", "4", "-c", "8192",
+                     "--flash-attn", "on", "--metrics", "-ngl", "99", "--no-warmup",
+                     "--seed", "0", "--spec-type", spec]
+        if kv_unified:
+            self.args += ["--kv-unified"]
+        if spec == "draft-mtp":
+            self.args += ["--spec-draft-n-max", "2"]
+        self.args += extra
+        self.env = dict(os.environ)
+        self.env["CUDA_VISIBLE_DEVICES"] = "3"
+        self.env.update(env_extra)
+
+    def __enter__(self):
+        self.fh = open(self.log_path, "ab")
+        self.fh.write(("\n=== " + " ".join(self.args) + "\n=== env " +
+                       json.dumps({k: v for k, v in self.env.items()
+                                   if k.startswith("GGML") or k == "CUDA_VISIBLE_DEVICES"}) + "\n").encode())
+        self.fh.flush()
+        self.p = subprocess.Popen(self.args, stdout=self.fh, stderr=subprocess.STDOUT,
+                                  env=self.env, start_new_session=True)
+        print(f"[server] pid={self.p.pid} port={self.port} log={self.log_path}", flush=True)
+        deadline = time.time() + 600
+        while time.time() < deadline:
+            if self.p.poll() is not None:
+                raise RuntimeError(f"server died rc={self.p.returncode}, see {self.log_path}")
+            try:
+                if get(self.port, "/health").get("status") == "ok":
+                    print("[server] ready", flush=True)
+                    return self
+            except Exception:
+                time.sleep(1.0)
+        raise RuntimeError("server did not become healthy")
+
+    def __exit__(self, *a):
+        print(f"[server] stopping pid={self.p.pid}", flush=True)
+        try:
+            os.killpg(os.getpgid(self.p.pid), signal.SIGTERM)
+            self.p.wait(timeout=60)
+        except Exception:
+            try:
+                os.killpg(os.getpgid(self.p.pid), signal.SIGKILL)
+            except Exception:
+                pass
+        self.fh.close()
+
+
+def run_concurrent(port, names, n_predict):
+    barrier = threading.Barrier(len(names))
+    out = {}
+
+    def work(name):
+        barrier.wait()
+        out[name] = completion(port, PROMPTS[name], n_predict)
+
+    ts = [threading.Thread(target=work, args=(n,)) for n in names]
+    t0 = time.time()
+    for t in ts:
+        t.start()
+    for t in ts:
+        t.join()
+    return out, time.time() - t0
+
+
+def first_diff(a, b):
+    for i, (x, y) in enumerate(zip(a, b)):
+        if x != y:
+            return i
+    return None if len(a) == len(b) else min(len(a), len(b))
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--reference")
+    ap.add_argument("--label", required=True)
+    ap.add_argument("--port", type=int, default=9601)
+    ap.add_argument("--binary", required=True)
+    ap.add_argument("--spec", default="none")
+    ap.add_argument("--n-predict", type=int, default=512)
+    ap.add_argument("--repeats", type=int, default=3)
+    ap.add_argument("--env", action="append", default=[])
+    ap.add_argument("--extra", action="append", default=[])
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--no-kv-unified", action="store_true")
+    a = ap.parse_args()
+
+    env_extra = dict(kv.split("=", 1) for kv in a.env)
+    res = {"label": a.label, "spec": a.spec, "n_predict": a.n_predict,
+           "env": env_extra, "extra": a.extra, "binary": a.binary,
+           "kv_unified": not a.no_kv_unified}
+
+    with Server(a.port, a.binary, a.extra, env_extra, a.out + ".server.log", a.spec,
+                kv_unified=not a.no_kv_unified) as s:
+        solo = completion(a.port, PROMPTS["P0"], a.n_predict)
+        ref = json.load(open(a.reference))["tokens"] if a.reference else solo["tokens"]
+        res["solo_first_diff"] = first_diff(ref, solo["tokens"])
+        res["reference"] = a.reference
+        res["solo"] = {"n_tokens": len(ref), "tok_per_s": solo["timings"]["predicted_per_second"],
+                       "text_sha": None}
+        # solo repeat, to prove solo itself is stable
+        solo2 = completion(a.port, PROMPTS["P0"], a.n_predict)
+        res["solo_repeat_first_diff"] = first_diff(ref, solo2["tokens"])
+        res["rounds"] = []
+        for r in range(a.repeats):
+            outs, wall = run_concurrent(a.port, ["P0", "P1", "P2", "P3"], a.n_predict)
+            p0 = outs["P0"]["tokens"]
+            fd = first_diff(ref, p0)
+            agg = sum(outs[n]["timings"]["predicted_per_second"] for n in outs)
+            row = {"round": r, "first_diff": fd, "n_tokens": len(p0),
+                   "identical": fd is None, "wall_s": wall,
+                   "p0_tok_per_s": outs["P0"]["timings"]["predicted_per_second"],
+                   "aggregate_tok_per_s": agg,
+                   "per_req_n": {n: len(outs[n]["tokens"]) for n in outs}, "p0_tokens": p0}
+            res["rounds"].append(row)
+            print(f"[round {r}] first_diff={fd} identical={fd is None} wall={wall:.1f}s agg={agg:.1f} tok/s", flush=True)
+        with urllib.request.urlopen(f"http://127.0.0.1:{a.port}/metrics") as response:
+            res["metrics"] = response.read().decode()
+        with open(a.out + ".p0_solo.json", "w") as f:
+            json.dump({"tokens": ref, "content": solo["content"]}, f)
+
+    with open(a.out, "w") as f:
+        json.dump(res, f, indent=2)
+    print(json.dumps({k: v for k, v in res.items() if k != "rounds"}, indent=2), flush=True)
+    print(json.dumps(res["rounds"], indent=2), flush=True)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp
new file mode 100644
index 00000000000..151463cd768
--- /dev/null
+++ b/scripts/batchinv/probe.cpp
@@ -0,0 +1,368 @@
+// Locate the first graph op whose sequence-0 output changes when the decode batch
+// holds four sequences instead of one. Prompt KV for seq 0 is built identically in
+// both phases, so the only difference is the width of the final decode ubatch.
+#include "llama.h"
+#include "ggml.h"
+#include "ggml-backend.h"
+
+#include <cstdio>
+#include <cmath>
+#include <cstdlib>
+#include <cstring>
+#include <cstdint>
+#include <fstream>
+#include <sstream>
+#include <string>
+#include <vector>
+
+struct node_rec {
+    std::string name;
+    std::string op;
+    std::string tname;
+    int64_t ne[4];
+    int64_t gdn_tokens = 0, gdn_seqs = 0;
+    size_t esize = 0;          // bytes per element, 0 = not byte comparable
+    std::vector<uint8_t> data; // empty when skipped
+    bool contiguous = false;
+    uint64_t hash = 0;
+};
+
+static bool g_record = false;
+static std::vector<node_rec> * g_sink = nullptr;
+
+static uint64_t fnv1a(const uint8_t * p, size_t n) {
+    uint64_t h = 1469598103934665603ULL;
+    for (size_t i = 0; i < n; ++i) { h ^= p[i]; h *= 1099511628211ULL; }
+    return h;
+}
+
+static bool eval_cb(struct ggml_tensor * t, bool ask, void * /*ud*/) {
+    if (!g_record) return false;
+    if (ask)       return true;
+
+    node_rec r;
+    r.name = ggml_get_name(t);
+    r.tname = ggml_type_name(t->type);
+    r.op   = t->op == GGML_OP_NONE ? "LEAF" : ggml_op_name(t->op);
+    if (t->op == GGML_OP_UNARY)  r.op = std::string("UNARY_")  + ggml_unary_op_name(ggml_get_unary_op(t));
+    if (t->op == GGML_OP_GLU)    r.op = std::string("GLU_")    + ggml_glu_op_name(ggml_get_glu_op(t));
+    for (int i = 0; i < 4; ++i) r.ne[i] = t->ne[i];
+    r.contiguous = ggml_is_contiguous(t);
+    if (r.name == "linear_attn_out-0") {
+        fprintf(stderr, "linear_attn_out-0: weight=%s input=[%lld,%lld,%lld,%lld]\n",
+            ggml_type_name(t->src[0]->type), (long long)t->src[1]->ne[0],
+            (long long)t->src[1]->ne[1], (long long)t->src[1]->ne[2], (long long)t->src[1]->ne[3]);
+    }
+    if (t->op == GGML_OP_GATED_DELTA_NET) {
+        r.gdn_tokens = t->src[2]->ne[2];
+        r.gdn_seqs = t->src[2]->ne[3];
+    }
+
+    const size_t nbytes = ggml_nbytes(t);
+    if (r.contiguous && ggml_blck_size(t->type) == 1 && nbytes <= (256u << 20)) {
+        r.esize = ggml_type_size(t->type);
+        r.data.resize(nbytes);
+        ggml_backend_tensor_get(t, r.data.data(), 0, nbytes);
+        if (t->op == GGML_OP_FLASH_ATTN_EXT) {
+            for (size_t i = 0; i < nbytes/sizeof(float); ++i) {
+                float v; memcpy(&v, r.data.data() + i*sizeof(float), sizeof(float));
+                if (!std::isfinite(v)) { fprintf(stderr, "nonfinite attention: %s\n", t->name); exit(5); }
+            }
+        }
+        r.hash = fnv1a(r.data.data(), nbytes);
+        if (nbytes > (64u << 20)) { r.data.clear(); } // keep the hash only for the big ones
+    }
+    g_sink->push_back(std::move(r));
+    return true;
+}
+
+static std::string slurp(const char * path) {
+    std::ifstream f(path);
+    std::stringstream ss; ss << f.rdbuf();
+    std::string s = ss.str();
+    while (!s.empty() && (s.back() == '\n' || s.back() == '\r')) s.pop_back();
+    return s;
+}
+
+static std::vector<llama_token> tokenize(const llama_vocab * v, const std::string & s) {
+    std::vector<llama_token> out(s.size() + 16);
+    int n = llama_tokenize(v, s.c_str(), (int) s.size(), out.data(), (int) out.size(), true, false);
+    if (n < 0) { out.resize(-n); n = llama_tokenize(v, s.c_str(), (int) s.size(), out.data(), (int) out.size(), true, false); }
+    out.resize(n);
+    return out;
+}
+
+struct batch_holder {
+    std::vector<llama_token>   tok;
+    std::vector<llama_pos>     pos;
+    std::vector<int32_t>       nsid;
+    std::vector<llama_seq_id>  sid;
+    std::vector<llama_seq_id*> sidp;
+    std::vector<int8_t>        out;
+    llama_batch get() {
+        sidp.resize(tok.size());
+        for (size_t i = 0; i < tok.size(); ++i) sidp[i] = &sid[i];
+        llama_batch b{};
+        b.n_tokens = (int32_t) tok.size();
+        b.token = tok.data(); b.pos = pos.data(); b.n_seq_id = nsid.data();
+        b.seq_id = sidp.data(); b.logits = out.data();
+        return b;
+    }
+};
+
+static llama_token greedy(llama_context * ctx, int32_t i, int n_vocab) {
+    const float * l = llama_get_logits_ith(ctx, i);
+    for (int k = 0; k < n_vocab; ++k) {
+        if (!std::isfinite(l[k])) { fprintf(stderr, "nonfinite logits at %d\n", k); exit(4); }
+    }
+    int best = 0;
+    for (int k = 1; k < n_vocab; ++k) if (l[k] > l[best]) best = k;
+    return best;
+}
+
+// Feed a prompt as one decode call for one sequence, return the greedy next token.
+static llama_token feed(llama_context * ctx, const std::vector<llama_token> & p, llama_seq_id seq, int n_vocab) {
+    batch_holder h;
+    for (size_t i = 0; i < p.size(); ++i) {
+        h.tok.push_back(p[i]); h.pos.push_back((llama_pos) i);
+        h.nsid.push_back(1);   h.sid.push_back(seq);
+        h.out.push_back(i + 1 == p.size());
+    }
+    llama_batch b = h.get();
+    if (llama_decode(ctx, b) != 0) { fprintf(stderr, "decode failed\n"); exit(1); }
+    return greedy(ctx, (int32_t) p.size() - 1, n_vocab);
+}
+
+int main(int argc, char ** argv) {
+    const bool prefill = getenv("PROBE_PREFILL") != nullptr;
+    const char * model_path = argv[1];
+    const int  n_seqs       = argc > 2 ? atoi(argv[2]) : 4;    // width of the probed decode batch
+    const char * out_path   = argc > 3 ? argv[3] : nullptr;
+    std::vector<std::string> prompts;
+    for (int i = 4; i < argc; ++i) prompts.push_back(slurp(argv[i]));
+
+    llama_backend_init();
+    llama_model_params mp = llama_model_default_params();
+    mp.n_gpu_layers = 99;
+    llama_model * model = llama_model_load_from_file(model_path, mp);
+    if (!model) { fprintf(stderr, "model load failed\n"); return 1; }
+    const llama_vocab * vocab = llama_model_get_vocab(model);
+    const int n_vocab = llama_vocab_n_tokens(vocab);
+
+    std::vector<std::vector<llama_token>> ptok;
+    for (auto & s : prompts) ptok.push_back(tokenize(vocab, s));
+    for (size_t i = 0; i < ptok.size(); ++i) fprintf(stderr, "prompt %zu: %zu tokens\n", i, ptok[i].size());
+
+    auto make_ctx = [&]() {
+        llama_context_params cp = llama_context_default_params();
+        cp.n_ctx = 8192; cp.n_batch = 2048; cp.n_ubatch = 512;
+        if (prefill) { cp.n_ubatch = 2048; }
+        cp.n_seq_max = 4; cp.kv_unified = true;
+        cp.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_ENABLED;
+        cp.cb_eval = eval_cb; cp.cb_eval_user_data = nullptr;
+        cp.no_perf = true;
+        return llama_init_from_model(model, cp);
+    };
+
+    std::vector<node_rec> rec_a, rec_b;
+    llama_token first_tok[4] = {0, 0, 0, 0};
+
+    // Phase A: decode ubatch width 1. PROBE_A_FILL controls how many sequences are
+    // already in the shared KV cache, which is what sets K->ne[1] for attention.
+    const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1;
+    // PROBE_A_PERM reorders which prompt goes into which sequence in phase A. With the same
+    // multiset of prompts the cache keeps its length but the masked cells hold different data.
+    int a_perm[4] = {0, 1, 2, 3};
+    if (const char * perm = getenv("PROBE_A_PERM")) {
+        for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0';
+    }
+    {
+        llama_context * ctx = make_ctx();
+        g_sink = &rec_a; g_record = prefill;
+        for (int s = 0; s < a_fill; ++s) {
+            const llama_token t = feed(ctx, ptok[a_perm[s]], s, n_vocab);
+            if (a_perm[s] == 0) first_tok[0] = t;
+        }
+        batch_holder h;
+        h.tok = {first_tok[0]}; h.pos = {(llama_pos) ptok[0].size()};
+        h.nsid = {1}; h.sid = {0}; h.out = {1};
+        llama_batch b = h.get();
+        g_sink = &rec_a; g_record = !prefill;
+        if (llama_decode(ctx, b) != 0) { fprintf(stderr, "A decode failed\n"); return 1; }
+        g_record = false;
+        llama_free(ctx);
+    }
+
+    // Phase B: same seq-0 prompt KV, then a decode ubatch holding n_seqs tokens.
+    {
+        llama_context * ctx = make_ctx();
+        if (prefill) {
+            batch_holder h;
+            for (int seq = 0; seq < n_seqs; ++seq) {
+                for (size_t i = 0; i < ptok[0].size(); ++i) {
+                    h.tok.push_back(ptok[seq][i%ptok[seq].size()]); h.pos.push_back(i);
+                    h.nsid.push_back(1); h.sid.push_back(seq); h.out.push_back(i+1 == ptok[0].size());
+                }
+            }
+            auto b = h.get();
+            g_sink = &rec_b; g_record = true;
+            if (llama_decode(ctx, b) != 0) { return 6; }
+            g_record = false;
+            for (int seq = 0; seq < n_seqs; ++seq) {
+                first_tok[seq] = greedy(ctx, (seq+1)*ptok[0].size()-1, n_vocab);
+            }
+        } else for (int k = 0; k < n_seqs; ++k) {
+            const int s = getenv("PROBE_B_REVERSE") ? n_seqs - 1 - k : k;
+            first_tok[s] = feed(ctx, ptok[s], s, n_vocab);
+        }
+        if (getenv("PROBE_RESTORE")) {
+            std::vector<uint8_t> state(llama_state_seq_get_size(ctx, 0));
+            if (llama_state_seq_get_data(ctx, state.data(), state.size(), 0) != state.size()) { return 2; }
+            llama_memory_seq_rm(llama_get_memory(ctx), 0, -1, -1);
+            llama_memory_seq_rm(llama_get_memory(ctx), 1, -1, -1);
+            if (llama_state_seq_set_data(ctx, state.data(), state.size(), 0) != state.size()) { return 3; }
+            first_tok[1] = feed(ctx, ptok[1], 1, n_vocab);
+        }
+        if (first_tok[0] != 0 && rec_a.size()) {}
+        batch_holder h;
+        for (int s = 0; s < n_seqs; ++s) {
+            h.tok.push_back(first_tok[s]); h.pos.push_back((llama_pos) ptok[s].size());
+            h.nsid.push_back(1); h.sid.push_back(s); h.out.push_back(1);
+        }
+        llama_batch b = h.get();
+        g_sink = &rec_b; g_record = !prefill;
+        if (!prefill && llama_decode(ctx, b) != 0) { fprintf(stderr, "B decode failed\n"); return 1; }
+        g_record = false;
+        llama_free(ctx);
+    }
+
+    // Optional: keep decoding and report the first step at which seq 0's token differs.
+    const int n_steps = getenv("PROBE_STEPS") ? atoi(getenv("PROBE_STEPS")) : 0;
+    int first_bad_step = -1;
+    if (n_steps > 0) {
+        std::vector<llama_token> tok_a, tok_b;
+        for (int phase = 0; phase < 2; ++phase) {
+            const int fill = phase == 0 ? a_fill : n_seqs;
+            const int width = phase == 0 ? 1 : n_seqs;
+            std::vector<llama_token> & out = phase == 0 ? tok_a : tok_b;
+            llama_context * ctx = make_ctx();
+            std::vector<llama_token> next(4, 0);
+            std::vector<llama_pos>   pos(4, 0);
+            for (int s = 0; s < fill; ++s) {
+                const int p = phase == 0 ? a_perm[s] : s;
+                next[s] = feed(ctx, ptok[p], s, n_vocab);
+                pos[s]  = (llama_pos) ptok[p].size();
+            }
+            for (int step = 0; step < n_steps; ++step) {
+                batch_holder h;
+                for (int s = 0; s < width; ++s) {
+                    h.tok.push_back(next[s]); h.pos.push_back(pos[s]);
+                    h.nsid.push_back(1); h.sid.push_back(s); h.out.push_back(1);
+                }
+                llama_batch b = h.get();
+                if (llama_decode(ctx, b) != 0) { fprintf(stderr, "step decode failed\n"); exit(1); }
+                out.push_back(next[0]);
+                for (int s = 0; s < width; ++s) { next[s] = greedy(ctx, s, n_vocab); pos[s] += 1; }
+            }
+            llama_free(ctx);
+        }
+        for (int i = 0; i < n_steps; ++i) {
+            if (tok_a[i] != tok_b[i]) { first_bad_step = i; break; }
+        }
+        fprintf(stderr, "steps: %d  first differing step: %d\n", n_steps, first_bad_step);
+    }
+
+    fprintf(stderr, "nodes: A=%zu B=%zu  first tokens: %d %d %d %d\n",
+            rec_a.size(), rec_b.size(), first_tok[0], first_tok[1], first_tok[2], first_tok[3]);
+
+    // Walk both node lists in order and compare seq 0's slice.
+    FILE * out = out_path ? fopen(out_path, "w") : stdout;
+    fprintf(out, "{\"n_seqs\":%d,\"first_bad_step\":%d,\"nodes_a\":%zu,\"nodes_b\":%zu,\"diffs\":[", n_seqs, first_bad_step, rec_a.size(), rec_b.size());
+    size_t n = rec_a.size() < rec_b.size() ? rec_a.size() : rec_b.size();
+    int emitted = 0;
+    for (size_t i = 0; i < n; ++i) {
+        const node_rec & A = rec_a[i];
+        const node_rec & B = rec_b[i];
+        const char * verdict = nullptr;
+        double max_abs = 0.0;
+        size_t ndiff = 0, ncmp = 0;
+
+        if (A.name != B.name || A.op != B.op) {
+            verdict = "misaligned";
+        } else if (A.op == "GATED_DELTA_NET" && A.gdn_tokens == B.gdn_tokens &&
+                !A.data.empty() && !B.data.empty()) {
+            // Packed GDN outputs put ALL token outputs before ALL sequence states.
+            // Sequence 0's state therefore moves when the number of sequences changes.
+            const size_t output = A.ne[0]*A.gdn_tokens;
+            const size_t state = A.ne[0]*A.ne[1]/A.gdn_seqs - output;
+            for (size_t k = 0; k < output + state; ++k) {
+                const size_t ia = k < output ? k : A.gdn_seqs*output + k-output;
+                const size_t ib = k < output ? k : B.gdn_seqs*output + k-output;
+                float va, vb;
+                memcpy(&va, A.data.data()+ia*4, 4); memcpy(&vb, B.data.data()+ib*4, 4);
+                ++ncmp;
+                if (memcmp(&va, &vb, 4)) {
+                    ++ndiff;
+                    if (std::abs(double(va)-vb) > max_abs) { max_abs = std::abs(double(va)-vb); }
+                }
+            }
+            verdict = ndiff ? "row-differs" : nullptr;
+        } else if (A.esize == 0 || B.esize == 0 || A.esize != B.esize) {
+            verdict = "skipped";
+        } else {
+            int tdim = -1; bool same = true;
+            for (int d = 0; d < 4; ++d) {
+                if (A.ne[d] == B.ne[d]) continue;
+                same = false;
+                if (B.ne[d] == n_seqs*A.ne[d] && tdim < 0) tdim = d; else { tdim = -2; break; }
+            }
+            if (tdim == -2) {
+                verdict = "shape-incomparable";
+            } else if (same) {
+                verdict = (A.hash == B.hash) ? nullptr : "whole-tensor-differs";
+            } else if (A.data.empty() || B.data.empty()) {
+                verdict = "too-large";
+            } else {
+                // compare element (.., i_tdim = 0, ..) across all other indices
+                int64_t st[4] = {1, A.ne[0], A.ne[0]*A.ne[1], A.ne[0]*A.ne[1]*A.ne[2]};
+                int64_t stb[4] = {1, B.ne[0], B.ne[0]*B.ne[1], B.ne[0]*B.ne[1]*B.ne[2]};
+                for (int64_t i3 = 0; i3 < A.ne[3]; ++i3)
+                for (int64_t i2 = 0; i2 < A.ne[2]; ++i2)
+                for (int64_t i1 = 0; i1 < A.ne[1]; ++i1)
+                for (int64_t i0 = 0; i0 < A.ne[0]; ++i0) {
+                    int64_t idx[4] = {i0, i1, i2, i3};
+
+                    size_t oa = 0, ob = 0;
+                    for (int d = 0; d < 4; ++d) { oa += idx[d]*st[d]; ob += idx[d]*stb[d]; }
+                    ncmp++;
+                    const uint8_t * pa = A.data.data() + oa*A.esize;
+                    const uint8_t * pb = B.data.data() + ob*B.esize;
+                    if (memcmp(pa, pb, A.esize) != 0) {
+                        ndiff++;
+                        if (A.esize == 4) {
+                            float fa, fb; memcpy(&fa, pa, 4); memcpy(&fb, pb, 4);
+                            double d2 = fa - fb; if (d2 < 0) d2 = -d2;
+                            if (d2 > max_abs) max_abs = d2;
+                        }
+                    }
+                }
+                verdict = ndiff ? "row-differs" : nullptr;
+            }
+        }
+        {
+            if (emitted++) fprintf(out, ",");
+            fprintf(out, "\n{\"i\":%zu,\"name\":\"%s\",\"op\":\"%s\",\"ne_a\":[%lld,%lld,%lld,%lld],"
+                         "\"ne_b\":[%lld,%lld,%lld,%lld],\"type\":\"%s\",\"verdict\":\"%s\",\"ndiff\":%zu,\"ncmp\":%zu,\"max_abs\":%.6g}",
+                    i, A.name.c_str(), A.op.c_str(),
+                    (long long)A.ne[0],(long long)A.ne[1],(long long)A.ne[2],(long long)A.ne[3],
+                    (long long)B.ne[0],(long long)B.ne[1],(long long)B.ne[2],(long long)B.ne[3],
+                    A.tname.c_str(), verdict ? verdict : "same", ndiff, ncmp, max_abs);
+        }
+    }
+    fprintf(out, "\n]}\n");
+    if (out_path) fclose(out);
+
+    llama_model_free(model);
+    llama_backend_free();
+    return 0;
+}
diff --git a/scripts/batchinv/prompts.py b/scripts/batchinv/prompts.py
new file mode 100644
index 00000000000..860b65fe08d
--- /dev/null
+++ b/scripts/batchinv/prompts.py
@@ -0,0 +1,53 @@
+# Four distinct prompts, each about 300 tokens of raw text (no chat template).
+_BODIES = {
+"P0": """The history of numerical computing is a history of compromises between speed and exactness.
+Early machines used fixed point arithmetic because it was cheap, and programmers carried scaling
+factors in their heads. Floating point hardware moved the bookkeeping into silicon, but it did not
+remove the compromise, it only hid it. Addition of floating point numbers is commutative but it is
+not associative, so the order in which a long sum is accumulated changes the last few bits of the
+result. On a single processor that order is fixed by the program text and nobody notices. On a
+parallel processor the order is fixed by how the work was divided, and the division is chosen for
+speed, not for reproducibility. A reduction split across two warps sums a different set of partial
+products than the same reduction split across four warps, and the two answers differ in the low
+bits. Nothing is wrong with either answer. Both are within a fraction of an ulp of the exact value.
+The trouble begins when a downstream decision is discrete. A comparison, a rounding to an integer,
+or the selection of the largest element of a vector turns a difference of one bit into a difference
+of one branch, and from there the two computations walk away from each other and never come back.
+Explain, carefully and at length, why this matters for a system that serves many users at once,
+and what an engineer would have to give up to make the answer depend only on the request and not
+on what else the machine happened to be doing at the time. Discuss the cost.""",
+"P1": """Consider a public library that lends physical books and must decide how many copies of a
+popular title to buy. The librarian has a fixed budget, a waiting list that grows and shrinks, and
+a shelf that is already full. Every copy purchased shortens the queue for that title and lengthens
+the queue for every other title, because the money and the shelf space are shared. The obvious
+policy, buy copies of whatever has the longest queue, is unstable, because a title that briefly
+becomes fashionable will absorb the whole budget and then sit unread for a decade. A better policy
+has to weigh how long the demand is likely to last against how long the book will remain useful,
+and it has to do this with almost no information. Describe in detail how you would design such a
+policy, what data you would collect, how you would test it without harming readers, and how you
+would know whether it was working. Consider what happens when the budget is cut in half without
+warning, when a title is suddenly assigned as required reading by a local school, and when the
+shelf itself must shrink because the building is being renovated. Explain the tradeoffs plainly.""",
+"P2": """A small coastal town has one bridge to the mainland and it is failing. The engineers say it
+has perhaps eight years left. Replacing it costs more than the town has ever spent on anything.
+Repairing it buys maybe four years and costs a third as much, and the repair work closes the bridge
+for two months in the summer, which is when the town earns most of its money. Doing nothing is
+free until the day it is not. The town council is split, the ferry operator has opinions, and the
+regional government will match funds only for a replacement, only if construction begins within
+three years, and only if the town covers the first quarter of the cost itself. Write a long and
+careful analysis of the options available to the council. Identify the assumptions that matter
+most, the ones where being wrong changes the recommendation, and say how the council could cheaply
+find out whether those assumptions hold. Then give a recommendation and state honestly what would
+have to be true for the recommendation to be wrong. Do not hedge. Commit to an answer at the end.""",
+"P3": """Describe the process by which a large body of water freezes over in winter, beginning with
+the surface layer and working downward, and explain why the ice floats rather than sinking, why a
+deep lake takes much longer to freeze than a shallow one of the same surface area, and why the
+temperature at the bottom of a frozen lake settles near four degrees Celsius rather than at zero.
+Then explain what this means for the animals that live there, how fish survive a winter under a
+solid lid, why a heavy snowfall on top of the ice can be more dangerous to them than the cold
+itself, and what happens in the spring when the whole column overturns. Use plain language and
+avoid equations. Where a common explanation is wrong or incomplete, say so and give the better one.
+Be thorough. Assume the reader is curious and patient but has no training in physics or biology,
+and would rather understand one thing properly than be told five things quickly.""",
+}
+PROMPTS = {k: " ".join(v.split()) for k, v in _BODIES.items()}
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 8fca8e1bc0e..71fce85bca7 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -468,6 +468,7 @@ void llm_graph_input_attn_no_cache::set_input(const llama_ubatch * ubatch) {
 }
 
 void llm_graph_input_attn_kv::set_input(const llama_ubatch * ubatch) {
+    if (self_pages && self_pages->buffer) { mctx->set_input_pages(self_pages, ubatch); }
     mctx->set_input_k_idxs(self_k_idxs, ubatch);
     mctx->set_input_v_idxs(self_v_idxs, ubatch);
 
@@ -1084,6 +1085,7 @@ void llm_graph_input_attn_cross::set_input(const llama_ubatch * ubatch) {
 }
 
 void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) {
+    if (inp_attn->self_pages) { mctx->get_attn()->set_input_pages(inp_attn->self_pages, ubatch); }
     mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch);
     mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch);
 
@@ -2547,7 +2549,8 @@ ggml_tensor * llm_graph_context::build_attn_mha(
          ggml_tensor * sinks,
          ggml_tensor * v_mla,
                float   kq_scale,
-                 int   il) const {
+                 int   il,
+         ggml_tensor * pages) const {
     const bool v_trans = v->nb[1] > v->nb[2];
 
     // split the batch into streams if needed
@@ -2580,6 +2583,7 @@ ggml_tensor * llm_graph_context::build_attn_mha(
 
         cur = ggml_flash_attn_ext(ctx0, q, k, v, kq_mask, kq_scale, hparams.f_max_alibi_bias,
                                   hparams.attn_soft_cap ? hparams.f_attn_logit_softcapping : 0.0f);
+        cur->src[5] = pages;
         res->add_fused_node({LLM_FUSED_OP_FLASH_ATTN, cur, il});
 
         ggml_flash_attn_ext_add_sinks(cur, sinks);
@@ -2769,6 +2773,8 @@ static std::unique_ptr<llm_graph_input_attn_kv> build_attn_inp_kv_impl(
         inp->self_kq_mask_cnv = inp->self_kq_mask;
     }
 
+    inp->self_pages = mctx_cur->build_input_pages(ctx0, ubatch);
+    GGML_ASSERT(!inp->self_pages || (cparams.flash_attn && cparams.causal_attn));
     inp->self_k_rot = mctx_cur->build_input_k_rot(ctx0);
     inp->self_v_rot = mctx_cur->build_input_v_rot(ctx0);
 
@@ -2831,7 +2837,7 @@ ggml_tensor * llm_graph_context::build_attn(
     ggml_tensor * k = mctx_cur->get_k(ctx0, il);
     ggml_tensor * v = mctx_cur->get_v(ctx0, il);
 
-    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il);
+    ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il, inp->self_pages);
     cb(cur, "kqv_out", il);
 
     if (inp->self_v_rot) {
diff --git a/src/llama-graph.h b/src/llama-graph.h
index b388e028cb5..26f2169532d 100644
--- a/src/llama-graph.h
+++ b/src/llama-graph.h
@@ -319,6 +319,7 @@ class llm_graph_input_attn_no_cache : public llm_graph_input_i {
 
 class llm_graph_input_attn_kv : public llm_graph_input_i {
 public:
+    ggml_tensor * self_pages = nullptr; // I32 [1 + physical pages, n_tokens]
     llm_graph_input_attn_kv(
             const llama_hparams & hparams,
             const llama_cparams & cparams,
@@ -1172,7 +1173,8 @@ struct llm_graph_context {
             ggml_tensor * sinks,   // [n_head_q]
             ggml_tensor * v_mla,   // [n_embd_head_v_mla, n_embd_head_v, n_head_v]
                   float   kq_scale,
-                    int   il) const;
+                    int   il,
+            ggml_tensor * pages = nullptr) const;
 
     llm_graph_input_attn_no_cache * build_attn_inp_no_cache() const;
 
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index ec0f5a75314..df643a047a9 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -84,6 +84,14 @@ llama_kv_cache::llama_kv_cache(
     v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),
     v_cells(*v_cells_impl) {
 
+    const char * exact_env = getenv("LLAMA_EXACT_CONCURRENCY");
+    exact_pages = exact_env && atoi(exact_env) != 0;
+    if (exact_pages) {
+        GGML_ASSERT(unified && offload && !v_trans && n_swa == 0);
+        GGML_ASSERT(type_k == GGML_TYPE_F16 && type_v == GGML_TYPE_F16);
+        GGML_ASSERT(kv_size % exact_page_size == 0);
+    }
+
     // shared cells view the source cache's K/V tensors, so the cell count
     // follows the source allocation: a fitted target can be smaller than the
     // draft default and oversized views would overflow the source tensors
@@ -447,6 +455,7 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
 }
 
 void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+    GGML_ASSERT(!exact_pages || seq_id_src == seq_id_dst);
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
@@ -566,6 +575,7 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
 }
 
 void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
+    GGML_ASSERT(!exact_pages || shift == 0);
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
@@ -616,6 +626,7 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll
 }
 
 void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
+    GGML_ASSERT(!exact_pages || d == 1);
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
@@ -961,6 +972,48 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
         }
     }
 
+    if (exact_pages) {
+        // Reconstruct page ownership from live cells. Empty pages are immediately reusable;
+        // prepare() can roll back its speculative allocations without a second metadata log.
+        const auto & cells = v_cells[0];
+        using page_key = std::pair<llama_seq_id, llama_pos>;
+        std::map<page_key, uint32_t> pages;
+        std::vector<bool> occupied(cells.size()/exact_page_size, false);
+        std::vector<bool> assigned(cells.size(), false);
+        for (uint32_t i = 0; i < cells.size(); ++i) {
+            if (cells.is_empty(i)) { continue; }
+            GGML_ASSERT(cells.seq_count(i) == 1);
+            const auto pos = cells.pos_get(i);
+            GGML_ASSERT(pos >= 0 && uint32_t(pos)%exact_page_size == i%exact_page_size);
+            const page_key key {cells.seq_get(i), pos/exact_page_size};
+            auto ins = pages.emplace(key, i/exact_page_size);
+            GGML_ASSERT(ins.first->second == i/exact_page_size);
+            occupied[i/exact_page_size] = true;
+        }
+        slot_info res {0, 0, {0}, {{}}};
+        for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+            GGML_ASSERT(ubatch.n_seq_id[i] == 1 && ubatch.pos[i] >= 0);
+            const page_key key {ubatch.seq_id[i][0], ubatch.pos[i]/exact_page_size};
+            auto it = pages.find(key);
+            if (it == pages.end()) {
+                // Round-robin free-page search deliberately permits nonmonotonic physical order.
+                uint32_t page = v_heads[0]/exact_page_size;
+                uint32_t tested = 0;
+                while (tested < occupied.size() && occupied[page%occupied.size()]) { ++page; ++tested; }
+                if (tested == occupied.size()) { return {}; }
+                page %= occupied.size();
+                occupied[page] = true;
+                it = pages.emplace(key, page).first;
+            }
+            const uint32_t idx = it->second*exact_page_size + ubatch.pos[i]%exact_page_size;
+            if (!cells.is_empty(idx) || assigned[idx]) { return {}; }
+            assigned[idx] = true;
+            res.idxs[0].push_back(idx);
+        }
+        if (cont && !res.is_contiguous()) { return {}; }
+        return res;
+    }
+
     uint32_t n_tokens = ubatch.n_tokens;
     uint32_t n_seqs   = 1;
 
@@ -1232,7 +1285,50 @@ const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const {
     return v_cells[seq_to_stream[seq_id]];
 }
 
+ggml_tensor * llama_kv_cache::build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const {
+    if (!exact_pages) { return nullptr; }
+    auto * pages = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1 + get_size()/exact_page_size, ubatch.n_tokens);
+    ggml_set_input(pages);
+    ggml_set_name(pages, "attn_logical_pages");
+    return pages;
+}
+
+void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const {
+    GGML_ASSERT(exact_pages && dst->ne[1] == ubatch->n_tokens);
+    std::map<llama_seq_id, std::map<llama_pos, uint32_t>> pages;
+    const auto & cells = v_cells[0];
+    for (uint32_t i = 0; i < cells.size(); ++i) {
+        if (!cells.is_empty(i)) {
+            GGML_ASSERT(cells.seq_count(i) == 1);
+            pages[cells.seq_get(i)][cells.pos_get(i)/exact_page_size] = i/exact_page_size;
+        }
+    }
+    std::vector<int32_t> data(ggml_nelements(dst), -1);
+    for (uint32_t i = 0; i < ubatch->n_tokens; ++i) {
+        GGML_ASSERT(ubatch->n_seq_id[i] == 1);
+        auto * row = data.data() + i*dst->ne[0];
+        row[0] = 0;
+        for (const auto & page : pages[ubatch->seq_id[i][0]]) {
+            // Exclude wholly future pages even when prefill includes later query rows.
+            if (page.first*exact_page_size > uint32_t(ubatch->pos[i])) { break; }
+            row[++row[0]] = page.second;
+        }
+    }
+    ggml_backend_tensor_set(dst, data.data(), 0, data.size()*sizeof(int32_t));
+}
+
+ggml_tensor * llama_kv_cache_context::build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const {
+    return kv->build_input_pages(ctx, ubatch);
+}
+
+void llama_kv_cache_context::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const {
+    kv->set_input_pages(dst, ubatch);
+}
+
 uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
+    // The physical view spans the pool. The page map, independently padded per query,
+    // is the only loop bound for exact attention; neighbours cannot extend that loop.
+    if (exact_pages) { return get_size(); }
     uint32_t result = 0;
 
     // pad the n_kv value so that the graph remains constant across batches and can be reused
@@ -2283,6 +2379,7 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
             GGML_ASSERT(cells.seq_has(idx, dest_seq_id));
         }
     } else {
+        GGML_ASSERT(!exact_pages && "exact mode supports per-sequence restore only");
         // whole KV cache restore
 
         if (cell_count > cells.size()) {
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index 6cb6dbd2f98..fa257422f02 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -171,6 +171,8 @@ class llama_kv_cache : public llama_memory_i {
     //
 
     uint32_t get_n_kv(const slot_info & sinfo) const;
+    ggml_tensor * build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const;
+    void set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const;
 
     // get views of the current state of the cache
     ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const;
@@ -235,6 +237,9 @@ class llama_kv_cache : public llama_memory_i {
         std::vector<ggml_tensor *> v_stream;
     };
 
+    static constexpr uint32_t exact_page_size = 256;
+    bool exact_pages = false;
+
     bool v_trans = true;  // the value tensor is transposed
 
     const uint32_t n_seq_max = 1;
@@ -365,6 +370,8 @@ class llama_kv_cache_context : public llama_memory_context_i {
     //
 
     uint32_t get_n_kv() const;
+    ggml_tensor * build_input_pages(ggml_context * ctx, const llama_ubatch & ubatch) const;
+    void set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const;
 
     ggml_type type_k() const;
     ggml_type type_v() const;
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 53e93a1448d..a7511396791 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -7193,6 +7193,41 @@ struct test_flash_attn_ext : public test_case {
     }
 };
 
+// Same mathematical attention as the CPU mask reference, but visit nonadjacent pages
+// in a different order. Covers a partial tail and different page counts per query.
+struct test_flash_attn_ext_pages : public test_flash_attn_ext {
+    test_flash_attn_ext_pages(int64_t batch) :
+        test_flash_attn_ext(256, 256, 2, {8, 1}, 1024, batch) {}
+
+    std::string vars() override { return test_flash_attn_ext::vars() + ",exact_pages=1"; }
+
+    ggml_tensor * build_graph(ggml_context * ctx) override {
+        auto * out = test_flash_attn_ext::build_graph(ctx);
+        out->src[5] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 5, nb);
+        ggml_set_name(out->src[5], "pages");
+        return out;
+    }
+
+    void initialize_tensors(ggml_context * ctx) override {
+        test_flash_attn_ext::initialize_tensors(ctx);
+        auto * pages = ggml_get_tensor(ctx, "pages");
+        auto * mask = ggml_get_tensor(ctx, "m");
+        std::vector<int32_t> ids(5*nb, -1);
+        std::vector<ggml_fp16_t> values(1024*nb, ggml_fp32_to_fp16(-INFINITY));
+        for (int64_t q = 0; q < nb; ++q) {
+            ids[5*q] = q%2 ? 1 : 2;
+            ids[5*q + 1] = 2;
+            ids[5*q + 2] = 0;
+            for (int j = 0; j < 256; ++j) { values[1024*q + 512 + j] = ggml_fp32_to_fp16(0.0f); }
+            if (q%2 == 0) {
+                for (int j = 0; j < 17; ++j) { values[1024*q + j] = ggml_fp32_to_fp16(0.0f); }
+            }
+        }
+        ggml_backend_tensor_set(pages, ids.data(), 0, ids.size()*sizeof(int32_t));
+        ggml_backend_tensor_set(mask, values.data(), 0, values.size()*sizeof(ggml_fp16_t));
+    }
+};
+
 // GGML_OP_CROSS_ENTROPY_LOSS
 struct test_cross_entropy_loss : public test_case {
     const ggml_type type;
@@ -9170,6 +9205,14 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
+    // Shared weights over sequence planes, as in a recurrent-model output projection.
+    for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0}) {
+        for (int n : {1, 17, 307}) {
+            test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {4, 1}));
+            test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {1, 4}));
+        }
+    }
+
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1}));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1}));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1}));
@@ -9937,6 +9980,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     }
 
     // mixed quant and Q1_0 test cases
+    for (int64_t batch : {1, 4, 12}) {
+        test_cases.emplace_back(new test_flash_attn_ext_pages(batch));
+    }
     test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q8_0, GGML_TYPE_Q4_0));
     test_cases.emplace_back(new test_flash_attn_ext(64, 64, 4, {1, 1}, 128, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_F16));
     test_cases.emplace_back(new test_flash_attn_ext(72, 72, 4, {1, 1}, 96, 2, true, false, 0, 0, GGML_PREC_F32, GGML_TYPE_Q4_0, GGML_TYPE_Q8_0));

From a94f76feaa9e7862ec2768e17d6a1b75a21c71e3 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 5 Sep 2026 23:20:47 +0000
Subject: [PATCH 04/33] batch: keep a prompt ubatch to one sequence under
 LLAMA_EXACT_CONCURRENCY

The recurrent half of a hybrid model is not invariant to the shape of the ubatch. With the
attention half made exact, a prompt processed in ubatches it shares with other sequences'
prompt tokens still leaves a different gated delta net state than the same prompt processed
alone: the first node to show it is the layer 1 recurrent state, 2.2e5 of 5.2e5 elements,
max 7.3e-4, and over a 512 token generation it flips a token at step 79.

split_equal grows an optional cap on the number of sequence sets per ubatch. The hybrid
memory passes 1 when the mode is on and some sequence contributes more than one token to
the batch, which is the prompt phase. A plain decode step, one token per sequence, is
already exact under the gather and stays batched, so the cost falls on prompt processing
only.
---
 src/llama-batch.cpp         | 21 ++++++++++++++++++++-
 src/llama-batch.h           |  8 +++++++-
 src/llama-impl.cpp          | 11 +++++++++++
 src/llama-impl.h            |  6 ++++++
 src/llama-memory-hybrid.cpp |  9 ++++++++-
 5 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 2b98a552f48..50f70bb0ff3 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -507,7 +507,21 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) {
     return ubatch_add(idxs, idxs.size(), false);
 }
 
-llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail) {
+bool llama_batch_allocr::has_multi_token_seq() const {
+    std::vector<uint32_t> n_per_seq(n_seq_max, 0);
+
+    for (int32_t i = 0; i < batch.n_tokens; ++i) {
+        for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
+            if (++n_per_seq[batch.seq_id[i][s]] > 1) {
+                return true;
+            }
+        }
+    }
+
+    return false;
+}
+
+llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t n_seqs_max) {
     if (sequential && has_cpl) {
         LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__);
 
@@ -547,6 +561,11 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
             if (cur_seq_set.size() > n_ubatch) {
                 break;
             }
+
+            // [TAG_EXACT_CONCURRENCY]
+            if (n_seqs_max > 0 && cur_seq_set.size() >= n_seqs_max) {
+                break;
+            }
         }
     }
 
diff --git a/src/llama-batch.h b/src/llama-batch.h
index a3d1889d4a0..d354c442d03 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -105,7 +105,13 @@ class llama_batch_allocr {
     // make ubatches of equal-length sequences sets
     // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
     // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
-    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail);
+    // n_seqs_max  = maximum sequence sets per ubatch, 0 = no limit
+    //               [TAG_EXACT_CONCURRENCY] passing 1 keeps a ubatch to a single sequence
+    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t n_seqs_max = 0);
+
+    // [TAG_EXACT_CONCURRENCY] true if some sequence contributes more than one token to the batch,
+    // i.e. this is not a plain one-token-per-sequence decode step
+    bool has_multi_token_seq() const;
 
     // sequence-set-wise split - each ubatch contains a single sequence-set
     llama_ubatch split_seq(uint32_t n_ubatch);
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index b3a94b946d2..bad0e55237a 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -6,6 +6,7 @@
 #include <cinttypes>
 #include <climits>
 #include <cstdarg>
+#include <cstdlib>
 #include <cstring>
 #include <vector>
 #include <sstream>
@@ -169,3 +170,13 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i) {
             return gguf_data_to_str(type, gguf_get_val_data(ctx_gguf, i), 0);
     }
 }
+
+// [TAG_EXACT_CONCURRENCY]
+bool llama_exact_concurrency() {
+    static const bool enabled = []() {
+        const char * val = getenv("LLAMA_EXACT_CONCURRENCY");
+        return val && atoi(val) != 0;
+    }();
+
+    return enabled;
+}
diff --git a/src/llama-impl.h b/src/llama-impl.h
index 4988b06d2ca..9b64431fedd 100644
--- a/src/llama-impl.h
+++ b/src/llama-impl.h
@@ -103,3 +103,9 @@ std::string llama_format_tensor_shape(const std::vector<int64_t> & ne);
 std::string llama_format_tensor_shape(const struct ggml_tensor * t);
 
 std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i);
+
+// [TAG_EXACT_CONCURRENCY]
+// opt-in mode under which a sequence's attention depends only on its own cells, in position order,
+// so that its output does not change when other sequences share the KV cache. Off by default.
+// Reads the same LLAMA_EXACT_CONCURRENCY variable as the paged KV cache and the CUDA backend.
+bool llama_exact_concurrency();
diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index 42c7381a9e6..ba54ab12923 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -86,7 +86,14 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
                 //   so that the rollback snapshots remain valid
                 const uint32_t n_rs_seq = mem_recr->n_rs_seq;
 
-                ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
+                // [TAG_EXACT_CONCURRENCY] the recurrent half of a hybrid model is not invariant to
+                // the shape of the ubatch: a prompt processed next to other sequences' prompt tokens
+                // leaves a different gated delta net state than the same prompt processed alone.
+                // Keeping such a ubatch to a single sequence removes that. A plain decode step, one
+                // token per sequence, is already exact and stays batched.
+                const uint32_t n_seqs_max = llama_exact_concurrency() && balloc.has_multi_token_seq() ? 1 : 0;
+
+                ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, n_seqs_max);
             }
 
             if (ubatch.n_tokens == 0) {

From 07d82f022d44d8339bd0270f1aedcf821c1b7e79 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 5 Sep 2026 23:31:58 +0000
Subject: [PATCH 05/33] cuda: let exact mode bound the column policy when
 prompt ubatches are per sequence

---
 ggml/src/ggml-cuda/ggml-cuda.cu | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 251758c60ee..2b1754baea6 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1854,9 +1854,13 @@ int ggml_cuda_batch_invariant() {
 
 int ggml_cuda_batch_invariant_max_cols() {
     static const int max_cols = []() {
-        if (ggml_cuda_exact_concurrency()) { return 0; }
+        // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill
+        // matmul shapes match its solo run, so exact mode no longer needs the column policy to be
+        // unbounded there. Honour an explicit bound when one is set; default to unbounded.
         const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
-        return val ? atoi(val) : 0;
+        if (val) { return atoi(val); }
+        if (ggml_cuda_exact_concurrency()) { return 0; }
+        return 0;
     }();
     return max_cols;
 }

From 4814a26852444bc8200038f6fa18827b874bffcb Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 02:10:23 +0000
Subject: [PATCH 06/33] cuda: give every MUL_MAT_ID token the single-token
 configuration

A mixture-of-experts decode groups the ubatch's tokens by the expert they routed
to, so the column count of an expert's matmul, the rows the per-expert copy
gathers and the width the activations are quantized at all depend on what the
other tokens in the ubatch picked. The quantized path makes it visible: at one
token per ubatch MUL_MAT_ID runs mul_mat_vec_q at ncols_dst 1, four warps
dividing the K loop and a shared memory reduction across them, and at more than
one token it runs the dedicated MoE kernel, one warp per token with a warp only
reduction. Whether those two agree bit for bit depends on the quantization type
and on K. Measured on a B200 they agree for Q4_K and Q5_K up to K 2048 and
disagree for Q6_K and Q8_0 from K 512, which is why on
Qwen3.6-35B-A3B-UD-Q4_K_XL the Q4_K gate and up projections of every layer
matched and the three Q6_K down projections, layers 34, 38 and 39, did not.

Under GGML_CUDA_BATCH_INVARIANT the node is now computed one token at a time.
Each call then sees the shapes a batch of one has, whatever the neighbours routed
to, which covers the ids variants of MMVQ, MMQ and MMF and the sorted per expert
fallback with one change. GGML_CUDA_BATCH_INVARIANT_MAX_COLS bounds it the way it
bounds the MUL_MAT column split. The CUDA graph fallback check is evaluated
against the single-token path as well, since that is what a split node runs.

On the 35B the sequence-0 slice of one decode step goes from 195 of 3727 nodes
differing between a one token and a four token ubatch to 0, and a standalone
MUL_MAT_ID probe over Q4_K, Q5_K, Q6_K, Q8_0, Q4_0, Q3_K, Q2_K, MXFP4, F16 and
BF16 at K 512, 2048 and 4096 goes to 0 at every token count up to 8.
---
 ggml/src/ggml-cuda/ggml-cuda.cu | 78 +++++++++++++++++++++++++++++++--
 tests/test-backend-ops.cpp      | 11 +++++
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 2b1754baea6..0995acdfd6a 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -2030,6 +2030,25 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
     GGML_ABORT("fatal error");
 }
 
+// [TAG_BATCH_INVARIANT]
+// True when the batch-invariant policy computes this MUL_MAT_ID one token at a time.
+// Every expert product then reduces the way it would in a batch of one, whatever the
+// rest of the ubatch routed to.
+static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) {
+    if (!ggml_cuda_batch_invariant()) {
+        return false;
+    }
+    const int64_t ntokens = dst->ne[2];
+    if (ntokens <= 1) {
+        return false;
+    }
+    const int max_cols = ggml_cuda_batch_invariant_max_cols();
+    if (max_cols > 0 && ntokens > max_cols) {
+        return false;
+    }
+    return true;
+}
+
 // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization
 // [TAG_MUL_MAT_ID_CUDA_GRAPHS]
 static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int cc) {
@@ -2040,9 +2059,13 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
         return true;
     }
 
-    if (dst->ne[2] <= MMVQ_MAX_BATCH_SIZE) {
+    // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path
+    // that decides whether the stream is synchronized is the single-token one.
+    const int64_t ntokens = ggml_cuda_mul_mat_id_splits_tokens(dst) ? 1 : dst->ne[2];
+
+    if (ntokens <= MMVQ_MAX_BATCH_SIZE) {
         if (ggml_is_quantized(src0->type)) {
-            if (dst->ne[2] <= get_mmvq_mmid_max_batch(src0->type, cc)) {
+            if (ntokens <= get_mmvq_mmid_max_batch(src0->type, cc)) {
                 return false;
             }
         } else if (GGML_CUDA_CC_IS_AMD(cc)) {
@@ -2050,17 +2073,57 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
         }
     }
 
-    if (ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[2], /*n_experts=*/src0->ne[2])) {
+    if (ggml_cuda_should_use_mmq(src0->type, cc, ntokens, /*n_experts=*/src0->ne[2])) {
         return false;
     }
 
-    if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) {
+    if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, ntokens, /*mul_mat_id=*/true)) {
         return false;
     }
 
     return true;
 }
 
+static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
+
+// [TAG_BATCH_INVARIANT]
+// Recompute dst one token at a time. Every implementation below groups the ubatch's tokens
+// by the expert they routed to, so the column count of an expert's matmul, the tokens the
+// per-expert copy gathers and the width the activations are quantized at all depend on what
+// the other tokens in the ubatch picked. Handing each token its own call removes that: the
+// callee sees the shapes a batch of one has, whatever the neighbours did.
+static void ggml_cuda_mul_mat_id_split_tokens(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
+    const ggml_tensor * src1 = dst->src[1];
+    const ggml_tensor * ids  = dst->src[2];
+
+    const int64_t ntokens = dst->ne[2];
+
+    for (int64_t i = 0; i < ntokens; ++i) {
+        ggml_tensor src1_token = *src1;
+        ggml_tensor ids_token  = *ids;
+        ggml_tensor dst_token  = *dst;
+
+        // src1 is [ne10, ne11, ntokens], one expert list per token in ids [n_expert_used, ntokens]
+        src1_token.ne[2] = 1;
+        src1_token.nb[3] = src1_token.nb[2];
+        src1_token.data  = (char *) src1->data + i*src1->nb[2];
+
+        ids_token.ne[1] = 1;
+        ids_token.nb[2] = ids_token.nb[1];
+        ids_token.nb[3] = ids_token.nb[1];
+        ids_token.data  = (char *) ids->data + i*ids->nb[1];
+
+        dst_token.ne[2] = 1;
+        dst_token.nb[3] = dst_token.nb[2];
+        dst_token.data  = (char *) dst->data + i*dst->nb[2];
+
+        dst_token.src[1] = &src1_token;
+        dst_token.src[2] = &ids_token;
+
+        ggml_cuda_mul_mat_id(ctx, &dst_token);
+    }
+}
+
 static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src0 = dst->src[0];
     const ggml_tensor * src1 = dst->src[1];
@@ -2073,6 +2136,13 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
 
     const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
 
+    // [TAG_BATCH_INVARIANT]
+    if (ggml_cuda_mul_mat_id_splits_tokens(dst)) {
+        GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1);
+        ggml_cuda_mul_mat_id_split_tokens(ctx, dst);
+        return;
+    }
+
     // [TAG_MUL_MAT_ID_CUDA_GRAPHS]
     if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
         static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE);
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index a7511396791..685193be77d 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -9213,6 +9213,17 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
+    // Mixture-of-experts projections at the token counts a decode ubatch forms. The gate and up
+    // projections broadcast one activation row over the expert list, the down projection carries
+    // one row per expert, and the mixed quantization of a real MoE gguf puts different types on
+    // the two. 17 tokens is past the width the exact-concurrency policy pins.
+    for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16}) {
+        for (int n : {1, 2, 4, 8, 17}) {
+            test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, true,   512, n, 2048));
+            test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, false, 2048, n,  512));
+        }
+    }
+
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q4_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1}));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_Q8_0, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1}));
     test_cases.emplace_back(new test_mul_mat(GGML_TYPE_MXFP4, GGML_TYPE_F32, 2880, 32, 2880, {1, 1}, {1, 1}));

From 7b0d7edb9eccadde314cffdd7445d5c9aefb8e10 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 02:10:34 +0000
Subject: [PATCH 07/33] cuda: leave the top-k routing unfused under
 GGML_CUDA_BATCH_INVARIANT

ggml_cuda_check_fusion_memory_ranges accepts the top-k MoE subgraph through an
explicit ggml_nrows(node) == 1 exception, which skips the aliasing test on the
grounds that each row is read entirely before it is written. With more than one
token the generic overlap test runs instead and refuses. So a request decoding
alone computes its routing weights with the fused warp local top-k kernel and the
same request decoding next to three others computes them with the softmax,
argsort, get_rows, sum, clamp and divide chain. Two algorithms for one set of
routing weights is exactly the batch dependence this knob removes, and the same
reason mul_mat plus GLU fusion is already off here.

A node probe cannot see this, which is worth recording: registering an eval
callback disables fusion, so both sides of the comparison take the unfused chain
and agree. It only appears when the two are compared without one. On
Qwen3.6-35B-A3B, with every node of one decode step already byte-identical,
sequence 0's logits differed in 248319 of 248320 entries from the first decode
step, by up to 2.6e-1, and the greedy stream flipped a token at step 47. The dense
4B is unaffected either way, and disabling every CUDA fusion removes it, which is
what named the fused op.

With the routing left unfused under the knob, 512 decode steps of sequence 0
alone against sequence 0 next to three neighbours are byte-identical on the 35B,
and every cell of the server matrix reads identical.
---
 ggml/src/ggml-cuda/ggml-cuda.cu | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 0995acdfd6a..30c33b23eaf 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -3528,9 +3528,14 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
         }
     }
 
-    //topk-moe
-    if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX ||
-            cgraph->nodes[i]->op == GGML_OP_ARGSORT) {
+    // topk-moe
+    // [TAG_BATCH_INVARIANT] The routing fusion passes its memory-range check only when the ubatch
+    // holds one token, so a request decoding alone picks the fused warp-local top-k kernel and the
+    // same request decoding next to neighbours picks the softmax, argsort and normalize chain.
+    // Two algorithms for one set of routing weights is the batch dependence this mode removes.
+    if (!ggml_cuda_batch_invariant() &&
+            (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX ||
+             cgraph->nodes[i]->op == GGML_OP_ARGSORT)) {
         ggml_cuda_topk_moe_args args;
         const bool              can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args);
         std::vector<ggml_op>    ops;

From 65860ea386d10a8bb7c662fdc72ce0d7b689b757 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 05:43:07 +0000
Subject: [PATCH 08/33] cuda: default exact mode to a bound that covers the
 speculative verify batch

Exact mode still defaults the column policy to unbounded, so every measurement
that recovered the prefill cost had to set GGML_CUDA_BATCH_INVARIANT_MAX_COLS by
hand, and the value everything was measured at, 8, does not cover a speculative
verify ubatch by construction: with --parallel 4 and --spec-type draft-mtp
--spec-draft-n-max 2 a verify ubatch holds one accepted token plus two drafts per
slot, up to 12 columns, and above the bound neither the MUL_MAT column split nor
the MUL_MAT_ID per token split fires.

Default the bound to 16 in exact mode instead, which covers four slots at up to
three tokens each. An explicitly set GGML_CUDA_BATCH_INVARIANT_MAX_COLS still
wins, in either mode, so a deployment with more slots or a wider draft can raise
it.

Measured on one B200, Qwen3.6-35B-A3B-UD-Q4_K_XL and Qwen3.5-4B-UD-Q4_K_XL,
--parallel 4 --kv-unified -c 8192 --flash-attn on -ngl 99 --seed 0, greedy
sampling, 512 predicted tokens, P0 solo twice then P0 concurrent with P1 to P3,
three rounds per cell. Every P0 completion was byte identical to its solo
reference and every solo repeat matched:

  35B, draft-mtp n-max 2, MAX_COLS=16                       identical x3
  35B, draft-mtp n-max 2, MAX_COLS=16, PREEMPT_EVERY=64     identical x3, 98/98 parks
  35B, draft-mtp n-max 4, MAX_COLS=8  (verify up to 20)     identical x3
  35B, draft-mtp n-max 4, MAX_COLS=32                       identical x3
  4B,  draft-mtp n-max 2, MAX_COLS=16, PREEMPT_EVERY=64     identical x3, 98/98 parks
  35B, draft-mtp n-max 2, new default, no env var           identical x3
  35B, spec off, new default, no env var                    identical x3

The 35B MTP cell at 16 reproduces the acceptance counters of the same cell at 8
exactly, 4091 of 6103 draft tokens, so raising the bound does not perturb the
generation. Cost, three interleaved MAX_COLS=8 against MAX_COLS=16 pairs on the
35B with speculation on, medians: solo decode 25.42 against 25.37 tok/s, four
chat aggregate decode 54.55 against 54.23 tok/s, four chat wall 19.98 against
20.15 s. Four interleaved pairs with speculation off, where a decode ubatch is
four columns wide and both bounds must behave identically, medians 31.39 against
31.40 tok/s solo and 111.41 against 111.41 tok/s aggregate. All within the run to
run spread.

test-backend-ops test -b CUDA0 -o MUL_MAT_ID,MUL_MAT under
LLAMA_EXACT_CONCURRENCY=1 GGML_CUDA_BATCH_INVARIANT=2 MAX_COLS=16:
2136/2136 passed, 2/2 backends.
---
 ggml/src/ggml-cuda/ggml-cuda.cu | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 30c33b23eaf..4adee44ca61 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1856,10 +1856,15 @@ int ggml_cuda_batch_invariant_max_cols() {
     static const int max_cols = []() {
         // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill
         // matmul shapes match its solo run, so exact mode no longer needs the column policy to be
-        // unbounded there. Honour an explicit bound when one is set; default to unbounded.
+        // unbounded there. An explicit bound always wins, in either mode.
         const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
         if (val) { return atoi(val); }
-        if (ggml_cuda_exact_concurrency()) { return 0; }
+        // Exact mode then only has to cover the widest ubatch a decode step can build: one column
+        // per slot, times one plus the number of speculative draft tokens carried with it. 16
+        // covers the default four slots at up to three tokens each, which is what
+        // --spec-type draft-mtp --spec-draft-n-max 2 produces. More slots, or a wider draft, need
+        // the bound set explicitly; above it the column split does not fire.
+        if (ggml_cuda_exact_concurrency()) { return 16; }
         return 0;
     }();
     return max_cols;

From 3cc003acf55f8d8195e50f805dede52fbb808051 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 06:41:33 +0000
Subject: [PATCH 09/33] kv-cache: refuse exact mode when a KV layer is not on
 the CUDA backend

The mode was gated on unified && offload && !v_trans && n_swa == 0 and never on
where the attention layers actually run. Only the CUDA FLASH_ATTN_EXT reads
src[5]; the CPU, Metal, Vulkan, SYCL, OpenCL and CANN kernels ignore it. So with
-ngl 0, a partial -ngl, or a non-CUDA GPU the pool was still paged and the page
table was still attached, but attention traversed physical cell order: the output
stayed correct and neighbour and relocation independence were silently lost while
the mode reported itself as on.

Check the placement where the cache is built instead. Every KV layer must be
offloaded and its device must belong to the CUDA family backend, which is also
built as ROCm and MUSA and carries the same paged kernel; anything else fails the
load naming the layer and the backend it landed on. As a second line, the
FLASH_ATTN_EXT of every backend that would ignore the page table now refuses an
op with src[5] set, so a scheduler decision made after the load cannot route it
somewhere that walks the pool in physical order. The CPU is deliberately left
accepting it and says so in place: it is the reference test-backend-ops compares
the paged CUDA kernel against, and that test builds a mask which selects exactly
the listed cells.

The four remaining preconditions were one bare GGML_ASSERT each, so a quantized
KV cache, an SWA model, a transposed V cache or a -c that is not a multiple of
256 aborted at model load without naming which one failed. Each now logs what it
needs and which flag sets it, and the load returns an error the way every other
KV cache failure does. The context size check also moved after the shared-source
override, so it tests the size the cache is actually built at.

  -ngl 10 with LLAMA_EXACT_CONCURRENCY=1:
  llama_kv_cache: LLAMA_EXACT_CONCURRENCY is set but layer 0 keeps its KV cache
  on CPU, which has no paged attention: every layer must be offloaded to the CUDA
  backend (pass -ngl to offload all layers and do not pass --no-kv-offload)
---
 ggml/src/ggml-cann/ggml-cann.cpp        |  5 ++
 ggml/src/ggml-cpu/ggml-cpu.cpp          |  7 +++
 ggml/src/ggml-metal/ggml-metal-device.m |  5 ++
 ggml/src/ggml-opencl/ggml-opencl.cpp    |  5 ++
 ggml/src/ggml-sycl/ggml-sycl.cpp        |  4 +-
 ggml/src/ggml-vulkan/ggml-vulkan.cpp    |  5 ++
 src/llama-kv-cache.cpp                  | 62 +++++++++++++++++++++++--
 7 files changed, 87 insertions(+), 6 deletions(-)

diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp
index 5e5541aac94..0e901ed0160 100644
--- a/ggml/src/ggml-cann/ggml-cann.cpp
+++ b/ggml/src/ggml-cann/ggml-cann.cpp
@@ -2656,6 +2656,11 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
             return true;
         case GGML_OP_FLASH_ATTN_EXT:
             {
+                // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
+                // the CUDA backend reads
+                if (op->src[5]) {
+                    return false;
+                }
 #ifdef ASCEND_310P
                 // FA not support on 310p device
                 return false;
diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp
index 8cece71f186..a548b33bd71 100644
--- a/ggml/src/ggml-cpu/ggml-cpu.cpp
+++ b/ggml/src/ggml-cpu/ggml-cpu.cpp
@@ -474,6 +474,13 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_SSM_SCAN:
             return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1;
+        // [TAG_EXACT_CONCURRENCY] note: GGML_OP_FLASH_ATTN_EXT with src[5] set, the
+        // exact-concurrency page table, is deliberately still accepted here. The CPU ignores the
+        // page table and attends in physical cell order, which is why every other backend refuses
+        // it, but the CPU is also the reference that test-backend-ops compares the paged CUDA
+        // kernel against, and that test builds a mask which selects exactly the listed cells. A KV
+        // cache layer cannot reach the CPU under the mode anyway: llama_kv_cache refuses to
+        // construct unless every KV layer is on the CUDA backend.
         default:
             return true;
     }
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index 19c57820e85..90873f2fab0 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -1592,6 +1592,11 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_ROLL:
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_FLASH_ATTN_EXT:
+            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the
+            // CUDA backend reads; walking the pool in physical order here would be silently wrong
+            if (op->src[5] != NULL) {
+                return false;
+            }
             // for new head sizes, add checks here
             if (op->src[0]->ne[0] != 32 &&
                 op->src[0]->ne[0] != 40 &&
diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp
index 64f3325b2a5..effd11714f9 100644
--- a/ggml/src/ggml-opencl/ggml-opencl.cpp
+++ b/ggml/src/ggml-opencl/ggml-opencl.cpp
@@ -7842,6 +7842,11 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
         case GGML_OP_MEAN:
             return op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_FLASH_ATTN_EXT: {
+            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the
+            // CUDA backend reads
+            if (op->src[5]) {
+                return false;
+            }
             // The E17 compilers segfault while building FA kernels, skip E17 for now
             if (adreno_e17_compiler_quirks(backend_ctx)) {
                 return false;
diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp
index 0573643d834..69a344ab790 100644
--- a/ggml/src/ggml-sycl/ggml-sycl.cpp
+++ b/ggml/src/ggml-sycl/ggml-sycl.cpp
@@ -6342,7 +6342,9 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_SOLVE_TRI:
             return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K;
         case GGML_OP_FLASH_ATTN_EXT:
-            return ggml_sycl_flash_attn_ext_supported(device, op);
+            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the
+            // CUDA backend reads
+            return op->src[5] == nullptr && ggml_sycl_flash_attn_ext_supported(device, op);
         default:
             return false;
     }
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index c1d86aaac5c..4a2347219b3 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -18192,6 +18192,11 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
             }
         case GGML_OP_FLASH_ATTN_EXT:
             {
+                // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
+                // the CUDA backend reads
+                if (op->src[5]) {
+                    return false;
+                }
                 bool coopmat2 = device->coopmat2;
                 uint32_t HSK = op->src[1]->ne[0];
                 uint32_t HSV = op->src[2]->ne[0];
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index df643a047a9..dcaf7bf687c 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -61,6 +61,29 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) {
 // llama_kv_cache
 //
 
+// [TAG_EXACT_CONCURRENCY]
+// The paged attention specialization that reads the logical page table lives in the CUDA backend
+// sources, which are also built as the ROCm and MUSA backends. Every other backend ignores src[5]
+// and walks the pool in physical cell order, so a KV layer placed there would silently lose the
+// guarantee the mode exists to provide.
+static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) {
+    if (!dev) {
+        return false;
+    }
+
+    ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(dev);
+    if (!reg) {
+        return false;
+    }
+
+    const char * name = ggml_backend_reg_name(reg);
+    if (!name) {
+        return false;
+    }
+
+    return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0;
+}
+
 llama_kv_cache::llama_kv_cache(
         const llama_model & model,
         const llama_hparams & hparams,
@@ -86,11 +109,6 @@ llama_kv_cache::llama_kv_cache(
 
     const char * exact_env = getenv("LLAMA_EXACT_CONCURRENCY");
     exact_pages = exact_env && atoi(exact_env) != 0;
-    if (exact_pages) {
-        GGML_ASSERT(unified && offload && !v_trans && n_swa == 0);
-        GGML_ASSERT(type_k == GGML_TYPE_F16 && type_v == GGML_TYPE_F16);
-        GGML_ASSERT(kv_size % exact_page_size == 0);
-    }
 
     // shared cells view the source cache's K/V tensors, so the cell count
     // follows the source allocation: a fitted target can be smaller than the
@@ -105,6 +123,30 @@ llama_kv_cache::llama_kv_cache(
 
     GGML_ASSERT(kv_size % n_pad == 0);
 
+    // [TAG_EXACT_CONCURRENCY]
+    // Every one of these is reachable from the command line, so report which one failed by name
+    // instead of aborting on a bare assert that only prints a file and a line.
+    if (exact_pages) {
+        const char * unsupported = nullptr;
+
+        if (!unified) {
+            unsupported = "it needs a unified KV cache (pass --kv-unified)";
+        } else if (v_trans) {
+            unsupported = "it needs a non-transposed V cache (pass --flash-attn on)";
+        } else if (n_swa != 0) {
+            unsupported = "the paged pool does not support sliding window attention";
+        } else if (type_k != GGML_TYPE_F16 || type_v != GGML_TYPE_F16) {
+            unsupported = "it needs an F16 KV cache (do not pass --cache-type-k or --cache-type-v)";
+        } else if (kv_size % exact_page_size != 0) {
+            unsupported = "the context size must be a multiple of 256 (pass -c as a multiple of 256)";
+        }
+
+        if (unsupported) {
+            LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s\n", __func__, unsupported);
+            throw std::runtime_error("exact concurrency: unsupported KV cache configuration");
+        }
+    }
+
     const uint32_t n_layer = hparams.n_layer_all;
 
     // define a comparator for the buft -> ctx map to ensure that the order is well-defined:
@@ -228,6 +270,16 @@ llama_kv_cache::llama_kv_cache(
 
         LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name);
 
+        // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order while
+        // the mode still reports itself as on, so refuse the load instead
+        if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) {
+            LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, "
+                    "which has no paged attention: every layer must be offloaded to the CUDA backend "
+                    "(pass -ngl to offload all layers and do not pass --no-kv-offload)\n",
+                    __func__, il, dev_name);
+            throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend");
+        }
+
         ggml_context * ctx = ctx_for_buft(buft);
         if (!ctx) {
             throw std::runtime_error("failed to create ggml context for kv cache");

From 66583931d928a472dcbbdc6a6c94a6e8fe5b52e2 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 06:43:01 +0000
Subject: [PATCH 10/33] kv-cache: report the cache transformations exact mode
 cannot do

get_can_shift() still returned true under exact_pages, so --context-shift and
--cache-reuse N passed every startup capability gate and then aborted the whole
process on the first seq_add with a nonzero shift. Both are opt-in flags, so the
default was safe, but llama-server accepted them silently and died on the first
request that needed them. The server already disables both at load for a cache
that cannot shift, with a warning, so returning false there reuses that path:

  srv    load_model: ctx_shift is not supported by this context, it will be disabled

The remaining aborts are reachable the same way, from one request parameter or
one API call, and abort() is not an acceptable answer to either in a network
server. seq_cp between two different sequences, a nonzero seq_add and a seq_div
now log which transformation was refused and on which sequence and return without
touching the cells, and the whole-context branch of state_read_meta() logs and
returns false the way the two failure paths next to it already do, so
llama_state_load_file() reports a recoverable error through an API that is
designed for one instead of killing the process.

The page invariant is protected exactly as before: none of these paths can now
run and leave a cell outside the page its position belongs to.
---
 src/llama-kv-cache.cpp | 42 ++++++++++++++++++++++++++++++++++++++----
 1 file changed, 38 insertions(+), 4 deletions(-)

diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index dcaf7bf687c..b39018c3ac4 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -507,7 +507,14 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
 }
 
 void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
-    GGML_ASSERT(!exact_pages || seq_id_src == seq_id_dst);
+    // [TAG_EXACT_CONCURRENCY] a page belongs to one (sequence, position/256) pair, so two sequences
+    // cannot share physical cells. Refuse the copy rather than abort the process: this is reachable
+    // from a request parameter.
+    if (exact_pages && seq_id_src != seq_id_dst) {
+        LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so cells cannot be shared between "
+                "sequences: ignoring the copy from seq %d to seq %d\n", __func__, seq_id_src, seq_id_dst);
+        return;
+    }
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
@@ -627,7 +634,14 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
 }
 
 void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
-    GGML_ASSERT(!exact_pages || shift == 0);
+    // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so a
+    // shift would have to move the cells too. get_can_shift() reports this so that --context-shift
+    // and --cache-reuse are turned off at load; this is the guard for the library API.
+    if (exact_pages && shift != 0) {
+        LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so positions cannot be shifted: "
+                "ignoring the shift of %d on seq %d\n", __func__, shift, seq_id);
+        return;
+    }
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
@@ -678,7 +692,13 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll
 }
 
 void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
-    GGML_ASSERT(!exact_pages || d == 1);
+    // [TAG_EXACT_CONCURRENCY] same reason as seq_add: the offset inside a page is derived from the
+    // position, so dividing the positions would leave every cell in the wrong slot.
+    if (exact_pages && d != 1) {
+        LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so positions cannot be divided: "
+                "ignoring the division by %d on seq %d\n", __func__, d, seq_id);
+        return;
+    }
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
@@ -1276,6 +1296,13 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
 }
 
 bool llama_kv_cache::get_can_shift() const {
+    // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so the
+    // paged pool cannot shift positions. Reporting it here is what makes the server disable
+    // --context-shift and --cache-reuse at load, with a warning, instead of accepting both and
+    // failing on the first request that needs them.
+    if (exact_pages) {
+        return false;
+    }
     // Step35 uses per-layer RoPE dims; K-shift assumes a single global n_rot.
     if (model.arch == LLM_ARCH_STEP35) {
         return false;
@@ -2431,9 +2458,16 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
             GGML_ASSERT(cells.seq_has(idx, dest_seq_id));
         }
     } else {
-        GGML_ASSERT(!exact_pages && "exact mode supports per-sequence restore only");
         // whole KV cache restore
 
+        // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical
+        // index, which the paged pool owns. Report it like every other failure in this function.
+        if (exact_pages) {
+            LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state "
+                    "restore only\n", __func__);
+            return false;
+        }
+
         if (cell_count > cells.size()) {
             LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__);
             return false;

From b81d8e7018bbc45f1a7c4ef93e92db69e3a2d62a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 06:44:24 +0000
Subject: [PATCH 11/33] batch: isolate only the prompt sequences, in every
 memory type

Two things were wrong with the one-sequence prompt ubatch rule.

It only existed in the hybrid memory. A dense transformer with a unified cache
takes split_simple, which packs every sequence's prompt tokens into one ubatch,
so its prefill matmuls ran at a width the solo run never sees and, with the
column policy bounded, produced K and V the solo run never produces. Exact mode
was therefore not exact by construction on dense models, which is most of what it
will be pointed at. A pure recurrent model still called the three-argument
split_equal for the same reason the hybrid one no longer does. Both now take the
rule: llama_kv_cache::init_batch keeps split_simple for a plain decode step and
switches to the sequence-set split when a prompt is present, and
llama_memory_recurrent::init_batch passes the flag through.

The rule itself then serialized more than it had to. has_multi_token_seq()
scanned the whole original batch and ignored used[], so it stayed true after the
prompt had been consumed, and the n_seqs_max cap it fed capped every sequence set
including one-token decode sets. One prompt chunk plus three decodes therefore
became four single-sequence ubatches and the three chats decoded one at a time
for the whole prefill, which contradicts the comment saying a plain decode step
stays batched. The predicate now skips used[] tokens, and the cap became an
isolate_multi_token_seqs flag: a sequence set with more than one token left to
place takes a ubatch of its own, sets with one token left keep grouping. One
prompt next to three decodes now costs one extra ubatch, not three.

The KV cache constructor also read getenv("LLAMA_EXACT_CONCURRENCY") directly
while llama_exact_concurrency() and ggml_cuda_exact_concurrency() each cache the
first value they see, so a process that created one context with the knob unset
and then set it got a paged cache on top of a dispatcher still in default mode.
It now reads the same cached value as the other two.
---
 src/llama-batch.cpp            | 40 +++++++++++++++++++++++++++++-----
 src/llama-batch.h              | 14 +++++++-----
 src/llama-kv-cache.cpp         | 16 +++++++++++---
 src/llama-memory-hybrid.cpp    |  6 ++---
 src/llama-memory-recurrent.cpp |  6 ++++-
 5 files changed, 63 insertions(+), 19 deletions(-)

diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 50f70bb0ff3..4b73ab2478b 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -511,6 +511,11 @@ bool llama_batch_allocr::has_multi_token_seq() const {
     std::vector<uint32_t> n_per_seq(n_seq_max, 0);
 
     for (int32_t i = 0; i < batch.n_tokens; ++i) {
+        // tokens already placed in an earlier ubatch do not make the rest of the batch a prompt
+        if (used[i]) {
+            continue;
+        }
+
         for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
             if (++n_per_seq[batch.seq_id[i][s]] > 1) {
                 return true;
@@ -521,7 +526,7 @@ bool llama_batch_allocr::has_multi_token_seq() const {
     return false;
 }
 
-llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t n_seqs_max) {
+llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, bool isolate_multi_token_seqs) {
     if (sequential && has_cpl) {
         LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__);
 
@@ -554,6 +559,34 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
         }
 
         if (add) {
+            // [TAG_EXACT_CONCURRENCY] a sequence set that still has more than one token to place is
+            // a prompt, and a prompt shares its arithmetic with whatever else is in the ubatch, so
+            // give it a ubatch of its own. Sets with one token left are a plain decode step, which
+            // is already exact, so keep grouping those: isolating them too would make one prompt
+            // serialize every concurrent decode for the whole of the prefill.
+            if (isolate_multi_token_seqs) {
+                uint32_t n_left = 0;
+
+                for (const auto idx : seq_set_map[seq_set[i]]) {
+                    if (!used[idx]) {
+                        ++n_left;
+                    }
+                }
+
+                if (n_left > 1) {
+                    if (!cur_seq_set.empty()) {
+                        // let the sets already taken have this ubatch; the prompt gets the next one
+                        break;
+                    }
+
+                    cur_seq_set.push_back(seq_set[i]);
+
+                    last_seq_id = batch.seq_id[i][0];
+
+                    break;
+                }
+            }
+
             cur_seq_set.push_back(seq_set[i]);
 
             last_seq_id = batch.seq_id[i][0];
@@ -561,11 +594,6 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
             if (cur_seq_set.size() > n_ubatch) {
                 break;
             }
-
-            // [TAG_EXACT_CONCURRENCY]
-            if (n_seqs_max > 0 && cur_seq_set.size() >= n_seqs_max) {
-                break;
-            }
         }
     }
 
diff --git a/src/llama-batch.h b/src/llama-batch.h
index d354c442d03..4bd2aa98f9f 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -105,12 +105,14 @@ class llama_batch_allocr {
     // make ubatches of equal-length sequences sets
     // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
     // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
-    // n_seqs_max  = maximum sequence sets per ubatch, 0 = no limit
-    //               [TAG_EXACT_CONCURRENCY] passing 1 keeps a ubatch to a single sequence
-    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t n_seqs_max = 0);
-
-    // [TAG_EXACT_CONCURRENCY] true if some sequence contributes more than one token to the batch,
-    // i.e. this is not a plain one-token-per-sequence decode step
+    // isolate_multi_token_seqs = [TAG_EXACT_CONCURRENCY] a sequence set with more than one token
+    //               left to place is given a ubatch of its own; sets with a single token left are
+    //               still grouped together, so a prompt next to three decodes costs one extra
+    //               ubatch and does not serialize the three decodes
+    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, bool isolate_multi_token_seqs = false);
+
+    // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than one token left to place,
+    // i.e. what remains of the batch is not a plain one-token-per-sequence decode step
     bool has_multi_token_seq() const;
 
     // sequence-set-wise split - each ubatch contains a single sequence-set
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index b39018c3ac4..16c04fb6e3b 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -107,8 +107,10 @@ llama_kv_cache::llama_kv_cache(
     v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),
     v_cells(*v_cells_impl) {
 
-    const char * exact_env = getenv("LLAMA_EXACT_CONCURRENCY");
-    exact_pages = exact_env && atoi(exact_env) != 0;
+    // [TAG_EXACT_CONCURRENCY] read the knob through the one cached reader that the graph and the
+    // CUDA dispatcher also use, so a process that sets it between two context creations cannot end
+    // up with a paged cache on top of a dispatcher that is still in default mode
+    exact_pages = llama_exact_concurrency();
 
     // shared cells view the source cache's K/V tensors, so the cell count
     // follows the source allocation: a fitted target can be smaller than the
@@ -791,7 +793,15 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
 
         std::vector<llama_ubatch> ubatches;
         while (true) {
-            auto ubatch = n_stream == 1 ? balloc.split_simple(n_ubatch) : balloc.split_equal(n_ubatch, true, 0);
+            // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt tokens into one
+            // ubatch, so a sequence's prefill would run at a width its solo run never sees. Take
+            // the sequence-set split instead, which can give each prompt a ubatch of its own; a
+            // plain decode step has nothing to isolate and keeps taking split_simple.
+            const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq();
+
+            auto ubatch = n_stream == 1 && !isolate
+                ? balloc.split_simple(n_ubatch)
+                : balloc.split_equal(n_ubatch, n_stream > 1, 0, isolate);
 
             if (ubatch.n_tokens == 0) {
                 break;
diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index ba54ab12923..4ebd476aa8b 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -89,11 +89,11 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
                 // [TAG_EXACT_CONCURRENCY] the recurrent half of a hybrid model is not invariant to
                 // the shape of the ubatch: a prompt processed next to other sequences' prompt tokens
                 // leaves a different gated delta net state than the same prompt processed alone.
-                // Keeping such a ubatch to a single sequence removes that. A plain decode step, one
+                // Giving such a sequence a ubatch of its own removes that. A plain decode step, one
                 // token per sequence, is already exact and stays batched.
-                const uint32_t n_seqs_max = llama_exact_concurrency() && balloc.has_multi_token_seq() ? 1 : 0;
+                const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq();
 
-                ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, n_seqs_max);
+                ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate);
             }
 
             if (ubatch.n_tokens == 0) {
diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp
index e2990972ef7..f639a25c5df 100644
--- a/src/llama-memory-recurrent.cpp
+++ b/src/llama-memory-recurrent.cpp
@@ -431,7 +431,11 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr &
                 // [TAG_RECURRENT_ROLLBACK_SPLITS]
                 // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
                 //   so that the rollback snapshots remain valid
-                ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0);
+                // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: a recurrent state that a
+                // prompt leaves behind depends on what shared its ubatch, so isolate the prompts
+                const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq();
+
+                ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate);
             }
 
             if (ubatch.n_tokens == 0) {

From c7027d60dd8c48a0b72d11a9ba07c27e7b8c9504 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 06:49:37 +0000
Subject: [PATCH 12/33] cuda: derive the exact mode column bound from the
 decode width

The fixed default of 16 silently turns exact mode off above 16 columns: MUL_MAT,
MUL_MAT_ID and the per-row attention split all fall back to the neighbour
dependent batched path. --parallel 6 --spec-draft-n-max 2 gives 18 columns and
--parallel 8 gives 24; both are ordinary server configurations and neither said
anything. The source comment documented the cliff, nothing at runtime did.

The bound only ever had to cover the widest ubatch a decode step can build, since
a prompt ubatch holds one sequence and gets its exactness from that. So let the
caller report that width. ggml_backend_cuda_set_exact_decode_width(), reachable
directly or through ggml_backend_reg_get_proc_address(), takes one column per slot
times one plus the draft length, and exact mode defaults the bound to it. common
computes it from n_parallel and the speculative type and reports it before the
warmup, which is the first graph any of these tools computes, and refuses at
startup an explicitly set GGML_CUDA_BATCH_INVARIANT_MAX_COLS that is smaller:

  GGML_CUDA_BATCH_INVARIANT_MAX_COLS is 8 but LLAMA_EXACT_CONCURRENCY needs at
  least 12 to cover a decode step of 4 slots, above which a matmul is left batched
  and its rows depend on the other rows in the ubatch. Raise it to 12, set it to 0
  for no bound, or unset it to let it default to 12.

--parallel 4 with speculation off now defaults to 4 rather than 16 and with
--spec-type draft-mtp --spec-draft-n-max 2 to 12 rather than 16, which is the same
guarantee over a narrower range of shapes: a decode ubatch of that model cannot be
wider than that, and everything above it is a prompt.

When nothing reported a width, the default stays 16 and the dispatcher warns once
per process the first time a MUL_MAT or MUL_MAT_ID above the bound is left
unsplit, naming both numbers. It deliberately stays quiet once a width is known,
because then the only batches above the bound are prompt ubatches and warning on
those would fire on every prefill for a case that is working as intended.
---
 common/arg.cpp                  |  7 +++
 common/common.cpp               | 82 +++++++++++++++++++++++++++++++++
 common/common.h                 | 13 ++++++
 ggml/include/ggml-cuda.h        |  9 ++++
 ggml/src/ggml-cuda/ggml-cuda.cu | 77 +++++++++++++++++++++++++------
 5 files changed, 175 insertions(+), 13 deletions(-)

diff --git a/common/arg.cpp b/common/arg.cpp
index 5bfa4adcdf0..e4f2e8f2517 100644
--- a/common/arg.cpp
+++ b/common/arg.cpp
@@ -1304,6 +1304,13 @@ bool common_params_parse(int argc, char ** argv, common_params & params, llama_e
             exit(0);
         }
         params.lr.init();
+
+        // [TAG_EXACT_CONCURRENCY] refuse a column bound that cannot cover a decode step before
+        // anything is loaded, rather than running with the guarantee quietly switched off
+        if (!common_exact_concurrency_init(ctx_arg.params)) {
+            ctx_arg.params = params_org;
+            return false;
+        }
     } catch (const std::invalid_argument & ex) {
         fprintf(stderr, "%s\n", ex.what());
         ctx_arg.params = params_org;
diff --git a/common/common.cpp b/common/common.cpp
index 3d54bd6002d..d2beebc8e2f 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1,4 +1,5 @@
 #include "ggml.h"
+#include "ggml-backend.h"
 #include "gguf.h"
 
 #include "build-info.h"
@@ -1433,6 +1434,82 @@ std::vector<llama_adapter_lora_ptr> & common_init_result::lora() {
     return pimpl->lora;
 }
 
+// [TAG_EXACT_CONCURRENCY]
+bool common_exact_concurrency() {
+    static const bool enabled = []() {
+        const char * val = getenv("LLAMA_EXACT_CONCURRENCY");
+        return val && atoi(val) != 0;
+    }();
+
+    return enabled;
+}
+
+// [TAG_EXACT_CONCURRENCY]
+int common_exact_decode_width(const common_params & params) {
+    const int n_slots = std::max(1, params.n_parallel);
+
+    // the draft tokens a slot carries into the verify ubatch alongside its accepted token
+    int n_draft = 0;
+
+    for (const auto type : params.speculative.types) {
+        switch (type) {
+            case COMMON_SPECULATIVE_TYPE_NONE:
+                break;
+            case COMMON_SPECULATIVE_TYPE_NGRAM_MOD:
+                n_draft = std::max(n_draft, params.speculative.ngram_mod.n_max);
+                break;
+            case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
+                n_draft = std::max(n_draft, (int) params.speculative.ngram_simple.size_m);
+                break;
+            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K:
+                n_draft = std::max(n_draft, (int) params.speculative.ngram_map_k.size_m);
+                break;
+            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V:
+                n_draft = std::max(n_draft, (int) params.speculative.ngram_map_k4v.size_m);
+                break;
+            default:
+                n_draft = std::max(n_draft, params.speculative.draft.n_max);
+                break;
+        }
+    }
+
+    return n_slots*(1 + std::max(0, n_draft));
+}
+
+// [TAG_EXACT_CONCURRENCY]
+bool common_exact_concurrency_init(const common_params & params) {
+    if (!common_exact_concurrency()) {
+        return true;
+    }
+
+    const int n_cols = common_exact_decode_width(params);
+
+    const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
+    if (bound) {
+        const int max_cols = atoi(bound);
+        if (max_cols > 0 && max_cols < n_cols) {
+            COM_ERR("GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at "
+                    "least %d to cover a decode step of %d slots, above which a matmul is left "
+                    "batched and its rows depend on the other rows in the ubatch. Raise it to %d, "
+                    "set it to 0 for no bound, or unset it to let it default to %d.\n",
+                    max_cols, n_cols, std::max(1, params.n_parallel), n_cols, n_cols);
+            return false;
+        }
+    }
+
+    // the CUDA backend may not be present or may be loaded dynamically, so go through the registry
+    for (size_t i = 0; i < ggml_backend_reg_count(); ++i) {
+        ggml_backend_reg_t reg = ggml_backend_reg_get(i);
+
+        auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width");
+        if (fn) {
+            fn(n_cols);
+        }
+    }
+
+    return true;
+}
+
 common_init_result_ptr common_init_from_params(common_params & params, bool model_only) {
     common_init_result_ptr res(new common_init_result(params, model_only));
 
@@ -1454,6 +1531,11 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
 
     const llama_vocab * vocab = llama_model_get_vocab(model);
 
+    // [TAG_EXACT_CONCURRENCY] before the warmup, which is the first graph this process computes
+    if (!common_exact_concurrency_init(params)) {
+        return res;
+    }
+
     if (params.ctx_shift && !llama_memory_can_shift(llama_get_memory(lctx))) {
         COM_WRN("%s", "KV cache shifting is not supported for this context, disabling KV cache shifting\n");
         params.ctx_shift = false;
diff --git a/common/common.h b/common/common.h
index c99269f9a96..2be2fab6a8b 100644
--- a/common/common.h
+++ b/common/common.h
@@ -931,6 +931,19 @@ using common_init_result_ptr = std::unique_ptr<common_init_result>;
 
 common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false);
 
+// [TAG_EXACT_CONCURRENCY]
+// true when LLAMA_EXACT_CONCURRENCY is set for this process
+bool common_exact_concurrency();
+
+// the widest ubatch a decode step can build with these parameters: one column per slot, times one
+// plus the number of speculative draft tokens carried with it. Under exact mode this is what the
+// CUDA column policy has to cover, and what its default bound is derived from.
+int common_exact_decode_width(const common_params & params);
+
+// report that width to the CUDA backend, and refuse an explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS
+// that is smaller than it. Returns false if the configuration must not run.
+bool common_exact_concurrency_init(const common_params & params);
+
 struct llama_model_params   common_model_params_to_llama  (      common_params & params);
 struct llama_context_params common_context_params_to_llama(const common_params & params);
 
diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h
index 1cd81eeaebc..c3dd87c97b7 100644
--- a/ggml/include/ggml-cuda.h
+++ b/ggml/include/ggml-cuda.h
@@ -38,6 +38,15 @@ GGML_BACKEND_API void ggml_backend_cuda_get_device_description(int device, char
 GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total);
 
 GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size);
+
+// [TAG_EXACT_CONCURRENCY]
+// Report the widest ubatch a decode step of this process can build: one column per slot, times one
+// plus the number of speculative draft tokens carried with it. Under LLAMA_EXACT_CONCURRENCY the
+// column policy then defaults to that width instead of a fixed number, so --parallel or a wider
+// draft cannot silently push a decode above the bound and leave it batched. An explicitly set
+// GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins. Call before the first graph is computed. Also
+// available through ggml_backend_reg_get_proc_address().
+GGML_BACKEND_API void ggml_backend_cuda_set_exact_decode_width(int n_cols);
 GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer);
 
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_cuda_reg(void);
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 4adee44ca61..c8f43bae382 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1852,22 +1852,67 @@ int ggml_cuda_batch_invariant() {
     return mode;
 }
 
+// [TAG_EXACT_CONCURRENCY] the widest decode ubatch the caller says it can build, 0 if it never said
+static std::atomic<int> g_exact_decode_width{0};
+
+void ggml_backend_cuda_set_exact_decode_width(int n_cols) {
+    g_exact_decode_width.store(n_cols > 0 ? n_cols : 0, std::memory_order_relaxed);
+}
+
 int ggml_cuda_batch_invariant_max_cols() {
-    static const int max_cols = []() {
-        // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill
-        // matmul shapes match its solo run, so exact mode no longer needs the column policy to be
-        // unbounded there. An explicit bound always wins, in either mode.
+    // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill
+    // matmul shapes match its solo run, so exact mode no longer needs the column policy to be
+    // unbounded there. An explicit bound always wins, in either mode.
+    static const int explicit_cols = []() {
         const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
-        if (val) { return atoi(val); }
-        // Exact mode then only has to cover the widest ubatch a decode step can build: one column
-        // per slot, times one plus the number of speculative draft tokens carried with it. 16
-        // covers the default four slots at up to three tokens each, which is what
-        // --spec-type draft-mtp --spec-draft-n-max 2 produces. More slots, or a wider draft, need
-        // the bound set explicitly; above it the column split does not fire.
-        if (ggml_cuda_exact_concurrency()) { return 16; }
-        return 0;
+        return val ? atoi(val) : -1;
     }();
-    return max_cols;
+
+    if (explicit_cols >= 0) {
+        return explicit_cols;
+    }
+
+    if (!ggml_cuda_exact_concurrency()) {
+        return 0;
+    }
+
+    // Exact mode only has to cover the widest ubatch a decode step can build: one column per slot,
+    // times one plus the number of speculative draft tokens carried with it. Use that width when
+    // the caller reported it through ggml_backend_cuda_set_exact_decode_width(). Nothing reported
+    // it, so fall back to 16, which covers four slots at up to three tokens each, which is what
+    // --parallel 4 --spec-type draft-mtp --spec-draft-n-max 2 produces. Above the bound the column
+    // split does not fire, and ggml_cuda_warn_above_exact_bound() says so once.
+    const int width = g_exact_decode_width.load(std::memory_order_relaxed);
+
+    return width > 0 ? width : 16;
+}
+
+// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so its rows depend on the
+// other rows in it and the mode does not hold for that op. Say so once, rather than never.
+//
+// Only when nothing reported a decode width. When one was reported the bound is derived from it, so
+// the only batches above the bound are prompt ubatches, and those hold a single sequence under this
+// mode: their exactness comes from that, not from the column policy, and leaving them batched is
+// the whole point of having a bound at all. Warning on those would be crying wolf on every prefill.
+static void ggml_cuda_warn_above_exact_bound(const char * op, int64_t ncols, int max_cols) {
+    if (!ggml_cuda_exact_concurrency()) {
+        return;
+    }
+
+    if (g_exact_decode_width.load(std::memory_order_relaxed) > 0) {
+        return;
+    }
+
+    static std::atomic_flag warned = ATOMIC_FLAG_INIT;
+    if (warned.test_and_set(std::memory_order_relaxed)) {
+        return;
+    }
+
+    GGML_LOG_WARN("%s: LLAMA_EXACT_CONCURRENCY is set, but this %s is %d columns wide while "
+            "GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d, so it is left batched and its result depends "
+            "on the other columns in the ubatch. Raise the bound, set it to 0 for no bound, or call "
+            "ggml_backend_cuda_set_exact_decode_width() with the widest decode this process builds. "
+            "Reported once.\n", __func__, op, (int) ncols, max_cols);
 }
 
 enum ggml_cuda_mm_path {
@@ -1952,6 +1997,7 @@ static bool ggml_cuda_mul_mat_split_columns(
     }
     const int max_cols = ggml_cuda_batch_invariant_max_cols();
     if (max_cols > 0 && ncols_dst > max_cols) {
+        ggml_cuda_warn_above_exact_bound("MUL_MAT", ncols_dst, max_cols);
         return false;
     }
 
@@ -2049,6 +2095,7 @@ static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) {
     }
     const int max_cols = ggml_cuda_batch_invariant_max_cols();
     if (max_cols > 0 && ntokens > max_cols) {
+        ggml_cuda_warn_above_exact_bound("MUL_MAT_ID", ntokens, max_cols);
         return false;
     }
     return true;
@@ -5733,6 +5780,10 @@ static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, con
     if (strcmp(name, "ggml_backend_get_features") == 0) {
         return (void *)ggml_backend_cuda_get_features;
     }
+    // [TAG_EXACT_CONCURRENCY]
+    if (strcmp(name, "ggml_backend_cuda_set_exact_decode_width") == 0) {
+        return (void *)ggml_backend_cuda_set_exact_decode_width;
+    }
     return nullptr;
 }
 

From 9cd1222437a91a4750651801b5ef146819791c77 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 06:50:19 +0000
Subject: [PATCH 13/33] graph: refuse exact mode on the V-less attention
 layouts

Page tables are wired into llm_graph_input_attn_kv only. llm_graph_input_attn_k
has no self_pages member and its build_attn calls build_attn_mha without a pages
argument, as do the DeepSeek sparse and sliding window variants. A model on one of
those layouts still got its cells placed in pages by the allocator and then
attended in physical order after a park and a restore, so the mode reported itself
as on and lost the one invariant it exists for. That is the same silent failure
the CUDA placement gate was added to stop, so answer it the same way: log which
layout it is and fail the context.

Rejecting is the smaller correct change of the two. Wiring self_pages into
llm_graph_input_attn_k is four lines and looks tempting, but it fixes one of the
four V-less input classes and DeepSeek 3.2 uses two of them: its sparse layers
build their mask from a top-k selection and would stay unpaged, leaving the model
half paged, which is worse than refused. None of these architectures was measured
here, and the paged kernel also requires 256-dimensional K and V heads, which none
of them was checked against.

Reaches the user through the path llama_init_from_model already has for a context
that cannot be built:

  llm_graph_reject_exact_concurrency: LLAMA_EXACT_CONCURRENCY is set, but this
  model uses the V-less KV (attn_k) attention layout, which carries no page table
  and would attend in physical cell order
  llama_init_from_model: failed to initialize the context: exact concurrency:
  unsupported attention layout
---
 src/llama-graph.cpp | 31 +++++++++++++++++++++++++++++++
 1 file changed, 31 insertions(+)

diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index 71fce85bca7..f8f5c83a2b6 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -21,11 +21,36 @@
 #include <cstring>
 #include <numeric>
 #include <sstream>
+#include <stdexcept>
 #include <string>
 #include <unordered_set>
 
 // dedup helpers
 
+// [TAG_EXACT_CONCURRENCY]
+// The page table is wired into llm_graph_input_attn_kv only. The V-less layouts build their
+// attention without one, so a model on one of those would get its cells placed in pages by the
+// allocator and then attend in physical cell order anyway: the mode would report itself as on and
+// lose the one invariant it exists for, which is the same silent failure the CUDA placement gate
+// was added to stop. Refuse the context instead.
+//
+// Rejecting is the smaller correct change here. Wiring self_pages into llm_graph_input_attn_k alone
+// is four lines, but it fixes only one of the four V-less input classes, and DeepSeek 3.2 uses two
+// of them: its sparse layers rewrite the mask from a top-k selection and would still be unpaged, so
+// the model would end up half paged, which is worse than refused. None of these architectures was
+// measured, and the paged kernel additionally requires 256-dimensional K and V heads.
+static void llm_graph_reject_exact_concurrency(const char * layout) {
+    if (!llama_exact_concurrency()) {
+        return;
+    }
+
+    LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, but this model uses the %s attention "
+            "layout, which carries no page table and would attend in physical cell order\n",
+            __func__, layout);
+
+    throw std::runtime_error("exact concurrency: unsupported attention layout");
+}
+
 static ggml_tensor * build_attn_inp_kq_mask(
         ggml_context * ctx,
         const llama_kv_cache_context * mctx,
@@ -2871,6 +2896,8 @@ static std::unique_ptr<llm_graph_input_attn_k> build_attn_inp_k_impl(
     const llama_cparams & cparams,
     const llama_kv_cache_context * mctx_cur) {
 
+    llm_graph_reject_exact_concurrency("V-less KV (attn_k)");
+
     auto inp = std::make_unique<llm_graph_input_attn_k>(hparams, cparams, mctx_cur);
 
     {
@@ -3247,6 +3274,8 @@ static std::unique_ptr<llm_graph_input_attn_k_dsa> build_attn_inp_k_dsa_impl(
     const llama_cparams & cparams,
     const llama_kv_cache_dsa_context * mctx_cur) {
 
+    llm_graph_reject_exact_concurrency("sparse V-less KV (attn_k_dsa)");
+
     auto inp = std::make_unique<llm_graph_input_attn_k_dsa>(hparams, cparams, mctx_cur);
 
     {
@@ -3364,6 +3393,8 @@ llm_graph_input_attn_kv_iswa * llm_graph_context::build_attn_inp_kv_iswa() const
 llm_graph_input_attn_k_iswa * llm_graph_context::build_attn_inp_k_iswa() const {
     const auto * mctx_cur = static_cast<const llama_kv_cache_iswa_context *>(mctx);
 
+    llm_graph_reject_exact_concurrency("V-less sliding window KV (attn_k_iswa)");
+
     auto inp = std::make_unique<llm_graph_input_attn_k_iswa>(hparams, cparams, mctx_cur);
 
     {

From b5e9ebdf2c6b42df8a241d321f876ddb3ebddb43 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 06:50:35 +0000
Subject: [PATCH 14/33] batchinv: stop the harness from certifying a reduced
 run

run_concurrent lost worker exceptions. A Python thread exception only prints a
traceback, join() returns, and the partial dict was returned, so a round where
P1 to P3 failed and P0 succeeded still classified P0 as identical and aggregated
throughput over whichever requests happened to survive. The evidence harness could
certify a solo run as a clean four-way concurrency result. Exceptions are now
collected under a lock and raised after join, the barrier is aborted so the other
workers do not block on one that will never fill, and every expected name has to
be present before the result is returned. bench.py imports the same helper, so its
throughput number is covered too.

Server.__enter__ raised after Popen had already started the server, and Python
does not call __exit__ when __enter__ raises, so a server that started but never
reported healthy kept the GPU, the port and the log handle. The health wait is now
wrapped and tears the server down before re-raising. __exit__ also waits after the
SIGKILL path instead of leaving a zombie, and says in place that it is POSIX only.

The run record and the server log header captured only variables starting with
GGML plus CUDA_VISIBLE_DEVICES, so an inherited LLAMA_EXACT_CONCURRENCY was
invisible in both and a run intended as the mode-off reference could silently have
been an exact-mode run while the JSON said "env": {}. That is the baseline the
whole divergence claim rests on. Both now record the environment the server
actually inherited, from an explicit allowlist that includes
LLAMA_EXACT_CONCURRENCY, GGML_CUDA_BATCH_INVARIANT,
GGML_CUDA_BATCH_INVARIANT_MAX_COLS, LLAMA_SERVER_PREEMPT_EVERY and
CUDA_VISIBLE_DEVICES, along with the resolved model path and the full server
command line. What the run asked for is kept separately as env_requested.

UNSLOTH_WORKSPACE was read at import, so both tools raised KeyError before argparse
ran and even --help failed. The model path is resolved when the server arguments
are built and raises a named RuntimeError.

The README still said the mode forces GGML_CUDA_BATCH_INVARIANT=2 with no column
limit including during prefill, which stopped being true two commits before this
branch. It now states the bound, where its default comes from, that an explicit
value below the decode width is refused at startup, and the load-time refusals for
placement and the V-less layouts.
---
 scripts/batchinv/README.md     |  31 +++++++--
 scripts/batchinv/divergence.py | 118 ++++++++++++++++++++++++---------
 2 files changed, 113 insertions(+), 36 deletions(-)

diff --git a/scripts/batchinv/README.md b/scripts/batchinv/README.md
index 62a51487047..82b769b6d04 100644
--- a/scripts/batchinv/README.md
+++ b/scripts/batchinv/README.md
@@ -1,18 +1,31 @@
 # Exact concurrency experiment p
 
 Opt in before loading the model with `LLAMA_EXACT_CONCURRENCY=1`. This also forces
-`GGML_CUDA_BATCH_INVARIANT=2` with no column limit, including during prefill.
+`GGML_CUDA_BATCH_INVARIANT=2` and gives `GGML_CUDA_BATCH_INVARIANT_MAX_COLS` a
+default. The column policy only has to cover the widest ubatch a decode step can
+build, one column per slot times one plus the draft length, because a prompt
+ubatch is kept to one sequence and gets its exactness from that instead. Tools
+built on `common` report that width, so the default is `--parallel` times one plus
+`--spec-draft-n-max`, and an explicitly set `GGML_CUDA_BATCH_INVARIANT_MAX_COLS`
+smaller than it is refused at startup. Nothing reported a width, the default is 16
+and the dispatcher warns once the first time a `MUL_MAT` or `MUL_MAT_ID` above the
+bound is left unsplit. Set the variable to `0` for no bound; above the bound the
+column policy does not fire, including during prefill.
 
 The experimental policy supports unified, offloaded F16 K/V, causal flash attention,
 256-dimensional K and V heads, no attention soft cap, and no sliding window.
 Shared-weight matmuls over multiple sequence planes are normalized to one plane
 before the inherited selective column dispatcher. Without this, the recurrent
 output projection bypasses batch invariance during concurrent prefill.
-It is measured on text prompts with Qwen3.5-4B on one B200. Context shifting,
-position division, cross-sequence prefix copies, shared-prefix input tokens, and
-whole-context state loading are unsupported. Per-sequence state save and restore
-is supported. Unsupported cache transformations assert instead of silently
-violating the page invariant.
+It is measured on text prompts with Qwen3.5-4B on one B200. Every KV layer has to
+be on the CUDA backend, since no other backend reads the page table; a partial or
+absent offload fails the load naming the layer. The V-less attention layouts have
+no page table either, and a model on one of those is refused at context creation.
+Context shifting, position division, cross-sequence prefix copies, shared-prefix
+input tokens, and whole-context state loading are unsupported. Per-sequence state
+save and restore is supported. Unsupported cache transformations are refused with
+a logged error and leave the cells untouched; `--context-shift` and
+`--cache-reuse` are reported as unsupported at load and disabled there.
 
 The allocator owns pages of 256 cells on behalf of one (sequence, position/256).
 Position modulo 256 fixes the cell offset. Empty pages remain in the unified pool
@@ -50,3 +63,9 @@ reference. `bench.py --modes 0,1 --pairs 3` measures default off against exact m
 on, with 256 predicted tokens. Set `UNSLOTH_WORKSPACE` to the model parent workspace
 and `LD_LIBRARY_PATH` to this build's bin directory. The harness uses GPU 3;
 select a port in 9601-9610 explicitly.
+
+A concurrent request that fails now fails the run instead of being dropped from
+the result, and every run records the environment the server actually inherited,
+including `LLAMA_EXACT_CONCURRENCY`, under `env` in the JSON and in the server log
+header, so a run labelled as the mode-off reference can be checked rather than
+trusted. Teardown is POSIX only.
diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py
index e451afca0e1..3e554f465b7 100644
--- a/scripts/batchinv/divergence.py
+++ b/scripts/batchinv/divergence.py
@@ -5,8 +5,22 @@
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 from prompts import PROMPTS
 
-WS = os.environ["UNSLOTH_WORKSPACE"]
-MODEL = f"{WS}/models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf"
+# Environment recorded with every run. LLAMA_EXACT_CONCURRENCY inherited from the shell is what
+# decides whether a run labelled as the mode-off reference actually was one, so it is not optional.
+RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT",
+                "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY",
+                "CUDA_VISIBLE_DEVICES")
+
+MODEL_REL = "models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf"
+
+
+def model_path():
+    """Resolved when the server args are built, so --help works without the variable set."""
+    ws = os.environ.get("UNSLOTH_WORKSPACE")
+    if not ws:
+        raise RuntimeError("UNSLOTH_WORKSPACE is not set; it must point at the workspace holding "
+                           + MODEL_REL)
+    return os.path.join(ws, MODEL_REL)
 
 
 def post(port, path, payload, timeout=1800):
@@ -34,7 +48,7 @@ def completion(port, prompt, n_predict):
 class Server:
     def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=True):
         self.port, self.log_path = port, log_path
-        self.args = [binary, "-m", MODEL, "--port", str(port), "--host", "127.0.0.1",
+        self.args = [binary, "-m", model_path(), "--port", str(port), "--host", "127.0.0.1",
                      "--parallel", "4", "-c", "8192",
                      "--flash-attn", "on", "--metrics", "-ngl", "99", "--no-warmup",
                      "--seed", "0", "--spec-type", spec]
@@ -46,48 +60,78 @@ def __init__(self, port, binary, extra, env_extra, log_path, spec, kv_unified=Tr
         self.env = dict(os.environ)
         self.env["CUDA_VISIBLE_DEVICES"] = "3"
         self.env.update(env_extra)
+        # what the server will actually see, not what this run meant to set
+        self.env_resolved = {k: self.env[k] for k in RECORDED_ENV if k in self.env}
+        self.p = None
+        self.fh = None
 
     def __enter__(self):
         self.fh = open(self.log_path, "ab")
         self.fh.write(("\n=== " + " ".join(self.args) + "\n=== env " +
-                       json.dumps({k: v for k, v in self.env.items()
-                                   if k.startswith("GGML") or k == "CUDA_VISIBLE_DEVICES"}) + "\n").encode())
+                       json.dumps(self.env_resolved) + "\n").encode())
         self.fh.flush()
         self.p = subprocess.Popen(self.args, stdout=self.fh, stderr=subprocess.STDOUT,
                                   env=self.env, start_new_session=True)
         print(f"[server] pid={self.p.pid} port={self.port} log={self.log_path}", flush=True)
-        deadline = time.time() + 600
-        while time.time() < deadline:
-            if self.p.poll() is not None:
-                raise RuntimeError(f"server died rc={self.p.returncode}, see {self.log_path}")
-            try:
-                if get(self.port, "/health").get("status") == "ok":
-                    print("[server] ready", flush=True)
-                    return self
-            except Exception:
-                time.sleep(1.0)
-        raise RuntimeError("server did not become healthy")
+        try:
+            deadline = time.time() + 600
+            while time.time() < deadline:
+                if self.p.poll() is not None:
+                    raise RuntimeError(f"server died rc={self.p.returncode}, see {self.log_path}")
+                try:
+                    if get(self.port, "/health").get("status") == "ok":
+                        print("[server] ready", flush=True)
+                        return self
+                except Exception:
+                    time.sleep(1.0)
+            raise RuntimeError("server did not become healthy")
+        except BaseException:
+            # __exit__ is not called when __enter__ raises, so a server that started but never
+            # reported healthy would keep the GPU, the port and the log handle
+            self.__exit__(None, None, None)
+            raise
 
     def __exit__(self, *a):
-        print(f"[server] stopping pid={self.p.pid}", flush=True)
-        try:
-            os.killpg(os.getpgid(self.p.pid), signal.SIGTERM)
-            self.p.wait(timeout=60)
-        except Exception:
+        # note: POSIX only. On Windows this needs CREATE_NEW_PROCESS_GROUP at Popen and
+        # terminate()/kill() here; the runs this harness backs are Linux only.
+        if self.p is not None:
+            print(f"[server] stopping pid={self.p.pid}", flush=True)
             try:
-                os.killpg(os.getpgid(self.p.pid), signal.SIGKILL)
+                os.killpg(os.getpgid(self.p.pid), signal.SIGTERM)
+                self.p.wait(timeout=60)
             except Exception:
-                pass
-        self.fh.close()
+                try:
+                    os.killpg(os.getpgid(self.p.pid), signal.SIGKILL)
+                except Exception:
+                    pass
+                try:
+                    self.p.wait(timeout=60)
+                except Exception:
+                    pass
+            self.p = None
+        if self.fh is not None:
+            self.fh.close()
+            self.fh = None
 
 
 def run_concurrent(port, names, n_predict):
     barrier = threading.Barrier(len(names))
+    lock = threading.Lock()
     out = {}
+    errors = []
 
     def work(name):
-        barrier.wait()
-        out[name] = completion(port, PROMPTS[name], n_predict)
+        try:
+            barrier.wait()
+            res = completion(port, PROMPTS[name], n_predict)
+        except BaseException as e:
+            with lock:
+                errors.append((name, e))
+            # release the others rather than let them block on a barrier that will never fill
+            barrier.abort()
+            return
+        with lock:
+            out[name] = res
 
     ts = [threading.Thread(target=work, args=(n,)) for n in names]
     t0 = time.time()
@@ -95,7 +139,18 @@ def work(name):
         t.start()
     for t in ts:
         t.join()
-    return out, time.time() - t0
+    wall = time.time() - t0
+
+    # a thread exception used to only print a traceback, so a run where P1..P3 failed and P0
+    # succeeded was still reported as a clean four-way concurrency result
+    if errors:
+        raise RuntimeError("concurrent requests failed: " +
+                           "; ".join(f"{n}: {type(e).__name__}: {e}" for n, e in errors))
+    missing = set(names) - set(out)
+    if missing:
+        raise RuntimeError(f"concurrent requests produced no result for {sorted(missing)}")
+
+    return out, wall
 
 
 def first_diff(a, b):
@@ -121,12 +176,15 @@ def main():
     a = ap.parse_args()
 
     env_extra = dict(kv.split("=", 1) for kv in a.env)
+    server = Server(a.port, a.binary, a.extra, env_extra, a.out + ".server.log", a.spec,
+                    kv_unified=not a.no_kv_unified)
     res = {"label": a.label, "spec": a.spec, "n_predict": a.n_predict,
-           "env": env_extra, "extra": a.extra, "binary": a.binary,
+           "env_requested": env_extra, "env": server.env_resolved,
+           "model": server.args[2], "args": server.args,
+           "extra": a.extra, "binary": a.binary,
            "kv_unified": not a.no_kv_unified}
 
-    with Server(a.port, a.binary, a.extra, env_extra, a.out + ".server.log", a.spec,
-                kv_unified=not a.no_kv_unified) as s:
+    with server as s:
         solo = completion(a.port, PROMPTS["P0"], a.n_predict)
         ref = json.load(open(a.reference))["tokens"] if a.reference else solo["tokens"]
         res["solo_first_diff"] = first_diff(ref, solo["tokens"])

From 712bee75c69120f869b9ae68e45ffe53978522e6 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 08:34:48 +0000
Subject: [PATCH 15/33] kv-cache: maintain page ownership instead of rebuilding
 it per ubatch

find_slot() and set_input_pages() each rebuilt the (sequence, logical page) to
physical page map from a scan of every live cell into an std::map, so the paged
allocator did O(cells log pages) work twice per ubatch, twice per decode step,
and the CPU cost grew with the size of the pool rather than with the number of
pages in it.

Keep the ownership in a flat vector with one entry per physical page, claimed in
apply_ubatch() as cells are placed and marked dirty by the paths that remove them,
seq_rm(), seq_keep() and clear(). prepare() snapshots it alongside the cells so
that undoing a speculative placement puts back what the allocator knew rather than
forcing a rebuild. Both readers now build their lookup from one entry per page:
32 entries at -c 8192 and 256 at -c 65536, against 8192 and 65536 cells. find_slot
also stops allocating a cells-sized bitmap per call. Nothing about the placement
policy changes, and the derived-from-live-cells rebuild is still there and still
authoritative: LLAMA_KV_CACHE_DEBUG=1 runs it on every read and asserts that the
maintained ownership says exactly what the cells say.

Four chats, 937 token prompts, 1536 tokens each, ignore_eos, speculation off,
--parallel 4 --kv-unified --flash-attn on -ngl 99, three interleaved pairs on one
B200, medians of four-chat aggregate decode tok/s:

  -c 8192    off 151.03   exact 136.70   0.905      before
  -c 8192    off 150.83   exact 137.94   0.915      after
  -c 65536   off 158.34   exact 136.29   0.861      before
  -c 65536   off 159.48   exact 143.62   0.901      after

So at 32 pages it is worth about a point, and at 256 pages it is worth four, which
is what a cost that followed the cell count and now follows the page count should
look like. With LLAMA_KV_CACHE_DEBUG=1 and LLAMA_SERVER_PREEMPT_EVERY=32, which
parks and restores every slot every 32 tokens and so exercises every path that
marks the ownership dirty, no assert fires and P0 stays byte identical to its solo
reference.
---
 scripts/batchinv/divergence.py |   2 +-
 src/llama-kv-cache.cpp         | 149 +++++++++++++++++++++++++++------
 src/llama-kv-cache.h           |  22 +++++
 3 files changed, 147 insertions(+), 26 deletions(-)

diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py
index 3e554f465b7..ec2c3a00467 100644
--- a/scripts/batchinv/divergence.py
+++ b/scripts/batchinv/divergence.py
@@ -9,7 +9,7 @@
 # decides whether a run labelled as the mode-off reference actually was one, so it is not optional.
 RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT",
                 "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY",
-                "CUDA_VISIBLE_DEVICES")
+                "LLAMA_KV_CACHE_DEBUG", "LLAMA_BATCH_DEBUG", "CUDA_VISIBLE_DEVICES")
 
 MODEL_REL = "models/Qwen3.5-4B-MTP-GGUF/Qwen3.5-4B-UD-Q4_K_XL.gguf"
 
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 16c04fb6e3b..06729158133 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -11,6 +11,7 @@
 #include <cstring>
 #include <limits>
 #include <map>
+#include <set>
 #include <stdexcept>
 
 static bool ggml_is_power_of_2(int n) {
@@ -426,7 +427,76 @@ llama_kv_cache::llama_kv_cache(
     debug = LLAMA_KV_CACHE_DEBUG ? atoi(LLAMA_KV_CACHE_DEBUG) : 0;
 }
 
+// [TAG_EXACT_CONCURRENCY]
+void llama_kv_cache::exact_pages_rebuild() const {
+    const auto & cells = v_cells[0];
+
+    exact_page_owner.assign(cells.size()/exact_page_size, exact_page{});
+
+    for (uint32_t i = 0; i < cells.size(); ++i) {
+        if (cells.is_empty(i)) {
+            continue;
+        }
+
+        GGML_ASSERT(cells.seq_count(i) == 1);
+
+        const auto pos = cells.pos_get(i);
+
+        GGML_ASSERT(pos >= 0 && uint32_t(pos)%exact_page_size == i%exact_page_size);
+
+        const exact_page cur { cells.seq_get(i), llama_pos(pos/(llama_pos) exact_page_size) };
+
+        auto & owner = exact_page_owner[i/exact_page_size];
+
+        GGML_ASSERT(owner.seq < 0 || (owner.seq == cur.seq && owner.lpg == cur.lpg));
+
+        owner = cur;
+    }
+
+    exact_page_owner_dirty = false;
+}
+
+// [TAG_EXACT_CONCURRENCY]
+void llama_kv_cache::exact_pages_sync() const {
+    if (exact_page_owner_dirty) {
+        exact_pages_rebuild();
+
+        return;
+    }
+
+    if (debug > 0) {
+        // the incrementally maintained ownership has to say what the cells say
+        const auto kept = exact_page_owner;
+
+        exact_pages_rebuild();
+
+        GGML_ASSERT(kept.size() == exact_page_owner.size());
+
+        for (size_t p = 0; p < kept.size(); ++p) {
+            GGML_ASSERT(kept[p].seq == exact_page_owner[p].seq && kept[p].lpg == exact_page_owner[p].lpg);
+        }
+    }
+}
+
+// [TAG_EXACT_CONCURRENCY]
+void llama_kv_cache::exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos) {
+    if (exact_page_owner_dirty || exact_page_owner.empty()) {
+        // the next sync rebuilds from the cells anyway
+        return;
+    }
+
+    const exact_page cur { seq, llama_pos(pos/(llama_pos) exact_page_size) };
+
+    auto & owner = exact_page_owner[idx/exact_page_size];
+
+    GGML_ASSERT(owner.seq < 0 || (owner.seq == cur.seq && owner.lpg == cur.lpg));
+
+    owner = cur;
+}
+
 void llama_kv_cache::clear(bool data) {
+    exact_page_owner_dirty = true;
+
     for (uint32_t s = 0; s < n_stream; ++s) {
         v_cells[s].reset();
         v_heads[s] = 0;
@@ -445,6 +515,9 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
         return true;
     }
 
+    // [TAG_EXACT_CONCURRENCY] a removal can empty a page, which only the cells know
+    exact_page_owner_dirty = true;
+
     // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG]
     GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()));
 
@@ -614,6 +687,9 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
         return;
     }
 
+    // [TAG_EXACT_CONCURRENCY] as in seq_rm, this can empty pages
+    exact_page_owner_dirty = true;
+
     GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
 
     auto & cells = v_cells[seq_to_stream[seq_id]];
@@ -848,6 +924,10 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_
         std::vector<uint32_t> v_heads_old; // old positions of the heads, before placing the ubatch
 
         std::vector<llama_kv_cells> v_cells; // copy of the old cells, before placing the ubatch
+
+        // [TAG_EXACT_CONCURRENCY] page ownership before placing the ubatch, so that undoing the
+        // speculative placement does not force a rebuild from every cell on the next ubatch
+        std::vector<exact_page> exact_page_owner_old;
     };
 
     // remember the old state of the cells so we can restore it in the end
@@ -868,7 +948,7 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_
 
         // store the old state of the cells in the recovery stack
         {
-            state_t state = { sinfo_new, v_heads, {} };
+            state_t state = { sinfo_new, v_heads, {}, exact_page_owner };
 
             for (uint32_t s = 0; s < sinfo_new.n_stream(); ++s) {
                 auto & cells = v_cells[sinfo_new.strm[s]];
@@ -896,6 +976,13 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_
             cells.set(sinfo.idxs[s], it->v_cells[s]);
             head = it->v_heads_old[s];
         }
+
+        // [TAG_EXACT_CONCURRENCY] the speculative placements are being undone behind the
+        // allocator's back. Put back what it knew before, unless something during the placement
+        // removed cells as well, in which case only the cells can say what is left.
+        if (!exact_page_owner_dirty) {
+            exact_page_owner = it->exact_page_owner_old;
+        }
     }
 
     if (!success) {
@@ -1055,23 +1142,26 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
     }
 
     if (exact_pages) {
-        // Reconstruct page ownership from live cells. Empty pages are immediately reusable;
-        // prepare() can roll back its speculative allocations without a second metadata log.
+        // Page ownership is maintained as cells are placed and invalidated when they are removed,
+        // so the allocator reads one entry per physical page rather than scanning every cell. The
+        // claims this call makes are local: prepare() can still roll back its speculative
+        // placements, and empty pages stay immediately reusable.
         const auto & cells = v_cells[0];
+
+        exact_pages_sync();
+
         using page_key = std::pair<llama_seq_id, llama_pos>;
+
+        std::vector<exact_page> owner = exact_page_owner;
         std::map<page_key, uint32_t> pages;
-        std::vector<bool> occupied(cells.size()/exact_page_size, false);
-        std::vector<bool> assigned(cells.size(), false);
-        for (uint32_t i = 0; i < cells.size(); ++i) {
-            if (cells.is_empty(i)) { continue; }
-            GGML_ASSERT(cells.seq_count(i) == 1);
-            const auto pos = cells.pos_get(i);
-            GGML_ASSERT(pos >= 0 && uint32_t(pos)%exact_page_size == i%exact_page_size);
-            const page_key key {cells.seq_get(i), pos/exact_page_size};
-            auto ins = pages.emplace(key, i/exact_page_size);
-            GGML_ASSERT(ins.first->second == i/exact_page_size);
-            occupied[i/exact_page_size] = true;
+
+        for (uint32_t p = 0; p < owner.size(); ++p) {
+            if (owner[p].seq >= 0) {
+                pages.emplace(page_key {owner[p].seq, owner[p].lpg}, p);
+            }
         }
+
+        std::set<uint32_t> assigned;
         slot_info res {0, 0, {0}, {{}}};
         for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
             GGML_ASSERT(ubatch.n_seq_id[i] == 1 && ubatch.pos[i] >= 0);
@@ -1081,15 +1171,14 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
                 // Round-robin free-page search deliberately permits nonmonotonic physical order.
                 uint32_t page = v_heads[0]/exact_page_size;
                 uint32_t tested = 0;
-                while (tested < occupied.size() && occupied[page%occupied.size()]) { ++page; ++tested; }
-                if (tested == occupied.size()) { return {}; }
-                page %= occupied.size();
-                occupied[page] = true;
+                while (tested < owner.size() && owner[page%owner.size()].seq >= 0) { ++page; ++tested; }
+                if (tested == owner.size()) { return {}; }
+                page %= owner.size();
+                owner[page] = exact_page {key.first, key.second};
                 it = pages.emplace(key, page).first;
             }
             const uint32_t idx = it->second*exact_page_size + ubatch.pos[i]%exact_page_size;
-            if (!cells.is_empty(idx) || assigned[idx]) { return {}; }
-            assigned[idx] = true;
+            if (!cells.is_empty(idx) || !assigned.insert(idx).second) { return {}; }
             res.idxs[0].push_back(idx);
         }
         if (cont && !res.is_contiguous()) { return {}; }
@@ -1274,6 +1363,13 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
             for (int32_t s = 0; s < ubatch.n_seq_id[i]; s++) {
                 cells.seq_add(idx, ubatch.seq_id[i][s]);
             }
+
+            // [TAG_EXACT_CONCURRENCY] the page this cell belongs to is now owned by its sequence
+            if (exact_pages) {
+                GGML_ASSERT(ubatch.n_seq_id[i] == 1);
+
+                exact_pages_claim(idx, ubatch.seq_id[i][0], ubatch.pos[i]);
+            }
         }
     }
 
@@ -1384,12 +1480,15 @@ ggml_tensor * llama_kv_cache::build_input_pages(ggml_context * ctx, const llama_
 
 void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * ubatch) const {
     GGML_ASSERT(exact_pages && dst->ne[1] == ubatch->n_tokens);
+
+    // [TAG_EXACT_CONCURRENCY] one entry per physical page, not one per cell
+    exact_pages_sync();
+
     std::map<llama_seq_id, std::map<llama_pos, uint32_t>> pages;
-    const auto & cells = v_cells[0];
-    for (uint32_t i = 0; i < cells.size(); ++i) {
-        if (!cells.is_empty(i)) {
-            GGML_ASSERT(cells.seq_count(i) == 1);
-            pages[cells.seq_get(i)][cells.pos_get(i)/exact_page_size] = i/exact_page_size;
+    for (uint32_t p = 0; p < exact_page_owner.size(); ++p) {
+        const auto & owner = exact_page_owner[p];
+        if (owner.seq >= 0) {
+            pages[owner.seq][owner.lpg] = p;
         }
     }
     std::vector<int32_t> data(ggml_nelements(dst), -1);
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index fa257422f02..8228752d38b 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -240,6 +240,28 @@ class llama_kv_cache : public llama_memory_i {
     static constexpr uint32_t exact_page_size = 256;
     bool exact_pages = false;
 
+    // [TAG_EXACT_CONCURRENCY]
+    // Which (sequence, logical page) owns each physical page of the pool; seq < 0 means the page is
+    // free. Kept current as cells are placed, and marked dirty by the paths that remove cells, so
+    // that find_slot() and set_input_pages() read one entry per page instead of rebuilding the map
+    // from every live cell twice per ubatch. Mutable because set_input_pages() is const.
+    struct exact_page {
+        llama_seq_id seq = -1;
+        llama_pos    lpg = -1;
+    };
+
+    mutable std::vector<exact_page> exact_page_owner;
+    mutable bool                    exact_page_owner_dirty = true;
+
+    // bring exact_page_owner up to date; rebuilds only when a removal marked it dirty
+    void exact_pages_sync() const;
+
+    // recompute it from the live cells
+    void exact_pages_rebuild() const;
+
+    // record that a cell of (seq, pos) now lives at physical cell idx
+    void exact_pages_claim(uint32_t idx, llama_seq_id seq, llama_pos pos);
+
     bool v_trans = true;  // the value tensor is transposed
 
     const uint32_t n_seq_max = 1;

From a2f9c081f4a9fe3fa5bc368d6a0646fb6a083e5a Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 06:12:58 +0000
Subject: [PATCH 16/33] server: refuse n > 1 under exact concurrency instead of
 aborting in the cache

A review of #194 pointed at the new GGML_ASSERT in llama_kv_cache::seq_cp, and it
is right. Reproduced on this branch with the 4B on one B200:

  LLAMA_EXACT_CONCURRENCY=1, POST /completion {"n": 2}
    -> llama-kv-cache.cpp:458: GGML_ASSERT(!exact_pages || seq_id_src == seq_id_dst)
       failed, through server_context_impl::decode -> common_memory::seq_cp
    -> process aborted, the next request gets connection refused

With the mode off the same request is served normally, so this is reachable by any
client of an exact-mode server and takes every other request on the machine with
it. Two changes:

The server refuses the request. n_cmpl > 1 works by copying the parent's cells to
a second sequence id, and exact mode gives a page to one sequence, so there is
nowhere for that copy to land. Rejecting it where the task is built turns it into
a 400 with a reason. The check reads LLAMA_EXACT_CONCURRENCY from the environment,
the same way the KV cache, the batch splitter and the CUDA backend each do,
because the answer is needed before a context exists and the mode has no other
representation.

The cache stops aborting. seq_cp, seq_add and seq_div log an error and return
instead of asserting, so a caller this branch does not know about degrades to a
refused operation rather than killing the server. The guards also move below the
shared-cells early return, which the asserts sat above: a draft cache forwards
these calls and copies nothing of its own, and it should not be judged by a rule
about cells it does not own.

After: the n=2 request returns 400 on both /completion and /v1/completions, the
server stays up, and a following ordinary request returns 200.
---
 src/llama-kv-cache.cpp          | 55 +++++++++++++++++++--------------
 tools/server/server-context.cpp | 25 +++++++++++++++
 2 files changed, 57 insertions(+), 23 deletions(-)

diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 06729158133..8ecf701a557 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -582,19 +582,24 @@ bool llama_kv_cache::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
 }
 
 void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
-    // [TAG_EXACT_CONCURRENCY] a page belongs to one (sequence, position/256) pair, so two sequences
-    // cannot share physical cells. Refuse the copy rather than abort the process: this is reachable
-    // from a request parameter.
-    if (exact_pages && seq_id_src != seq_id_dst) {
-        LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so cells cannot be shared between "
-                "sequences: ignoring the copy from seq %d to seq %d\n", __func__, seq_id_src, seq_id_dst);
-        return;
-    }
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
     }
 
+    // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared
+    // between two of them. Refuse the operation rather than abort the process: a server
+    // rejects the request that would reach here (n_cmpl > 1), and any caller this does not
+    // cover degrades to a failed copy it can report instead of killing every other request
+    // on the machine. Placed after the shared-cells return so a draft cache, which copies
+    // nothing of its own, is unaffected.
+    if (exact_pages && seq_id_src != seq_id_dst) {
+        LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between "
+                        "sequences (%d -> %d); ignoring the copy\n",
+                        __func__, seq_id_src, seq_id_dst);
+        return;
+    }
+
     GGML_ASSERT(seq_id_src >= 0 && (size_t) seq_id_src < seq_to_stream.size());
     GGML_ASSERT(seq_id_dst >= 0 && (size_t) seq_id_dst < seq_to_stream.size());
 
@@ -712,19 +717,21 @@ void llama_kv_cache::seq_keep(llama_seq_id seq_id) {
 }
 
 void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
-    // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so a
-    // shift would have to move the cells too. get_can_shift() reports this so that --context-shift
-    // and --cache-reuse are turned off at load; this is the guard for the library API.
-    if (exact_pages && shift != 0) {
-        LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so positions cannot be shifted: "
-                "ignoring the shift of %d on seq %d\n", __func__, shift, seq_id);
-        return;
-    }
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
     }
 
+    // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo the
+    // page size, so shifting positions would put every cell of the sequence in the wrong
+    // place. Context shift is unsupported in exact mode; say so rather than abort.
+    if (exact_pages && shift != 0) {
+        LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions "
+                        "(seq %d, shift %d); ignoring the shift\n",
+                        __func__, seq_id, shift);
+        return;
+    }
+
     GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
     GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_add() is only supported for n_pos_per_embd() == 1");
 
@@ -770,18 +777,20 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll
 }
 
 void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
-    // [TAG_EXACT_CONCURRENCY] same reason as seq_add: the offset inside a page is derived from the
-    // position, so dividing the positions would leave every cell in the wrong slot.
-    if (exact_pages && d != 1) {
-        LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, so positions cannot be divided: "
-                "ignoring the division by %d on seq %d\n", __func__, d, seq_id);
-        return;
-    }
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
     }
 
+    // [TAG_EXACT_CONCURRENCY] same reason as seq_add: dividing positions breaks the
+    // identity between a cell's position and its offset inside its page.
+    if (exact_pages && d != 1) {
+        LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions "
+                        "(seq %d, d %d); ignoring the division\n",
+                        __func__, seq_id, d);
+        return;
+    }
+
     GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
     GGML_ASSERT(hparams.n_pos_per_embd() == 1 && "seq_div() is only supported for n_pos_per_embd() == 1");
 
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 6723c51397e..9f3d1c53eb5 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -37,6 +37,19 @@
 
 constexpr int HTTP_POLLING_SECONDS = 1;
 
+// [TAG_EXACT_CONCURRENCY] the knob is read from the environment by the KV cache, the batch
+// splitter and the CUDA backend independently, because it has to be answered before a
+// context exists. The server needs the same answer to refuse the one request shape the mode
+// cannot serve, so it reads it the same way rather than growing a public API for it.
+static bool server_exact_concurrency() {
+    static const bool enabled = []() {
+        const char * val = getenv("LLAMA_EXACT_CONCURRENCY");
+        return val && atoi(val) != 0;
+    }();
+
+    return enabled;
+}
+
 static common_speculative_output_limits server_output_limits(const common_params & params) {
     if (params.embedding ||
             (params.pooling_type != LLAMA_POOLING_TYPE_UNSPECIFIED && params.pooling_type != LLAMA_POOLING_TYPE_NONE)) {
@@ -4686,6 +4699,18 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
             task.params.oaicompat_cmpl_id = completion_id;
             task.params.oaicompat_model   = meta->model_name;
 
+            // [TAG_EXACT_CONCURRENCY] the children of an n_cmpl > 1 task are served by
+            // copying the parent's cells to another sequence id, and exact mode gives a KV
+            // page to one sequence, so there is nothing for that copy to land in. Refuse
+            // the request here, where it becomes a 400 the client can read, rather than
+            // letting it reach seq_cp with nothing to do.
+            if (task.params.n_cmpl > 1 && server_exact_concurrency()) {
+                throw std::runtime_error(
+                    "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each "
+                    "completion needs its own sequence, and in exact mode a KV page belongs "
+                    "to a single sequence. Send n separate requests, or unset the variable.");
+            }
+
             // prepare child tasks
             if (task.params.n_cmpl > 1) {
                 int n_children = task.params.n_cmpl - 1;

From 72aca44b5e6f7b8bb3a5bf33bcbad16790435971 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 08:47:59 +0000
Subject: [PATCH 17/33] memory: let the server ask how many cells an allocation
 takes

Under exact concurrency the KV cache hands out 256-cell pages, one to each
(sequence, position / 256) pair, so a sequence can hold up to 255 cells that
nobody else can be given. A preemption planner that counts tokens does not see
those cells: it believes there is room, never parks anybody, and the pool fills
until every request ends in the old context error.

llama_memory_i gains alloc_granularity(), defaulting to 1 so no module that
allocates a cell per token changes; the KV cache returns its page size under
exact mode and the hybrid memory forwards to its attention half.
llama_memory_alloc_granularity() exposes it. The server side of this, rounding
its planner figures by that value, follows separately.
---
 include/llama.h             | 8 ++++++++
 src/llama-context.cpp       | 8 ++++++++
 src/llama-kv-cache.cpp      | 7 +++++++
 src/llama-kv-cache.h        | 3 +++
 src/llama-memory-hybrid.cpp | 6 ++++++
 src/llama-memory-hybrid.h   | 2 ++
 src/llama-memory.h          | 9 +++++++++
 7 files changed, 43 insertions(+)

diff --git a/include/llama.h b/include/llama.h
index a04177f9f7d..43d79b40a2a 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -795,6 +795,14 @@ extern "C" {
     // Check if the memory supports shifting
     LLAMA_API bool llama_memory_can_shift(llama_memory_t mem);
 
+    // [TAG_EXACT_CONCURRENCY] Cells the memory allocates in one indivisible unit.
+    //
+    // 1 in every ordinary configuration. Larger where a mode places cells in blocks, and
+    // then a sequence of n tokens occupies round_up(n, granularity) cells. A caller that
+    // decides whether the pool has room by counting tokens has to round the same way, or it
+    // will believe there is space that cannot be handed out.
+    LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem);
+
     //
     // State / sessions
     //
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 66940d4fc61..3ab84de780c 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -4030,6 +4030,14 @@ bool llama_memory_can_shift(llama_memory_t mem) {
     return mem->get_can_shift();
 }
 
+uint32_t llama_memory_alloc_granularity(llama_memory_t mem) {
+    if (!mem) {
+        return 1;
+    }
+
+    return mem->alloc_granularity();
+}
+
 // llama state API
 
 // deprecated
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 8ecf701a557..dced34f6a14 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -1410,6 +1410,13 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
     }
 }
 
+uint32_t llama_kv_cache::alloc_granularity() const {
+    // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so a
+    // sequence holding n tokens holds round_up(n, exact_page_size) cells: its tail page is
+    // charged in full whether or not it is full.
+    return exact_pages ? exact_page_size : 1;
+}
+
 bool llama_kv_cache::get_can_shift() const {
     // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so the
     // paged pool cannot shift positions. Reporting it here is what makes the server disable
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index 8228752d38b..af3a04be39d 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -131,6 +131,9 @@ class llama_kv_cache : public llama_memory_i {
 
     bool get_can_shift() const override;
 
+    // [TAG_EXACT_CONCURRENCY] the page size under exact mode, 1 otherwise
+    uint32_t alloc_granularity() const override;
+
     void clear(bool data) override;
 
     bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;
diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index 4ebd476aa8b..48fdea8b49c 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -142,6 +142,12 @@ bool llama_memory_hybrid::get_can_shift() const {
     return mem_attn->get_can_shift();
 }
 
+uint32_t llama_memory_hybrid::alloc_granularity() const {
+    // the recurrent half holds one state per sequence rather than per token, so the
+    // attention half is the one whose cells a caller is planning capacity for
+    return mem_attn->alloc_granularity();
+}
+
 void llama_memory_hybrid::clear(bool data) {
     mem_attn->clear(data);
     mem_recr->clear(data);
diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h
index 484eafb7499..70ba19ca323 100644
--- a/src/llama-memory-hybrid.h
+++ b/src/llama-memory-hybrid.h
@@ -58,6 +58,8 @@ class llama_memory_hybrid : public llama_memory_i {
 
     bool get_can_shift() const override;
 
+    uint32_t alloc_granularity() const override;
+
     void clear(bool data) override;
 
     bool seq_rm  (llama_seq_id seq_id,                              llama_pos p0, llama_pos p1) override;
diff --git a/src/llama-memory.h b/src/llama-memory.h
index db825396645..51539a03919 100644
--- a/src/llama-memory.h
+++ b/src/llama-memory.h
@@ -100,6 +100,15 @@ struct llama_memory_i {
     // getters
     virtual bool get_can_shift() const = 0;
 
+    // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit.
+    //
+    // 1 for every module that allocates a cell per token, which is all of them unless a mode
+    // is on that allocates in larger blocks. Where it is larger, a sequence of n tokens
+    // occupies round_up(n, granularity) cells, and a caller that plans pool capacity by
+    // counting tokens will believe there is room that does not exist. Not pure, so a module
+    // that has never heard of this inherits the answer that has always been true of it.
+    virtual uint32_t alloc_granularity() const { return 1; }
+
     //
     // ops
     //

From da8556d316d2f7bb66b690d87d089019e5a720ce Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 09:23:52 +0000
Subject: [PATCH 18/33] server: plan the kv pool in cells rather than tokens

Exact mode and #184's preemption planner cannot both be right about how full the
pool is, and on this branch they are not. Reproduced with the 4B on one B200,
LLAMA_EXACT_CONCURRENCY=1, four chats with 1000-token prompts generating 2048
tokens each at --parallel 4 --kv-unified -c 8192 --spec-type draft-mtp
--spec-draft-n-max 2, no forced-park knob, three rounds:

  0 of 4, 0 of 4 and 1 of 4 completions, 12 "Context size has been exceeded",
  0 parks and 0 restores.

Nothing was ever parked. preempt_kv_used(), preempt_n_need() and
preempt_kv_reserve() count tokens, and exact mode's allocator hands out 256-cell
pages, one page to one (sequence, position / 256) pair. Four sequences can
therefore be holding up to 1020 cells that no other sequence can be given, and the
planner, seeing room in tokens that find_slot cannot find in pages, never reaches
the threshold that would park anybody. The retry ladder then halves n_batch to 1
and ends every request, which is the pre-#184 behaviour that preemption exists to
remove.

Ask the memory how it allocates instead of assuming. The server reads
llama_memory_alloc_granularity() once at load and rounds:

  preempt_kv_used()    charges every slot's tail page in full, because a page
                       belongs to one sequence however little of it is used
  preempt_n_need()     rounds what a resume must be given, since a restore takes
                       fresh pages
  preempt_kv_reserve() reserves the cells the next step ADDS rather than its
                       tokens, because on a rounded used figure a step is free
                       until it crosses a page boundary and costs a whole page
                       when it does, and that crossing is the only moment the pool
                       can run out
  preempt_n_margin()   rounds the spare cells up to a page, since a margin of
                       eight is no margin at all where a step can cost 256

With a granularity of 1 every one of these is the arithmetic it was, which is
pinned by static assertions on the two rounding helpers rather than left to be
read: at 1 they are the identity, so nothing changes with the mode off.

After, same configuration and three rounds: 4 of 4 completions each round, 3, 4
and 4 parks and the same number of restores, no context errors, and P0
byte-identical to its solo run in all three. The same run with exact mode off is
also 4 of 4 with 3, 2 and 2 parks, and its planner figures are still the token
counts they always were.

LLAMA_SERVER_PREEMPT_GRANULARITY overrides the figure the memory reports. It is a
test knob, next to LLAMA_SERVER_PREEMPT_EVERY: the paged attention kernel needs a
head size of 256, which the harness model does not have, so this is the only way
to reach the paged arithmetic from tools/server/tests. The new test drives two
slots over a 256-cell pool at a granularity of 64 and asserts every figure the
planner logs is a whole number of blocks. Counting tokens the same run logs
kv 119/256 and wanted 249; counting cells it logs kv 64/256 and wanted 256.
---
 tools/server/server-context.cpp         | 117 ++++++++++++++++++++++--
 tools/server/tests/unit/test_preempt.py |  44 +++++++++
 2 files changed, 152 insertions(+), 9 deletions(-)

diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 9f3d1c53eb5..947ec514772 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -93,6 +93,45 @@ constexpr int32_t PREEMPT_N_STARVED  = 3;  // preemptions after which a slot is
 constexpr int32_t PREEMPT_N_FAIL_MAX = 8;  // failed restores before the slot is given up on
 constexpr int64_t PREEMPT_FAIL_US    = 60ll * 1000 * 1000;  // ... and only after this long parked
 
+// [TAG_EXACT_CONCURRENCY] The planner above counts cells, not tokens, because the two are not
+// the same number under every mode. llama_memory_alloc_granularity() reports how many cells the
+// pool hands out at a time: 1 in every ordinary configuration, and the exact concurrency page
+// size when that mode is on, where one page belongs to one (sequence, position / page) pair and
+// a sequence of n tokens therefore occupies round_up(n, page) cells. Four sequences can be
+// holding up to 4 * (page - 1) cells that nobody else can be given, and a planner counting
+// tokens sees room in the pool that find_slot cannot find in pages: it never reaches the
+// threshold that would park anybody, the retry ladder halves n_batch to 1, and every request
+// ends in the context error that preemption exists to remove.
+
+// cells a run of n_tokens occupies when the pool allocates g at a time
+static constexpr int32_t preempt_n_cells_g(int32_t n_tokens, int32_t g) {
+    return (g <= 1 || n_tokens <= 0) ? n_tokens : ((n_tokens + g - 1) / g) * g;
+}
+
+// cells a run of n_tokens has to be given for a step of n_step more: nothing until the step
+// crosses a page boundary, a whole page when it does
+static constexpr int32_t preempt_n_cells_step_g(int32_t n_tokens, int32_t n_step, int32_t g) {
+    return preempt_n_cells_g(n_tokens + n_step, g) - preempt_n_cells_g(n_tokens, g);
+}
+
+// At a granularity of 1 both are the identity, so every figure the planner computes is exactly
+// the arithmetic it did before it started asking the memory how it allocates, and nothing
+// changes in any configuration that does not page.
+static_assert(preempt_n_cells_g(0, 1) == 0 && preempt_n_cells_g(1, 1) == 1 &&
+              preempt_n_cells_g(8191, 1) == 8191 && preempt_n_cells_g(-3, 1) == -3,
+              "at a granularity of 1 a run of n tokens has to cost exactly n cells");
+static_assert(preempt_n_cells_step_g(0, 1, 1) == 1 && preempt_n_cells_step_g(8191, 1, 1) == 1 &&
+              preempt_n_cells_step_g(1000, 512, 1) == 512,
+              "at a granularity of 1 a step of n tokens has to cost exactly n cells");
+
+// and the page arithmetic itself, so the rounding cannot be changed by accident
+static_assert(preempt_n_cells_g(1, 256) == 256 && preempt_n_cells_g(256, 256) == 256 &&
+              preempt_n_cells_g(257, 256) == 512,
+              "a tail page is charged in full");
+static_assert(preempt_n_cells_step_g(255, 1, 256) == 0 && preempt_n_cells_step_g(256, 1, 256) == 256 &&
+              preempt_n_cells_step_g(256, 257, 256) == 512,
+              "a step is free until it crosses a page boundary and costs whole pages when it does");
+
 struct server_slot; // forward declaration
 
 struct server_batch {
@@ -1415,6 +1454,28 @@ struct server_context_impl {
             }
         }
 
+        // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates, rather than assume a cell per
+        // token. 1 in every ordinary configuration, so this changes nothing unless a mode that
+        // places cells in blocks is on.
+        {
+            preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt)));
+
+            // a test knob: the paged attention kernel only supports a head size of 256, so a
+            // harness model cannot turn exact concurrency on, and this is the only way to reach
+            // the paged arithmetic of the planner from the server tests
+            const char * LLAMA_SERVER_PREEMPT_GRANULARITY = getenv("LLAMA_SERVER_PREEMPT_GRANULARITY");
+
+            if (LLAMA_SERVER_PREEMPT_GRANULARITY) {
+                preempt_alloc_granularity = std::max(1, atoi(LLAMA_SERVER_PREEMPT_GRANULARITY));
+
+                SRV_WRN("LLAMA_SERVER_PREEMPT_GRANULARITY = %d (test knob: planning the kv pool in blocks of %d cells)\n",
+                        preempt_alloc_granularity, preempt_alloc_granularity);
+            } else if (preempt_alloc_granularity > 1) {
+                SRV_INF("preemption: the kv pool allocates %d cells at a time, planning in pages\n",
+                        preempt_alloc_granularity);
+            }
+        }
+
         {
             const char * LLAMA_SERVER_PREEMPT_EVERY = getenv("LLAMA_SERVER_PREEMPT_EVERY");
             preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0;
@@ -2868,6 +2929,30 @@ struct server_context_impl {
     // uninterrupted one is the preemption's fault and nothing else's.
     int32_t preempt_test_every = 0;
 
+    // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load from the
+    // memory itself: 1 in every ordinary configuration, the page size under exact concurrency.
+    // Everything below plans in cells because of it. LLAMA_SERVER_PREEMPT_GRANULARITY overrides
+    // it, which is how the harness reaches the paged arithmetic on a model whose head size the
+    // paged attention kernel does not support.
+    int32_t preempt_alloc_granularity = 1;
+
+    // cells a slot holding n_tokens actually occupies
+    int32_t preempt_n_cells(int32_t n_tokens) const {
+        return preempt_n_cells_g(n_tokens, preempt_alloc_granularity);
+    }
+
+    // cells a slot holding n_tokens has to be given for a step of n_step more
+    int32_t preempt_n_cells_step(int32_t n_tokens, int32_t n_step) const {
+        return preempt_n_cells_step_g(n_tokens, n_step, preempt_alloc_granularity);
+    }
+
+    // Cells kept spare on top of the reservation. A step that crosses a page boundary costs a
+    // whole page rather than a cell, so a margin of a few cells is no margin at all under a page
+    // allocator: round it up to one page. With a granularity of 1 this is PREEMPT_N_MARGIN.
+    int32_t preempt_n_margin() const {
+        return preempt_n_cells(PREEMPT_N_MARGIN);
+    }
+
     int32_t preempt_n_spec_max() const {
         return spec ? std::max(0, common_speculative_n_max(&params_base.speculative)) : 0;
     }
@@ -2906,7 +2991,10 @@ struct server_context_impl {
             res += std::max(1, std::min((int32_t) llama_n_batch(ctx_tgt), n_left));
         }
 
-        return res;
+        // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in
+        // full, so what the pool has to have free for this slot is the rounded figure. Under
+        // counting here is what admits a resume that find_slot then cannot satisfy.
+        return preempt_n_cells(res);
     }
 
     // Cells the pool is holding right now. A released slot keeps its prompt in the cache
@@ -2921,7 +3009,9 @@ struct server_context_impl {
                 continue; // parked: its cells are in host RAM, not in the pool
             }
 
-            res += slot.prompt.n_tokens();
+            // [TAG_EXACT_CONCURRENCY] the slot's tail page is charged in full: it belongs to
+            // this sequence and cannot be given to anybody else, however little of it is used
+            res += preempt_n_cells(slot.prompt.n_tokens());
         }
 
         return res;
@@ -2935,27 +3025,36 @@ struct server_context_impl {
         int32_t res     = 0;
         int32_t res_pmt = 0;
 
+        // [TAG_EXACT_CONCURRENCY] each slot reserves the cells its next step ADDS, not the
+        // tokens it adds. preempt_kv_used() already charges every slot's tail page in full, so
+        // with a granularity of 1 these are the same number and nothing changes; with a larger
+        // one the step is free until it crosses a page boundary and costs a whole page when it
+        // does. Reserving tokens on top of a rounded used figure would miss exactly that
+        // crossing, which is the only moment the pool can actually run out.
         for (const auto & slot : slots) {
+            const int32_t n_cur = slot.prompt.n_tokens();
+
             switch (slot.state) {
                 case SLOT_STATE_GENERATING:
                 case SLOT_STATE_DONE_PROMPT:
                     {
-                        res += 1 + n_spec;
+                        res += preempt_n_cells_step(n_cur, 1 + n_spec);
                     } break;
                 case SLOT_STATE_STARTED:
                 case SLOT_STATE_PROCESSING_PROMPT:
                     {
-                        const int32_t n_left = slot.task ? slot.task->n_tokens() - slot.prompt.n_tokens() : 0;
+                        const int32_t n_left = slot.task ? slot.task->n_tokens() - n_cur : 0;
 
-                        res_pmt += std::max(1, std::min(n_batch, n_left));
+                        res_pmt += preempt_n_cells_step(n_cur, std::max(1, std::min(n_batch, n_left)));
                     } break;
                 default:
                     break;
             }
         }
 
-        // one batch is all the prompt slots get between them, however many are waiting
-        return res + std::min(res_pmt, n_batch);
+        // one batch is all the prompt slots get between them, however many are waiting; in
+        // cells that batch can straddle one boundary more than it has tokens for
+        return res + std::min(res_pmt, preempt_n_cells(n_batch));
     }
 
     // Keep the slot that is furthest along -- it is the closest to finishing and to giving
@@ -3068,7 +3167,7 @@ struct server_context_impl {
             // continue, so give those cells up first - same call the KV-full path makes.
             for (;;) {
                 for (auto * slot : parked) {
-                    if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + PREEMPT_N_MARGIN <= n_cells) {
+                    if (preempt_kv_used() + preempt_kv_reserve() + preempt_n_need(*slot) + preempt_n_margin() <= n_cells) {
                         best = slot;
                         break;
                     }
@@ -3132,7 +3231,7 @@ struct server_context_impl {
         for (;;) {
             const int32_t n_used = preempt_kv_used() + preempt_kv_reserve();
 
-            if (n_used + PREEMPT_N_MARGIN <= n_cells) {
+            if (n_used + preempt_n_margin() <= n_cells) {
                 break;
             }
 
diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py
index 0da885bcafd..1e22d0a6ddb 100644
--- a/tools/server/tests/unit/test_preempt.py
+++ b/tools/server/tests/unit/test_preempt.py
@@ -1,4 +1,5 @@
 import os
+import re
 import time
 import tempfile
 import pytest
@@ -38,6 +39,7 @@ def create_server():
     os.close(fd)
     yield
     os.environ.pop("LLAMA_SERVER_PREEMPT_EVERY", None)
+    os.environ.pop("LLAMA_SERVER_PREEMPT_GRANULARITY", None)
     os.environ.pop("LLAMA_ARG_PREEMPT_RAM", None)
 
 
@@ -111,6 +113,48 @@ def test_two_slots_that_overflow_the_pool_together_both_finish():
         assert len(res.body["tokens"]) == n_predict
 
 
+def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages():
+    # A pool that hands out cells in blocks gives a whole block to one sequence, so a sequence
+    # of n tokens occupies round_up(n, block) cells and holds the rest of its tail block against
+    # everybody else. The planner has to count those cells: counting tokens, it sees room the
+    # allocator cannot find, never parks anybody, and the retry ladder ends every request.
+    #
+    # llama_memory_alloc_granularity() reports the block size, and the only mode that returns
+    # more than 1 today is exact concurrency, whose paged attention kernel needs a head size this
+    # model does not have. LLAMA_SERVER_PREEMPT_GRANULARITY injects the figure instead: what is
+    # under test is the server's arithmetic, which is the same at 64 as at 256.
+    global server
+    server.n_ctx = 256
+    os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64"
+    server.start()
+    log = LogReader(server.log_path)
+    assert "LLAMA_SERVER_PREEMPT_GRANULARITY = 64" in log.drain()
+
+    n_predict = 160
+    results = parallel_function_calls([
+        (_complete, (n_predict, "Once upon a time there was a brave knight who")),
+        (_complete, (n_predict, "The quick brown fox jumps over the lazy dog and")),
+    ])
+
+    text = log.drain()
+    assert "Context size has been exceeded" not in text
+    assert "preempted:" in text
+    assert "resumed after" in text
+
+    # every figure the planner logs is a whole number of blocks: "kv N/256" is what the pool is
+    # holding and "(wanted N)" is that plus what the next decode reserves. Counting tokens, both
+    # land wherever the sequences happen to be.
+    held   = [int(n) for n in re.findall(r"kv (\d+)/256", text)]
+    wanted = [int(n) for n in re.findall(r"\(wanted (\d+)\)", text)]
+    assert held and wanted, f"the planner logged no figures:\n{text}"
+    assert all(n % 64 == 0 for n in held + wanted), f"not whole blocks: {held} {wanted}"
+
+    for res in results:
+        assert res.status_code == 200
+        assert res.body["timings"]["predicted_n"] == n_predict
+        assert res.body["truncated"] is False
+        assert len(res.body["tokens"]) == n_predict
+
 
 _WORDS = (
     "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor "

From c6c3cb671ded0994d682003b534f8d6c6593ce3b Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 10:50:22 +0000
Subject: [PATCH 19/33] exact: group speculative verify batches and slice the
 column split

Under LLAMA_EXACT_CONCURRENCY a decode step with speculative drafts cost
about twice what the same step cost with the mode off, and the mode itself
was not where the time went.

The batch splitter isolated any sequence set with more than one token left
to place, on the reasoning that such a set is a prompt whose prefill would
otherwise share a ubatch with other sequences. A speculative verify batch
is such a set too: with two MTP drafts every slot brings three tokens, so
each slot's verify step became a ubatch of its own and every decode step ran
the whole graph once per slot.

The splitter now isolates by width. llama_set_exact_decode_tokens() tells
the library how many tokens one sequence contributes to a decode step, one
plus the draft length, and only a set with more tokens than that left to
place is a prompt. The server already derives the CUDA column bound from
the same figure, so a grouped verify batch is at most that wide and the
column policy keeps every column at its batch-of-one arithmetic. The
default of 1 is the previous behaviour.

On the CUDA side the column policy recomputed a batch above the bound one
column at a time, reading the weights once per column. A column's result
depends on the implementation and, for MMVQ, on the warp count of the
launch, and neither depends on the other columns in the launch, so the
split now runs in the widest slices whose configuration matches a batch of
one: a twelve-column verify batch is three MMVQ launches on a table whose
configuration holds up to four columns, rather than twelve. Mode 1 of
GGML_CUDA_BATCH_INVARIANT still recomputes one column at a time.

Qwen3.5-4B, four chats with two MTP drafts each, exact mode on: byte
identical to the solo run in every round, with and without a forced park
every 64 tokens, at about twice the previous aggregate rate.
Qwen3.6-35B-A3B, the same cell: identical in every round at 92 to 96 tok/s
against 55 before. MUL_MAT op tests 1217 of 1217 under the mode; server
preemption tests 7 of 7.
---
 common/common.cpp               |  3 ++
 ggml/src/ggml-cuda/ggml-cuda.cu | 53 ++++++++++++++++++++++++++++-----
 include/llama.h                 |  8 +++++
 src/llama-batch.cpp             | 22 +++++++-------
 src/llama-batch.h               | 20 +++++++------
 src/llama-impl.cpp              | 12 ++++++++
 src/llama-kv-cache.cpp          |  2 +-
 src/llama-memory-hybrid.cpp     |  2 +-
 src/llama-memory-recurrent.cpp  |  2 +-
 9 files changed, 94 insertions(+), 30 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index d2beebc8e2f..8a1c76793cc 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1484,6 +1484,9 @@ bool common_exact_concurrency_init(const common_params & params) {
 
     const int n_cols = common_exact_decode_width(params);
 
+    // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step is
+    llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel)));
+
     const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
     if (bound) {
         const int max_cols = atoi(bound);
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index c8f43bae382..2d3fc641ae6 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1962,7 +1962,35 @@ static ggml_cuda_mm_path ggml_cuda_mul_mat_path(
 
 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
 
-// Recompute dst one column at a time so that each column sees the batch-of-one configuration.
+// [TAG_BATCH_INVARIANT]
+// The widest slice of columns that can be recomputed in one launch while every column in it still
+// sums the way a batch of one would. A column's result depends on the implementation and, for
+// MMVQ, on the warp count of the launch, and neither depends on the values of the other columns
+// in the launch, so a slice as wide as the batch-of-one configuration reaches gives each of its
+// columns the batch-of-one value while reading the weights once for all of them instead of once
+// per column. A twelve-column speculative decode over a table whose configuration holds up to four
+// columns then costs three weight reads rather than twelve. Always below ncols_dst, so the
+// recursive call cannot land back here with the same shape.
+static int64_t ggml_cuda_mul_mat_invariant_width(
+        int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst,
+        ggml_cuda_mm_path path_one, int64_t ncols_dst) {
+    if (path_one != GGML_CUDA_MM_MMVF && path_one != GGML_CUDA_MM_MMVQ) {
+        return 1;
+    }
+    const int64_t widest = path_one == GGML_CUDA_MM_MMVF ? MMVF_MAX_BATCH_SIZE : MMVQ_MAX_BATCH_SIZE;
+    for (int64_t w = std::min<int64_t>(ncols_dst - 1, widest); w > 1; --w) {
+        if (ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, w) != path_one) {
+            continue;
+        }
+        if (path_one == GGML_CUDA_MM_MMVQ && !ggml_cuda_mmvq_matches_single_column(src0->type, cc, w)) {
+            continue;
+        }
+        return w;
+    }
+    return 1;
+}
+
+// Recompute dst in slices of columns so that each column sees the batch-of-one configuration.
 // Returns false when the batched launch already gives every column that same value.
 static bool ggml_cuda_mul_mat_split_columns(
         ggml_backend_cuda_context & ctx, int cc, int warp_size,
@@ -2001,6 +2029,9 @@ static bool ggml_cuda_mul_mat_split_columns(
         return false;
     }
 
+    // Mode 1 recomputes one column at a time. Mode 2 recomputes in the widest slices that keep the
+    // batch-of-one arithmetic, which is what the exact concurrency mode runs under.
+    int64_t width = 1;
     if (ggml_cuda_batch_invariant() >= 2) {
         const ggml_cuda_mm_path path_one     = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1);
         const ggml_cuda_mm_path path_batched = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ncols_dst);
@@ -2014,20 +2045,26 @@ static bool ggml_cuda_mul_mat_split_columns(
                 return false;
             }
         }
+        width = ggml_cuda_mul_mat_invariant_width(cc, warp_size, src0, src1, dst, path_one, ncols_dst);
+        if (width >= ncols_dst) {
+            width = 1;
+        }
     }
 
-    for (int64_t i = 0; i < ncols_dst; ++i) {
+    for (int64_t i = 0; i < ncols_dst; i += width) {
+        const int64_t n = std::min(width, ncols_dst - i);
+
         ggml_tensor src1_col = *src1;
         ggml_tensor dst_col  = *dst;
 
-        src1_col.ne[1] = 1;
-        src1_col.nb[2] = src1_col.nb[1];
-        src1_col.nb[3] = src1_col.nb[1];
+        src1_col.ne[1] = n;
+        src1_col.nb[2] = n*src1_col.nb[1];
+        src1_col.nb[3] = n*src1_col.nb[1];
         src1_col.data  = (char *) src1->data + i*src1->nb[1];
 
-        dst_col.ne[1] = 1;
-        dst_col.nb[2] = dst_col.nb[1];
-        dst_col.nb[3] = dst_col.nb[1];
+        dst_col.ne[1] = n;
+        dst_col.nb[2] = n*dst_col.nb[1];
+        dst_col.nb[3] = n*dst_col.nb[1];
         dst_col.data  = (char *) dst->data + i*dst->nb[1];
 
         ggml_cuda_mul_mat(ctx, src0, &src1_col, &dst_col);
diff --git a/include/llama.h b/include/llama.h
index 43d79b40a2a..0e44ff052c2 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -803,6 +803,14 @@ extern "C" {
     // will believe there is space that cannot be handed out.
     LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem);
 
+    // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or
+    // 1 plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence
+    // set with more tokens than this left to place is a prompt and is prefilled in a ubatch of its
+    // own; a set at or below it is a decode step and stays grouped with the other decodes, so a
+    // speculative verify batch is not run once per sequence. Process-wide, default 1.
+    LLAMA_API void     llama_set_exact_decode_tokens(uint32_t n_tokens);
+    LLAMA_API uint32_t llama_exact_decode_tokens(void);
+
     //
     // State / sessions
     //
diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 4b73ab2478b..4080ecd11d0 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -507,7 +507,7 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) {
     return ubatch_add(idxs, idxs.size(), false);
 }
 
-bool llama_batch_allocr::has_multi_token_seq() const {
+bool llama_batch_allocr::has_seq_wider_than(uint32_t n_tokens) const {
     std::vector<uint32_t> n_per_seq(n_seq_max, 0);
 
     for (int32_t i = 0; i < batch.n_tokens; ++i) {
@@ -517,7 +517,7 @@ bool llama_batch_allocr::has_multi_token_seq() const {
         }
 
         for (int32_t s = 0; s < batch.n_seq_id[i]; ++s) {
-            if (++n_per_seq[batch.seq_id[i][s]] > 1) {
+            if (++n_per_seq[batch.seq_id[i][s]] > n_tokens) {
                 return true;
             }
         }
@@ -526,7 +526,7 @@ bool llama_batch_allocr::has_multi_token_seq() const {
     return false;
 }
 
-llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, bool isolate_multi_token_seqs) {
+llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above) {
     if (sequential && has_cpl) {
         LLAMA_LOG_ERROR("%s: sequential split is not supported when there are coupled sequences in the input batch (you may need to use the -kvu flag)\n", __func__);
 
@@ -559,12 +559,14 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
         }
 
         if (add) {
-            // [TAG_EXACT_CONCURRENCY] a sequence set that still has more than one token to place is
-            // a prompt, and a prompt shares its arithmetic with whatever else is in the ubatch, so
-            // give it a ubatch of its own. Sets with one token left are a plain decode step, which
-            // is already exact, so keep grouping those: isolating them too would make one prompt
-            // serialize every concurrent decode for the whole of the prefill.
-            if (isolate_multi_token_seqs) {
+            // [TAG_EXACT_CONCURRENCY] a sequence set that still has more tokens to place than a
+            // decode step carries is a prompt, and a prompt shares its arithmetic with whatever
+            // else is in the ubatch, so give it a ubatch of its own. Sets at or below that width
+            // are decode steps, plain or speculative, whose columns the backend's column policy
+            // already keeps exact, so keep grouping those: isolating them too would make one
+            // prompt serialize every concurrent decode for the whole of the prefill, and would run
+            // a speculative verify step once per sequence.
+            if (isolate_seqs_above > 0) {
                 uint32_t n_left = 0;
 
                 for (const auto idx : seq_set_map[seq_set[i]]) {
@@ -573,7 +575,7 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
                     }
                 }
 
-                if (n_left > 1) {
+                if (n_left > isolate_seqs_above) {
                     if (!cur_seq_set.empty()) {
                         // let the sets already taken have this ubatch; the prompt gets the next one
                         break;
diff --git a/src/llama-batch.h b/src/llama-batch.h
index 4bd2aa98f9f..7b638b20d30 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -105,15 +105,17 @@ class llama_batch_allocr {
     // make ubatches of equal-length sequences sets
     // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
     // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
-    // isolate_multi_token_seqs = [TAG_EXACT_CONCURRENCY] a sequence set with more than one token
-    //               left to place is given a ubatch of its own; sets with a single token left are
-    //               still grouped together, so a prompt next to three decodes costs one extra
-    //               ubatch and does not serialize the three decodes
-    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, bool isolate_multi_token_seqs = false);
-
-    // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than one token left to place,
-    // i.e. what remains of the batch is not a plain one-token-per-sequence decode step
-    bool has_multi_token_seq() const;
+    // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this many
+    //               tokens left to place is a prompt and is given a ubatch of its own; sets at or
+    //               below it are decode steps (one token, or one plus the speculative drafts) and
+    //               stay grouped together, so a prompt next to three decodes costs one extra ubatch
+    //               and does not serialize the three decodes, and a speculative verify batch is not
+    //               run once per sequence
+    llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above = 0);
+
+    // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place,
+    // i.e. what remains of the batch holds a prompt rather than decode steps only
+    bool has_seq_wider_than(uint32_t n_tokens) const;
 
     // sequence-set-wise split - each ubatch contains a single sequence-set
     llama_ubatch split_seq(uint32_t n_ubatch);
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index bad0e55237a..8c95e842f99 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -5,6 +5,7 @@
 
 #include <cinttypes>
 #include <climits>
+#include <atomic>
 #include <cstdarg>
 #include <cstdlib>
 #include <cstring>
@@ -180,3 +181,14 @@ bool llama_exact_concurrency() {
 
     return enabled;
 }
+
+// [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h
+static std::atomic<uint32_t> g_exact_decode_tokens{1};
+
+void llama_set_exact_decode_tokens(uint32_t n_tokens) {
+    g_exact_decode_tokens.store(n_tokens > 0 ? n_tokens : 1, std::memory_order_relaxed);
+}
+
+uint32_t llama_exact_decode_tokens(void) {
+    return g_exact_decode_tokens.load(std::memory_order_relaxed);
+}
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index dced34f6a14..bf37cdd291f 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -882,7 +882,7 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
             // ubatch, so a sequence's prefill would run at a width its solo run never sees. Take
             // the sequence-set split instead, which can give each prompt a ubatch of its own; a
             // plain decode step has nothing to isolate and keeps taking split_simple.
-            const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq();
+            const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0;
 
             auto ubatch = n_stream == 1 && !isolate
                 ? balloc.split_simple(n_ubatch)
diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index 48fdea8b49c..7f502f3fa0c 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -91,7 +91,7 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
                 // leaves a different gated delta net state than the same prompt processed alone.
                 // Giving such a sequence a ubatch of its own removes that. A plain decode step, one
                 // token per sequence, is already exact and stays batched.
-                const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq();
+                const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0;
 
                 ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate);
             }
diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp
index f639a25c5df..8943faf316c 100644
--- a/src/llama-memory-recurrent.cpp
+++ b/src/llama-memory-recurrent.cpp
@@ -433,7 +433,7 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr &
                 //   so that the rollback snapshots remain valid
                 // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: a recurrent state that a
                 // prompt leaves behind depends on what shared its ubatch, so isolate the prompts
-                const bool isolate = llama_exact_concurrency() && balloc.has_multi_token_seq();
+                const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0;
 
                 ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate);
             }

From 379ca5d42a187c05dd4b663be749b088422fa033 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 11:27:32 +0000
Subject: [PATCH 20/33] cuda: run the single-token MUL_MAT_ID configuration
 over every token in one launch

Under the batch-invariant policy a MUL_MAT_ID with several tokens was
recomputed one token at a time, re-entering the op once per token, so a
decode step of four slots with two MTP drafts each cost twelve serial
expert launches per projection.

The single-column MMVQ kernel already carries a sample axis. With ids, a
sample is now a token: the wrapper launches the ncols_dst = 1 configuration
once with the tokens on the z axis, y advancing by a token per sample and
the expert index read per sample, so every (token, expert slot) block runs
exactly the instructions the token alone would run, on the same data, with
the same warp count, row split and K loop. Only the block indices differ.
The stock single-token launch has one sample and is unchanged, as is every
launch without the knob.

A quantized expert matrix takes this path for any token count under the
knob; other types still go one token at a time.

Single-op probe, token 0's slice against its own single-token run, knob on,
Q4_K, Q5_K, Q6_K and Q8_0 at K 2048 in both the gate-up and the down
layout and Q6_K at K 512 (the shape whose multi-token kernel differed),
2 to 12 tokens: 0 differing elements in every one of 99 comparisons.
---
 ggml/src/ggml-cuda/ggml-cuda.cu |  7 +++++++
 ggml/src/ggml-cuda/mmvq.cu      | 28 +++++++++++++++++++++++++---
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 2d3fc641ae6..e5e66f95265 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -2228,6 +2228,13 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
     // [TAG_BATCH_INVARIANT]
     if (ggml_cuda_mul_mat_id_splits_tokens(dst)) {
         GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1);
+        // A quantized expert matrix takes the single-token MMVQ path for every token count, and
+        // that path can put the tokens on its sample axis in one launch rather than being
+        // re-entered once per token. Anything else is still recomputed one token at a time.
+        if (ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+            ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst);
+            return;
+        }
         ggml_cuda_mul_mat_id_split_tokens(ctx, dst);
         return;
     }
diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu
index b14ef9681c5..b4e5196f172 100644
--- a/ggml/src/ggml-cuda/mmvq.cu
+++ b/ggml/src/ggml-cuda/mmvq.cu
@@ -592,9 +592,13 @@ static __global__ void mul_mat_vec_q(
     uint32_t sample_dst;
 
     ggml_cuda_pdl_sync();
-    channel_x  = ncols_dst == 1 && ids ? ids[channel_dst]                     : fastdiv(channel_dst, channel_ratio);
-    channel_y  = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst;
     sample_dst = blockIdx.z;
+    // [TAG_BATCH_INVARIANT] with ids, a sample is a token: the batch-invariant MUL_MAT_ID launch
+    // puts every token of the batch on the z axis of one single-column launch, so each (token,
+    // expert slot) block runs the exact single-token configuration. The stock single-token launch
+    // has one sample, where this indexing is ids[channel_dst] as before.
+    channel_x  = ncols_dst == 1 && ids ? ids[sample_dst*ids_stride + channel_dst] : fastdiv(channel_dst, channel_ratio);
+    channel_y  = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y)      : channel_dst;
 
     const uint32_t sample_x    = fastdiv(sample_dst, sample_ratio);
     const uint32_t sample_y    = sample_dst;
@@ -1281,7 +1285,14 @@ void ggml_cuda_mul_mat_vec_q(
     GGML_ASSERT(        nb0        == ts_dst);
     GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type));
 
-    GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE);
+    // [TAG_BATCH_INVARIANT] under the knob a MUL_MAT_ID with several tokens is computed as one
+    // launch of the single-token configuration with the tokens on the sample axis, so every
+    // (token, expert slot) block reduces exactly as the token alone would. The token count is
+    // then not bounded by the column templates.
+    const bool tokens_as_samples = ids && ne2 > 1 && ggml_cuda_batch_invariant();
+
+    GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE || tokens_as_samples);
+    GGML_ASSERT(!tokens_as_samples || !fusion);
 
     const float   * src1_d =       (const float   *) src1->data;
     const int32_t *  ids_d = ids ? (const int32_t *)  ids->data : nullptr;
@@ -1369,6 +1380,17 @@ void ggml_cuda_mul_mat_vec_q(
 
     const int64_t ids_stride = ids ? ids->nb[1] / ggml_type_size(ids->type) : 0;
 
+    if (tokens_as_samples) {
+        GGML_ASSERT(ne03 == 1 && ne13 == 1 && ne3 == 1);
+        // one column, one sample per token: y advances by s12 per token, dst by s2, x not at all
+        mul_mat_vec_q_switch_type(
+            src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00,
+            ne01,              1,             s01, stride_col_y,     stride_col_dst,
+            ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst,
+            1,                 ne2,           s03, s12,              s2,               ids_stride, stream);
+        return;
+    }
+
     mul_mat_vec_q_switch_type(
         src0->data, src0->type, src1_q8_1.get(), ids_d, fusion_local, dst_d, ne00,
         ne01,              ncols_dst,     s01, stride_col_y,     stride_col_dst,

From f4e45646d6c76992d9aa3730eb660fc9e97caba0 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 12:55:09 +0000
Subject: [PATCH 21/33] batch: group only sets with the same number of tokens
 left under exact mode

With prompts isolated by width, sets at or below the decode width were
grouped whatever their token counts, and the equal-length expansion then
placed a three-token verify step beside a two-token one as two tokens now
and one later. Attention and the gated delta net do not care, and the 4B
reads identical either way, but a memory that reduces over a chunk of
tokens, such as a chunked state space scan, would sum in a different
order than the solo run's single three-token ubatch.

A set now joins the ubatch only if it has as many tokens left as the first
set taken, so every set in a ubatch finishes in that ubatch and each
sequence's step has the shape it has alone. Sets with a different count
wait for a later ubatch. Speculative verify steps of equal width, the
common case, still share one ubatch.

4B, two MTP drafts, three rounds each: identical to the solo run with and
without a forced park every 64 tokens; speculation off with forced parks
identical; server preemption tests 7 of 7.
---
 src/llama-batch.cpp | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 4080ecd11d0..0c28cec32ba 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -537,6 +537,10 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
 
     llama_seq_id last_seq_id = -1;
 
+    // [TAG_EXACT_CONCURRENCY] tokens left in the first set taken, when isolating: only sets with
+    // the same count join it, so that every set in the ubatch finishes in this ubatch
+    uint32_t n_left_first = 0;
+
     // determine the non-overlapping sequence sets participating in this ubatch
     for (int32_t i = 0; i < batch.n_tokens; ++i) {
         if (used[i]) {
@@ -565,7 +569,12 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
             // are decode steps, plain or speculative, whose columns the backend's column policy
             // already keeps exact, so keep grouping those: isolating them too would make one
             // prompt serialize every concurrent decode for the whole of the prefill, and would run
-            // a speculative verify step once per sequence.
+            // a speculative verify step once per sequence. Grouped sets must have the same number
+            // of tokens left: the equal-length expansion below would otherwise place a three-token
+            // verify step beside a two-token one as two tokens now and one later, and a memory
+            // that reduces over a chunk of tokens (a chunked state space scan) would then sum in a
+            // different order than the solo run's single three-token ubatch. A set with a
+            // different count waits for a later ubatch.
             if (isolate_seqs_above > 0) {
                 uint32_t n_left = 0;
 
@@ -587,6 +596,12 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
 
                     break;
                 }
+
+                if (cur_seq_set.empty()) {
+                    n_left_first = n_left;
+                } else if (n_left != n_left_first) {
+                    continue;
+                }
             }
 
             cur_seq_set.push_back(seq_set[i]);

From 1a6f7da42092a7ec7192831e7ea8deb00b8f7dce Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 15:36:01 +0000
Subject: [PATCH 22/33] exact: refuse by name what the mode cannot run, before
 it runs

Three gaps found in review, all of the same kind: exact mode accepted an
input it could not implement and either ran something else or asserted.

A model whose K or V heads are not 256 wide passed every load-time check
while the paged attention kernel refuses such heads, so attention ran
unpaged on the CPU with the mode reporting itself as on. The load now
fails with the head widths in the message, as the other preconditions do.
MLA layouts are refused the same way.

A token carrying several sequence ids reached the page placement and hit
an assertion there. init_batch now refuses the batch with a logged error
and the failed-prepare status, for the unified cache and the hybrid one.

On a hybrid memory the attention half already refused a cross-sequence
copy, a position shift and a position division under the mode, but the
recurrent half went ahead, leaving the two halves describing different
states. The hybrid memory now refuses all three before either half is
touched.

The harness model (8-wide heads) is now refused at load under the mode
instead of loading unpaged; the 4B is unaffected (two MTP rounds
identical); server preemption tests 7 of 7.
---
 src/llama-batch.cpp         | 10 ++++++++++
 src/llama-batch.h           |  3 +++
 src/llama-kv-cache.cpp      | 17 +++++++++++++++++
 src/llama-memory-hybrid.cpp | 27 +++++++++++++++++++++++++++
 4 files changed, 57 insertions(+)

diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 0c28cec32ba..5db683db281 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -507,6 +507,16 @@ llama_ubatch llama_batch_allocr::split_simple(uint32_t n_ubatch) {
     return ubatch_add(idxs, idxs.size(), false);
 }
 
+bool llama_batch_allocr::has_shared_tokens() const {
+    for (int32_t i = 0; i < batch.n_tokens; ++i) {
+        if (batch.n_seq_id[i] > 1) {
+            return true;
+        }
+    }
+
+    return false;
+}
+
 bool llama_batch_allocr::has_seq_wider_than(uint32_t n_tokens) const {
     std::vector<uint32_t> n_per_seq(n_seq_max, 0);
 
diff --git a/src/llama-batch.h b/src/llama-batch.h
index 7b638b20d30..52a375ad49e 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -117,6 +117,9 @@ class llama_batch_allocr {
     // i.e. what remains of the batch holds a prompt rather than decode steps only
     bool has_seq_wider_than(uint32_t n_tokens) const;
 
+    // [TAG_EXACT_CONCURRENCY] true if some token carries more than one sequence id
+    bool has_shared_tokens() const;
+
     // sequence-set-wise split - each ubatch contains a single sequence-set
     llama_ubatch split_seq(uint32_t n_ubatch);
 
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index bf37cdd291f..1c31e57a761 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -275,6 +275,15 @@ llama_kv_cache::llama_kv_cache(
 
         // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order while
         // the mode still reports itself as on, so refuse the load instead
+        // [TAG_EXACT_CONCURRENCY] the paged attention kernel handles 256-wide K and V heads only;
+        // any other width would run unpaged on the CPU while the mode reports itself as on
+        if (exact_pages && (hparams.n_embd_head_k(il) != 256 || (!is_mla && hparams.n_embd_head_v(il) != 256) || is_mla)) {
+            LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d has %u-wide K heads and %u-wide V heads%s, "
+                    "and the paged attention kernel supports 256-wide K and V heads only\n",
+                    __func__, il, hparams.n_embd_head_k(il), hparams.n_embd_head_v(il), is_mla ? " (MLA)" : "");
+            throw std::runtime_error("exact concurrency: unsupported attention head size");
+        }
+
         if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) {
             LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, "
                     "which has no paged attention: every layer must be offloaded to the CUDA backend "
@@ -874,6 +883,14 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
     GGML_UNUSED(embd_all);
 
     do {
+        // [TAG_EXACT_CONCURRENCY] a token shared by several sequences would be one cell in a page
+        // that belongs to one sequence; the placement asserts on it later, so refuse it here
+        if (exact_pages && balloc.has_shared_tokens()) {
+            LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; "
+                    "give every token exactly one sequence id\n", __func__);
+            break;
+        }
+
         balloc.split_reset();
 
         std::vector<llama_ubatch> ubatches;
diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index 7f502f3fa0c..3fea4bb0ba3 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -66,6 +66,13 @@ llama_memory_hybrid::llama_memory_hybrid(
 
 llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) {
     do {
+        // [TAG_EXACT_CONCURRENCY] refused before the attention half asserts on it, see llama_kv_cache::init_batch
+        if (llama_exact_concurrency() && balloc.has_shared_tokens()) {
+            LLAMA_LOG_ERROR("%s: exact concurrency does not support tokens shared by several sequence ids; "
+                    "give every token exactly one sequence id\n", __func__);
+            break;
+        }
+
         balloc.split_reset();
 
         // follow the recurrent pattern for creating the ubatch splits
@@ -163,6 +170,14 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
 }
 
 void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
+    // [TAG_EXACT_CONCURRENCY] the attention half refuses this under exact mode; refuse it here
+    // before either half is touched, so the two halves cannot end up describing different states
+    if (llama_exact_concurrency() && seq_id_src != seq_id_dst) {
+        LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between sequences (%d -> %d); ignoring the copy\n",
+                __func__, seq_id_src, seq_id_dst);
+        return;
+    }
+
     mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1);
     mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1);
 }
@@ -173,11 +188,23 @@ void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) {
 }
 
 void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) {
+    if (llama_exact_concurrency() && shift != 0) {
+        LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions (seq %d, shift %d); ignoring the shift\n",
+                __func__, seq_id, shift);
+        return;
+    }
+
     mem_attn->seq_add(seq_id, p0, p1, shift);
     mem_recr->seq_add(seq_id, p0, p1, shift);
 }
 
 void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) {
+    if (llama_exact_concurrency() && d != 1) {
+        LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions (seq %d, d %d); ignoring the division\n",
+                __func__, seq_id, d);
+        return;
+    }
+
     mem_attn->seq_div(seq_id, p0, p1, d);
     mem_recr->seq_div(seq_id, p0, p1, d);
 }

From 0c9fc0ed8b03c7845c16feee12df7ed3c0db765e Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 16:48:25 +0000
Subject: [PATCH 23/33] exact: every context reports the widest decode step it
 can build

The column bound the CUDA backend splits at was reported by the server
only. A program using the library directly got a fallback of 16 columns
and, above it, a once-only warning that the mode did not hold for that
op; with 32 sequences in a step that is silent inexactness after one
line of log.

llama_set_exact_decode_width() is the one place the width is reported,
and it never lowers what was reported. Every llama_context reports its
own at creation, n_seq_max times the per-sequence width, so a decode of
any context stays within the bound its kernels split at whether or not
the caller knew there was one. The server's common path goes through
the same call, so the figure it reports is the one its contexts would.

Probe: 32 sequences on the 4B, greedy, each with its own copy of the
prompt. Before, the warning fires; after, it does not, and all 32 match
the solo run either way at this length.
---
 common/common.cpp     | 12 +++---------
 include/llama.h       |  9 +++++++++
 src/llama-context.cpp |  8 ++++++++
 src/llama-impl.cpp    | 29 +++++++++++++++++++++++++++++
 4 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 8a1c76793cc..abca3630e13 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1500,15 +1500,9 @@ bool common_exact_concurrency_init(const common_params & params) {
         }
     }
 
-    // the CUDA backend may not be present or may be loaded dynamically, so go through the registry
-    for (size_t i = 0; i < ggml_backend_reg_count(); ++i) {
-        ggml_backend_reg_t reg = ggml_backend_reg_get(i);
-
-        auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width");
-        if (fn) {
-            fn(n_cols);
-        }
-    }
+    // a context created later reports n_seq_max times the per-sequence width, which is this
+    // figure again; reporting it here as well covers a caller that decodes before that
+    llama_set_exact_decode_width((uint32_t) n_cols);
 
     return true;
 }
diff --git a/include/llama.h b/include/llama.h
index 0e44ff052c2..299249c5daf 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -811,6 +811,15 @@ extern "C" {
     LLAMA_API void     llama_set_exact_decode_tokens(uint32_t n_tokens);
     LLAMA_API uint32_t llama_exact_decode_tokens(void);
 
+    // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns: the
+    // sequences a context can hold times the tokens each contributes to a step. Every context
+    // reports its own at creation and a backend keeps the widest it has heard, so a decode of any
+    // context stays within the bound its kernels split at. A caller that builds wider steps than
+    // the contexts imply (a draft of its own, say) reports the width itself, before creating the
+    // context or before the first decode. Never lowers what was reported.
+    LLAMA_API void     llama_set_exact_decode_width(uint32_t n_cols);
+    LLAMA_API uint32_t llama_exact_decode_width(void);
+
     //
     // State / sessions
     //
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 3ab84de780c..e1840ac57d3 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -101,6 +101,14 @@ llama_context::llama_context(
         throw std::runtime_error("n_seq_max must be <= " + std::to_string(LLAMA_MAX_SEQ));
     }
 
+    // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build: one column per
+    // sequence, times the tokens a sequence contributes to a step. Reported so that a backend
+    // splitting columns for exactness covers it without the caller having to know the bound; a
+    // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width.
+    if (llama_exact_concurrency()) {
+        llama_set_exact_decode_width(cparams.n_seq_max * llama_exact_decode_tokens());
+    }
+
     cparams.n_rs_seq = params.n_rs_seq;
     if (cparams.n_rs_seq > 0 && !llm_arch_supports_rs_rollback(model.arch)) {
         LLAMA_LOG_DEBUG("%s: n_rs_seq=%u requested but model does not support recurrent partial rollback; clamping to 0\n",
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index 8c95e842f99..ad63fc5ea3b 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -1,5 +1,6 @@
 #include "llama-impl.h"
 
+#include "ggml-backend.h"
 #include "gguf.h"
 #include "llama.h"
 
@@ -192,3 +193,31 @@ void llama_set_exact_decode_tokens(uint32_t n_tokens) {
 uint32_t llama_exact_decode_tokens(void) {
     return g_exact_decode_tokens.load(std::memory_order_relaxed);
 }
+
+// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h. A backend that
+// splits columns to make a decode exact reads it through ggml_backend_cuda_set_exact_decode_width,
+// reached through the registry so that a backend that is absent or loaded late costs nothing.
+static std::atomic<uint32_t> g_exact_decode_width{0};
+
+void llama_set_exact_decode_width(uint32_t n_cols) {
+    uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed);
+
+    while (n_cols > cur) {
+        if (g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) {
+            for (size_t i = 0; i < ggml_backend_reg_count(); ++i) {
+                ggml_backend_reg_t reg = ggml_backend_reg_get(i);
+
+                auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width");
+                if (fn) {
+                    fn((int) n_cols);
+                }
+            }
+
+            return;
+        }
+    }
+}
+
+uint32_t llama_exact_decode_width(void) {
+    return g_exact_decode_width.load(std::memory_order_relaxed);
+}

From 46e7fa742fa7c78ad0bfb8c7247f5146b2a0c2e2 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 17:43:09 +0000
Subject: [PATCH 24/33] exact: four refusals and one width from review

The decode width per slot came from a switch over the speculation types
that did not know the ngram cache verifies eight tokens, so under that
type the per-sequence width was 4 rather than 9 and every verify group
was taken for a prompt and isolated. The width now comes from
common_speculative_n_max(), the same place the speculation code takes it.

DFlash drafting turns causal attention off on its draft context, and the
paged attention the mode runs on needs it; the next graph asserted. The
server refuses the combination at setup by name, and a context with a
cache under the mode refuses to turn causal attention off.

A library user setting GGML_CUDA_BATCH_INVARIANT_MAX_COLS below a
context's decode width was not caught: the explicit bound wins in the
backend and its warning is silent once a width was reported. The context
refuses to be created, the way the server's setup refuses.

A width reported before a backend was loaded never reached it, since the
setter only forwarded on a raise. The widest figure now goes to every
backend on every call, and every context reports at creation.
---
 common/common.cpp     | 38 +++++++++++++-------------------------
 src/llama-context.cpp | 24 +++++++++++++++++++++++-
 src/llama-impl.cpp    | 24 +++++++++++++-----------
 3 files changed, 49 insertions(+), 37 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index abca3630e13..7b6fc8ad2e2 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1448,32 +1448,11 @@ bool common_exact_concurrency() {
 int common_exact_decode_width(const common_params & params) {
     const int n_slots = std::max(1, params.n_parallel);
 
-    // the draft tokens a slot carries into the verify ubatch alongside its accepted token
-    int n_draft = 0;
+    // the draft tokens a slot carries into the verify ubatch alongside its accepted token, per
+    // speculation type, from the same place the speculation code takes its own width
+    const int n_draft = std::max(0, (int) common_speculative_n_max(&params.speculative));
 
-    for (const auto type : params.speculative.types) {
-        switch (type) {
-            case COMMON_SPECULATIVE_TYPE_NONE:
-                break;
-            case COMMON_SPECULATIVE_TYPE_NGRAM_MOD:
-                n_draft = std::max(n_draft, params.speculative.ngram_mod.n_max);
-                break;
-            case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE:
-                n_draft = std::max(n_draft, (int) params.speculative.ngram_simple.size_m);
-                break;
-            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K:
-                n_draft = std::max(n_draft, (int) params.speculative.ngram_map_k.size_m);
-                break;
-            case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V:
-                n_draft = std::max(n_draft, (int) params.speculative.ngram_map_k4v.size_m);
-                break;
-            default:
-                n_draft = std::max(n_draft, params.speculative.draft.n_max);
-                break;
-        }
-    }
-
-    return n_slots*(1 + std::max(0, n_draft));
+    return n_slots*(1 + n_draft);
 }
 
 // [TAG_EXACT_CONCURRENCY]
@@ -1482,6 +1461,15 @@ bool common_exact_concurrency_init(const common_params & params) {
         return true;
     }
 
+    // DFlash drafting turns causal attention off on its draft context, and the paged
+    // attention the mode runs on needs it; say so instead of asserting in the graph
+    for (const auto type : params.speculative.types) {
+        if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) {
+            COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash: it disables causal attention, which the paged attention needs\n");
+            return false;
+        }
+    }
+
     const int n_cols = common_exact_decode_width(params);
 
     // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step is
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index e1840ac57d3..c95493e06e6 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -106,7 +106,22 @@ llama_context::llama_context(
     // splitting columns for exactness covers it without the caller having to know the bound; a
     // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width.
     if (llama_exact_concurrency()) {
-        llama_set_exact_decode_width(cparams.n_seq_max * llama_exact_decode_tokens());
+        const uint32_t n_cols = cparams.n_seq_max * llama_exact_decode_tokens();
+
+        // an explicit column bound wins over the reported width in the backend, so one below
+        // this context's width would leave its decodes batched above the bound with the mode
+        // still reporting itself on; refuse it here, the way the server's setup does
+        if (const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS")) {
+            const int max_cols = atoi(bound);
+
+            if (max_cols > 0 && (uint32_t) max_cols < n_cols) {
+                LLAMA_LOG_ERROR("%s: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at least %u to cover a decode step of %u sequences; raise it, set it to 0 for no bound, or unset it\n",
+                        __func__, max_cols, n_cols, cparams.n_seq_max);
+                throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width");
+            }
+        }
+
+        llama_set_exact_decode_width(n_cols);
     }
 
     cparams.n_rs_seq = params.n_rs_seq;
@@ -1196,6 +1211,13 @@ void llama_context::set_causal_attn(bool value) {
         return;
     }
 
+    // [TAG_EXACT_CONCURRENCY] the paged attention the mode runs on is causal; a context with a
+    // cache under the mode keeps causal attention rather than asserting in the next graph
+    if (!value && memory && llama_exact_concurrency()) {
+        LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so causal attention cannot be turned off; the change is refused\n", __func__);
+        return;
+    }
+
     cparams.causal_attn = value;
 
     sched_need_reserve = true;
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index ad63fc5ea3b..f8ce14014f1 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -202,18 +202,20 @@ static std::atomic<uint32_t> g_exact_decode_width{0};
 void llama_set_exact_decode_width(uint32_t n_cols) {
     uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed);
 
-    while (n_cols > cur) {
-        if (g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) {
-            for (size_t i = 0; i < ggml_backend_reg_count(); ++i) {
-                ggml_backend_reg_t reg = ggml_backend_reg_get(i);
-
-                auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width");
-                if (fn) {
-                    fn((int) n_cols);
-                }
-            }
+    while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) {
+    }
+
+    // The widest figure so far goes to every backend on every call, not only when it grew: a
+    // width reported before a backend was loaded would otherwise never reach it, and every
+    // context reports at creation, by which time the backends are there.
+    const uint32_t widest = g_exact_decode_width.load(std::memory_order_relaxed);
+
+    for (size_t i = 0; i < ggml_backend_reg_count(); ++i) {
+        ggml_backend_reg_t reg = ggml_backend_reg_get(i);
 
-            return;
+        auto * fn = (void (*)(int)) ggml_backend_reg_get_proc_address(reg, "ggml_backend_cuda_set_exact_decode_width");
+        if (fn) {
+            fn((int) widest);
         }
     }
 }

From eae462431761ee3ac5774873eae2c99e43bda6df Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 18:13:27 +0000
Subject: [PATCH 25/33] exact concurrency: refuse a non-causal context with a
 cache at creation, keep a page boundary per prompt slot in the reserve cap

A context created with LLAMA_ATTENTION_TYPE_NON_CAUSAL under the mode used to
pass creation and assert on its first graph; llama_context now refuses it right
after the memory is created, so callers get an error from llama_init_from_model
instead of an abort. Contexts without a cache are unaffected.

The reserve cap gave every waiting prompt slot one batch between them, rounded
to a page once; under page allocation each prompt slot can cross a boundary of
its own within that batch, so the cap now keeps one boundary per prompt slot.
---
 src/llama-context.cpp           |  7 +++++++
 tools/server/server-context.cpp | 13 ++++++++++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index c95493e06e6..f3b1df04c48 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -416,6 +416,13 @@ llama_context::llama_context(
         };
 
         memory.reset(model.create_memory(params_mem, cparams));
+
+        // [TAG_EXACT_CONCURRENCY] the paged attention the mode runs on is causal; a context
+        // created non-causal with a cache would assert on its first graph, so it is refused here
+        if (llama_exact_concurrency() && memory && !cparams.causal_attn) {
+            LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so it cannot be created with non-causal attention\n", __func__);
+            throw std::runtime_error("exact concurrency: non-causal attention is not supported with a KV cache");
+        }
     }
 
     // init backends
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index 56b48cca60c..f61cb2406ad 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -3166,7 +3166,18 @@ struct server_context_impl {
 
         // one batch is all the prompt slots get between them, however many are waiting; in
         // cells that batch can straddle one boundary more than it has tokens for
-        return res + std::min(res_pmt, preempt_n_cells(n_batch));
+        // one batch is all the prompt slots get between them, however many are waiting; under
+        // page allocation each of them can still cross a page boundary of its own within that
+        // batch, so the cap keeps one boundary per prompt slot on top of the batch
+        int32_t n_pmt = 0;
+
+        for (const auto & slot : slots) {
+            if (slot.state == SLOT_STATE_STARTED || slot.state == SLOT_STATE_PROCESSING_PROMPT) {
+                n_pmt++;
+            }
+        }
+
+        return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1));
     }
 
     // Keep the slot that is furthest along -- it is the closest to finishing and to giving

From dbd82ca5ee65c21a2006b10a847875056957a787 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 18:24:14 +0000
Subject: [PATCH 26/33] exact concurrency: the decode width of every context
 follows the token figure

A context reported n_seq_max times the tokens per sequence once, at creation.
Raising the process-wide token figure afterwards widened the decode step of
every earlier context while the width they reported stayed put, so a context
created under a narrower figure could batch above the bound it reported. Each
context now reports its sequence count; the widest count seen times the current
token figure is re-reported whenever the figure changes.
---
 src/llama-context.cpp |  4 +++-
 src/llama-impl.cpp    | 21 +++++++++++++++++++++
 src/llama-impl.h      |  4 ++++
 3 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index f3b1df04c48..d14123779a5 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -105,6 +105,8 @@ llama_context::llama_context(
     // sequence, times the tokens a sequence contributes to a step. Reported so that a backend
     // splitting columns for exactness covers it without the caller having to know the bound; a
     // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width.
+    // The sequence count is what is reported: the tokens figure can be raised later for the
+    // whole process, and the width then follows it for this context too.
     if (llama_exact_concurrency()) {
         const uint32_t n_cols = cparams.n_seq_max * llama_exact_decode_tokens();
 
@@ -121,7 +123,7 @@ llama_context::llama_context(
             }
         }
 
-        llama_set_exact_decode_width(n_cols);
+        llama_exact_report_n_seq(cparams.n_seq_max);
     }
 
     cparams.n_rs_seq = params.n_rs_seq;
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index f8ce14014f1..51c50898bdb 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -186,8 +186,29 @@ bool llama_exact_concurrency() {
 // [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h
 static std::atomic<uint32_t> g_exact_decode_tokens{1};
 
+// the most sequences any context so far was created with. The tokens figure is process
+// wide, so raising it widens the decode step of every context that already exists; the
+// width those contexts reported at creation is re-reported here with the new figure, or a
+// context created under a narrower figure would batch above the bound it reported.
+static std::atomic<uint32_t> g_exact_max_n_seq{0};
+
+void llama_exact_report_n_seq(uint32_t n_seq) {
+    uint32_t cur = g_exact_max_n_seq.load(std::memory_order_relaxed);
+
+    while (n_seq > cur && !g_exact_max_n_seq.compare_exchange_weak(cur, n_seq, std::memory_order_relaxed)) {
+    }
+
+    llama_set_exact_decode_width(g_exact_max_n_seq.load(std::memory_order_relaxed) * llama_exact_decode_tokens());
+}
+
 void llama_set_exact_decode_tokens(uint32_t n_tokens) {
     g_exact_decode_tokens.store(n_tokens > 0 ? n_tokens : 1, std::memory_order_relaxed);
+
+    const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed);
+
+    if (n_seq > 0) {
+        llama_set_exact_decode_width(n_seq * llama_exact_decode_tokens());
+    }
 }
 
 uint32_t llama_exact_decode_tokens(void) {
diff --git a/src/llama-impl.h b/src/llama-impl.h
index 9b64431fedd..3c720cec150 100644
--- a/src/llama-impl.h
+++ b/src/llama-impl.h
@@ -109,3 +109,7 @@ std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i);
 // so that its output does not change when other sequences share the KV cache. Off by default.
 // Reads the same LLAMA_EXACT_CONCURRENCY variable as the paged KV cache and the CUDA backend.
 bool llama_exact_concurrency();
+
+// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so that the
+// decode width every context needs is known to the backend and follows llama_set_exact_decode_tokens
+void llama_exact_report_n_seq(uint32_t n_seq);

From d4e3fc8a6995fee7fa81b414ad4beb4b20e6aad8 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 18:58:00 +0000
Subject: [PATCH 27/33] exact concurrency: equal-count grouping stays on for
 recurrent and hybrid memories; a width the explicit bound cannot cover is
 refused

The recurrent and hybrid memories passed the isolation figure to split_equal
only when a prompt was in the batch, so a three-token verify step could be
placed beside a two-token one as two tokens now and one later, and a memory
that reduces over a chunk of tokens would sum in an order the solo run never
had. The figure is passed whenever the mode is on.

An explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS wins in the backend, so a
width raised past it after a context exists would leave decodes batched
above the bound. llama_set_exact_decode_width and llama_set_exact_decode_tokens
now return false and change nothing when the bound cannot cover the width;
the context constructor and common's setup treat that as the error it is.
---
 common/common.cpp              | 16 +++++++----
 include/llama.h                | 13 ++++++---
 src/llama-context.cpp          |  4 ++-
 src/llama-impl.cpp             | 50 +++++++++++++++++++++++++++++-----
 src/llama-impl.h               |  2 +-
 src/llama-memory-hybrid.cpp    |  4 ++-
 src/llama-memory-recurrent.cpp |  8 ++++--
 7 files changed, 75 insertions(+), 22 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 7b6fc8ad2e2..8847feaf9b2 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1472,9 +1472,6 @@ bool common_exact_concurrency_init(const common_params & params) {
 
     const int n_cols = common_exact_decode_width(params);
 
-    // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step is
-    llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel)));
-
     const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
     if (bound) {
         const int max_cols = atoi(bound);
@@ -1488,9 +1485,16 @@ bool common_exact_concurrency_init(const common_params & params) {
         }
     }
 
-    // a context created later reports n_seq_max times the per-sequence width, which is this
-    // figure again; reporting it here as well covers a caller that decodes before that
-    llama_set_exact_decode_width((uint32_t) n_cols);
+    // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode
+    // step is; a context created later reports n_seq_max times that figure, which is n_cols
+    // again, and reporting n_cols here as well covers a caller that decodes before that. Both
+    // refuse a width the explicit bound above cannot cover, which the check above already
+    // caught for this process; contexts created earlier by the caller are covered here.
+    if (!llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))) ||
+        !llama_set_exact_decode_width((uint32_t) n_cols)) {
+        COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: the decode width could not be reported, see the error above\n");
+        return false;
+    }
 
     return true;
 }
diff --git a/include/llama.h b/include/llama.h
index 299249c5daf..1608be8a6a0 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -807,8 +807,11 @@ extern "C" {
     // 1 plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence
     // set with more tokens than this left to place is a prompt and is prefilled in a ubatch of its
     // own; a set at or below it is a decode step and stays grouped with the other decodes, so a
-    // speculative verify batch is not run once per sequence. Process-wide, default 1.
-    LLAMA_API void     llama_set_exact_decode_tokens(uint32_t n_tokens);
+    // speculative verify batch is not run once per sequence. Process-wide, default 1. Raising it
+    // widens the decode step of every context that exists, and their width is re-reported with
+    // it; false, and no change, when an explicit column bound given to the backend cannot cover
+    // that width (see llama_set_exact_decode_width).
+    LLAMA_API bool     llama_set_exact_decode_tokens(uint32_t n_tokens);
     LLAMA_API uint32_t llama_exact_decode_tokens(void);
 
     // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns: the
@@ -816,8 +819,10 @@ extern "C" {
     // reports its own at creation and a backend keeps the widest it has heard, so a decode of any
     // context stays within the bound its kernels split at. A caller that builds wider steps than
     // the contexts imply (a draft of its own, say) reports the width itself, before creating the
-    // context or before the first decode. Never lowers what was reported.
-    LLAMA_API void     llama_set_exact_decode_width(uint32_t n_cols);
+    // context or before the first decode. Never lowers what was reported. Returns false, and
+    // reports nothing, when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is set to a positive figure below
+    // the width: that bound wins in the backend, so decodes above it would be left batched.
+    LLAMA_API bool     llama_set_exact_decode_width(uint32_t n_cols);
     LLAMA_API uint32_t llama_exact_decode_width(void);
 
     //
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index d14123779a5..e8ccf985477 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -123,7 +123,9 @@ llama_context::llama_context(
             }
         }
 
-        llama_exact_report_n_seq(cparams.n_seq_max);
+        if (!llama_exact_report_n_seq(cparams.n_seq_max)) {
+            throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width");
+        }
     }
 
     cparams.n_rs_seq = params.n_rs_seq;
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index 51c50898bdb..20739bb7c96 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -192,23 +192,35 @@ static std::atomic<uint32_t> g_exact_decode_tokens{1};
 // context created under a narrower figure would batch above the bound it reported.
 static std::atomic<uint32_t> g_exact_max_n_seq{0};
 
-void llama_exact_report_n_seq(uint32_t n_seq) {
+bool llama_exact_report_n_seq(uint32_t n_seq) {
+    const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed));
+
+    if (!llama_set_exact_decode_width(n_seq_max * llama_exact_decode_tokens())) {
+        return false;
+    }
+
     uint32_t cur = g_exact_max_n_seq.load(std::memory_order_relaxed);
 
     while (n_seq > cur && !g_exact_max_n_seq.compare_exchange_weak(cur, n_seq, std::memory_order_relaxed)) {
     }
 
-    llama_set_exact_decode_width(g_exact_max_n_seq.load(std::memory_order_relaxed) * llama_exact_decode_tokens());
+    return true;
 }
 
-void llama_set_exact_decode_tokens(uint32_t n_tokens) {
-    g_exact_decode_tokens.store(n_tokens > 0 ? n_tokens : 1, std::memory_order_relaxed);
+bool llama_set_exact_decode_tokens(uint32_t n_tokens) {
+    n_tokens = n_tokens > 0 ? n_tokens : 1;
 
+    // every context that exists widens with the figure, so the width they will need is
+    // reported first; a figure the explicit bound cannot cover leaves the old one in place
     const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed);
 
-    if (n_seq > 0) {
-        llama_set_exact_decode_width(n_seq * llama_exact_decode_tokens());
+    if (n_seq > 0 && !llama_set_exact_decode_width(n_seq * n_tokens)) {
+        return false;
     }
+
+    g_exact_decode_tokens.store(n_tokens, std::memory_order_relaxed);
+
+    return true;
 }
 
 uint32_t llama_exact_decode_tokens(void) {
@@ -220,7 +232,29 @@ uint32_t llama_exact_decode_tokens(void) {
 // reached through the registry so that a backend that is absent or loaded late costs nothing.
 static std::atomic<uint32_t> g_exact_decode_width{0};
 
-void llama_set_exact_decode_width(uint32_t n_cols) {
+// an explicit column bound given to the CUDA backend wins over the reported width there, so a
+// width above it would leave decodes batched past the bound with the mode still reporting itself
+// on; a width the bound does not cover is refused instead of stored
+static bool llama_exact_width_within_explicit_bound(uint32_t n_cols) {
+    static const int explicit_cols = []() {
+        const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
+        return val ? atoi(val) : -1;
+    }();
+
+    if (explicit_cols > 0 && (uint32_t) explicit_cols < n_cols) {
+        LLAMA_LOG_ERROR("%s: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at least %u columns for the decode step just requested; raise it, set it to 0 for no bound, or unset it\n",
+                __func__, explicit_cols, n_cols);
+        return false;
+    }
+
+    return true;
+}
+
+bool llama_set_exact_decode_width(uint32_t n_cols) {
+    if (!llama_exact_width_within_explicit_bound(n_cols)) {
+        return false;
+    }
+
     uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed);
 
     while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) {
@@ -239,6 +273,8 @@ void llama_set_exact_decode_width(uint32_t n_cols) {
             fn((int) widest);
         }
     }
+
+    return true;
 }
 
 uint32_t llama_exact_decode_width(void) {
diff --git a/src/llama-impl.h b/src/llama-impl.h
index 3c720cec150..de5c6a2d216 100644
--- a/src/llama-impl.h
+++ b/src/llama-impl.h
@@ -112,4 +112,4 @@ bool llama_exact_concurrency();
 
 // [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so that the
 // decode width every context needs is known to the backend and follows llama_set_exact_decode_tokens
-void llama_exact_report_n_seq(uint32_t n_seq);
+bool llama_exact_report_n_seq(uint32_t n_seq);
diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index 3fea4bb0ba3..e8d80c770ba 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -98,7 +98,9 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
                 // leaves a different gated delta net state than the same prompt processed alone.
                 // Giving such a sequence a ubatch of its own removes that. A plain decode step, one
                 // token per sequence, is already exact and stays batched.
-                const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0;
+                // The figure is passed whenever the mode is on, not only when a prompt is present:
+                // it also keeps sets of unequal token counts apart (see llama_batch_allocr::split_equal).
+                const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0;
 
                 ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate);
             }
diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp
index 8943faf316c..61463c72964 100644
--- a/src/llama-memory-recurrent.cpp
+++ b/src/llama-memory-recurrent.cpp
@@ -432,8 +432,12 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr &
                 // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
                 //   so that the rollback snapshots remain valid
                 // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: a recurrent state that a
-                // prompt leaves behind depends on what shared its ubatch, so isolate the prompts
-                const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0;
+                // prompt leaves behind depends on what shared its ubatch, so isolate the prompts.
+                // The figure is passed whenever the mode is on, not only when a prompt is present:
+                // it also keeps sets of unequal token counts apart, and a three-token verify step
+                // placed beside a two-token one as two now and one later would be reduced in
+                // chunks the solo run never had.
+                const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0;
 
                 ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate);
             }

From bf00ac38a0af5c26c8156b654967dad131d649dc Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 19:30:46 +0000
Subject: [PATCH 28/33] exact concurrency: soft-capped attention refused at
 load; width reports serialised and monotonic in the backend; the token figure
 never lowered

The paged attention kernel has no soft-capped variant and asserted on its
first call; a model with attn_soft_cap is refused when the cache is created,
with the other unsupported layouts.

Two contexts created at once could hand the backend a narrower width after
a wider one: the report is now made under a lock, and the CUDA setter keeps
the widest figure it has heard whatever the order.

A narrower context set up after a speculative one lowered the process-wide
token figure and turned the existing context's verify steps into prompts,
serialising them; the figure is never lowered now, as the width never was.
---
 ggml/src/ggml-cuda/ggml-cuda.cu |  6 +++++-
 include/llama.h                 |  3 ++-
 src/llama-impl.cpp              | 12 ++++++++++++
 src/llama-kv-cache.cpp          |  8 ++++++++
 4 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index e5e66f95265..8c5b2a408c3 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1856,7 +1856,11 @@ int ggml_cuda_batch_invariant() {
 static std::atomic<int> g_exact_decode_width{0};
 
 void ggml_backend_cuda_set_exact_decode_width(int n_cols) {
-    g_exact_decode_width.store(n_cols > 0 ? n_cols : 0, std::memory_order_relaxed);
+    // monotonic: the widest figure ever reported stays, whatever order the reports arrive in
+    int cur = g_exact_decode_width.load(std::memory_order_relaxed);
+
+    while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) {
+    }
 }
 
 int ggml_cuda_batch_invariant_max_cols() {
diff --git a/include/llama.h b/include/llama.h
index 1608be8a6a0..09c331ff3e1 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -810,7 +810,8 @@ extern "C" {
     // speculative verify batch is not run once per sequence. Process-wide, default 1. Raising it
     // widens the decode step of every context that exists, and their width is re-reported with
     // it; false, and no change, when an explicit column bound given to the backend cannot cover
-    // that width (see llama_set_exact_decode_width).
+    // that width (see llama_set_exact_decode_width). Never lowers what was set: a narrower context
+    // set up later must not turn an existing context's verify steps into prompts.
     LLAMA_API bool     llama_set_exact_decode_tokens(uint32_t n_tokens);
     LLAMA_API uint32_t llama_exact_decode_tokens(void);
 
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index 20739bb7c96..2a33c57e869 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -7,6 +7,7 @@
 #include <cinttypes>
 #include <climits>
 #include <atomic>
+#include <mutex>
 #include <cstdarg>
 #include <cstdlib>
 #include <cstring>
@@ -210,6 +211,12 @@ bool llama_exact_report_n_seq(uint32_t n_seq) {
 bool llama_set_exact_decode_tokens(uint32_t n_tokens) {
     n_tokens = n_tokens > 0 ? n_tokens : 1;
 
+    // never lowered: a narrower context set up later would otherwise turn the verify steps of
+    // an existing speculative context into prompts and serialise them
+    if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) {
+        return true;
+    }
+
     // every context that exists widens with the figure, so the width they will need is
     // reported first; a figure the explicit bound cannot cover leaves the old one in place
     const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed);
@@ -255,6 +262,11 @@ bool llama_set_exact_decode_width(uint32_t n_cols) {
         return false;
     }
 
+    // one reporter at a time: the widest figure is read and handed to the backends below as
+    // one step, so a narrower report cannot overtake a wider one on its way to a backend
+    static std::mutex mutex;
+    std::lock_guard<std::mutex> lock(mutex);
+
     uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed);
 
     while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) {
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 1c31e57a761..bb947c9ed09 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -284,6 +284,14 @@ llama_kv_cache::llama_kv_cache(
             throw std::runtime_error("exact concurrency: unsupported attention head size");
         }
 
+        // [TAG_EXACT_CONCURRENCY] the paged attention kernel has no soft-capped variant and would
+        // assert on its first call, so a soft-capped model is refused at load instead
+        if (exact_pages && hparams.attn_soft_cap) {
+            LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but this model soft-caps its attention logits (%.1f), "
+                    "which the paged attention kernel does not apply\n", __func__, hparams.f_attn_logit_softcapping);
+            throw std::runtime_error("exact concurrency: attention soft cap is not supported");
+        }
+
         if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) {
             LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, "
                     "which has no paged attention: every layer must be offloaded to the CUDA backend "

From 98fe86dfc0a42940906f57ea6a27cf9bdc3907ba Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 19:56:18 +0000
Subject: [PATCH 29/33] exact concurrency: the setup runs before anything is
 loaded; one lock for the token figure, the sequence count and the width

common_init_from_params() checked the mode after the context existed, so a
caller that skipped common_params_parse() could be handed a live context
under a bound the setup had just refused. The check moved to the front of
the init result's constructor, ahead of the fitting contexts and the model
load; on failure nothing is loaded.

The token figure, the widest sequence count and the width moved as three
separate atomics, so a context reporting its count while the figure changed
could leave the backend with a width that covered neither. One recursive
lock now spans every transition.
---
 common/common.cpp  | 14 +++++++++-----
 src/llama-impl.cpp | 15 +++++++++++----
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 8847feaf9b2..966620c7910 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1290,6 +1290,15 @@ struct common_init_result::impl {
 
 common_init_result::common_init_result(common_params & params, bool model_only) :
     pimpl(new impl{}) {
+    // [TAG_EXACT_CONCURRENCY] before any context exists, the fitting ones included: the
+    // per-sequence figure and the column bound are checked against the explicit bound first,
+    // so a context is never created under a figure the bound does not cover. A caller that
+    // skipped common_params_parse() gets the same check here; on failure nothing is loaded.
+    if (!model_only && !common_exact_concurrency_init(params)) {
+        COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to load the model, see the error above\n");
+        return;
+    }
+
     auto mparams = common_model_params_to_llama(params);
     auto cparams = common_context_params_to_llama(params);
 
@@ -1520,11 +1529,6 @@ common_init_result_ptr common_init_from_params(common_params & params, bool mode
 
     const llama_vocab * vocab = llama_model_get_vocab(model);
 
-    // [TAG_EXACT_CONCURRENCY] before the warmup, which is the first graph this process computes
-    if (!common_exact_concurrency_init(params)) {
-        return res;
-    }
-
     if (params.ctx_shift && !llama_memory_can_shift(llama_get_memory(lctx))) {
         COM_WRN("%s", "KV cache shifting is not supported for this context, disabling KV cache shifting\n");
         params.ctx_shift = false;
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index 2a33c57e869..266a14dd4d4 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -187,6 +187,12 @@ bool llama_exact_concurrency() {
 // [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h
 static std::atomic<uint32_t> g_exact_decode_tokens{1};
 
+// one lock for the token figure, the sequence count and the width: the three move together
+// (a context reports its count and the width that follows; a new token figure re-reports the
+// width for every count seen), and a report interleaved with a change of figure could leave
+// the backend with a width that covers neither. Recursive, since the setters call each other.
+static std::recursive_mutex g_exact_mutex;
+
 // the most sequences any context so far was created with. The tokens figure is process
 // wide, so raising it widens the decode step of every context that already exists; the
 // width those contexts reported at creation is re-reported here with the new figure, or a
@@ -194,6 +200,8 @@ static std::atomic<uint32_t> g_exact_decode_tokens{1};
 static std::atomic<uint32_t> g_exact_max_n_seq{0};
 
 bool llama_exact_report_n_seq(uint32_t n_seq) {
+    std::lock_guard<std::recursive_mutex> lock(g_exact_mutex);
+
     const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed));
 
     if (!llama_set_exact_decode_width(n_seq_max * llama_exact_decode_tokens())) {
@@ -211,6 +219,8 @@ bool llama_exact_report_n_seq(uint32_t n_seq) {
 bool llama_set_exact_decode_tokens(uint32_t n_tokens) {
     n_tokens = n_tokens > 0 ? n_tokens : 1;
 
+    std::lock_guard<std::recursive_mutex> lock(g_exact_mutex);
+
     // never lowered: a narrower context set up later would otherwise turn the verify steps of
     // an existing speculative context into prompts and serialise them
     if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) {
@@ -262,10 +272,7 @@ bool llama_set_exact_decode_width(uint32_t n_cols) {
         return false;
     }
 
-    // one reporter at a time: the widest figure is read and handed to the backends below as
-    // one step, so a narrower report cannot overtake a wider one on its way to a backend
-    static std::mutex mutex;
-    std::lock_guard<std::mutex> lock(mutex);
+    std::lock_guard<std::recursive_mutex> lock(g_exact_mutex);
 
     uint32_t cur = g_exact_decode_width.load(std::memory_order_relaxed);
 

From 77d318522e4b24d3f8e2f74cce1d4fdc52455a1c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 20:58:28 +0000
Subject: [PATCH 30/33] exact concurrency: an isolated ubatch takes only sets
 that finish in it; one bound check at context creation

Under isolation the equal-count guard chose which sets join a ubatch, but
the expansion could still cut them all part way when the sets together
exceeded n_ubatch, the chunking the guard exists to prevent. A set that
would not finish in the ubatch waits for the next one.

The context constructor checked the explicit column bound itself and then
again through the report; the report's refusal is the error now.
---
 src/llama-batch.cpp   |  4 ++++
 src/llama-context.cpp | 14 +-------------
 2 files changed, 5 insertions(+), 13 deletions(-)

diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index 5db683db281..cc73d83963c 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -611,6 +611,10 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
                     n_left_first = n_left;
                 } else if (n_left != n_left_first) {
                     continue;
+                } else if ((cur_seq_set.size() + 1) * n_left_first > n_ubatch) {
+                    // one more set would not finish in this ubatch: the expansion below would
+                    // then cut every set part way, the chunking the guard exists to prevent
+                    break;
                 }
             }
 
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index e8ccf985477..68caf978337 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -108,21 +108,9 @@ llama_context::llama_context(
     // The sequence count is what is reported: the tokens figure can be raised later for the
     // whole process, and the width then follows it for this context too.
     if (llama_exact_concurrency()) {
-        const uint32_t n_cols = cparams.n_seq_max * llama_exact_decode_tokens();
-
         // an explicit column bound wins over the reported width in the backend, so one below
         // this context's width would leave its decodes batched above the bound with the mode
-        // still reporting itself on; refuse it here, the way the server's setup does
-        if (const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS")) {
-            const int max_cols = atoi(bound);
-
-            if (max_cols > 0 && (uint32_t) max_cols < n_cols) {
-                LLAMA_LOG_ERROR("%s: GGML_CUDA_BATCH_INVARIANT_MAX_COLS is %d but LLAMA_EXACT_CONCURRENCY needs at least %u to cover a decode step of %u sequences; raise it, set it to 0 for no bound, or unset it\n",
-                        __func__, max_cols, n_cols, cparams.n_seq_max);
-                throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width");
-            }
-        }
-
+        // still reporting itself on; the report refuses that, and the refusal is an error here
         if (!llama_exact_report_n_seq(cparams.n_seq_max)) {
             throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width");
         }

From 2f2258dc0faa6b39be80fc20d0c1b45adecfa3fd Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 21:35:50 +0000
Subject: [PATCH 31/33] exact concurrency: refuse a whole-context restore
 before it clears the cache, publish the width once construction succeeds,
 refuse the page table on every backend that ignores it

A whole-context restore under LLAMA_EXACT_CONCURRENCY was refused inside
state_read_meta(), after which the generic restore path cleared the live
cache. It is refused at the top of state_read_data() and
llama_kv_cache::state_read() now, before a byte is read, so
llama_state_set_data() returns 0 and the sequences are untouched.

The context reported its sequence count at the front of the constructor,
so a construction that failed later on left a width behind that no
context needed. The count is checked against the explicit column bound
early and reported at the end, once nothing can fail any more.

WebGPU, ExecuTorch, Hexagon, OpenVINO and RPC advertised FLASH_ATTN_EXT
with the page table in src[5] that only the CUDA kernels read; they
refuse it like CANN, Metal, OpenCL, SYCL and Vulkan already did.

DSpark runs on the DFlash implementation and turns causal attention off
on its draft context, so it is refused alongside DFlash. The decode width
is computed in 64 bits and refused above INT32_MAX instead of wrapping.
The probe's PROBE_A_PERM keeps prompt 0 on sequence 0, which phase B
compares against.
---
 common/common.cpp                        | 23 ++++++++++++----
 ggml/src/ggml-et/ggml-et.cpp             |  6 ++++
 ggml/src/ggml-hexagon/ggml-hexagon.cpp   |  4 ++-
 ggml/src/ggml-openvino/ggml-openvino.cpp |  5 ++++
 ggml/src/ggml-rpc/ggml-rpc.cpp           |  6 +++-
 ggml/src/ggml-webgpu/ggml-webgpu.cpp     |  7 +++++
 scripts/batchinv/probe.cpp               |  4 +++
 src/llama-context.cpp                    | 19 ++++++++++++-
 src/llama-impl.cpp                       | 35 ++++++++++++++++++++++--
 src/llama-impl.h                         |  4 +++
 src/llama-kv-cache.cpp                   | 17 +++++++-----
 11 files changed, 112 insertions(+), 18 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index 966620c7910..04ce8744359 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1455,13 +1455,17 @@ bool common_exact_concurrency() {
 
 // [TAG_EXACT_CONCURRENCY]
 int common_exact_decode_width(const common_params & params) {
-    const int n_slots = std::max(1, params.n_parallel);
+    const int64_t n_slots = std::max(1, params.n_parallel);
 
     // the draft tokens a slot carries into the verify ubatch alongside its accepted token, per
     // speculation type, from the same place the speculation code takes its own width
-    const int n_draft = std::max(0, (int) common_speculative_n_max(&params.speculative));
+    const int64_t n_draft = std::max(0, (int) common_speculative_n_max(&params.speculative));
 
-    return n_slots*(1 + n_draft);
+    // the product is what a backend is asked to split columns by, as an int; one that does not
+    // fit is reported as such rather than wrapped
+    const int64_t n_cols = n_slots*(1 + n_draft);
+
+    return n_cols > INT32_MAX ? -1 : (int) n_cols;
 }
 
 // [TAG_EXACT_CONCURRENCY]
@@ -1471,16 +1475,23 @@ bool common_exact_concurrency_init(const common_params & params) {
     }
 
     // DFlash drafting turns causal attention off on its draft context, and the paged
-    // attention the mode runs on needs it; say so instead of asserting in the graph
+    // attention the mode runs on needs it; say so instead of asserting in the graph. DSpark
+    // is the same implementation under another name, so it is refused with it.
     for (const auto type : params.speculative.types) {
-        if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) {
-            COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash: it disables causal attention, which the paged attention needs\n");
+        if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) {
+            COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash or draft-dspark: both disable causal attention on the draft, which the paged attention needs\n");
             return false;
         }
     }
 
     const int n_cols = common_exact_decode_width(params);
 
+    if (n_cols < 0) {
+        COM_ERR("LLAMA_EXACT_CONCURRENCY: a decode step of %d slots with %d draft tokens each is too wide to report\n",
+                std::max(1, params.n_parallel), std::max(0, (int) common_speculative_n_max(&params.speculative)));
+        return false;
+    }
+
     const char * bound = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
     if (bound) {
         const int max_cols = atoi(bound);
diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp
index b87b189a57a..a3792f3852b 100644
--- a/ggml/src/ggml-et/ggml-et.cpp
+++ b/ggml/src/ggml-et/ggml-et.cpp
@@ -1266,6 +1266,12 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                         (op->src[1]->ne[1] % op->src[4]->ne[1] == 0);
             break;
         case GGML_OP_FLASH_ATTN_EXT:
+            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
+            // the CUDA backend reads
+            if (op->src[5]) {
+                supported = false;
+                break;
+            }
             if (op->type == GGML_TYPE_F32 && op->src[0] && op->src[0]->type == GGML_TYPE_F32 && op->src[1] &&
                 (op->src[1]->type == GGML_TYPE_F32 || op->src[1]->type == GGML_TYPE_F16) && op->src[2] &&
                 (op->src[2]->type == GGML_TYPE_F32 || op->src[2]->type == GGML_TYPE_F16) && op->src[4] == nullptr &&
diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
index e8a5009b381..aa20083ec05 100644
--- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp
+++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
@@ -4157,7 +4157,9 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
             break;
 
         case GGML_OP_FLASH_ATTN_EXT:
-            supp = ggml_hexagon_supported_flash_attn_ext(sess, op);
+            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
+            // the CUDA backend reads
+            supp = op->src[5] == nullptr && ggml_hexagon_supported_flash_attn_ext(sess, op);
             break;
 
         case GGML_OP_SET_ROWS:
diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp
index e299e16c778..dfc9f90926f 100644
--- a/ggml/src/ggml-openvino/ggml-openvino.cpp
+++ b/ggml/src/ggml-openvino/ggml-openvino.cpp
@@ -1128,6 +1128,11 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
         break;
     }
     case GGML_OP_FLASH_ATTN_EXT: {
+        // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
+        // the CUDA backend reads
+        if (op->src[5]) {
+            return true;
+        }
         float scale = 1.0f;
         float max_bias = 0.0f;
         float logit_softcap = 0.0f;
diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp
index 69a8a08ae17..6fb8851a904 100644
--- a/ggml/src/ggml-rpc/ggml-rpc.cpp
+++ b/ggml/src/ggml-rpc/ggml-rpc.cpp
@@ -1915,7 +1915,11 @@ static ggml_backend_buffer_type_t ggml_backend_rpc_device_get_buffer_type(ggml_b
 
 static bool ggml_backend_rpc_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
     GGML_UNUSED(dev);
-    GGML_UNUSED(op);
+    // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the
+    // CUDA backend reads; the remote end is not asked, so it is not claimed here
+    if (op->op == GGML_OP_FLASH_ATTN_EXT && op->src[5]) {
+        return false;
+    }
     //TODO: call the remote backend and cache the results
     return true;
 }
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index 2434848a55a..70462a97f3c 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -4408,6 +4408,13 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
             break;
         case GGML_OP_FLASH_ATTN_EXT:
             {
+                // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
+                // the CUDA backend reads
+                if (op->src[5]) {
+                    supports_op = false;
+                    break;
+                }
+
                 // conservative support checks for whether the more resource-intensive shader paths
                 // can be used, to avoid cases where flash_attn is assigned to the CPU later on
                 supports_op = src0->type == GGML_TYPE_F32 &&
diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp
index 151463cd768..c4e478f97c7 100644
--- a/scripts/batchinv/probe.cpp
+++ b/scripts/batchinv/probe.cpp
@@ -172,9 +172,13 @@ int main(int argc, char ** argv) {
     const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1;
     // PROBE_A_PERM reorders which prompt goes into which sequence in phase A. With the same
     // multiset of prompts the cache keeps its length but the masked cells hold different data.
+    // Phase B decodes prompt 0's first token on sequence 0, so the permutation may only move
+    // the neighbours: sequence 0 keeps prompt 0, or the two phases would compare different
+    // sequences.
     int a_perm[4] = {0, 1, 2, 3};
     if (const char * perm = getenv("PROBE_A_PERM")) {
         for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0';
+        if (a_perm[0] != 0) { fprintf(stderr, "PROBE_A_PERM must keep prompt 0 on sequence 0\n"); return 1; }
     }
     {
         llama_context * ctx = make_ctx();
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 68caf978337..8d35922a47b 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -107,11 +107,14 @@ llama_context::llama_context(
     // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width.
     // The sequence count is what is reported: the tokens figure can be raised later for the
     // whole process, and the width then follows it for this context too.
+    // Checked here and reported at the end of the constructor: the count is process-wide
+    // state that outlives a context, so a construction that fails later on, an unsupported
+    // cache layout say, must not leave a width behind that no context needs.
     if (llama_exact_concurrency()) {
         // an explicit column bound wins over the reported width in the backend, so one below
         // this context's width would leave its decodes batched above the bound with the mode
         // still reporting itself on; the report refuses that, and the refusal is an error here
-        if (!llama_exact_report_n_seq(cparams.n_seq_max)) {
+        if (!llama_exact_check_n_seq(cparams.n_seq_max)) {
             throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width");
         }
     }
@@ -498,6 +501,13 @@ llama_context::llama_context(
             sampling.token_ids_full_vocab[i] = i;
         }
     }
+
+    // [TAG_EXACT_CONCURRENCY] nothing above can fail any more, so the width this context
+    // needs is published now; checked against the explicit bound at the top, so this
+    // cannot refuse unless the bound moved underneath it, which is an error all the same
+    if (llama_exact_concurrency() && !llama_exact_report_n_seq(cparams.n_seq_max)) {
+        throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width");
+    }
 }
 
 llama_context::~llama_context() {
@@ -3255,6 +3265,13 @@ size_t llama_context::state_write_data(llama_io_write_i & io) {
 }
 
 size_t llama_context::state_read_data(llama_io_read_i & io) {
+    // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical
+    // index, which the paged pool owns. Refused here, before anything is parsed, so that the
+    // cache the caller has is left as it was: the generic restore path clears it on failure.
+    if (memory && memory->alloc_granularity() > 1) {
+        throw std::runtime_error("whole-context restore is not supported with LLAMA_EXACT_CONCURRENCY, restore per sequence");
+    }
+
     LLAMA_LOG_DEBUG("%s: reading state\n", __func__);
 
     // read model info
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index 266a14dd4d4..0b218bc64e7 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -199,12 +199,41 @@ static std::recursive_mutex g_exact_mutex;
 // context created under a narrower figure would batch above the bound it reported.
 static std::atomic<uint32_t> g_exact_max_n_seq{0};
 
+static bool llama_exact_width_within_explicit_bound(uint32_t n_cols);
+
+// the width is sequences times tokens, handed to a backend as an int; a product that does not
+// fit is refused rather than wrapped
+static bool llama_exact_width_of(uint32_t n_seq, uint32_t n_tokens, uint32_t & n_cols) {
+    const uint64_t w = (uint64_t) n_seq * (uint64_t) n_tokens;
+
+    if (w > (uint64_t) INT32_MAX) {
+        LLAMA_LOG_ERROR("%s: a decode step of %u sequences with %u tokens each is too wide to report\n", __func__, n_seq, n_tokens);
+        return false;
+    }
+
+    n_cols = (uint32_t) w;
+
+    return true;
+}
+
+bool llama_exact_check_n_seq(uint32_t n_seq) {
+    std::lock_guard<std::recursive_mutex> lock(g_exact_mutex);
+
+    const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed));
+
+    uint32_t n_cols = 0;
+
+    return llama_exact_width_of(n_seq_max, llama_exact_decode_tokens(), n_cols) && llama_exact_width_within_explicit_bound(n_cols);
+}
+
 bool llama_exact_report_n_seq(uint32_t n_seq) {
     std::lock_guard<std::recursive_mutex> lock(g_exact_mutex);
 
     const uint32_t n_seq_max = std::max(n_seq, g_exact_max_n_seq.load(std::memory_order_relaxed));
 
-    if (!llama_set_exact_decode_width(n_seq_max * llama_exact_decode_tokens())) {
+    uint32_t n_cols = 0;
+
+    if (!llama_exact_width_of(n_seq_max, llama_exact_decode_tokens(), n_cols) || !llama_set_exact_decode_width(n_cols)) {
         return false;
     }
 
@@ -231,7 +260,9 @@ bool llama_set_exact_decode_tokens(uint32_t n_tokens) {
     // reported first; a figure the explicit bound cannot cover leaves the old one in place
     const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed);
 
-    if (n_seq > 0 && !llama_set_exact_decode_width(n_seq * n_tokens)) {
+    uint32_t n_cols = 0;
+
+    if (n_seq > 0 && (!llama_exact_width_of(n_seq, n_tokens, n_cols) || !llama_set_exact_decode_width(n_cols))) {
         return false;
     }
 
diff --git a/src/llama-impl.h b/src/llama-impl.h
index de5c6a2d216..65d56a51d1d 100644
--- a/src/llama-impl.h
+++ b/src/llama-impl.h
@@ -113,3 +113,7 @@ bool llama_exact_concurrency();
 // [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so that the
 // decode width every context needs is known to the backend and follows llama_set_exact_decode_tokens
 bool llama_exact_report_n_seq(uint32_t n_seq);
+
+// the same check without the report: whether a context of n_seq sequences could be reported
+// under the explicit column bound, for a constructor that may still fail after asking
+bool llama_exact_check_n_seq(uint32_t n_seq);
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index bb947c9ed09..c384e46d520 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -2362,6 +2362,14 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla
 }
 
 void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
+    // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical
+    // index, which the paged pool owns. Refused before a byte is read, so that the failure
+    // path below, which clears the cache, is never entered for it.
+    if (exact_pages && seq_id == -1) {
+        LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state restore only\n", __func__);
+        throw std::runtime_error("whole-cache restore is not supported with LLAMA_EXACT_CONCURRENCY");
+    }
+
     // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS]
     if (other) {
         return;
@@ -2610,13 +2618,8 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32
     } else {
         // whole KV cache restore
 
-        // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical
-        // index, which the paged pool owns. Report it like every other failure in this function.
-        if (exact_pages) {
-            LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state "
-                    "restore only\n", __func__);
-            return false;
-        }
+        // [TAG_EXACT_CONCURRENCY] refused at the top of state_read(), before anything is read
+        GGML_ASSERT(!exact_pages);
 
         if (cell_count > cells.size()) {
             LLAMA_LOG_ERROR("%s: not enough cells in kv cache\n", __func__);

From 918a8bf4f022023953c4e3f2da9cef279d33251c Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sun, 6 Sep 2026 23:33:32 +0000
Subject: [PATCH 32/33] exact concurrency: ask the device whether it can run
 the paged attention, not only which backend it is

The KV cache accepted a layer on any CUDA, ROCm or MUSA device by the
registry name alone. A build without the flash attention kernels, or a
device and head shape they do not cover, would then have the scheduler
hand the paged op to the CPU, which accepts the page table as the
reference for test-backend-ops and ignores it, and the mode would report
itself on while attending in physical order. The constructor now builds
the attention op the way the graph does, page table attached, at the
widths of a decode step, a verify step and a prompt chunk, and asks the
device; a refusal is a load error naming the layer and the types.
---
 src/llama-kv-cache.cpp | 67 ++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 67 insertions(+)

diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index c384e46d520..31154639ad9 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -85,6 +85,61 @@ static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) {
     return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0;
 }
 
+// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a
+// layer of this shape. The registry name says which backends carry the kernels; it does not
+// say the build has them (FLASH_ATTN_AVAILABLE), nor that the device's architecture, the
+// head width and the K/V types land on a kernel. Where they do not, the scheduler would hand
+// the op to the CPU, which accepts the page table as the reference for test-backend-ops and
+// ignores it, and the mode would report itself on while attending in physical order. So the
+// op is built the way the graph builds it, at the widths a decode step, a verify step and a
+// prompt chunk use, and the device is asked.
+static bool llama_dev_supports_paged_attn(
+        ggml_backend_dev_t dev,
+        ggml_type type_k, ggml_type type_v,
+        uint32_t n_embd_head_k, uint32_t n_embd_head_v,
+        uint32_t n_head, uint32_t n_head_kv,
+        uint32_t n_cells, uint32_t page_size) {
+    if (!llama_dev_has_paged_attn(dev)) {
+        return false;
+    }
+
+    ggml_init_params ip = {
+        /*.mem_size   =*/ ggml_tensor_overhead()*16 + ggml_graph_overhead(),
+        /*.mem_buffer =*/ nullptr,
+        /*.no_alloc   =*/ true,
+    };
+
+    ggml_context * ctx = ggml_init(ip);
+    if (!ctx) {
+        return false;
+    }
+
+    bool res = true;
+
+    const int64_t n_kv = page_size;
+
+    for (const int64_t n_tokens : { (int64_t) 1, (int64_t) 4, (int64_t) 16, (int64_t) 512 }) {
+        ggml_tensor * q = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_embd_head_k, n_tokens, n_head,    1);
+        ggml_tensor * k = ggml_new_tensor_4d(ctx, type_k,        n_embd_head_k, n_kv,     n_head_kv, 1);
+        ggml_tensor * v = ggml_new_tensor_4d(ctx, type_v,        n_embd_head_v, n_kv,     n_head_kv, 1);
+        ggml_tensor * m = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_tokens, 1, 1);
+
+        ggml_tensor * op = ggml_flash_attn_ext(ctx, q, k, v, m, 1.0f/sqrtf((float) n_embd_head_k), 0.0f, 0.0f);
+        ggml_flash_attn_ext_set_prec(op, GGML_PREC_F32);
+
+        op->src[5] = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, 1 + n_cells/page_size, n_tokens);
+
+        if (!ggml_backend_dev_supports_op(dev, op)) {
+            res = false;
+            break;
+        }
+    }
+
+    ggml_free(ctx);
+
+    return res;
+}
+
 llama_kv_cache::llama_kv_cache(
         const llama_model & model,
         const llama_hparams & hparams,
@@ -300,6 +355,18 @@ llama_kv_cache::llama_kv_cache(
             throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend");
         }
 
+        // [TAG_EXACT_CONCURRENCY] the backend is the right one; ask it whether this layer's
+        // attention, with the page table attached, lands on one of its kernels at all
+        if (exact_pages && !llama_dev_supports_paged_attn(model.dev_layer(il), type_k, type_v,
+                    hparams.n_embd_head_k(il), hparams.n_embd_head_v(il),
+                    hparams.n_head(il), hparams.n_head_kv(il), kv_size, exact_page_size)) {
+            LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but %s cannot run the paged attention for layer %d "
+                    "(K %s, V %s, %u-wide heads): the build or the device has no flash attention kernel for it, "
+                    "and the op would fall to the CPU, which ignores the page table\n",
+                    __func__, dev_name, il, ggml_type_name(type_k), ggml_type_name(type_v), hparams.n_embd_head_k(il));
+            throw std::runtime_error("exact concurrency: the device cannot run the paged attention");
+        }
+
         ggml_context * ctx = ctx_for_buft(buft);
         if (!ctx) {
             throw std::runtime_error("failed to create ggml context for kv cache");

From b4b0f9bd71f1b411df4a3a00d373140272be9d91 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 7 Sep 2026 04:06:23 +0000
Subject: [PATCH 33/33] exact concurrency: shorter comments

Comment-only pass over the PR's diff: collapse the long explanations to one or two lines each and drop the ones the code already says.
---
 .github/workflows/unsloth-pin-preflight.yml |  57 ++--
 .github/workflows/unsloth-pr-set-lint.yml   |   6 +-
 .github/workflows/unsloth-prebuilt.yml      |  23 +-
 common/common.cpp                           |  27 +-
 common/common.h                             |  12 +-
 ggml/include/ggml-cuda.h                    |  12 +-
 ggml/src/ggml-cann/ggml-cann.cpp            |   3 +-
 ggml/src/ggml-cpu/ggml-cpu.cpp              |  11 +-
 ggml/src/ggml-cuda/common.cuh               |   4 +-
 ggml/src/ggml-cuda/fattn-common.cuh         |   9 +-
 ggml/src/ggml-cuda/fattn-vec.cuh            |   4 +-
 ggml/src/ggml-cuda/fattn.cu                 |  10 +-
 ggml/src/ggml-cuda/ggml-cuda.cu             | 111 +++----
 ggml/src/ggml-cuda/mmvq.cu                  |  16 +-
 ggml/src/ggml-cuda/mmvq.cuh                 |   5 +-
 ggml/src/ggml-et/ggml-et.cpp                |   3 +-
 ggml/src/ggml-hexagon/ggml-hexagon.cpp      |   3 +-
 ggml/src/ggml-metal/ggml-metal-device.m     |   4 +-
 ggml/src/ggml-opencl/ggml-opencl.cpp        |   3 +-
 ggml/src/ggml-openvino/ggml-openvino.cpp    |   3 +-
 ggml/src/ggml-rpc/ggml-rpc.cpp              |   4 +-
 ggml/src/ggml-sycl/ggml-sycl.cpp            |   3 +-
 ggml/src/ggml-vulkan/ggml-vulkan.cpp        |   3 +-
 ggml/src/ggml-webgpu/ggml-webgpu.cpp        |   3 +-
 include/llama.h                             |  36 +--
 scripts/batchinv/divergence.py              |  12 +-
 scripts/batchinv/probe.cpp                  |  30 +-
 scripts/batchinv/prompts.py                 |   2 +-
 scripts/unsloth/additive_merge.py           |  43 +--
 scripts/unsloth/feature_matrix.py           |  14 +-
 scripts/unsloth/pin_contract.py             |  36 +--
 scripts/unsloth/test_additive_merge.py      |   6 +-
 scripts/unsloth/test_pin_contract.py        |  11 +-
 src/llama-batch.cpp                         |  23 +-
 src/llama-batch.h                           |   9 +-
 src/llama-context.cpp                       |  44 ++-
 src/llama-graph.cpp                         |  17 +-
 src/llama-impl.cpp                          |  40 +--
 src/llama-impl.h                            |  14 +-
 src/llama-kv-cache.cpp                      | 110 +++----
 src/llama-kv-cache.h                        |   9 +-
 src/llama-memory-hybrid.cpp                 |  19 +-
 src/llama-memory-recurrent.cpp              |  10 +-
 src/llama-memory.h                          |  11 +-
 tests/test-backend-ops.cpp                  |  13 +-
 tests/test-state-restore-fragmented.cpp     |   4 +-
 tools/server/server-context.cpp             | 333 +++++++-------------
 tools/server/tests/unit/test_preempt.py     | 111 +++----
 48 files changed, 497 insertions(+), 799 deletions(-)

diff --git a/.github/workflows/unsloth-pin-preflight.yml b/.github/workflows/unsloth-pin-preflight.yml
index f390f381e6f..494bca228c8 100644
--- a/.github/workflows/unsloth-pin-preflight.yml
+++ b/.github/workflows/unsloth-pin-preflight.yml
@@ -28,16 +28,10 @@ permissions:
   contents: write
   issues: write
 
-# Two runs of the same ref probe the same pins against the same base, so the
-# second adds nothing and just competes for runners. On 08-04 a dispatch and the
-# schedule sat queued together for an hour. Newest wins: it sees the newest
-# pr-set.json.
-#
-# Per ref, though, not globally. This file also runs on any push that touches
-# pr-set.json, so with one shared group a push to a second branch cancelled the
-# first branch's run: observed on 09-03, where the run that would have said
-# whether a repin fixed the nightly was cancelled by an unrelated branch, and
-# the PR was left showing the failure from before the fix.
+# Two runs of the same ref probe the same pins against the same base, so the second only
+# competes for runners; newest wins, since it sees the newest pr-set.json. Per ref, not
+# globally: this also runs on any push touching pr-set.json, and with one shared group a push
+# to a second branch cancelled the first branch's run (09-03).
 concurrency:
   group: unsloth-pin-preflight-${{ github.ref }}
   cancel-in-progress: true
@@ -56,10 +50,8 @@ jobs:
         id: p
         run: |
           set -uo pipefail
-          # Everything below reports through `status`/`details`, so a death
-          # anywhere else leaves both empty and the alert blank: a red X on a
-          # scheduled run nobody opens. Report the abort through the same
-          # channel as a finding, so the repin bot sees a failure either way.
+          # everything below reports through `status`/`details`, so a death anywhere else
+          # leaves the alert blank; report the abort through the same channel as a finding
           trap 'rc=$?; if [ "$rc" != 0 ]; then {
                   echo "status=failure"
                   echo "details<<ALERT_EOF"
@@ -161,15 +153,10 @@ jobs:
               continue
             fi
             FILES="$(git diff --name-only --diff-filter=U | sed 's/^/  /')"
-            # `|| true` is load-bearing, not tidying. GitHub runs a `run:` block
-            # under `bash -e` whatever this script's own `set` line says, `head`
-            # closes the pipe after 20 lines, and pipefail then makes the whole
-            # assignment fail. So on 09-03 the step died right here, on the first
-            # real conflict, with `grep: write error: Broken pipe` and no alert:
-            # the one path this job exists to report was the one it could not
-            # survive. It only fires when the conflict diff is bigger than the
-            # 64 KiB pipe buffer, since a smaller one is written before `head`
-            # ever closes it, which is why most conflicts got reported fine.
+            # `|| true` is load-bearing: a `run:` block is `bash -e` whatever this script
+            # says, `head` closes the pipe after 20 lines, and pipefail then fails the whole
+            # assignment. On 09-03 that killed the step on the first real conflict with
+            # `grep: write error: Broken pipe`. Only fires above the 64 KiB pipe buffer.
             HUNKS="$(git diff --diff-filter=U -U0 2>/dev/null | grep -E '^\+|^-' | grep -vE '^(\+\+\+|---)' | head -20 || true)"
             git merge --abort 2>/dev/null
             PROBLEMS="${PROBLEMS}- \`${SRC}#${NUM}\` (\`${SHA:0:10}\`) does not merge onto \`${BASE}\` + the pins before it.\n\n  Conflicting files:\n\n\`\`\`\n${FILES}\n\`\`\`\n\n  <details><summary>conflict hunks</summary>\n\n\`\`\`diff\n${HUNKS}\n\`\`\`\n\n  </details>\n"
@@ -193,11 +180,9 @@ jobs:
               PROBLEMS="${PROBLEMS}- the merged tree builds, but \`scripts/unsloth/merge_checks.py\` found a resolution that is silently wrong. See the run log for file and line.\n"
             fi
 
-            # The other half of that question. merge_checks.py asks whether the
-            # tree contains something wrong; this asks whether it still contains
-            # what each pin carries. A pin that has rotted into a no-op, or an
-            # arch registration a resolution quietly dropped, is invisible to
-            # every other check here and to the compiler.
+            # the other half: merge_checks.py asks whether the tree contains something
+            # wrong, this asks whether it still contains what each pin carries. A pin rotted
+            # into a no-op is invisible to every other check here and to the compiler.
             if ! python3 ../scripts/unsloth/pin_contract.py --root . --base "$BASE" \
                  --pr-set ../scripts/unsloth/pr-set.json --report "${RUNNER_TEMP}/pin_contract.json" ; then
               PROBLEMS="${PROBLEMS}- the merged tree is missing code a pin carries. See the run log for the pin and file.\n"
@@ -207,12 +192,10 @@ jobs:
               PROBLEMS="${PROBLEMS}- pins upstream has taken over, safe to delete from \`pr-set.json\`:\n\n\`\`\`\n${NOTES}\n\`\`\`\n"
             fi
 
-            # A clean merge is not a compiling tree. On 09-03 ggml-org#27754
-            # merged with no conflicts at all and did not compile: upstream had
-            # added a parameter to build_attn_mha and the pin's new
-            # build_attn_sparse still called the old signature. Nothing above
-            # can see that. CPU only and the `llama` target only, which is where
-            # that translation unit lives; 59s cold at -j4 with no ccache.
+            # a clean merge is not a compiling tree: on 09-03 ggml-org#27754 merged with no
+            # conflicts and did not compile, upstream having added a parameter to
+            # build_attn_mha that the pin's build_attn_sparse still called without. CPU only,
+            # 59s cold at -j4 with no ccache.
             GATE_OK=1
             if ! cmake -B "${RUNNER_TEMP}/gate" -DCMAKE_BUILD_TYPE=Release \
                    -DGGML_CUDA=OFF -DLLAMA_BUILD_TESTS=ON -DLLAMA_BUILD_SERVER=OFF \
@@ -223,10 +206,8 @@ jobs:
               PROBLEMS="${PROBLEMS}- the pins merge cleanly and the merged tree does not compile. See the run log for the file and line; this is the failure that only shows up in the CUDA leg once the nightly has fanned out.\n"
             fi
 
-            # The last question, and the only one that needs a binary: does each
-            # feature we ship still work. Everything above is about the source.
-            # CPU only, because no runner in this pipeline has a GPU -- see the
-            # note in feature_matrix.py about what that does and does not prove.
+            # the only question that needs a binary: does each feature we ship still work.
+            # CPU only, since no runner here has a GPU; see the note in feature_matrix.py.
             if [ -n "$GATE_OK" ]; then
               if ! python3 ../scripts/unsloth/feature_matrix.py \
                    --build-dir "${RUNNER_TEMP}/gate" \
diff --git a/.github/workflows/unsloth-pr-set-lint.yml b/.github/workflows/unsloth-pr-set-lint.yml
index 8a898539f02..99b06b74f34 100644
--- a/.github/workflows/unsloth-pr-set-lint.yml
+++ b/.github/workflows/unsloth-pr-set-lint.yml
@@ -128,9 +128,9 @@ jobs:
           done
           exit "$fail"
 
-      # A pin nobody decided about is the failure this whole file exists to stop.
-      # Being in `unchecked` with a reason is a fine answer; being in neither map
-      # is how DiffusionGemma went five weeks with no coverage and no record of it.
+      # a pin nobody decided about is the failure this file exists to stop: being in
+      # `unchecked` with a reason is fine, being in neither map is how DiffusionGemma went
+      # five weeks with no coverage and no record of it
       - name: Every pin is either checked or knowingly unchecked
         run: |
           set -euo pipefail
diff --git a/.github/workflows/unsloth-prebuilt.yml b/.github/workflows/unsloth-prebuilt.yml
index 834ff544909..353d88417ca 100644
--- a/.github/workflows/unsloth-prebuilt.yml
+++ b/.github/workflows/unsloth-prebuilt.yml
@@ -282,8 +282,8 @@ jobs:
           # .github/workflows, which upstream history routinely does).
           if [ "$EXISTS" != "true" ] || [ "${{ github.event_name }}" = "workflow_dispatch" ]; then
             git remote add upstream https://github.com/ggml-org/llama.cpp.git
-            # The upstream checkout below takes scripts/unsloth/ away. Copy the
-            # whole dir out, not file by file: see the note above the step.
+            # the upstream checkout below takes scripts/unsloth/ away; copy the whole dir
+            # out, not file by file, see the note above the step
             cp -r scripts/unsloth "${RUNNER_TEMP}/us"
             ADDITIVE_MERGE="${RUNNER_TEMP}/us/additive_merge.py"
             if [ "$(jq length <<<"$PRS")" != 0 ]; then
@@ -446,8 +446,7 @@ jobs:
 
       # A bad pin resolution can still build fine, so it must be caught before the source artifact ships. See merge_checks.py.
       # Its own step, not more script in `resolve`: GitHub caps one workflow string at 21000 chars and that step is near it. See check_workflow_scalars.py.
-      # That is also why `resolve` copies all of scripts/unsloth/ to ${RUNNER_TEMP}/us in one line rather than one cp per script: every check added
-      # here would otherwise cost another line inside the capped block, and going over silently disables the whole workflow.
+      # That is also why `resolve` copies all of scripts/unsloth/ to ${RUNNER_TEMP}/us in one line rather than one cp per script: another line inside the capped block per check would eventually go over, which silently disables the whole workflow.
       - name: Check the merged tree for silently wrong resolutions
         if: ${{ env.MERGED_PINS == '1' }}
         run: |
@@ -457,13 +456,11 @@ jobs:
             exit 1
           fi
 
-      # merge_checks.py asserts the ABSENCE of two known-bad shapes. This asserts the PRESENCE of what each pin carries, which is a different question and
-      # the one that goes unanswered when a pin rots into a no-op or a resolution quietly drops an arch registration. Free, so it runs before the compile gate.
+      # merge_checks.py asserts the ABSENCE of two known-bad shapes; this asserts the PRESENCE of what each pin carries, the question that goes unanswered when a pin rots into a no-op. Free, so it runs before the compile gate.
       - name: Check every pin still contributes what it carries
         if: ${{ env.MERGED_PINS == '1' }}
-        # Through env, never interpolated into the script: `prs` carries PR
-        # titles, which are third-party text, and `${{ }}` pastes them into the
-        # shell source before bash ever sees it.
+        # through env, never interpolated: `prs` carries PR titles, which are third-party
+        # text, and `${{ }}` pastes them into the shell source before bash sees it
         env:
           PRS: ${{ steps.r.outputs.prs }}
           BASE: ${{ steps.r.outputs.base }}
@@ -475,12 +472,8 @@ jobs:
             exit 1
           fi
 
-      # The gap this closes, observed 09-03: ggml-org#27754 merged with zero conflicts and did not compile, because upstream had added a parameter to
-      # build_attn_mha and the pin's new build_attn_sparse still called the old signature. Nothing before this point can see that, and without it the release
-      # dies in the CUDA leg after the 38-job fan-out. CPU only: a cold `llama` build took 59s at -j4 with no ccache, against 20-60 minutes for a CUDA build.
-      # mtmd is in the gate because `llama` alone is not enough: observed 09-04, ggml-org#25731 built `llama` clean while tools/mtmd did not compile at all,
-      # upstream having made mtmd_image_preprocessor::preprocess const while the pin's Inkling subclass stayed non-const, so it overrode nothing and the
-      # vision and audio towers were abstract. Every vision pin lands in mtmd, so a gate that skips it cannot see the whole class.
+      # The gap this closes, observed 09-03: ggml-org#27754 merged with zero conflicts and did not compile, upstream having added a parameter to build_attn_mha that the pin's build_attn_sparse still called without. Without this the release dies in the CUDA leg after the 38-job fan-out. CPU only: a cold `llama` build took 59s at -j4, against 20-60 minutes for CUDA.
+      # mtmd is in the gate because `llama` alone is not enough: on 09-04 ggml-org#25731 built `llama` clean while tools/mtmd did not compile at all, upstream having made mtmd_image_preprocessor::preprocess const while the pin's Inkling subclass stayed non-const. Every vision pin lands in mtmd.
       - name: Compile gate (CPU, llama and mtmd targets)
         if: ${{ env.MERGED_PINS == '1' }}
         run: |
diff --git a/common/common.cpp b/common/common.cpp
index 04ce8744359..b74bee734ae 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -1290,10 +1290,9 @@ struct common_init_result::impl {
 
 common_init_result::common_init_result(common_params & params, bool model_only) :
     pimpl(new impl{}) {
-    // [TAG_EXACT_CONCURRENCY] before any context exists, the fitting ones included: the
-    // per-sequence figure and the column bound are checked against the explicit bound first,
-    // so a context is never created under a figure the bound does not cover. A caller that
-    // skipped common_params_parse() gets the same check here; on failure nothing is loaded.
+    // [TAG_EXACT_CONCURRENCY] before any context exists, so one is never created under a figure
+    // the explicit bound does not cover; this also covers a caller that skipped
+    // common_params_parse(). On failure nothing is loaded.
     if (!model_only && !common_exact_concurrency_init(params)) {
         COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: refusing to load the model, see the error above\n");
         return;
@@ -1457,12 +1456,11 @@ bool common_exact_concurrency() {
 int common_exact_decode_width(const common_params & params) {
     const int64_t n_slots = std::max(1, params.n_parallel);
 
-    // the draft tokens a slot carries into the verify ubatch alongside its accepted token, per
-    // speculation type, from the same place the speculation code takes its own width
+    // draft tokens a slot carries into the verify ubatch, from the same place the speculation
+    // code takes its own width
     const int64_t n_draft = std::max(0, (int) common_speculative_n_max(&params.speculative));
 
-    // the product is what a backend is asked to split columns by, as an int; one that does not
-    // fit is reported as such rather than wrapped
+    // the product is handed to a backend as an int; one that overflows is reported, not wrapped
     const int64_t n_cols = n_slots*(1 + n_draft);
 
     return n_cols > INT32_MAX ? -1 : (int) n_cols;
@@ -1474,9 +1472,8 @@ bool common_exact_concurrency_init(const common_params & params) {
         return true;
     }
 
-    // DFlash drafting turns causal attention off on its draft context, and the paged
-    // attention the mode runs on needs it; say so instead of asserting in the graph. DSpark
-    // is the same implementation under another name, so it is refused with it.
+    // DFlash drafting turns causal attention off on its draft context, which the paged attention
+    // needs; say so instead of asserting in the graph. DSpark is the same implementation.
     for (const auto type : params.speculative.types) {
         if (type == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) {
             COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY does not support --spec-type draft-dflash or draft-dspark: both disable causal attention on the draft, which the paged attention needs\n");
@@ -1505,11 +1502,9 @@ bool common_exact_concurrency_init(const common_params & params) {
         }
     }
 
-    // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode
-    // step is; a context created later reports n_seq_max times that figure, which is n_cols
-    // again, and reporting n_cols here as well covers a caller that decodes before that. Both
-    // refuse a width the explicit bound above cannot cover, which the check above already
-    // caught for this process; contexts created earlier by the caller are covered here.
+    // the batch splitter isolates prompts by width, so tell it how wide one sequence's decode step
+    // is. A context created later reports n_seq_max times that, which is n_cols again; reporting
+    // n_cols here too covers a caller that decodes first, or contexts it created earlier.
     if (!llama_set_exact_decode_tokens((uint32_t) (n_cols / std::max(1, params.n_parallel))) ||
         !llama_set_exact_decode_width((uint32_t) n_cols)) {
         COM_ERR("%s", "LLAMA_EXACT_CONCURRENCY: the decode width could not be reported, see the error above\n");
diff --git a/common/common.h b/common/common.h
index 2be2fab6a8b..4616b73153e 100644
--- a/common/common.h
+++ b/common/common.h
@@ -931,17 +931,15 @@ using common_init_result_ptr = std::unique_ptr<common_init_result>;
 
 common_init_result_ptr common_init_from_params(common_params & params, bool model_only = false);
 
-// [TAG_EXACT_CONCURRENCY]
-// true when LLAMA_EXACT_CONCURRENCY is set for this process
+// [TAG_EXACT_CONCURRENCY] true when LLAMA_EXACT_CONCURRENCY is set for this process
 bool common_exact_concurrency();
 
-// the widest ubatch a decode step can build with these parameters: one column per slot, times one
-// plus the number of speculative draft tokens carried with it. Under exact mode this is what the
-// CUDA column policy has to cover, and what its default bound is derived from.
+// the widest ubatch a decode step can build here: one column per slot times one plus its draft
+// tokens. Under exact mode the CUDA column policy has to cover this, and derives its bound from it.
 int common_exact_decode_width(const common_params & params);
 
-// report that width to the CUDA backend, and refuse an explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS
-// that is smaller than it. Returns false if the configuration must not run.
+// report that width to the CUDA backend, refusing a smaller explicit
+// GGML_CUDA_BATCH_INVARIANT_MAX_COLS; false if the configuration must not run
 bool common_exact_concurrency_init(const common_params & params);
 
 struct llama_model_params   common_model_params_to_llama  (      common_params & params);
diff --git a/ggml/include/ggml-cuda.h b/ggml/include/ggml-cuda.h
index c3dd87c97b7..07131df327c 100644
--- a/ggml/include/ggml-cuda.h
+++ b/ggml/include/ggml-cuda.h
@@ -39,13 +39,11 @@ GGML_BACKEND_API void ggml_backend_cuda_get_device_memory(int device, size_t * f
 
 GGML_BACKEND_API bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size);
 
-// [TAG_EXACT_CONCURRENCY]
-// Report the widest ubatch a decode step of this process can build: one column per slot, times one
-// plus the number of speculative draft tokens carried with it. Under LLAMA_EXACT_CONCURRENCY the
-// column policy then defaults to that width instead of a fixed number, so --parallel or a wider
-// draft cannot silently push a decode above the bound and leave it batched. An explicitly set
-// GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins. Call before the first graph is computed. Also
-// available through ggml_backend_reg_get_proc_address().
+// [TAG_EXACT_CONCURRENCY] report the widest ubatch a decode step of this process can build: one
+// column per slot times one plus its draft tokens. Under LLAMA_EXACT_CONCURRENCY the column policy
+// defaults to that instead of a fixed number, so --parallel or a wider draft cannot silently push a
+// decode above the bound. An explicit GGML_CUDA_BATCH_INVARIANT_MAX_COLS still wins. Call before
+// the first graph is computed; also available through ggml_backend_reg_get_proc_address().
 GGML_BACKEND_API void ggml_backend_cuda_set_exact_decode_width(int n_cols);
 GGML_BACKEND_API void ggml_backend_cuda_unregister_host_buffer(void * buffer);
 
diff --git a/ggml/src/ggml-cann/ggml-cann.cpp b/ggml/src/ggml-cann/ggml-cann.cpp
index 0e901ed0160..9b9213df856 100644
--- a/ggml/src/ggml-cann/ggml-cann.cpp
+++ b/ggml/src/ggml-cann/ggml-cann.cpp
@@ -2656,8 +2656,7 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
             return true;
         case GGML_OP_FLASH_ATTN_EXT:
             {
-                // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
-                // the CUDA backend reads
+                // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads
                 if (op->src[5]) {
                     return false;
                 }
diff --git a/ggml/src/ggml-cpu/ggml-cpu.cpp b/ggml/src/ggml-cpu/ggml-cpu.cpp
index a548b33bd71..8bed84e44bb 100644
--- a/ggml/src/ggml-cpu/ggml-cpu.cpp
+++ b/ggml/src/ggml-cpu/ggml-cpu.cpp
@@ -474,13 +474,10 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_SSM_SCAN:
             return ggml_get_op_params_i32(op, 0) == 1 || op->src[3]->ne[0] == 1;
-        // [TAG_EXACT_CONCURRENCY] note: GGML_OP_FLASH_ATTN_EXT with src[5] set, the
-        // exact-concurrency page table, is deliberately still accepted here. The CPU ignores the
-        // page table and attends in physical cell order, which is why every other backend refuses
-        // it, but the CPU is also the reference that test-backend-ops compares the paged CUDA
-        // kernel against, and that test builds a mask which selects exactly the listed cells. A KV
-        // cache layer cannot reach the CPU under the mode anyway: llama_kv_cache refuses to
-        // construct unless every KV layer is on the CUDA backend.
+        // [TAG_EXACT_CONCURRENCY] note: FLASH_ATTN_EXT with src[5], the page table, is deliberately
+        // still accepted. The CPU ignores it and attends in physical order, but it is also the
+        // reference test-backend-ops compares the paged CUDA kernel against, and that test's mask
+        // selects exactly the listed cells. A KV layer cannot reach the CPU under the mode anyway.
         default:
             return true;
     }
diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
index 728aa08dcb5..d749b42ee48 100644
--- a/ggml/src/ggml-cuda/common.cuh
+++ b/ggml/src/ggml-cuda/common.cuh
@@ -51,8 +51,8 @@
 #define GGML_CUDA_CC_DP4A            610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products
 // [TAG_BATCH_INVARIANT] 0 = off, 1 = split every batched matmul, 2 = split only where it changes bits
 int ggml_cuda_batch_invariant();
-// Widest batch the split is applied to, 0 = no bound. Prompt-sized batches cost far more to
-// split than decode-sized ones, and only prompt-phase invariance is given up by bounding it.
+// widest batch the split applies to, 0 = no bound; bounding it gives up prompt-phase invariance
+// only, and prompt-sized batches cost far more to split
 int ggml_cuda_batch_invariant_max_cols();
 
 #define GGML_CUDA_CC_VOLTA           700
diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh
index f6aa5b03ec1..10f28eb229b 100644
--- a/ggml/src/ggml-cuda/fattn-common.cuh
+++ b/ggml/src/ggml-cuda/fattn-common.cuh
@@ -1091,8 +1091,8 @@ void launch_fattn(
     // Optional optimization where the mask is scanned to determine whether part of the calculation can be skipped.
     // Only worth the overhead if there is at lease one FATTN_KQ_STRIDE x FATTN_KQ_STRIDE square to be skipped or
     //     multiple sequences of possibly different lengths.
-    // [TAG_BATCH_INVARIANT] Without this scan the KV loop runs to K->ne[1], which grows with the
-    // other sequences sharing the cache. Scanning the mask bounds it by the sequence's own extent.
+    // [TAG_BATCH_INVARIANT] without this scan the KV loop runs to K->ne[1], which grows with the
+    // other sequences sharing the cache; scanning the mask bounds it by the sequence's own extent
     const bool batch_invariant_KV_max = ggml_cuda_batch_invariant() != 0;
     if (!dst->src[5] && mask && K->ne[1] % FATTN_KQ_STRIDE == 0 && (Q->ne[1] >= 1024 || Q->ne[3] > 1 || batch_invariant_KV_max)) {
         const int64_t s31 = mask->nb[1] / sizeof(half2);
@@ -1152,9 +1152,8 @@ void launch_fattn(
             dst_tmp_meta.alloc((size_t(blocks_num.x) * ncols * (2 + DV/2)));
         }
     } else if (dst->src[5] || ggml_cuda_batch_invariant()) {
-        // [TAG_BATCH_INVARIANT] How the KV cache is split between blocks, and therefore the order
-        // in which the partial attention results are combined, follows K->ne[1]. That length grows
-        // with the other sequences sharing the cache, so pin the split to a single block per tile.
+        // [TAG_BATCH_INVARIANT] the KV split between blocks, and so the order the partials combine
+        // in, follows K->ne[1], which grows with the other sequences: pin it to one block per tile
         parallel_blocks = 1;
 
         blocks_num.x = ntiles_x;
diff --git a/ggml/src/ggml-cuda/fattn-vec.cuh b/ggml/src/ggml-cuda/fattn-vec.cuh
index f402795942c..4005d28879e 100644
--- a/ggml/src/ggml-cuda/fattn-vec.cuh
+++ b/ggml/src/ggml-cuda/fattn-vec.cuh
@@ -247,8 +247,8 @@ static __global__ void flash_attn_ext_vec(
 #endif // V_DOT2_F32_F16_AVAILABLE
     }
 
-    // In the paged specialization KV_max carries [count, physical page IDs...] per query.
-    // The loop and each warp's recurrence follow logical positions, never physical addresses.
+    // in the paged specialization KV_max carries [count, physical page IDs...] per query; the loop
+    // and each warp's recurrence follow logical positions, never physical addresses
     static_assert(!paged || ncols == 1, "paged attention has one query per block");
     const int * pages = paged ? KV_max + (sequence*int(ne01.z) + ic0)*(1 + ne11/FATTN_KQ_STRIDE) : nullptr;
     const int k_VKQ_max = paged ? pages[0]*FATTN_KQ_STRIDE : (KV_max ? KV_max[sequence*gridDim.x + blockIdx.x] : ne11);
diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu
index eff18212272..915c2d04da8 100644
--- a/ggml/src/ggml-cuda/fattn.cu
+++ b/ggml/src/ggml-cuda/fattn.cu
@@ -457,9 +457,8 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const
     // 192 satisfies % 64 == 0 but has no vec instance (DKQ != DV); force it onto the MMA path.
     const bool can_use_vector_kernel = Q->ne[0] <= 256 && Q->ne[0] % 64 == 0 && Q->ne[0] != 192 && K->ne[1] % FATTN_KQ_STRIDE == 0;
 
-    // [TAG_BATCH_INVARIANT] Every choice below switches on Q->ne[1] or on K->ne[1], and both
-    // grow with the other sequences in the batch and in the shared KV cache. Pin the kernel a
-    // batch of one would use so a request is never moved onto a different algorithm by its neighbours.
+    // [TAG_BATCH_INVARIANT] every choice below switches on Q->ne[1] or K->ne[1], both of which grow
+    // with the other sequences, so pin the kernel a batch of one would use
     if (ggml_cuda_batch_invariant() && can_use_vector_kernel && Q->ne[1] == 1) {
         return BEST_FATTN_KERNEL_VEC;
     }
@@ -592,7 +591,7 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
         return;
     }
 
-    // [TAG_BATCH_INVARIANT] Attend one query row at a time, as a batch of one would.
+    // [TAG_BATCH_INVARIANT] attend one query row at a time, as a batch of one would
     const int fattn_max_cols = ggml_cuda_batch_invariant_max_cols();
     if (ggml_cuda_batch_invariant() && dst->src[0]->ne[1] > 1 && dst->src[0]->ne[3] == 1 &&
             (fattn_max_cols <= 0 || dst->src[0]->ne[1] <= fattn_max_cols)) {
@@ -606,8 +605,7 @@ void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst
 
             ggml_tensor mask_row;
             ggml_tensor dst_row = *dst;
-            // ne[2] keeps running to the end of dst so that the scratch space for F16 copies of
-            // K and V, which is placed right behind dst, is still put in the same place.
+            // ne[2] runs to the end of dst so the F16 K/V scratch behind dst stays in place
             dst_row.ne[2] = dst->ne[2] - i;
             dst_row.data  = (char *) dst->data + i*dst->nb[2];
             dst_row.src[0] = &Q_row;
diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu
index 8c5b2a408c3..76fa4dc66c6 100644
--- a/ggml/src/ggml-cuda/ggml-cuda.cu
+++ b/ggml/src/ggml-cuda/ggml-cuda.cu
@@ -1758,8 +1758,8 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up,
 }
 
 static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) {
-    // [TAG_BATCH_INVARIANT] mul_mat+GLU is only fused for a single destination column, so
-    // leaving it on would give a solo request a different code path from a batched one.
+    // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving
+    // it on would give a solo request a different code path from a batched one
     if (ggml_cuda_batch_invariant()) {
         return false;
     }
@@ -1791,8 +1791,8 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) {
 }
 
 static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
-    // [TAG_BATCH_INVARIANT] mul_mat+GLU is only fused for a single destination column, so
-    // leaving it on would give a solo request a different code path from a batched one.
+    // [TAG_BATCH_INVARIANT] mul_mat+GLU is fused for a single destination column only, so leaving
+    // it on would give a solo request a different code path from a batched one
     if (ggml_cuda_batch_invariant()) {
         return false;
     }
@@ -1825,16 +1825,11 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) {
     return use_mul_mat_vec_q;
 }
 
-// [TAG_BATCH_INVARIANT]
-// The number of tokens in a batch picks both the matmul implementation below and, inside
-// several of them, how the K loop is divided between threads. Both change the order in
-// which the partial products of one destination element are summed, so the same request
-// produces different bits depending on how many other requests decode alongside it.
-//
-// GGML_CUDA_BATCH_INVARIANT removes that dependency:
-//   1 - compute every destination column on its own, exactly as a batch of one would.
-//   2 - split off only the columns whose batch-of-one configuration differs from the
-//       batched one, leaving the already invariant matmuls batched.
+// [TAG_BATCH_INVARIANT] the token count picks the matmul implementation and how its K loop is
+// divided between threads, both of which change the summation order, so the same request produces
+// different bits depending on how many others decode alongside it. GGML_CUDA_BATCH_INVARIANT:
+//   1 - compute every destination column on its own, exactly as a batch of one would
+//   2 - split off only the columns whose batch-of-one configuration differs from the batched one
 static bool ggml_cuda_exact_concurrency() {
     static const bool exact = []() {
         const char * value = getenv("LLAMA_EXACT_CONCURRENCY");
@@ -1864,9 +1859,8 @@ void ggml_backend_cuda_set_exact_decode_width(int n_cols) {
 }
 
 int ggml_cuda_batch_invariant_max_cols() {
-    // [TAG_EXACT_CONCURRENCY] With prompt ubatches kept to one sequence, a sequence's prefill
-    // matmul shapes match its solo run, so exact mode no longer needs the column policy to be
-    // unbounded there. An explicit bound always wins, in either mode.
+    // [TAG_EXACT_CONCURRENCY] prompt ubatches hold one sequence, so a prefill already matches its
+    // solo run and needs no unbounded column policy. An explicit bound always wins.
     static const int explicit_cols = []() {
         const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
         return val ? atoi(val) : -1;
@@ -1880,24 +1874,19 @@ int ggml_cuda_batch_invariant_max_cols() {
         return 0;
     }
 
-    // Exact mode only has to cover the widest ubatch a decode step can build: one column per slot,
-    // times one plus the number of speculative draft tokens carried with it. Use that width when
-    // the caller reported it through ggml_backend_cuda_set_exact_decode_width(). Nothing reported
-    // it, so fall back to 16, which covers four slots at up to three tokens each, which is what
-    // --parallel 4 --spec-type draft-mtp --spec-draft-n-max 2 produces. Above the bound the column
-    // split does not fire, and ggml_cuda_warn_above_exact_bound() says so once.
+    // the widest ubatch a decode step can build: one column per slot times one plus its draft
+    // tokens, as reported by ggml_backend_cuda_set_exact_decode_width(). Failing a report, 16,
+    // which covers --parallel 4 --spec-type draft-mtp --spec-draft-n-max 2. Above the bound the
+    // column split does not fire and ggml_cuda_warn_above_exact_bound() says so once.
     const int width = g_exact_decode_width.load(std::memory_order_relaxed);
 
     return width > 0 ? width : 16;
 }
 
-// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so its rows depend on the
-// other rows in it and the mode does not hold for that op. Say so once, rather than never.
-//
-// Only when nothing reported a decode width. When one was reported the bound is derived from it, so
-// the only batches above the bound are prompt ubatches, and those hold a single sequence under this
-// mode: their exactness comes from that, not from the column policy, and leaving them batched is
-// the whole point of having a bound at all. Warning on those would be crying wolf on every prefill.
+// [TAG_EXACT_CONCURRENCY] a batch wider than the bound is left batched, so the mode does not hold
+// for that op; say so once. Only when nothing reported a decode width: a reported one makes the
+// batches above the bound prompt ubatches, which are exact by holding a single sequence, so
+// warning on those would be crying wolf on every prefill.
 static void ggml_cuda_warn_above_exact_bound(const char * op, int64_t ncols, int max_cols) {
     if (!ggml_cuda_exact_concurrency()) {
         return;
@@ -1929,7 +1918,7 @@ enum ggml_cuda_mm_path {
     GGML_CUDA_MM_CUBLAS,
 };
 
-// The implementation ggml_cuda_mul_mat would pick for a batch of ne11 columns.
+// the implementation ggml_cuda_mul_mat would pick for a batch of ne11 columns
 static ggml_cuda_mm_path ggml_cuda_mul_mat_path(
         int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst, int64_t ne11) {
     // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q.
@@ -1966,14 +1955,10 @@ static ggml_cuda_mm_path ggml_cuda_mul_mat_path(
 
 static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst);
 
-// [TAG_BATCH_INVARIANT]
-// The widest slice of columns that can be recomputed in one launch while every column in it still
-// sums the way a batch of one would. A column's result depends on the implementation and, for
-// MMVQ, on the warp count of the launch, and neither depends on the values of the other columns
-// in the launch, so a slice as wide as the batch-of-one configuration reaches gives each of its
-// columns the batch-of-one value while reading the weights once for all of them instead of once
-// per column. A twelve-column speculative decode over a table whose configuration holds up to four
-// columns then costs three weight reads rather than twelve. Always below ncols_dst, so the
+// [TAG_BATCH_INVARIANT] the widest slice of columns that can be recomputed in one launch while
+// every column still sums as a batch of one would. A column's result depends on the implementation
+// and, for MMVQ, the launch's warp count, never on the other columns, so a slice this wide reads
+// the weights once for all of them instead of once per column. Always below ncols_dst, so the
 // recursive call cannot land back here with the same shape.
 static int64_t ggml_cuda_mul_mat_invariant_width(
         int cc, int warp_size, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst,
@@ -1994,14 +1979,13 @@ static int64_t ggml_cuda_mul_mat_invariant_width(
     return 1;
 }
 
-// Recompute dst in slices of columns so that each column sees the batch-of-one configuration.
-// Returns false when the batched launch already gives every column that same value.
+// recompute dst in slices of columns so each column sees the batch-of-one configuration; false
+// when the batched launch already gives every column that same value
 static bool ggml_cuda_mul_mat_split_columns(
         ggml_backend_cuda_context & ctx, int cc, int warp_size,
         const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
-    // Recurrent-model output projections broadcast one weight matrix over sequence
-    // planes. These are token projections too, even though ne[2] or ne[3] is > 1.
-    // Normalize each plane before applying the existing selective column policy.
+    // recurrent output projections broadcast one weight matrix over sequence planes: these are
+    // token projections too, so normalize each plane before applying the column policy
     if (ggml_cuda_exact_concurrency() && src0->ne[2] == 1 && src0->ne[3] == 1 &&
             (dst->ne[2] > 1 || dst->ne[3] > 1) &&
             src1->ne[2] == dst->ne[2] && src1->ne[3] == dst->ne[3]) {
@@ -2023,7 +2007,7 @@ static bool ggml_cuda_mul_mat_split_columns(
     if (ncols_dst <= 1 || src1->ne[1] != ncols_dst) {
         return false;
     }
-    // Only the token dimension is split, batched matmuls (attention) keep their shape.
+    // only the token dimension is split; batched matmuls (attention) keep their shape
     if (src1->ne[2] != 1 || src1->ne[3] != 1 || dst->ne[2] != 1 || dst->ne[3] != 1) {
         return false;
     }
@@ -2033,14 +2017,14 @@ static bool ggml_cuda_mul_mat_split_columns(
         return false;
     }
 
-    // Mode 1 recomputes one column at a time. Mode 2 recomputes in the widest slices that keep the
-    // batch-of-one arithmetic, which is what the exact concurrency mode runs under.
+    // mode 1 recomputes one column at a time; mode 2, which exact concurrency runs under, uses the
+    // widest slices that keep the batch-of-one arithmetic
     int64_t width = 1;
     if (ggml_cuda_batch_invariant() >= 2) {
         const ggml_cuda_mm_path path_one     = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, 1);
         const ggml_cuda_mm_path path_batched = ggml_cuda_mul_mat_path(cc, warp_size, src0, src1, dst, ncols_dst);
         if (path_one == path_batched) {
-            // Same implementation, but it still has to sum each destination element in the same order.
+            // same implementation, but it still has to sum in the same order
             if (path_batched == GGML_CUDA_MM_MMVF) {
                 return false; // the block size follows K alone
             }
@@ -2122,10 +2106,8 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor
     GGML_ABORT("fatal error");
 }
 
-// [TAG_BATCH_INVARIANT]
-// True when the batch-invariant policy computes this MUL_MAT_ID one token at a time.
-// Every expert product then reduces the way it would in a batch of one, whatever the
-// rest of the ubatch routed to.
+// [TAG_BATCH_INVARIANT] true when the policy computes this MUL_MAT_ID one token at a time, so
+// every expert product reduces as it would in a batch of one
 static bool ggml_cuda_mul_mat_id_splits_tokens(const ggml_tensor * dst) {
     if (!ggml_cuda_batch_invariant()) {
         return false;
@@ -2152,8 +2134,8 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
         return true;
     }
 
-    // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path
-    // that decides whether the stream is synchronized is the single-token one.
+    // [TAG_BATCH_INVARIANT] a split node runs as ntokens single-token calls, so the path that
+    // decides whether the stream is synchronized is the single-token one
     const int64_t ntokens = ggml_cuda_mul_mat_id_splits_tokens(dst) ? 1 : dst->ne[2];
 
     if (ntokens <= MMVQ_MAX_BATCH_SIZE) {
@@ -2179,12 +2161,9 @@ static bool ggml_cuda_mul_mat_id_needs_sync(const ggml_tensor * dst, const int c
 
 static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst);
 
-// [TAG_BATCH_INVARIANT]
-// Recompute dst one token at a time. Every implementation below groups the ubatch's tokens
-// by the expert they routed to, so the column count of an expert's matmul, the tokens the
-// per-expert copy gathers and the width the activations are quantized at all depend on what
-// the other tokens in the ubatch picked. Handing each token its own call removes that: the
-// callee sees the shapes a batch of one has, whatever the neighbours did.
+// [TAG_BATCH_INVARIANT] recompute dst one token at a time. Every implementation below groups the
+// ubatch's tokens by the expert they routed to, so shapes depend on what the other tokens picked;
+// one call per token makes the callee see the shapes a batch of one has.
 static void ggml_cuda_mul_mat_id_split_tokens(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
     const ggml_tensor * src1 = dst->src[1];
     const ggml_tensor * ids  = dst->src[2];
@@ -2232,9 +2211,8 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor *
     // [TAG_BATCH_INVARIANT]
     if (ggml_cuda_mul_mat_id_splits_tokens(dst)) {
         GGML_ASSERT(ne3 == 1 && src1->ne[3] == 1 && ids->ne[2] == 1 && ids->ne[3] == 1);
-        // A quantized expert matrix takes the single-token MMVQ path for every token count, and
-        // that path can put the tokens on its sample axis in one launch rather than being
-        // re-entered once per token. Anything else is still recomputed one token at a time.
+        // a quantized expert matrix takes the single-token MMVQ path at every token count, and that
+        // path can put the tokens on its sample axis in one launch; anything else goes token by token
         if (ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
             ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst);
             return;
@@ -3629,10 +3607,9 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph
     }
 
     // topk-moe
-    // [TAG_BATCH_INVARIANT] The routing fusion passes its memory-range check only when the ubatch
-    // holds one token, so a request decoding alone picks the fused warp-local top-k kernel and the
-    // same request decoding next to neighbours picks the softmax, argsort and normalize chain.
-    // Two algorithms for one set of routing weights is the batch dependence this mode removes.
+    // [TAG_BATCH_INVARIANT] the routing fusion passes its memory-range check only for a one-token
+    // ubatch, so a solo request takes the fused top-k kernel and a batched one takes the softmax,
+    // argsort and normalize chain: two algorithms for one set of routing weights
     if (!ggml_cuda_batch_invariant() &&
             (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX ||
              cgraph->nodes[i]->op == GGML_OP_ARGSORT)) {
diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu
index b4e5196f172..30f37c7722e 100644
--- a/ggml/src/ggml-cuda/mmvq.cu
+++ b/ggml/src/ggml-cuda/mmvq.cu
@@ -551,8 +551,8 @@ bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t n
         // There nwarps also depends on the K loop trip count, which the caller does not pass in.
         return ncols_dst == 1;
     }
-    // blocks_per_iter, which is what assigns K blocks to threads, is proportional to nwarps.
-    // rows_per_cuda_block only changes which rows a block owns, not the order within a row.
+    // blocks_per_iter, which assigns K blocks to threads, is proportional to nwarps;
+    // rows_per_cuda_block only changes which rows a block owns, not the order within a row
     return calc_nwarps(type, 1, table_id) == calc_nwarps(type, (int) ncols_dst, table_id);
 }
 
@@ -594,9 +594,8 @@ static __global__ void mul_mat_vec_q(
     ggml_cuda_pdl_sync();
     sample_dst = blockIdx.z;
     // [TAG_BATCH_INVARIANT] with ids, a sample is a token: the batch-invariant MUL_MAT_ID launch
-    // puts every token of the batch on the z axis of one single-column launch, so each (token,
-    // expert slot) block runs the exact single-token configuration. The stock single-token launch
-    // has one sample, where this indexing is ids[channel_dst] as before.
+    // puts every token on the z axis of one single-column launch, so each (token, expert slot)
+    // block runs the single-token configuration. The stock launch has one sample, as before.
     channel_x  = ncols_dst == 1 && ids ? ids[sample_dst*ids_stride + channel_dst] : fastdiv(channel_dst, channel_ratio);
     channel_y  = ncols_dst == 1 && ids ? fastmodulo(channel_dst, nchannels_y)      : channel_dst;
 
@@ -1285,10 +1284,9 @@ void ggml_cuda_mul_mat_vec_q(
     GGML_ASSERT(        nb0        == ts_dst);
     GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type));
 
-    // [TAG_BATCH_INVARIANT] under the knob a MUL_MAT_ID with several tokens is computed as one
-    // launch of the single-token configuration with the tokens on the sample axis, so every
-    // (token, expert slot) block reduces exactly as the token alone would. The token count is
-    // then not bounded by the column templates.
+    // [TAG_BATCH_INVARIANT] a multi-token MUL_MAT_ID becomes one launch of the single-token
+    // configuration with the tokens on the sample axis, so every (token, expert slot) block
+    // reduces as the token alone would and the count is not bounded by the column templates
     const bool tokens_as_samples = ids && ne2 > 1 && ggml_cuda_batch_invariant();
 
     GGML_ASSERT(!ids || ne12 <= MMVQ_MAX_BATCH_SIZE || tokens_as_samples);
diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh
index 61a88b851ec..67d69b1415d 100644
--- a/ggml/src/ggml-cuda/mmvq.cuh
+++ b/ggml/src/ggml-cuda/mmvq.cuh
@@ -4,9 +4,8 @@
 
 bool ggml_cuda_should_use_mmvq(enum ggml_type type, int cc, int64_t ne11);
 
-// [TAG_BATCH_INVARIANT]
-// True when an MMVQ launch of ncols_dst columns sums each destination element in the same
-// order as a launch of a single column, i.e. when the column count leaves nwarps unchanged.
+// [TAG_BATCH_INVARIANT] true when an MMVQ launch of ncols_dst columns sums each destination element
+// in the same order as a single-column launch, i.e. when the column count leaves nwarps unchanged
 bool ggml_cuda_mmvq_matches_single_column(enum ggml_type type, int cc, int64_t ncols_dst);
 
 // Returns the maximum batch size for which MMVQ should be used for MUL_MAT_ID,
diff --git a/ggml/src/ggml-et/ggml-et.cpp b/ggml/src/ggml-et/ggml-et.cpp
index a3792f3852b..b787108a7aa 100644
--- a/ggml/src/ggml-et/ggml-et.cpp
+++ b/ggml/src/ggml-et/ggml-et.cpp
@@ -1266,8 +1266,7 @@ static bool ggml_backend_et_device_supports_op(ggml_backend_dev_t dev, const ggm
                         (op->src[1]->ne[1] % op->src[4]->ne[1] == 0);
             break;
         case GGML_OP_FLASH_ATTN_EXT:
-            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
-            // the CUDA backend reads
+            // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads
             if (op->src[5]) {
                 supported = false;
                 break;
diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
index aa20083ec05..574276fed15 100644
--- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp
+++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp
@@ -4157,8 +4157,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons
             break;
 
         case GGML_OP_FLASH_ATTN_EXT:
-            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
-            // the CUDA backend reads
+            // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads
             supp = op->src[5] == nullptr && ggml_hexagon_supported_flash_attn_ext(sess, op);
             break;
 
diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m
index 90873f2fab0..e5fc6a8063b 100644
--- a/ggml/src/ggml-metal/ggml-metal-device.m
+++ b/ggml/src/ggml-metal/ggml-metal-device.m
@@ -1592,8 +1592,8 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
         case GGML_OP_ROLL:
             return ggml_is_contiguous(op->src[0]);
         case GGML_OP_FLASH_ATTN_EXT:
-            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the
-            // CUDA backend reads; walking the pool in physical order here would be silently wrong
+            // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads;
+            // walking the pool in physical order here would be silently wrong
             if (op->src[5] != NULL) {
                 return false;
             }
diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp
index effd11714f9..9c2362c8c4b 100644
--- a/ggml/src/ggml-opencl/ggml-opencl.cpp
+++ b/ggml/src/ggml-opencl/ggml-opencl.cpp
@@ -7842,8 +7842,7 @@ static bool ggml_opencl_supports_op(ggml_backend_dev_t dev, const struct ggml_te
         case GGML_OP_MEAN:
             return op->src[0]->type == GGML_TYPE_F32;
         case GGML_OP_FLASH_ATTN_EXT: {
-            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the
-            // CUDA backend reads
+            // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads
             if (op->src[5]) {
                 return false;
             }
diff --git a/ggml/src/ggml-openvino/ggml-openvino.cpp b/ggml/src/ggml-openvino/ggml-openvino.cpp
index dfc9f90926f..1cec1583d51 100644
--- a/ggml/src/ggml-openvino/ggml-openvino.cpp
+++ b/ggml/src/ggml-openvino/ggml-openvino.cpp
@@ -1128,8 +1128,7 @@ static bool is_op_unsupported_case(const ggml_tensor * op) {
         break;
     }
     case GGML_OP_FLASH_ATTN_EXT: {
-        // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
-        // the CUDA backend reads
+        // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads
         if (op->src[5]) {
             return true;
         }
diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp
index 6fb8851a904..51b552acf75 100644
--- a/ggml/src/ggml-rpc/ggml-rpc.cpp
+++ b/ggml/src/ggml-rpc/ggml-rpc.cpp
@@ -1915,8 +1915,8 @@ static ggml_backend_buffer_type_t ggml_backend_rpc_device_get_buffer_type(ggml_b
 
 static bool ggml_backend_rpc_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) {
     GGML_UNUSED(dev);
-    // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the
-    // CUDA backend reads; the remote end is not asked, so it is not claimed here
+    // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads;
+    // the remote end is not asked, so it is not claimed here
     if (op->op == GGML_OP_FLASH_ATTN_EXT && op->src[5]) {
         return false;
     }
diff --git a/ggml/src/ggml-sycl/ggml-sycl.cpp b/ggml/src/ggml-sycl/ggml-sycl.cpp
index 69a344ab790..a31a6a41ca8 100644
--- a/ggml/src/ggml-sycl/ggml-sycl.cpp
+++ b/ggml/src/ggml-sycl/ggml-sycl.cpp
@@ -6342,8 +6342,7 @@ static bool do_ggml_backend_sycl_device_supports_op(ggml_backend_dev_t dev, cons
         case GGML_OP_SOLVE_TRI:
             return op->src[0]->ne[0] <= SYCL_SOLVE_TRI_MAX_N && op->src[1]->ne[0] <= SYCL_SOLVE_TRI_MAX_K;
         case GGML_OP_FLASH_ATTN_EXT:
-            // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only the
-            // CUDA backend reads
+            // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads
             return op->src[5] == nullptr && ggml_sycl_flash_attn_ext_supported(device, op);
         default:
             return false;
diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
index 4a2347219b3..eb855f24d17 100644
--- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp
+++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp
@@ -18192,8 +18192,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
             }
         case GGML_OP_FLASH_ATTN_EXT:
             {
-                // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
-                // the CUDA backend reads
+                // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads
                 if (op->src[5]) {
                     return false;
                 }
diff --git a/ggml/src/ggml-webgpu/ggml-webgpu.cpp b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
index 70462a97f3c..26717d8804d 100644
--- a/ggml/src/ggml-webgpu/ggml-webgpu.cpp
+++ b/ggml/src/ggml-webgpu/ggml-webgpu.cpp
@@ -4408,8 +4408,7 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
             break;
         case GGML_OP_FLASH_ATTN_EXT:
             {
-                // [TAG_EXACT_CONCURRENCY] src[5] is the exact-concurrency page table, which only
-                // the CUDA backend reads
+                // [TAG_EXACT_CONCURRENCY] src[5] is the page table, which only the CUDA backend reads
                 if (op->src[5]) {
                     supports_op = false;
                     break;
diff --git a/include/llama.h b/include/llama.h
index 09c331ff3e1..077f17ecabf 100644
--- a/include/llama.h
+++ b/include/llama.h
@@ -795,34 +795,26 @@ extern "C" {
     // Check if the memory supports shifting
     LLAMA_API bool llama_memory_can_shift(llama_memory_t mem);
 
-    // [TAG_EXACT_CONCURRENCY] Cells the memory allocates in one indivisible unit.
-    //
-    // 1 in every ordinary configuration. Larger where a mode places cells in blocks, and
-    // then a sequence of n tokens occupies round_up(n, granularity) cells. A caller that
-    // decides whether the pool has room by counting tokens has to round the same way, or it
-    // will believe there is space that cannot be handed out.
+    // [TAG_EXACT_CONCURRENCY] cells the memory allocates in one indivisible unit: 1 ordinarily,
+    // larger where a mode places cells in blocks, and then n tokens occupy round_up(n, granularity)
+    // cells. A caller deciding whether the pool has room must round the same way.
     LLAMA_API uint32_t llama_memory_alloc_granularity(llama_memory_t mem);
 
-    // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or
-    // 1 plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence
-    // set with more tokens than this left to place is a prompt and is prefilled in a ubatch of its
-    // own; a set at or below it is a decode step and stays grouped with the other decodes, so a
-    // speculative verify batch is not run once per sequence. Process-wide, default 1. Raising it
-    // widens the decode step of every context that exists, and their width is re-reported with
-    // it; false, and no change, when an explicit column bound given to the backend cannot cover
-    // that width (see llama_set_exact_decode_width). Never lowers what was set: a narrower context
-    // set up later must not turn an existing context's verify steps into prompts.
+    // [TAG_EXACT_CONCURRENCY] the most tokens one sequence contributes to a decode step: 1, or 1
+    // plus the draft length under speculative decoding. Under LLAMA_EXACT_CONCURRENCY a sequence
+    // set with more left to place is a prompt and is prefilled in a ubatch of its own; one at or
+    // below stays grouped with the other decodes. Process-wide, default 1, never lowered. Raising
+    // it widens every existing context's decode step and re-reports their width; returns false,
+    // and changes nothing, when an explicit column bound cannot cover that width.
     LLAMA_API bool     llama_set_exact_decode_tokens(uint32_t n_tokens);
     LLAMA_API uint32_t llama_exact_decode_tokens(void);
 
     // [TAG_EXACT_CONCURRENCY] the widest decode ubatch this process can build, in columns: the
-    // sequences a context can hold times the tokens each contributes to a step. Every context
-    // reports its own at creation and a backend keeps the widest it has heard, so a decode of any
-    // context stays within the bound its kernels split at. A caller that builds wider steps than
-    // the contexts imply (a draft of its own, say) reports the width itself, before creating the
-    // context or before the first decode. Never lowers what was reported. Returns false, and
-    // reports nothing, when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is set to a positive figure below
-    // the width: that bound wins in the backend, so decodes above it would be left batched.
+    // sequences a context holds times the tokens each contributes. Every context reports its own
+    // at creation and a backend keeps the widest it has heard. A caller that builds wider steps
+    // reports the width itself, before the context or the first decode. Never lowered. Returns
+    // false, reporting nothing, when GGML_CUDA_BATCH_INVARIANT_MAX_COLS is positive and below the
+    // width: that bound wins in the backend, so decodes above it would be left batched.
     LLAMA_API bool     llama_set_exact_decode_width(uint32_t n_cols);
     LLAMA_API uint32_t llama_exact_decode_width(void);
 
diff --git a/scripts/batchinv/divergence.py b/scripts/batchinv/divergence.py
index ec2c3a00467..198c2bdf881 100644
--- a/scripts/batchinv/divergence.py
+++ b/scripts/batchinv/divergence.py
@@ -5,8 +5,8 @@
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 from prompts import PROMPTS
 
-# Environment recorded with every run. LLAMA_EXACT_CONCURRENCY inherited from the shell is what
-# decides whether a run labelled as the mode-off reference actually was one, so it is not optional.
+# recorded with every run; LLAMA_EXACT_CONCURRENCY inherited from the shell decides whether a run
+# labelled as the mode-off reference actually was one, so it is not optional
 RECORDED_ENV = ("LLAMA_EXACT_CONCURRENCY", "GGML_CUDA_BATCH_INVARIANT",
                 "GGML_CUDA_BATCH_INVARIANT_MAX_COLS", "LLAMA_SERVER_PREEMPT_EVERY",
                 "LLAMA_KV_CACHE_DEBUG", "LLAMA_BATCH_DEBUG", "CUDA_VISIBLE_DEVICES")
@@ -86,14 +86,13 @@ def __enter__(self):
                     time.sleep(1.0)
             raise RuntimeError("server did not become healthy")
         except BaseException:
-            # __exit__ is not called when __enter__ raises, so a server that started but never
+            # __exit__ is not called when __enter__ raises, and a server that started but never
             # reported healthy would keep the GPU, the port and the log handle
             self.__exit__(None, None, None)
             raise
 
     def __exit__(self, *a):
-        # note: POSIX only. On Windows this needs CREATE_NEW_PROCESS_GROUP at Popen and
-        # terminate()/kill() here; the runs this harness backs are Linux only.
+        # note: POSIX only; Windows would need CREATE_NEW_PROCESS_GROUP at Popen
         if self.p is not None:
             print(f"[server] stopping pid={self.p.pid}", flush=True)
             try:
@@ -141,8 +140,7 @@ def work(name):
         t.join()
     wall = time.time() - t0
 
-    # a thread exception used to only print a traceback, so a run where P1..P3 failed and P0
-    # succeeded was still reported as a clean four-way concurrency result
+    # without this a run where P1..P3 failed and P0 succeeded reads as a clean four-way result
     if errors:
         raise RuntimeError("concurrent requests failed: " +
                            "; ".join(f"{n}: {type(e).__name__}: {e}" for n, e in errors))
diff --git a/scripts/batchinv/probe.cpp b/scripts/batchinv/probe.cpp
index c4e478f97c7..204379c423b 100644
--- a/scripts/batchinv/probe.cpp
+++ b/scripts/batchinv/probe.cpp
@@ -1,6 +1,6 @@
-// Locate the first graph op whose sequence-0 output changes when the decode batch
-// holds four sequences instead of one. Prompt KV for seq 0 is built identically in
-// both phases, so the only difference is the width of the final decode ubatch.
+// Locate the first graph op whose sequence-0 output changes when the decode batch holds four
+// sequences instead of one. Seq 0's prompt KV is identical in both phases, so the only
+// difference is the width of the final decode ubatch.
 #include "llama.h"
 #include "ggml.h"
 #include "ggml-backend.h"
@@ -120,7 +120,7 @@ static llama_token greedy(llama_context * ctx, int32_t i, int n_vocab) {
     return best;
 }
 
-// Feed a prompt as one decode call for one sequence, return the greedy next token.
+// feed a prompt as one decode call for one sequence, return the greedy next token
 static llama_token feed(llama_context * ctx, const std::vector<llama_token> & p, llama_seq_id seq, int n_vocab) {
     batch_holder h;
     for (size_t i = 0; i < p.size(); ++i) {
@@ -167,14 +167,12 @@ int main(int argc, char ** argv) {
     std::vector<node_rec> rec_a, rec_b;
     llama_token first_tok[4] = {0, 0, 0, 0};
 
-    // Phase A: decode ubatch width 1. PROBE_A_FILL controls how many sequences are
-    // already in the shared KV cache, which is what sets K->ne[1] for attention.
+    // phase A: decode ubatch width 1. PROBE_A_FILL is how many sequences are already in the
+    // shared KV cache, which sets K->ne[1] for attention.
     const int a_fill = getenv("PROBE_A_FILL") ? atoi(getenv("PROBE_A_FILL")) : 1;
-    // PROBE_A_PERM reorders which prompt goes into which sequence in phase A. With the same
-    // multiset of prompts the cache keeps its length but the masked cells hold different data.
-    // Phase B decodes prompt 0's first token on sequence 0, so the permutation may only move
-    // the neighbours: sequence 0 keeps prompt 0, or the two phases would compare different
-    // sequences.
+    // PROBE_A_PERM reorders which prompt goes into which sequence in phase A, keeping the cache
+    // length but changing what the masked cells hold. It may only move the neighbours: sequence 0
+    // keeps prompt 0, or the two phases would compare different sequences.
     int a_perm[4] = {0, 1, 2, 3};
     if (const char * perm = getenv("PROBE_A_PERM")) {
         for (int k = 0; k < 4 && perm[2*k]; ++k) a_perm[k] = perm[2*k] - '0';
@@ -197,7 +195,7 @@ int main(int argc, char ** argv) {
         llama_free(ctx);
     }
 
-    // Phase B: same seq-0 prompt KV, then a decode ubatch holding n_seqs tokens.
+    // phase B: same seq-0 prompt KV, then a decode ubatch holding n_seqs tokens
     {
         llama_context * ctx = make_ctx();
         if (prefill) {
@@ -240,7 +238,7 @@ int main(int argc, char ** argv) {
         llama_free(ctx);
     }
 
-    // Optional: keep decoding and report the first step at which seq 0's token differs.
+    // optional: keep decoding and report the first step at which seq 0's token differs
     const int n_steps = getenv("PROBE_STEPS") ? atoi(getenv("PROBE_STEPS")) : 0;
     int first_bad_step = -1;
     if (n_steps > 0) {
@@ -279,7 +277,7 @@ int main(int argc, char ** argv) {
     fprintf(stderr, "nodes: A=%zu B=%zu  first tokens: %d %d %d %d\n",
             rec_a.size(), rec_b.size(), first_tok[0], first_tok[1], first_tok[2], first_tok[3]);
 
-    // Walk both node lists in order and compare seq 0's slice.
+    // walk both node lists in order and compare seq 0's slice
     FILE * out = out_path ? fopen(out_path, "w") : stdout;
     fprintf(out, "{\"n_seqs\":%d,\"first_bad_step\":%d,\"nodes_a\":%zu,\"nodes_b\":%zu,\"diffs\":[", n_seqs, first_bad_step, rec_a.size(), rec_b.size());
     size_t n = rec_a.size() < rec_b.size() ? rec_a.size() : rec_b.size();
@@ -295,8 +293,8 @@ int main(int argc, char ** argv) {
             verdict = "misaligned";
         } else if (A.op == "GATED_DELTA_NET" && A.gdn_tokens == B.gdn_tokens &&
                 !A.data.empty() && !B.data.empty()) {
-            // Packed GDN outputs put ALL token outputs before ALL sequence states.
-            // Sequence 0's state therefore moves when the number of sequences changes.
+            // packed GDN outputs put all token outputs before all sequence states, so seq 0's
+            // state moves when the number of sequences changes
             const size_t output = A.ne[0]*A.gdn_tokens;
             const size_t state = A.ne[0]*A.ne[1]/A.gdn_seqs - output;
             for (size_t k = 0; k < output + state; ++k) {
diff --git a/scripts/batchinv/prompts.py b/scripts/batchinv/prompts.py
index 860b65fe08d..efabc448a9d 100644
--- a/scripts/batchinv/prompts.py
+++ b/scripts/batchinv/prompts.py
@@ -1,4 +1,4 @@
-# Four distinct prompts, each about 300 tokens of raw text (no chat template).
+# four distinct prompts, each about 300 tokens of raw text (no chat template)
 _BODIES = {
 "P0": """The history of numerical computing is a history of compromises between speed and exactness.
 Early machines used fixed point arithmetic because it was cheap, and programmers carried scaling
diff --git a/scripts/unsloth/additive_merge.py b/scripts/unsloth/additive_merge.py
index 5364dc1b7c1..d930158a542 100644
--- a/scripts/unsloth/additive_merge.py
+++ b/scripts/unsloth/additive_merge.py
@@ -94,18 +94,11 @@ def nonblank(lines: list[str]) -> list[str]:
     return [ln.strip() for ln in lines if ln.strip()]
 
 
-# A line that closes or opens a block and nothing else. Two INDEPENDENT case
-# arms in the same switch share these by construction -- `{`, `} break;`, `}`
-# are what a case arm is made of, not what makes it that case arm -- so finding
-# them on both sides says nothing about whether the two sides added the same
-# construct. Matching them as "shared" is what refused the real add/add of
-# PROJECTOR_TYPE_KIMIK3 next to PROJECTOR_TYPE_DEEPSEEK4V in tools/mtmd/clip.cpp
-# with "one change made twice: {, } break;", when the two arms had no line of
-# actual content in common.
-#
-# Deliberately narrow: braces, brackets, parens, semicolons and commas, around
-# at most one bare block-terminating keyword. `break;` matches, `return true;`
-# does not, and anything naming a type, a constant or a function does not.
+# A line that only opens or closes a block. Two independent case arms share these by
+# construction, so finding them on both sides says nothing about the two sides adding the
+# same construct: treating them as shared is what refused the real PROJECTOR_TYPE_KIMIK3 /
+# PROJECTOR_TYPE_DEEPSEEK4V add/add in tools/mtmd/clip.cpp. Deliberately narrow: brackets,
+# semicolons and commas around at most one bare block-terminating keyword.
 STRUCTURAL = re.compile(r"^[\s{}()\[\];,]*(?:break|continue|return|pass)?[\s{}()\[\];,]*$")
 
 
@@ -114,8 +107,7 @@ def identifying(lines: list[str]) -> set[str]:
     return {ln for ln in nonblank(lines) if not STRUCTURAL.match(ln)}
 
 
-# `case FOO:`, `case FOO :`, `default:`. A fallthrough label may carry no body
-# at all, which is the shape the nightly hits most often.
+# `case FOO:`, `case FOO :`, `default:`; a fallthrough label may carry no body at all
 CASE_LABEL = re.compile(r"^(?:case\s+[^:]+|default\s*):")
 
 
@@ -147,32 +139,23 @@ def resolve_region(ours: list[str], base: list[str], theirs: list[str]) -> list[
         return list(ours)
     ours_arms, theirs_arms = case_arms(ours), case_arms(theirs)
     if ours_arms and theirs_arms and ours_arms.isdisjoint(theirs_arms):
-        # Both sides added case arms, and not one label is on both sides. Two
-        # arms of the same switch labelled differently are two constructs, so
-        # any line they happen to share is body text, not a duplicate: the real
-        # tools/mtmd/clip.cpp collision has a KIMIK3 arm and a DEEPSEEK4V arm
-        # that both set `hparams.rope_theta = 10000.0f;`, and refusing on that
-        # coincidence is what the shared-line check is for, backwards.
-        #
-        # The same change made twice would keep its label, so it lands in the
-        # check below instead. This is the one place where a shared line is
-        # allowed, and it is allowed because the labels prove the arms are
-        # distinct -- a duplicated label would not even compile.
+        # Both sides added case arms and no label is on both, so they are two constructs
+        # and any line they share is body text: the real clip.cpp collision has arms that
+        # both set `hparams.rope_theta = 10000.0f;`. The same change made twice would keep
+        # its label and land in the check below, so this is the one place a shared line is
+        # allowed - a duplicated label would not even compile.
         return list(theirs) + list(ours)
     shared = identifying(ours) & identifying(theirs)
     if shared:
         # Overlapping content is the signature of one construct added twice,
         # not two independent additions. Unioning it would duplicate code.
-        # Scaffolding lines are excluded above, so what is left is content both
-        # sides genuinely wrote, which is the thing that makes this a duplicate.
+        # scaffolding is excluded above, so what is left is content both sides wrote
         raise Unresolvable(
             "both sides add the same line(s), so this is one change made twice: "
             + ", ".join(sorted(shared)[:3])
         )
     if not identifying(ours) or not identifying(theirs):
-        # Everything one side added is scaffolding, so there is no content to
-        # tell the two additions apart and the exclusion above has nothing left
-        # to work with. Refuse rather than union braces onto braces.
+        # one side is all scaffolding, so there is no content to tell the additions apart
         raise Unresolvable(
             "one side adds only block scaffolding, so the two additions cannot "
             "be told apart"
diff --git a/scripts/unsloth/feature_matrix.py b/scripts/unsloth/feature_matrix.py
index 00b392ca1ec..d8fad55d469 100644
--- a/scripts/unsloth/feature_matrix.py
+++ b/scripts/unsloth/feature_matrix.py
@@ -37,7 +37,7 @@
 import sys
 from pathlib import Path
 
-# Output that means "this did not run" from a process that exited 0.
+# output that means "this did not run" from a process that exited 0
 SKIP_RE = re.compile(r"\bSKIP\b|not supported|unsupported|no tests|0 tests", re.I)
 
 
@@ -78,7 +78,7 @@ def probe_arch(check: dict, b: Path, gpu: bool) -> str:
     rc, out = run([str(b / "test-llama-archs"), "-a", arch, "-s", "1234"], b, gpu)
     if rc != 0:
         raise Unproven(f"test-llama-archs -a {arch} exited {rc}")
-    # The arch's own rows, not the header and not another arch's.
+    # the arch's own rows, not the header and not another arch's
     rows = [ln for ln in out.splitlines() if ln.strip().startswith("|") and f"|{arch:>16}|" in ln
             or (ln.strip().startswith("|") and ln.split("|")[1].strip() == arch)]
     if not rows:
@@ -119,8 +119,8 @@ def probe_mtmd(check: dict, b: Path, gpu: bool) -> str:
     m = re.search(r"assertions\s*:\s*(\d+)", out)
     if not m or int(m.group(1)) == 0:
         raise Unproven("test_projector_registry ran no assertions; the filter matched nothing")
-    # The registry test walks the whole enum, so it proves the table is sound.
-    # That the specific projector is IN the enum is pin_contract.py's job.
+    # the test walks the whole enum, so it proves the table is sound; that this projector
+    # is IN the enum is pin_contract.py's job
     return f"projector registry intact over {m.group(1)} assertions"
 
 
@@ -173,8 +173,7 @@ def main() -> int:
             print(f"ok   {name}: " + "; ".join(r["evidence"] for r in entry["results"])
                   + (f"  [{len(entry['deferred'])} needs a GPU]" if entry["deferred"] else ""))
         else:
-            # Nothing was shown either way. Not a failure here, but it must not
-            # read as one of the ok lines.
+            # nothing shown either way: not a failure, but not an ok line either
             print(f"--   {name}: nothing provable without a GPU "
                   f"({len(entry['deferred'])} check(s) deferred)")
 
@@ -188,8 +187,7 @@ def main() -> int:
     if failed:
         print(f"\n{failed} feature(s) could not be shown to work", file=sys.stderr)
         return 1
-    # Say what was NOT proven in the same breath as what was. A run that only
-    # ever prints a success line teaches the reader that green means covered.
+    # say what was NOT proven alongside what was, or green starts to read as covered
     tail = f", {deferred} check(s) need a GPU and were not run" if deferred else ""
     print(f"\nall {len(report['features'])} features demonstrated"
           + (" on GPU" if args.gpu else " on CPU") + tail)
diff --git a/scripts/unsloth/pin_contract.py b/scripts/unsloth/pin_contract.py
index bf07d58408e..d6a60746098 100644
--- a/scripts/unsloth/pin_contract.py
+++ b/scripts/unsloth/pin_contract.py
@@ -57,34 +57,28 @@
     r"^https://github\.com/([^/]+)/llama\.cpp/pull/(\d+)/commits/([0-9a-f]{40})/?$"
 )
 
-# Identifier families that name a FEATURE. Deliberately not "every new symbol":
-# a helper function renamed by a later upstream commit is not a lost feature,
-# but a missing LLM_ARCH_ entry always is. These are the tables that decide
-# whether an architecture, an op, a projector or a quant type exists at all.
+# Identifier families that name a FEATURE, not every new symbol: a renamed helper is not a
+# lost feature, but a missing LLM_ARCH_ entry always is.
 SYMBOL_FAMILIES = (
     "LLM_ARCH_", "LLM_TENSOR_", "LLM_KV_", "LLM_TYPE_",
     "PROJECTOR_TYPE_", "GGML_OP_", "GGML_TYPE_", "LLAMA_FTYPE_",
 )
 SYMBOL_RE = re.compile(r"\b(?:" + "|".join(SYMBOL_FAMILIES) + r")[A-Z0-9_]+\b")
 
-# The subset that names a whole feature rather than one of its tensors. Used
-# only to keep --emit readable; the check itself uses all of SYMBOL_FAMILIES.
+# the subset naming a whole feature; only to keep --emit readable, the check uses them all
 HEADLINE = ("LLM_ARCH_", "GGML_OP_", "GGML_TYPE_", "PROJECTOR_TYPE_", "LLAMA_FTYPE_")
 
-# A line worth tracking for survival. Comments and short punctuation drift with
-# every reformat and would make the check noise; a substantial code line does
-# not move on its own.
+# a line worth tracking for survival: comments and short punctuation drift with every
+# reformat, a substantial code line does not move on its own
 TRIVIAL_RE = re.compile(r"^\s*(?://|/\*|\*|\*/|#\s|$)")
 MIN_LINE = 12
 
-# Comments are stripped before anything is read off a line. A pin that merely
-# NAMES an arch in a comment has not registered it, and holding the comment's
-# wording as a contract fails the moment upstream rewords it. Observed on
-# unslothai#70, whose comment mentions GGML_OP_SSM_SCAN to explain why it does
-# NOT use it.
+# Comments are stripped first: a pin that merely NAMES an arch in a comment has not
+# registered it, and holding the wording as a contract fails when upstream rewords it.
+# Observed on unslothai#70, whose comment mentions GGML_OP_SSM_SCAN to say it is not used.
 COMMENT_RE = re.compile(r"//.*$|/\*.*?\*/|(?<!\S)#(?!\s*(?:include|define|if|el|endif|pragma)).*$")
 
-# Binary and generated paths whose "lines" are meaningless.
+# binary and generated paths whose "lines" are meaningless
 SKIP_SUFFIXES = (".npy", ".png", ".jpg", ".gguf", ".bin", ".safetensors", ".ico", ".pdf")
 
 
@@ -170,9 +164,7 @@ def derive(pin: dict, base: str, cwd: Path) -> dict:
             if code:
                 symbols[cur].update(SYMBOL_RE.findall(code))
 
-    # Only symbols the base does not ALREADY have in that file are evidence of
-    # this pin. Upstream naming an arch in a file the pin also touches is not
-    # something the pin is owed.
+    # only symbols the base does not already have in that file are evidence of this pin
     new_symbols: dict[str, list[str]] = {}
     for path, names in symbols.items():
         fresh = sorted(n for n in names
@@ -306,9 +298,8 @@ def main() -> int:
 
         if args.emit:
             report["pins"].append(entry)
-            # Only the families that NAME a feature are printed. Every symbol
-            # is still checked; a new file legitimately contributes a hundred
-            # LLM_TENSOR_ names and listing them buries the one that matters.
+            # only feature-naming families are printed; all are still checked, but a
+            # hundred LLM_TENSOR_ names would bury the one that matters
             sym = sorted({s for v in contract["symbols"].values() for s in v
                           if s.startswith(HEADLINE)})
             print(f"{name:>18}  {entry['line_count']:>5} lines, "
@@ -345,8 +336,7 @@ def main() -> int:
     if args.emit:
         return 0
 
-    # Notices after the verdict lines, never mixed into them: "upstream took
-    # this, drop the entry" is housekeeping and must not read as a failure.
+    # notices after the verdict lines: housekeeping must not read as a failure
     for n in notices:
         print(f"note {n}")
     if failed:
diff --git a/scripts/unsloth/test_additive_merge.py b/scripts/unsloth/test_additive_merge.py
index 0f91e9afd12..dd2b22c6cac 100644
--- a/scripts/unsloth/test_additive_merge.py
+++ b/scripts/unsloth/test_additive_merge.py
@@ -97,8 +97,7 @@ def run(repo, *extra):
       reason.endswith('twice: log("same");'), reason)
 
 # --- 3b. two independent case arms: braces are shared, content is not -------
-# The real tools/mtmd/clip.cpp shape. Refusing this on `{` and `} break;` is
-# what took the 09-02 nightly's last pin down.
+# the real clip.cpp shape; refusing it on `{` and `} break;` took the 09-02 nightly down
 base = "switch (t) {\n}\n"
 ours = ("switch (t) {\n    case PROJECTOR_TYPE_KIMIK3:\n        {\n"
         "            builder = std::make_unique<clip_graph_kimik3>(ctx, img);\n"
@@ -115,8 +114,7 @@ def run(repo, *extra):
       txt.count("} break;") == 2 and txt.count("clip_graph_kimik3") == 1, txt)
 
 # --- 3b2. two case arms that share a body line, which is a coincidence ------
-# The clip.cpp shape after upstream landed DEEPSEEK4V: both arms set the same
-# rope_theta, and refusing on that is the shared-line check backwards.
+# clip.cpp after upstream landed DEEPSEEK4V: both arms set the same rope_theta
 base = "switch (t) {\n}\n"
 ours = ("switch (t) {\n    case PROJECTOR_TYPE_KIMIK3:\n        {\n"
         "            hparams.image_resize_algo = RESIZE_ALGO_BILINEAR;\n"
diff --git a/scripts/unsloth/test_pin_contract.py b/scripts/unsloth/test_pin_contract.py
index 3c7b58d62d5..e76586b1fc4 100644
--- a/scripts/unsloth/test_pin_contract.py
+++ b/scripts/unsloth/test_pin_contract.py
@@ -98,8 +98,7 @@ def run(repo, pr_set, *extra):
 check("intact merge reports no notices", rep["notices"] == [], rep)
 
 # --- 2. the arm is dropped from ONE file: a tree-wide grep would pass ------
-# The real shape: LLM_ARCH_INKLING survives in the enum and the dispatch arm
-# that makes it do anything is gone.
+# the real shape: LLM_ARCH_INKLING survives in the enum, the dispatch arm is gone
 repo, pr_set, sha = make_repo()
 p = repo / "src" / "llama-model.cpp"
 p.write_text(MODEL_CPP_BASE)
@@ -128,8 +127,8 @@ def run(repo, pr_set, *extra):
       any("do_the_banded_thing" in x for x in rep["pins"][0]["problems"]), rep)
 
 # --- 5. redundancy: the base already has everything the pin adds ----------
-# Built the way it happens for real: upstream lands the same work, so the base
-# tag has it and the pin is not an ancestor of anything.
+# as it happens for real: upstream lands the same work, so the base tag has it and
+# the pin is not an ancestor of anything
 d = Path(tempfile.mkdtemp(prefix="pc_"))
 git(d, "init", "-q", "-b", "main")
 (d / "src").mkdir()
@@ -163,8 +162,8 @@ def run(repo, pr_set, *extra):
       rep["pins"][0]["added_files"] == ["src/inkling.cpp"], rep)
 
 # --- 7. a comment is not a contract ---------------------------------------
-# unslothai#70 has a comment naming GGML_OP_SSM_SCAN to say it does NOT use it.
-# Holding comment wording would fail the moment upstream rewords it.
+# unslothai#70 names GGML_OP_SSM_SCAN in a comment to say it does NOT use it, and
+# holding that wording would fail the moment upstream rewords it
 repo, pr_set, sha = make_repo()
 git(repo, "checkout", "-q", "pin")
 (repo / "src" / "note.cpp").write_text(
diff --git a/src/llama-batch.cpp b/src/llama-batch.cpp
index cc73d83963c..4fc1c808631 100644
--- a/src/llama-batch.cpp
+++ b/src/llama-batch.cpp
@@ -573,18 +573,13 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
         }
 
         if (add) {
-            // [TAG_EXACT_CONCURRENCY] a sequence set that still has more tokens to place than a
-            // decode step carries is a prompt, and a prompt shares its arithmetic with whatever
-            // else is in the ubatch, so give it a ubatch of its own. Sets at or below that width
-            // are decode steps, plain or speculative, whose columns the backend's column policy
-            // already keeps exact, so keep grouping those: isolating them too would make one
-            // prompt serialize every concurrent decode for the whole of the prefill, and would run
-            // a speculative verify step once per sequence. Grouped sets must have the same number
-            // of tokens left: the equal-length expansion below would otherwise place a three-token
-            // verify step beside a two-token one as two tokens now and one later, and a memory
-            // that reduces over a chunk of tokens (a chunked state space scan) would then sum in a
-            // different order than the solo run's single three-token ubatch. A set with a
-            // different count waits for a later ubatch.
+            // [TAG_EXACT_CONCURRENCY] a set with more tokens left than a decode step carries is a
+            // prompt, and a prompt shares its arithmetic with whatever else is in the ubatch, so
+            // give it one of its own. Sets at or below that width are decode steps, kept exact by
+            // the backend's column policy, so keep grouping them or one prompt would serialize
+            // every concurrent decode. Grouped sets must have the same number of tokens left, or
+            // the equal-length expansion below would cut a longer set in two and a memory that
+            // reduces over a chunk of tokens would sum in a different order than the solo run.
             if (isolate_seqs_above > 0) {
                 uint32_t n_left = 0;
 
@@ -612,8 +607,8 @@ llama_ubatch llama_batch_allocr::split_equal(uint32_t n_ubatch, bool sequential,
                 } else if (n_left != n_left_first) {
                     continue;
                 } else if ((cur_seq_set.size() + 1) * n_left_first > n_ubatch) {
-                    // one more set would not finish in this ubatch: the expansion below would
-                    // then cut every set part way, the chunking the guard exists to prevent
+                    // one more set would not finish here, and the expansion below would then cut
+                    // every set part way: the chunking this guard exists to prevent
                     break;
                 }
             }
diff --git a/src/llama-batch.h b/src/llama-batch.h
index 52a375ad49e..f0ecc9d8407 100644
--- a/src/llama-batch.h
+++ b/src/llama-batch.h
@@ -105,12 +105,9 @@ class llama_batch_allocr {
     // make ubatches of equal-length sequences sets
     // if sequential == true, the tokens in the ubatch will have increasing sequential sequence ids
     // n_keep_tail = minimum trailing tokens of a seq that must land in the same ubatch
-    // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this many
-    //               tokens left to place is a prompt and is given a ubatch of its own; sets at or
-    //               below it are decode steps (one token, or one plus the speculative drafts) and
-    //               stay grouped together, so a prompt next to three decodes costs one extra ubatch
-    //               and does not serialize the three decodes, and a speculative verify batch is not
-    //               run once per sequence
+    // isolate_seqs_above = [TAG_EXACT_CONCURRENCY] when > 0, a sequence set with more than this
+    //               many tokens left to place is a prompt and gets a ubatch of its own; sets at or
+    //               below it are decode steps and stay grouped, so a prompt does not serialize them
     llama_ubatch split_equal(uint32_t n_ubatch, bool sequential, uint32_t n_keep_tail, uint32_t isolate_seqs_above = 0);
 
     // [TAG_EXACT_CONCURRENCY] true if some sequence still has more than n_tokens left to place,
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 8d35922a47b..16e090a0623 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -102,18 +102,14 @@ llama_context::llama_context(
     }
 
     // [TAG_EXACT_CONCURRENCY] the widest decode step this context can build: one column per
-    // sequence, times the tokens a sequence contributes to a step. Reported so that a backend
-    // splitting columns for exactness covers it without the caller having to know the bound; a
-    // caller that builds wider steps reports the width itself, see llama_set_exact_decode_width.
-    // The sequence count is what is reported: the tokens figure can be raised later for the
-    // whole process, and the width then follows it for this context too.
-    // Checked here and reported at the end of the constructor: the count is process-wide
-    // state that outlives a context, so a construction that fails later on, an unsupported
-    // cache layout say, must not leave a width behind that no context needs.
+    // sequence times the tokens a sequence contributes, reported so a backend splitting columns
+    // covers it (a caller that builds wider steps uses llama_set_exact_decode_width). The
+    // sequence count is what is reported, so a later rise in the tokens figure follows it here
+    // too. Checked now but reported at the end of the constructor, so a construction that fails
+    // later does not leave a width behind that no context needs.
     if (llama_exact_concurrency()) {
-        // an explicit column bound wins over the reported width in the backend, so one below
-        // this context's width would leave its decodes batched above the bound with the mode
-        // still reporting itself on; the report refuses that, and the refusal is an error here
+        // an explicit column bound wins in the backend, so one below this context's width would
+        // leave decodes batched above it; the report refuses that, and that is an error here
         if (!llama_exact_check_n_seq(cparams.n_seq_max)) {
             throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width");
         }
@@ -412,8 +408,8 @@ llama_context::llama_context(
 
         memory.reset(model.create_memory(params_mem, cparams));
 
-        // [TAG_EXACT_CONCURRENCY] the paged attention the mode runs on is causal; a context
-        // created non-causal with a cache would assert on its first graph, so it is refused here
+        // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a non-causal context with a
+        // cache would assert on its first graph
         if (llama_exact_concurrency() && memory && !cparams.causal_attn) {
             LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so it cannot be created with non-causal attention\n", __func__);
             throw std::runtime_error("exact concurrency: non-causal attention is not supported with a KV cache");
@@ -502,9 +498,8 @@ llama_context::llama_context(
         }
     }
 
-    // [TAG_EXACT_CONCURRENCY] nothing above can fail any more, so the width this context
-    // needs is published now; checked against the explicit bound at the top, so this
-    // cannot refuse unless the bound moved underneath it, which is an error all the same
+    // [TAG_EXACT_CONCURRENCY] nothing above can fail now, so publish the width; already checked
+    // against the explicit bound at the top, so a refusal here means the bound moved
     if (llama_exact_concurrency() && !llama_exact_report_n_seq(cparams.n_seq_max)) {
         throw std::runtime_error("exact concurrency: the explicit column bound is below this context's decode width");
     }
@@ -1220,8 +1215,8 @@ void llama_context::set_causal_attn(bool value) {
         return;
     }
 
-    // [TAG_EXACT_CONCURRENCY] the paged attention the mode runs on is causal; a context with a
-    // cache under the mode keeps causal attention rather than asserting in the next graph
+    // [TAG_EXACT_CONCURRENCY] the paged attention is causal, so a context with a cache keeps
+    // causal attention rather than asserting in the next graph
     if (!value && memory && llama_exact_concurrency()) {
         LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set and this context has a KV cache, so causal attention cannot be turned off; the change is refused\n", __func__);
         return;
@@ -2664,17 +2659,16 @@ class llama_io_read_host : public llama_io_read_i {
             }
             const size_t tensor_bytes = ggml_nbytes(tensor);
             auto * buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
-            // A fragmented sequence can require thousands of synchronous device
-            // transfers per layer. For bounded tensors, stage the tensor once and
-            // preserve every byte belonging to other sequences. Bound scratch RAM
-            // and leave ordinary contiguous transfers on their original fast path.
+            // a fragmented sequence can need thousands of synchronous device transfers per
+            // layer: stage a bounded tensor once instead, preserving other sequences' bytes and
+            // leaving ordinary contiguous transfers on their fast path
             if (end - i >= 64 && tensor_bytes <= 64 * 1024 * 1024 &&
                     !ggml_backend_buffer_is_host(buffer)) {
                 std::vector<uint8_t> staging;
                 try {
                     staging.resize(tensor_bytes);
                 } catch (const std::bad_alloc &) {
-                    // Fall back to the individual transfers below.
+                    // fall back to the individual transfers below
                 }
                 if (!staging.empty()) {
                     ggml_backend_tensor_get(tensor, staging.data(), 0, tensor_bytes);
@@ -3266,8 +3260,8 @@ size_t llama_context::state_write_data(llama_io_write_i & io) {
 
 size_t llama_context::state_read_data(llama_io_read_i & io) {
     // [TAG_EXACT_CONCURRENCY] a whole-context restore writes cells at their recorded physical
-    // index, which the paged pool owns. Refused here, before anything is parsed, so that the
-    // cache the caller has is left as it was: the generic restore path clears it on failure.
+    // index, which the paged pool owns. Refused before anything is parsed, so the caller's cache
+    // is left as it was: the generic restore path clears it on failure.
     if (memory && memory->alloc_granularity() > 1) {
         throw std::runtime_error("whole-context restore is not supported with LLAMA_EXACT_CONCURRENCY, restore per sequence");
     }
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index f8f5c83a2b6..dc02560fc66 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -27,18 +27,11 @@
 
 // dedup helpers
 
-// [TAG_EXACT_CONCURRENCY]
-// The page table is wired into llm_graph_input_attn_kv only. The V-less layouts build their
-// attention without one, so a model on one of those would get its cells placed in pages by the
-// allocator and then attend in physical cell order anyway: the mode would report itself as on and
-// lose the one invariant it exists for, which is the same silent failure the CUDA placement gate
-// was added to stop. Refuse the context instead.
-//
-// Rejecting is the smaller correct change here. Wiring self_pages into llm_graph_input_attn_k alone
-// is four lines, but it fixes only one of the four V-less input classes, and DeepSeek 3.2 uses two
-// of them: its sparse layers rewrite the mask from a top-k selection and would still be unpaged, so
-// the model would end up half paged, which is worse than refused. None of these architectures was
-// measured, and the paged kernel additionally requires 256-dimensional K and V heads.
+// [TAG_EXACT_CONCURRENCY] the page table is wired into llm_graph_input_attn_kv only, so a V-less
+// layout would have its cells placed in pages and then attend in physical order anyway, with the
+// mode reporting itself as on. Refuse the context instead: wiring self_pages into
+// llm_graph_input_attn_k alone fixes one of the four V-less classes, and DeepSeek 3.2 uses two of
+// them, so the model would end up half paged, which is worse than refused.
 static void llm_graph_reject_exact_concurrency(const char * layout) {
     if (!llama_exact_concurrency()) {
         return;
diff --git a/src/llama-impl.cpp b/src/llama-impl.cpp
index 0b218bc64e7..f577414a51c 100644
--- a/src/llama-impl.cpp
+++ b/src/llama-impl.cpp
@@ -187,22 +187,18 @@ bool llama_exact_concurrency() {
 // [TAG_EXACT_CONCURRENCY] tokens one sequence contributes to a decode step, see llama.h
 static std::atomic<uint32_t> g_exact_decode_tokens{1};
 
-// one lock for the token figure, the sequence count and the width: the three move together
-// (a context reports its count and the width that follows; a new token figure re-reports the
-// width for every count seen), and a report interleaved with a change of figure could leave
-// the backend with a width that covers neither. Recursive, since the setters call each other.
+// one lock for the token figure, the sequence count and the width, since the three move together
+// and a report interleaved with a change of figure could leave the backend with a width that
+// covers neither; recursive, since the setters call each other
 static std::recursive_mutex g_exact_mutex;
 
-// the most sequences any context so far was created with. The tokens figure is process
-// wide, so raising it widens the decode step of every context that already exists; the
-// width those contexts reported at creation is re-reported here with the new figure, or a
-// context created under a narrower figure would batch above the bound it reported.
+// the most sequences any context was created with. The tokens figure is process wide, so raising
+// it widens every existing context's decode step and their width is re-reported with it.
 static std::atomic<uint32_t> g_exact_max_n_seq{0};
 
 static bool llama_exact_width_within_explicit_bound(uint32_t n_cols);
 
-// the width is sequences times tokens, handed to a backend as an int; a product that does not
-// fit is refused rather than wrapped
+// sequences times tokens, handed to a backend as an int; a product that overflows is refused
 static bool llama_exact_width_of(uint32_t n_seq, uint32_t n_tokens, uint32_t & n_cols) {
     const uint64_t w = (uint64_t) n_seq * (uint64_t) n_tokens;
 
@@ -250,14 +246,14 @@ bool llama_set_exact_decode_tokens(uint32_t n_tokens) {
 
     std::lock_guard<std::recursive_mutex> lock(g_exact_mutex);
 
-    // never lowered: a narrower context set up later would otherwise turn the verify steps of
-    // an existing speculative context into prompts and serialise them
+    // never lowered: a narrower context set up later would turn an existing speculative
+    // context's verify steps into prompts and serialise them
     if (n_tokens <= g_exact_decode_tokens.load(std::memory_order_relaxed)) {
         return true;
     }
 
-    // every context that exists widens with the figure, so the width they will need is
-    // reported first; a figure the explicit bound cannot cover leaves the old one in place
+    // every context widens with the figure, so report the width first; one the explicit bound
+    // cannot cover leaves the old figure in place
     const uint32_t n_seq = g_exact_max_n_seq.load(std::memory_order_relaxed);
 
     uint32_t n_cols = 0;
@@ -275,14 +271,13 @@ uint32_t llama_exact_decode_tokens(void) {
     return g_exact_decode_tokens.load(std::memory_order_relaxed);
 }
 
-// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h. A backend that
-// splits columns to make a decode exact reads it through ggml_backend_cuda_set_exact_decode_width,
-// reached through the registry so that a backend that is absent or loaded late costs nothing.
+// [TAG_EXACT_CONCURRENCY] the widest decode ubatch reported so far, see llama.h. Backends read it
+// through ggml_backend_cuda_set_exact_decode_width, reached through the registry so an absent or
+// late-loaded backend costs nothing.
 static std::atomic<uint32_t> g_exact_decode_width{0};
 
-// an explicit column bound given to the CUDA backend wins over the reported width there, so a
-// width above it would leave decodes batched past the bound with the mode still reporting itself
-// on; a width the bound does not cover is refused instead of stored
+// an explicit column bound wins in the CUDA backend, so a width above it would leave decodes
+// batched past the bound: refuse such a width instead of storing it
 static bool llama_exact_width_within_explicit_bound(uint32_t n_cols) {
     static const int explicit_cols = []() {
         const char * val = getenv("GGML_CUDA_BATCH_INVARIANT_MAX_COLS");
@@ -310,9 +305,8 @@ bool llama_set_exact_decode_width(uint32_t n_cols) {
     while (n_cols > cur && !g_exact_decode_width.compare_exchange_weak(cur, n_cols, std::memory_order_relaxed)) {
     }
 
-    // The widest figure so far goes to every backend on every call, not only when it grew: a
-    // width reported before a backend was loaded would otherwise never reach it, and every
-    // context reports at creation, by which time the backends are there.
+    // the widest figure goes to every backend on every call, not only when it grew, or a width
+    // reported before a backend was loaded would never reach it
     const uint32_t widest = g_exact_decode_width.load(std::memory_order_relaxed);
 
     for (size_t i = 0; i < ggml_backend_reg_count(); ++i) {
diff --git a/src/llama-impl.h b/src/llama-impl.h
index 65d56a51d1d..d0ab2e5cf1c 100644
--- a/src/llama-impl.h
+++ b/src/llama-impl.h
@@ -104,16 +104,14 @@ std::string llama_format_tensor_shape(const struct ggml_tensor * t);
 
 std::string gguf_kv_to_str(const struct gguf_context * ctx_gguf, int i);
 
-// [TAG_EXACT_CONCURRENCY]
-// opt-in mode under which a sequence's attention depends only on its own cells, in position order,
-// so that its output does not change when other sequences share the KV cache. Off by default.
-// Reads the same LLAMA_EXACT_CONCURRENCY variable as the paged KV cache and the CUDA backend.
+// [TAG_EXACT_CONCURRENCY] opt-in mode under which a sequence's attention depends only on its own
+// cells, in position order, so its output does not change when others share the KV cache. Off by
+// default; reads the same variable as the paged KV cache and the CUDA backend.
 bool llama_exact_concurrency();
 
-// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so that the
-// decode width every context needs is known to the backend and follows llama_set_exact_decode_tokens
+// [TAG_EXACT_CONCURRENCY] a context reports how many sequences it was created with, so the backend
+// knows the width every context needs and it follows llama_set_exact_decode_tokens
 bool llama_exact_report_n_seq(uint32_t n_seq);
 
-// the same check without the report: whether a context of n_seq sequences could be reported
-// under the explicit column bound, for a constructor that may still fail after asking
+// the same check without the report, for a constructor that may still fail after asking
 bool llama_exact_check_n_seq(uint32_t n_seq);
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 31154639ad9..428fbc6b567 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -62,11 +62,9 @@ static void ggml_gen_hadamard(ggml_tensor * tensor) {
 // llama_kv_cache
 //
 
-// [TAG_EXACT_CONCURRENCY]
-// The paged attention specialization that reads the logical page table lives in the CUDA backend
-// sources, which are also built as the ROCm and MUSA backends. Every other backend ignores src[5]
-// and walks the pool in physical cell order, so a KV layer placed there would silently lose the
-// guarantee the mode exists to provide.
+// [TAG_EXACT_CONCURRENCY] the paged attention specialization lives in the CUDA sources, which are
+// also built as ROCm and MUSA. Every other backend ignores src[5] and walks the pool in physical
+// cell order, so a KV layer placed there would silently lose the mode's guarantee.
 static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) {
     if (!dev) {
         return false;
@@ -85,14 +83,11 @@ static bool llama_dev_has_paged_attn(ggml_backend_dev_t dev) {
     return strcmp(name, "CUDA") == 0 || strcmp(name, "ROCm") == 0 || strcmp(name, "MUSA") == 0;
 }
 
-// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a
-// layer of this shape. The registry name says which backends carry the kernels; it does not
-// say the build has them (FLASH_ATTN_AVAILABLE), nor that the device's architecture, the
-// head width and the K/V types land on a kernel. Where they do not, the scheduler would hand
-// the op to the CPU, which accepts the page table as the reference for test-backend-ops and
-// ignores it, and the mode would report itself on while attending in physical order. So the
-// op is built the way the graph builds it, at the widths a decode step, a verify step and a
-// prompt chunk use, and the device is asked.
+// [TAG_EXACT_CONCURRENCY] whether the device can actually run the paged attention op for a layer
+// of this shape. The registry name only says which backends carry the kernels, not that the build
+// has them or that this architecture, head width and K/V types land on one; otherwise the op
+// falls to the CPU, which ignores the page table. So build the op as the graph does, at the
+// widths a decode step, a verify step and a prompt chunk use, and ask the device.
 static bool llama_dev_supports_paged_attn(
         ggml_backend_dev_t dev,
         ggml_type type_k, ggml_type type_v,
@@ -163,9 +158,8 @@ llama_kv_cache::llama_kv_cache(
     v_cells_impl(other ? other->v_cells_impl : std::make_shared<llama_kv_cells_vec>()),
     v_cells(*v_cells_impl) {
 
-    // [TAG_EXACT_CONCURRENCY] read the knob through the one cached reader that the graph and the
-    // CUDA dispatcher also use, so a process that sets it between two context creations cannot end
-    // up with a paged cache on top of a dispatcher that is still in default mode
+    // [TAG_EXACT_CONCURRENCY] read the knob through the same cached reader the graph and the CUDA
+    // dispatcher use, so a mid-process change cannot leave the two disagreeing
     exact_pages = llama_exact_concurrency();
 
     // shared cells view the source cache's K/V tensors, so the cell count
@@ -181,9 +175,8 @@ llama_kv_cache::llama_kv_cache(
 
     GGML_ASSERT(kv_size % n_pad == 0);
 
-    // [TAG_EXACT_CONCURRENCY]
-    // Every one of these is reachable from the command line, so report which one failed by name
-    // instead of aborting on a bare assert that only prints a file and a line.
+    // [TAG_EXACT_CONCURRENCY] all of these are reachable from the command line, so name the one
+    // that failed instead of aborting on a bare assert
     if (exact_pages) {
         const char * unsupported = nullptr;
 
@@ -328,10 +321,8 @@ llama_kv_cache::llama_kv_cache(
 
         LLAMA_LOG_DEBUG("%s: layer %3d: dev = %s\n", __func__, il, dev_name);
 
-        // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order while
-        // the mode still reports itself as on, so refuse the load instead
-        // [TAG_EXACT_CONCURRENCY] the paged attention kernel handles 256-wide K and V heads only;
-        // any other width would run unpaged on the CPU while the mode reports itself as on
+        // [TAG_EXACT_CONCURRENCY] the paged kernel handles 256-wide K and V heads only; any other
+        // width would run unpaged while the mode reports itself as on
         if (exact_pages && (hparams.n_embd_head_k(il) != 256 || (!is_mla && hparams.n_embd_head_v(il) != 256) || is_mla)) {
             LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d has %u-wide K heads and %u-wide V heads%s, "
                     "and the paged attention kernel supports 256-wide K and V heads only\n",
@@ -339,14 +330,14 @@ llama_kv_cache::llama_kv_cache(
             throw std::runtime_error("exact concurrency: unsupported attention head size");
         }
 
-        // [TAG_EXACT_CONCURRENCY] the paged attention kernel has no soft-capped variant and would
-        // assert on its first call, so a soft-capped model is refused at load instead
+        // [TAG_EXACT_CONCURRENCY] the paged kernel has no soft-capped variant and would assert
         if (exact_pages && hparams.attn_soft_cap) {
             LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but this model soft-caps its attention logits (%.1f), "
                     "which the paged attention kernel does not apply\n", __func__, hparams.f_attn_logit_softcapping);
             throw std::runtime_error("exact concurrency: attention soft cap is not supported");
         }
 
+        // [TAG_EXACT_CONCURRENCY] a layer left anywhere else attends in physical cell order
         if (exact_pages && !(offload && llama_dev_has_paged_attn(model.dev_layer(il)))) {
             LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set but layer %d keeps its KV cache on %s, "
                     "which has no paged attention: every layer must be offloaded to the CUDA backend "
@@ -355,8 +346,8 @@ llama_kv_cache::llama_kv_cache(
             throw std::runtime_error("exact concurrency: KV cache layer is not on the CUDA backend");
         }
 
-        // [TAG_EXACT_CONCURRENCY] the backend is the right one; ask it whether this layer's
-        // attention, with the page table attached, lands on one of its kernels at all
+        // [TAG_EXACT_CONCURRENCY] right backend; ask whether this layer's attention, with the
+        // page table attached, lands on one of its kernels at all
         if (exact_pages && !llama_dev_supports_paged_attn(model.dev_layer(il), type_k, type_v,
                     hparams.n_embd_head_k(il), hparams.n_embd_head_v(il),
                     hparams.n_head(il), hparams.n_head_kv(il), kv_size, exact_page_size)) {
@@ -671,12 +662,9 @@ void llama_kv_cache::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, ll
         return;
     }
 
-    // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared
-    // between two of them. Refuse the operation rather than abort the process: a server
-    // rejects the request that would reach here (n_cmpl > 1), and any caller this does not
-    // cover degrades to a failed copy it can report instead of killing every other request
-    // on the machine. Placed after the shared-cells return so a draft cache, which copies
-    // nothing of its own, is unaffected.
+    // [TAG_EXACT_CONCURRENCY] a page belongs to one sequence, so cells cannot be shared between
+    // two. Refuse rather than abort the process, so an uncovered caller gets a failed copy it can
+    // report. After the shared-cells return, so a draft cache is unaffected.
     if (exact_pages && seq_id_src != seq_id_dst) {
         LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between "
                         "sequences (%d -> %d); ignoring the copy\n",
@@ -806,9 +794,8 @@ void llama_kv_cache::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, ll
         return;
     }
 
-    // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo the
-    // page size, so shifting positions would put every cell of the sequence in the wrong
-    // place. Context shift is unsupported in exact mode; say so rather than abort.
+    // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo the page size,
+    // so shifting positions would misplace every cell; say so rather than abort
     if (exact_pages && shift != 0) {
         LLAMA_LOG_ERROR("%s: exact concurrency does not support shifting positions "
                         "(seq %d, shift %d); ignoring the shift\n",
@@ -866,8 +853,7 @@ void llama_kv_cache::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, in
         return;
     }
 
-    // [TAG_EXACT_CONCURRENCY] same reason as seq_add: dividing positions breaks the
-    // identity between a cell's position and its offset inside its page.
+    // [TAG_EXACT_CONCURRENCY] as in seq_add: dividing positions breaks the position/offset identity
     if (exact_pages && d != 1) {
         LLAMA_LOG_ERROR("%s: exact concurrency does not support dividing positions "
                         "(seq %d, d %d); ignoring the division\n",
@@ -970,10 +956,9 @@ llama_memory_context_ptr llama_kv_cache::init_batch(
 
         std::vector<llama_ubatch> ubatches;
         while (true) {
-            // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt tokens into one
-            // ubatch, so a sequence's prefill would run at a width its solo run never sees. Take
-            // the sequence-set split instead, which can give each prompt a ubatch of its own; a
-            // plain decode step has nothing to isolate and keeps taking split_simple.
+            // [TAG_EXACT_CONCURRENCY] split_simple packs every sequence's prompt into one ubatch,
+            // so a prefill would run at a width its solo run never sees. The sequence-set split
+            // gives each prompt its own ubatch; a plain decode step keeps taking split_simple.
             const uint32_t isolate = llama_exact_concurrency() && balloc.has_seq_wider_than(llama_exact_decode_tokens()) ? llama_exact_decode_tokens() : 0;
 
             auto ubatch = n_stream == 1 && !isolate
@@ -1026,8 +1011,8 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_
 
         std::vector<llama_kv_cells> v_cells; // copy of the old cells, before placing the ubatch
 
-        // [TAG_EXACT_CONCURRENCY] page ownership before placing the ubatch, so that undoing the
-        // speculative placement does not force a rebuild from every cell on the next ubatch
+        // [TAG_EXACT_CONCURRENCY] page ownership before the ubatch, so undoing a speculative
+        // placement does not force a rebuild from every cell
         std::vector<exact_page> exact_page_owner_old;
     };
 
@@ -1078,9 +1063,8 @@ llama_kv_cache::slot_info_vec_t llama_kv_cache::prepare(const std::vector<llama_
             head = it->v_heads_old[s];
         }
 
-        // [TAG_EXACT_CONCURRENCY] the speculative placements are being undone behind the
-        // allocator's back. Put back what it knew before, unless something during the placement
-        // removed cells as well, in which case only the cells can say what is left.
+        // [TAG_EXACT_CONCURRENCY] put back what the allocator knew, unless the placement also
+        // removed cells, in which case only the cells can say what is left
         if (!exact_page_owner_dirty) {
             exact_page_owner = it->exact_page_owner_old;
         }
@@ -1243,10 +1227,8 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
     }
 
     if (exact_pages) {
-        // Page ownership is maintained as cells are placed and invalidated when they are removed,
-        // so the allocator reads one entry per physical page rather than scanning every cell. The
-        // claims this call makes are local: prepare() can still roll back its speculative
-        // placements, and empty pages stay immediately reusable.
+        // ownership is maintained as cells are placed, so this reads one entry per page rather
+        // than scanning every cell; the claims are local and prepare() can still roll them back
         const auto & cells = v_cells[0];
 
         exact_pages_sync();
@@ -1269,7 +1251,7 @@ llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch,
             const page_key key {ubatch.seq_id[i][0], ubatch.pos[i]/exact_page_size};
             auto it = pages.find(key);
             if (it == pages.end()) {
-                // Round-robin free-page search deliberately permits nonmonotonic physical order.
+                // round-robin free-page search, deliberately nonmonotonic in physical order
                 uint32_t page = v_heads[0]/exact_page_size;
                 uint32_t tested = 0;
                 while (tested < owner.size() && owner[page%owner.size()].seq >= 0) { ++page; ++tested; }
@@ -1503,17 +1485,15 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch &
 }
 
 uint32_t llama_kv_cache::alloc_granularity() const {
-    // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so a
-    // sequence holding n tokens holds round_up(n, exact_page_size) cells: its tail page is
-    // charged in full whether or not it is full.
+    // [TAG_EXACT_CONCURRENCY] a page is given to one (sequence, position / page) pair, so n
+    // tokens hold round_up(n, exact_page_size) cells: the tail page is charged in full
     return exact_pages ? exact_page_size : 1;
 }
 
 bool llama_kv_cache::get_can_shift() const {
-    // [TAG_EXACT_CONCURRENCY] a cell's offset inside its page is its position modulo 256, so the
-    // paged pool cannot shift positions. Reporting it here is what makes the server disable
-    // --context-shift and --cache-reuse at load, with a warning, instead of accepting both and
-    // failing on the first request that needs them.
+    // [TAG_EXACT_CONCURRENCY] a cell's offset in its page is its position modulo 256, so the pool
+    // cannot shift positions. Reporting it here is what disables --context-shift and
+    // --cache-reuse at load rather than failing on the first request that needs them.
     if (exact_pages) {
         return false;
     }
@@ -1605,7 +1585,7 @@ void llama_kv_cache::set_input_pages(ggml_tensor * dst, const llama_ubatch * uba
         auto * row = data.data() + i*dst->ne[0];
         row[0] = 0;
         for (const auto & page : pages[ubatch->seq_id[i][0]]) {
-            // Exclude wholly future pages even when prefill includes later query rows.
+            // exclude wholly future pages even when prefill includes later query rows
             if (page.first*exact_page_size > uint32_t(ubatch->pos[i])) { break; }
             row[++row[0]] = page.second;
         }
@@ -1622,8 +1602,8 @@ void llama_kv_cache_context::set_input_pages(ggml_tensor * dst, const llama_ubat
 }
 
 uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
-    // The physical view spans the pool. The page map, independently padded per query,
-    // is the only loop bound for exact attention; neighbours cannot extend that loop.
+    // the physical view spans the pool; the per-query page map is the only loop bound for exact
+    // attention, so neighbours cannot extend it
     if (exact_pages) { return get_size(); }
     uint32_t result = 0;
 
@@ -2429,9 +2409,9 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla
 }
 
 void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) {
-    // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical
-    // index, which the paged pool owns. Refused before a byte is read, so that the failure
-    // path below, which clears the cache, is never entered for it.
+    // [TAG_EXACT_CONCURRENCY] a whole-cache restore writes cells at their recorded physical index,
+    // which the paged pool owns. Refused before a byte is read, so the clearing failure path below
+    // is never entered for it.
     if (exact_pages && seq_id == -1) {
         LLAMA_LOG_ERROR("%s: LLAMA_EXACT_CONCURRENCY is set, which supports per-sequence state restore only\n", __func__);
         throw std::runtime_error("whole-cache restore is not supported with LLAMA_EXACT_CONCURRENCY");
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index af3a04be39d..c139ede57e0 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -243,11 +243,10 @@ class llama_kv_cache : public llama_memory_i {
     static constexpr uint32_t exact_page_size = 256;
     bool exact_pages = false;
 
-    // [TAG_EXACT_CONCURRENCY]
-    // Which (sequence, logical page) owns each physical page of the pool; seq < 0 means the page is
-    // free. Kept current as cells are placed, and marked dirty by the paths that remove cells, so
-    // that find_slot() and set_input_pages() read one entry per page instead of rebuilding the map
-    // from every live cell twice per ubatch. Mutable because set_input_pages() is const.
+    // [TAG_EXACT_CONCURRENCY] which (sequence, logical page) owns each physical page; seq < 0 means
+    // free. Kept current as cells are placed and marked dirty by removals, so find_slot() and
+    // set_input_pages() read one entry per page instead of rebuilding from every live cell twice per
+    // ubatch. Mutable because set_input_pages() is const.
     struct exact_page {
         llama_seq_id seq = -1;
         llama_pos    lpg = -1;
diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp
index e8d80c770ba..d9e609fee20 100644
--- a/src/llama-memory-hybrid.cpp
+++ b/src/llama-memory-hybrid.cpp
@@ -93,13 +93,10 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba
                 //   so that the rollback snapshots remain valid
                 const uint32_t n_rs_seq = mem_recr->n_rs_seq;
 
-                // [TAG_EXACT_CONCURRENCY] the recurrent half of a hybrid model is not invariant to
-                // the shape of the ubatch: a prompt processed next to other sequences' prompt tokens
-                // leaves a different gated delta net state than the same prompt processed alone.
-                // Giving such a sequence a ubatch of its own removes that. A plain decode step, one
-                // token per sequence, is already exact and stays batched.
-                // The figure is passed whenever the mode is on, not only when a prompt is present:
-                // it also keeps sets of unequal token counts apart (see llama_batch_allocr::split_equal).
+                // [TAG_EXACT_CONCURRENCY] the recurrent half is not invariant to the ubatch shape:
+                // a prompt processed next to other prompts leaves a different gated delta net state,
+                // so it gets a ubatch of its own while plain decode steps stay batched. Passed
+                // whenever the mode is on, since it also keeps sets of unequal token counts apart.
                 const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0;
 
                 ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate);
@@ -152,8 +149,8 @@ bool llama_memory_hybrid::get_can_shift() const {
 }
 
 uint32_t llama_memory_hybrid::alloc_granularity() const {
-    // the recurrent half holds one state per sequence rather than per token, so the
-    // attention half is the one whose cells a caller is planning capacity for
+    // the recurrent half holds one state per sequence, so the attention half is the one whose
+    // cells a caller is planning capacity for
     return mem_attn->alloc_granularity();
 }
 
@@ -172,8 +169,8 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1
 }
 
 void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) {
-    // [TAG_EXACT_CONCURRENCY] the attention half refuses this under exact mode; refuse it here
-    // before either half is touched, so the two halves cannot end up describing different states
+    // [TAG_EXACT_CONCURRENCY] the attention half refuses this, so refuse before either half is
+    // touched or the two could end up describing different states
     if (llama_exact_concurrency() && seq_id_src != seq_id_dst) {
         LLAMA_LOG_ERROR("%s: exact concurrency does not support copying cells between sequences (%d -> %d); ignoring the copy\n",
                 __func__, seq_id_src, seq_id_dst);
diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp
index 61463c72964..dadc0661344 100644
--- a/src/llama-memory-recurrent.cpp
+++ b/src/llama-memory-recurrent.cpp
@@ -431,12 +431,10 @@ llama_memory_context_ptr llama_memory_recurrent::init_batch(llama_batch_allocr &
                 // [TAG_RECURRENT_ROLLBACK_SPLITS]
                 // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch
                 //   so that the rollback snapshots remain valid
-                // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: a recurrent state that a
-                // prompt leaves behind depends on what shared its ubatch, so isolate the prompts.
-                // The figure is passed whenever the mode is on, not only when a prompt is present:
-                // it also keeps sets of unequal token counts apart, and a three-token verify step
-                // placed beside a two-token one as two now and one later would be reduced in
-                // chunks the solo run never had.
+                // [TAG_EXACT_CONCURRENCY] same rule as the hybrid memory: the state a prompt leaves
+                // behind depends on what shared its ubatch, so isolate prompts. Passed whenever the
+                // mode is on, since it also keeps sets of unequal token counts apart, which would
+                // otherwise be reduced in chunks the solo run never had.
                 const uint32_t isolate = llama_exact_concurrency() ? llama_exact_decode_tokens() : 0;
 
                 ubatch = balloc.split_equal(n_ubatch, true, n_rs_seq > 0 ? n_rs_seq + 1 : 0, isolate);
diff --git a/src/llama-memory.h b/src/llama-memory.h
index 51539a03919..c0bc2ad7e30 100644
--- a/src/llama-memory.h
+++ b/src/llama-memory.h
@@ -100,13 +100,10 @@ struct llama_memory_i {
     // getters
     virtual bool get_can_shift() const = 0;
 
-    // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit.
-    //
-    // 1 for every module that allocates a cell per token, which is all of them unless a mode
-    // is on that allocates in larger blocks. Where it is larger, a sequence of n tokens
-    // occupies round_up(n, granularity) cells, and a caller that plans pool capacity by
-    // counting tokens will believe there is room that does not exist. Not pure, so a module
-    // that has never heard of this inherits the answer that has always been true of it.
+    // [TAG_EXACT_CONCURRENCY] cells this module hands out in one indivisible unit: 1 unless a mode
+    // that allocates in larger blocks is on, and then n tokens occupy round_up(n, granularity)
+    // cells, so a caller planning capacity in tokens would see room that does not exist. Not pure,
+    // so a module that has never heard of this inherits the answer that was always true of it.
     virtual uint32_t alloc_granularity() const { return 1; }
 
     //
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 685193be77d..3f08a060ef1 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -7193,8 +7193,8 @@ struct test_flash_attn_ext : public test_case {
     }
 };
 
-// Same mathematical attention as the CPU mask reference, but visit nonadjacent pages
-// in a different order. Covers a partial tail and different page counts per query.
+// same attention as the CPU mask reference, but visiting nonadjacent pages in a different order;
+// covers a partial tail and different page counts per query
 struct test_flash_attn_ext_pages : public test_flash_attn_ext {
     test_flash_attn_ext_pages(int64_t batch) :
         test_flash_attn_ext(256, 256, 2, {8, 1}, 1024, batch) {}
@@ -9205,7 +9205,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
-    // Shared weights over sequence planes, as in a recurrent-model output projection.
+    // shared weights over sequence planes, as in a recurrent-model output projection
     for (ggml_type type : {GGML_TYPE_F32, GGML_TYPE_Q4_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0}) {
         for (int n : {1, 17, 307}) {
             test_cases.emplace_back(new test_mul_mat(type, GGML_TYPE_F32, 64, n, 256, {1, 1}, {4, 1}));
@@ -9213,10 +9213,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
         }
     }
 
-    // Mixture-of-experts projections at the token counts a decode ubatch forms. The gate and up
-    // projections broadcast one activation row over the expert list, the down projection carries
-    // one row per expert, and the mixed quantization of a real MoE gguf puts different types on
-    // the two. 17 tokens is past the width the exact-concurrency policy pins.
+    // MoE projections at the token counts a decode ubatch forms: gate and up broadcast one
+    // activation row over the expert list, down carries one row per expert, and a real MoE gguf
+    // puts different types on the two. 17 tokens is past the width exact concurrency pins.
     for (ggml_type type_a : {GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_Q8_0, GGML_TYPE_F16}) {
         for (int n : {1, 2, 4, 8, 17}) {
             test_cases.emplace_back(new test_mul_mat_id(type_a, GGML_TYPE_F32, 16, 8, true,   512, n, 2048));
diff --git a/tests/test-state-restore-fragmented.cpp b/tests/test-state-restore-fragmented.cpp
index 428a9252981..24b6359b05f 100644
--- a/tests/test-state-restore-fragmented.cpp
+++ b/tests/test-state-restore-fragmented.cpp
@@ -73,8 +73,8 @@ int main(int argc, char ** argv) {
     }
     fprintf(stderr, "%s : saved seq 1 state, %zu bytes\n", __func__, ncopy);
 
-    // A fragmented restore may stage a whole device tensor. Check every
-    // sequence byte-for-byte, including the neighbours that must be preserved.
+    // a fragmented restore may stage a whole device tensor, so check every sequence
+    // byte-for-byte, including the neighbours that must be preserved
     std::vector<std::vector<uint8_t>> before(params.n_parallel);
     for (int s = 0; s < params.n_parallel; ++s) {
         before[s].resize(llama_state_seq_get_size(ctx, s));
diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp
index ae47bec194c..6fe51f8cf1a 100644
--- a/tools/server/server-context.cpp
+++ b/tools/server/server-context.cpp
@@ -38,10 +38,8 @@
 
 constexpr int HTTP_POLLING_SECONDS = 1;
 
-// [TAG_EXACT_CONCURRENCY] the knob is read from the environment by the KV cache, the batch
-// splitter and the CUDA backend independently, because it has to be answered before a
-// context exists. The server needs the same answer to refuse the one request shape the mode
-// cannot serve, so it reads it the same way rather than growing a public API for it.
+// [TAG_EXACT_CONCURRENCY] read from the env like the KV cache, batch splitter and CUDA
+// backend do: the answer is needed before a context exists.
 static bool server_exact_concurrency() {
     static const bool enabled = []() {
         const char * val = getenv("LLAMA_EXACT_CONCURRENCY");
@@ -78,28 +76,18 @@ enum slot_state {
 
 // [TAG_PREEMPT] server-side request preemption
 //
-// With --kv-unified the cells are one pool shared by every slot, and each slot believes it
-// has all of them. When the pool fills, llama_decode returns 1, the retry ladder halves
-// n_batch down to 1, and the server ends EVERY conversation in flight with "Context size
-// has been exceeded" -- including the ones nowhere near their own limit. Upstream marks the
-// spot in decode(): "TODO: try to terminate only the largest active slot/sequence and
-// continue with the rest".
-//
-// Nothing is terminated here. The cells of one slot are taken back and given to it again
-// later: its sequence is copied to host RAM, its cells are released, and when the pool has
-// room the copy goes back and the slot carries on with the same sampler, the same generated
-// text and the same open stream. A streaming client sees a pause, not an error.
+// With --kv-unified one full pool ends EVERY conversation in flight with "Context size has
+// been exceeded", including the ones nowhere near their own limit. Instead of terminating,
+// one slot's sequence is copied to host RAM and its cells released; when there is room the
+// copy goes back and the slot carries on with the same sampler, text and open stream, so a
+// streaming client sees a pause rather than an error.
 constexpr int32_t PREEMPT_N_MARGIN   = 8;  // cells left spare on top of the reservation
 constexpr int32_t PREEMPT_N_STARVED  = 3;  // preemptions after which a slot is protected
 
-// [TAG_PREEMPT] The order parked slots come back in. Head of the line by park time, and nobody
-// passes a head that does not fit yet: the head keeps the room the pool frees until it fits, so
-// its wait is bounded by the slots ahead of it and not by how often a smaller slot can squeeze
-// in, grow, and be parked again. Simulated over 60 seeds at eight chats this cuts the longest
-// single wait by 2.5 to 3x for 0 to 3 percent of makespan at 8192 cells, and parks less often.
-// LLAMA_SERVER_PREEMPT_RESUME=pass keeps the previous order: most-preempted first, then longest
-// parked, and a smaller slot may pass a head that does not fit.
-// LLAMA_SERVER_PREEMPT_RESUME=head (the default) or pass; read once in load_model() and logged.
+// [TAG_PREEMPT] the order parked slots come back in. Default (LLAMA_SERVER_PREEMPT_RESUME=head,
+// read once in load_model()): head of the line by park time, and it keeps the room the pool
+// frees until it fits, bounding its wait by the slots ahead of it. =pass keeps the previous
+// order, most-preempted then longest parked, where a smaller slot may pass a head.
 static bool g_preempt_resume_head_of_line = true;
 
 static bool preempt_resume_head_of_line() {
@@ -109,30 +97,22 @@ constexpr int32_t PREEMPT_N_FAIL_MAX = 8;  // failed restores before the slot is
 constexpr int64_t PREEMPT_FAIL_US    = 60ll * 1000 * 1000;  // ... and only after this long parked
 constexpr int64_t PREEMPT_ROTATE_US  =  2ll * 1000 * 1000;  // a resident cycling through context shifts gives way to a parked head that has waited this long
 
-// [TAG_EXACT_CONCURRENCY] The planner above counts cells, not tokens, because the two are not
-// the same number under every mode. llama_memory_alloc_granularity() reports how many cells the
-// pool hands out at a time: 1 in every ordinary configuration, and the exact concurrency page
-// size when that mode is on, where one page belongs to one (sequence, position / page) pair and
-// a sequence of n tokens therefore occupies round_up(n, page) cells. Four sequences can be
-// holding up to 4 * (page - 1) cells that nobody else can be given, and a planner counting
-// tokens sees room in the pool that find_slot cannot find in pages: it never reaches the
-// threshold that would park anybody, the retry ladder halves n_batch to 1, and every request
-// ends in the context error that preemption exists to remove.
+// [TAG_EXACT_CONCURRENCY] the planner counts cells, not tokens: llama_memory_alloc_granularity()
+// is 1 ordinarily but the page size under exact concurrency, where a sequence of n tokens
+// occupies round_up(n, page) cells. A planner counting tokens would see room find_slot cannot
+// find in pages, never park anybody, and end every request in the context error instead.
 
 // cells a run of n_tokens occupies when the pool allocates g at a time
 static constexpr int32_t preempt_n_cells_g(int32_t n_tokens, int32_t g) {
     return (g <= 1 || n_tokens <= 0) ? n_tokens : ((n_tokens + g - 1) / g) * g;
 }
 
-// cells a run of n_tokens has to be given for a step of n_step more: nothing until the step
-// crosses a page boundary, a whole page when it does
+// cells a step of n_step more costs: nothing until it crosses a page boundary, a page when it does
 static constexpr int32_t preempt_n_cells_step_g(int32_t n_tokens, int32_t n_step, int32_t g) {
     return preempt_n_cells_g(n_tokens + n_step, g) - preempt_n_cells_g(n_tokens, g);
 }
 
-// At a granularity of 1 both are the identity, so every figure the planner computes is exactly
-// the arithmetic it did before it started asking the memory how it allocates, and nothing
-// changes in any configuration that does not page.
+// at a granularity of 1 both are the identity, so nothing changes in a configuration that does not page
 static_assert(preempt_n_cells_g(0, 1) == 0 && preempt_n_cells_g(1, 1) == 1 &&
               preempt_n_cells_g(8191, 1) == 8191 && preempt_n_cells_g(-3, 1) == -3,
               "at a granularity of 1 a run of n tokens has to cost exactly n cells");
@@ -140,7 +120,7 @@ static_assert(preempt_n_cells_step_g(0, 1, 1) == 1 && preempt_n_cells_step_g(819
               preempt_n_cells_step_g(1000, 512, 1) == 512,
               "at a granularity of 1 a step of n tokens has to cost exactly n cells");
 
-// and the page arithmetic itself, so the rounding cannot be changed by accident
+// and the page arithmetic itself, so the rounding cannot change by accident
 static_assert(preempt_n_cells_g(1, 256) == 256 && preempt_n_cells_g(256, 256) == 256 &&
               preempt_n_cells_g(257, 256) == 512,
               "a tail page is charged in full");
@@ -380,16 +360,13 @@ struct server_slot {
         prompt.clear();
     }
 
-    // [TAG_PREEMPT] state of a slot whose cells were taken back
-    //
-    // Only the KV cells leave. The task, the sampler, the generated text and the position
-    // the stream has reached stay on the slot, so a resume is a memcpy and not a new
-    // request: no retokenisation, no replayed prompt, no seam in the output.
+    // [TAG_PREEMPT] state of a slot whose cells were taken back. Only the KV cells leave;
+    // the task, sampler, generated text and stream position stay, so a resume is a memcpy.
     slot_state           state_before_preempt = SLOT_STATE_IDLE;
     std::vector<uint8_t> preempt_state_tgt;
     std::vector<uint8_t> preempt_state_dft;
     int32_t              n_preempt      = 0;   // times the CURRENT task has been preempted
-    int32_t              n_ctx_shift    = 0;   // context shifts the CURRENT task has made: it is at the pool's limit and cycling
+    int32_t              n_ctx_shift    = 0;   // context shifts it has made: it is at the pool's limit and cycling
     int32_t              n_preempt_fail = 0;   // consecutive failed restores
     int64_t              t_preempt_us   = 0;   // when it was parked
 
@@ -438,10 +415,8 @@ struct server_slot {
             return false;
         }
 
-        // The draft is a prediction, not a result, so it goes with the cells. Preemption
-        // runs before the batch is built, so spec_i_batch is empty and prompt.tokens already
-        // holds exactly the tokens the state above covers -- including the rollback done by
-        // the checkpoint path when a draft was only partially accepted.
+        // the draft is a prediction, not a result, so it goes with the cells; prompt.tokens
+        // already covers exactly what the state above holds
         spec_draft.clear();
         spec_i_batch.clear();
         spec_ckpt.clear();
@@ -449,8 +424,7 @@ struct server_slot {
 
         i_batch = -1;
 
-        // note: prompt.tokens is deliberately kept. It is the mirror of the state just
-        //       copied out, and the resume needs it to know how many cells to ask for.
+        // note: prompt.tokens is deliberately kept - the resume sizes its request from it
         mem.seq_rm(id, -1, -1);
 
         state_before_preempt = state;
@@ -487,10 +461,8 @@ struct server_slot {
 
         state = state_before_preempt;
 
-        // same call the DONE_PROMPT -> GENERATING transition makes; for MTP it only checks
-        // that the draft context is where it should be, which the restore above ensures.
-        // A slot parked while still processing its prompt makes that transition itself
-        // once the prompt is done.
+        // same call the DONE_PROMPT -> GENERATING transition makes; a slot parked mid-prompt
+        // makes that transition itself once the prompt is done
         if (state == SLOT_STATE_GENERATING && can_speculate()) {
             common_speculative_begin(spec, id, prompt.tokens.get_text_tokens());
         }
@@ -498,12 +470,9 @@ struct server_slot {
         return true;
     }
 
-    // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds for this sequence.
-    // For a batch that is given up after it was built: the tokens added for this slot
-    // that were never decoded come off, the sampled token stays in `sampled` and goes into
-    // the next batch the way it went into this one, and a draft is a prediction that goes
-    // with them. A chunk that failed to decode left nothing in the cache, so the cache is
-    // the boundary.
+    // [TAG_PREEMPT] bring prompt.tokens back to what the cache holds, for a batch given up
+    // after it was built: never-decoded tokens and the draft come off, `sampled` is kept for
+    // the next batch. A failed chunk left nothing in the cache, so the cache is the boundary.
     void rewind_to_cache() {
         const int32_t n_cached = llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), id) + 1;
 
@@ -511,8 +480,7 @@ struct server_slot {
             prompt.tokens.keep_first(n_cached);
         }
 
-        // a prompt whose last chunk was in the batch was marked done when the chunk was
-        // built; the chunk never ran, so the prompt is not done
+        // the last chunk was marked done when it was built but never ran, so it is not done
         if (state == SLOT_STATE_DONE_PROMPT && task && prompt.n_tokens() < task->n_tokens()) {
             state = SLOT_STATE_PROCESSING_PROMPT;
         }
@@ -743,9 +711,8 @@ struct server_slot {
 
             t_last_used = ggml_time_us();
 
-            // [TAG_PREEMPT] the cells are already gone (a cancelled or failed slot can be
-            // released while parked), so the mirror of them must not outlive them: the next
-            // task on this slot would otherwise take a prefix match against an empty cache
+            // [TAG_PREEMPT] the cells are already gone, so the mirror of them must not outlive
+            // them: the next task would take a prefix match against an empty cache
             if (state == SLOT_STATE_PREEMPTED) {
                 preempt_state_free();
                 prompt_clear();
@@ -1499,15 +1466,13 @@ struct server_context_impl {
             }
         }
 
-        // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates, rather than assume a cell per
-        // token. 1 in every ordinary configuration, so this changes nothing unless a mode that
-        // places cells in blocks is on.
+        // [TAG_EXACT_CONCURRENCY] ask the cache how it allocates rather than assume a cell per
+        // token; 1 unless a mode that places cells in blocks is on
         {
             preempt_alloc_granularity = (int32_t) std::max(1u, llama_memory_alloc_granularity(llama_get_memory(ctx_tgt)));
 
-            // a test knob: the paged attention kernel only supports a head size of 256, so a
-            // harness model cannot turn exact concurrency on, and this is the only way to reach
-            // the paged arithmetic of the planner from the server tests
+            // test knob: the paged kernel needs a head size of 256, so a harness model cannot
+            // turn exact concurrency on and this is the only way to reach the paged arithmetic
             const char * LLAMA_SERVER_PREEMPT_GRANULARITY = getenv("LLAMA_SERVER_PREEMPT_GRANULARITY");
 
             if (LLAMA_SERVER_PREEMPT_GRANULARITY) {
@@ -1537,9 +1502,8 @@ struct server_context_impl {
             preempt_test_every = LLAMA_SERVER_PREEMPT_EVERY ? atoi(LLAMA_SERVER_PREEMPT_EVERY) : 0;
 
             // LLAMA_SERVER_PREEMPT_POLICY: which non-leader the planner parks, for comparing
-            // policies against each other on the same workload. smallest (the default and the
-            // shipped one), largest, youngest (the most recent task, as vLLM's scheduler
-            // preempts), oldest. The leader is kept and the starvation guard applies under all.
+            // policies on the same workload: smallest (default), largest, youngest, oldest.
+            // The leader is kept and the starvation guard applies under all.
             const char * LLAMA_SERVER_PREEMPT_POLICY = getenv("LLAMA_SERVER_PREEMPT_POLICY");
             preempt_test_policy = LLAMA_SERVER_PREEMPT_POLICY ? LLAMA_SERVER_PREEMPT_POLICY : "smallest";
 
@@ -2954,9 +2918,8 @@ struct server_context_impl {
 
     void abort_all_slots(const std::string & reason) {
         for (auto & slot : slots) {
-            // [TAG_PREEMPT] a parked slot took no part in what failed: its sequence is in
-            // host RAM, not in the cache, and it comes back when there is room, the same as
-            // in the decode error sweep
+            // [TAG_PREEMPT] a parked slot took no part in what failed and comes back when
+            // there is room, the same as in the decode error sweep
             if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) {
                 send_error(slot, reason, ERROR_TYPE_SERVER);
                 slot.release();
@@ -3000,18 +2963,14 @@ struct server_context_impl {
 
 
     // LLAMA_SERVER_PREEMPT_EVERY=N preempts every generating slot every N generated tokens,
-    // whether or not the pool is under pressure. It exists to answer the only question that
-    // matters about a resume: with one request on an idle server the batch has the same
-    // shape at every step, so a preempted continuation that is not byte-identical to an
-    // uninterrupted one is the preemption's fault and nothing else's.
+    // under pressure or not: on an idle server the batch shape is fixed, so a continuation
+    // that is not byte-identical to an uninterrupted one is the preemption's fault.
     int32_t preempt_test_every = 0;
     std::string preempt_test_policy = "smallest"; // LLAMA_SERVER_PREEMPT_POLICY, see load_model
 
-    // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load from the
-    // memory itself: 1 in every ordinary configuration, the page size under exact concurrency.
-    // Everything below plans in cells because of it. LLAMA_SERVER_PREEMPT_GRANULARITY overrides
-    // it, which is how the harness reaches the paged arithmetic on a model whose head size the
-    // paged attention kernel does not support.
+    // [TAG_EXACT_CONCURRENCY] cells the pool hands out at a time, read once at load: 1
+    // ordinarily, the page size under exact concurrency. Everything below plans in cells
+    // because of it. LLAMA_SERVER_PREEMPT_GRANULARITY overrides it for the harness.
     int32_t preempt_alloc_granularity = 1;
 
     // cells a slot holding n_tokens actually occupies
@@ -3024,16 +2983,14 @@ struct server_context_impl {
         return preempt_n_cells_step_g(n_tokens, n_step, preempt_alloc_granularity);
     }
 
-    // Cells kept spare on top of the reservation. A step that crosses a page boundary costs a
-    // whole page rather than a cell, so a margin of a few cells is no margin at all under a page
-    // allocator: round it up to one page. With a granularity of 1 this is PREEMPT_N_MARGIN.
+    // cells kept spare on top of the reservation, rounded up to a page since a boundary
+    // crossing costs a whole one; PREEMPT_N_MARGIN at a granularity of 1
     int32_t preempt_n_margin() const {
         return preempt_n_cells(PREEMPT_N_MARGIN);
     }
 
-    // env: LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, so
-    // the KV-full retry ladder and its last resort are the only thing between a full pool
-    // and the context error
+    // LLAMA_SERVER_PREEMPT_PLANNER=off (test knob): no parking ahead of the decode, leaving
+    // only the KV-full retry ladder and its last resort
     bool preempt_planner_off = false;
 
     // set by preempt_last_resort(): the batch being decoded was given up, stop the chunk loop
@@ -3043,8 +3000,8 @@ struct server_context_impl {
         return spec ? std::max(0, common_speculative_n_max(&params_base.speculative)) : 0;
     }
 
-    // draft tokens this slot's next step can actually carry: the configured maximum, cut to
-    // what its context and its prediction budget leave, the way get_n_draft_max() cuts it
+    // draft tokens the next step can carry: the maximum cut to what context and prediction
+    // budget leave, the way get_n_draft_max() cuts it
     int32_t preempt_n_spec(const server_slot & slot) const {
         int32_t res = preempt_n_spec_max();
 
@@ -3083,10 +3040,8 @@ struct server_context_impl {
         return preempt_ram_used() + slot.preempt_state_required() <= budget;
     }
 
-    // the same for a rotation: the parked head is restored on the pass that parks the
-    // resident, so its bytes are on their way out and are not held against the resident.
-    // A budget that holds one sequence but not two would otherwise refuse every rotation
-    // and leave the head parked for as long as the resident cares to generate.
+    // the same for a rotation: the head is restored on the pass that parks the resident, so
+    // its bytes are on their way out and a one-sequence budget still allows the swap
     bool preempt_fits_budget_for_rotation(const server_slot & slot, const server_slot & head) const {
         if (params_base.preempt_ram_mib < 0) {
             return true;
@@ -3106,9 +3061,8 @@ struct server_context_impl {
         if (slot.state_before_preempt == SLOT_STATE_GENERATING) {
             res += 1 + preempt_n_spec(slot);
         } else {
-            // a slot just given a task still mirrors the previous request's prompt; the batch
-            // builder keeps the prefix the two share and drops the rest, so what it holds and
-            // what it is about to ask for both count from that prefix, not from the old prompt
+            // a slot just given a task still mirrors the previous prompt; the batch builder
+            // keeps only the shared prefix, so count from that prefix
             if (slot.state == SLOT_STATE_STARTED && slot.task) {
                 res = (int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens);
             }
@@ -3119,21 +3073,17 @@ struct server_context_impl {
         }
 
         // [TAG_EXACT_CONCURRENCY] a restore takes fresh pages and its tail page is charged in
-        // full, so what the pool has to have free for this slot is the rounded figure. Under
-        // counting here is what admits a resume that find_slot then cannot satisfy.
+        // full; undercounting here admits a resume that find_slot cannot satisfy
         return preempt_n_cells(res);
     }
 
-    // Cells the pool is holding right now. A released slot keeps its prompt in the cache
-    // for the next request to reuse as a prefix, so idle slots count too: the first version
-    // of this counted only the running ones, decided a pool holding 8185 cached cells was
-    // empty, and every resume failed against a cache that was actually full.
+    // cells the pool is holding right now. A released slot keeps its prompt in the cache as a
+    // prefix for the next request, so idle slots count too or a full pool looks empty.
     int32_t preempt_kv_used() const {
         int32_t res = 0;
 
-        // n_cmpl > 1: the parent and its children share the prompt's cells through seq_cp, so
-        // the prompt is charged once per family, to whichever resident member comes first;
-        // the others are charged only what they generated on top of it
+        // n_cmpl > 1: a family shares the prompt's cells through seq_cp, so the prompt is
+        // charged once, to the first resident member; the others only for what they added
         std::vector<int> charged;
 
         for (const auto & slot : slots) {
@@ -3141,12 +3091,11 @@ struct server_context_impl {
                 continue; // parked: its cells are in host RAM, not in the pool
             }
 
-            // [TAG_EXACT_CONCURRENCY] the slot's tail page is charged in full: it belongs to
-            // this sequence and cannot be given to anybody else, however little of it is used
+            // [TAG_EXACT_CONCURRENCY] the tail page is charged in full: it cannot be given to
+            // anybody else, however little of it is used
 
-            // a child waiting for its parent's prompt does not share anything yet: until
-            // copy_state_to() runs it still holds whatever the previous request left in its
-            // cells, so it is charged that on its own, outside the family
+            // a child waiting for its parent shares nothing until copy_state_to() runs, so it
+            // is charged its own stale cells, outside the family
             if (slot.state == SLOT_STATE_WAIT_OTHER) {
                 res += preempt_n_cells(slot.prompt.n_tokens());
                 continue;
@@ -3163,9 +3112,8 @@ struct server_context_impl {
                 charged.push_back(family);
             }
 
-            // a slot just given a task still mirrors the previous request's prompt until the
-            // batch builder keeps the prefix the two share and drops the rest; what stays is
-            // the prefix, so that is what the pool holds for it
+            // a slot just given a task keeps only the prefix it shares with the new prompt,
+            // so that is what the pool holds for it
             if (slot.state == SLOT_STATE_STARTED && slot.task) {
                 res += preempt_n_cells((int32_t) slot.prompt.tokens.get_common_prefix(slot.task->tokens));
                 continue;
@@ -3185,11 +3133,8 @@ struct server_context_impl {
         int32_t res_pmt = 0;
 
         // [TAG_EXACT_CONCURRENCY] each slot reserves the cells its next step ADDS, not the
-        // tokens it adds. preempt_kv_used() already charges every slot's tail page in full, so
-        // with a granularity of 1 these are the same number and nothing changes; with a larger
-        // one the step is free until it crosses a page boundary and costs a whole page when it
-        // does. Reserving tokens on top of a rounded used figure would miss exactly that
-        // crossing, which is the only moment the pool can actually run out.
+        // tokens: preempt_kv_used() already rounds up the tail page, so reserving tokens on
+        // top of it would miss the boundary crossing, the only moment the pool can run out
         for (const auto & slot : slots) {
             const int32_t n_cur = slot.prompt.n_tokens();
 
@@ -3211,11 +3156,9 @@ struct server_context_impl {
             }
         }
 
-        // one batch is all the prompt slots get between them, however many are waiting; in
-        // cells that batch can straddle one boundary more than it has tokens for
         // one batch is all the prompt slots get between them, however many are waiting; under
-        // page allocation each of them can still cross a page boundary of its own within that
-        // batch, so the cap keeps one boundary per prompt slot on top of the batch
+        // page allocation each can still cross a boundary of its own, so the cap allows one
+        // boundary per prompt slot on top of the batch
         int32_t n_pmt = 0;
 
         for (const auto & slot : slots) {
@@ -3227,16 +3170,10 @@ struct server_context_impl {
         return res + std::min(res_pmt, preempt_n_cells(n_batch) + std::max(0, n_pmt - 1) * (preempt_alloc_granularity - 1));
     }
 
-    // Keep the slot that is furthest along -- it is the closest to finishing and to giving
-    // its cells back -- and among the rest prefer one that has not been preempted
-    // PREEMPT_N_STARVED times already, then the smallest.
-    // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt
-    // until the batch builder keeps the prefix the two share and drops the rest (see the
-    // SLOT_STATE_STARTED block of update_slots). Parked as it is, it would be copied out,
-    // charged and sized by the old prompt, and a short unrelated request could exceed the
-    // budget or stay parked for room it will never use. Keeping only the shared prefix now
-    // is what the batch builder does anyway; the chunk reuse it can add on top is given up
-    // for a slot the planner has to touch, which is rare.
+    // [TAG_PREEMPT] a slot just given a task still mirrors the previous request's prompt until
+    // the batch builder drops it (see the SLOT_STATE_STARTED block of update_slots). Parked as
+    // it is, it would be copied out and sized by the old prompt. Keeping only the shared prefix
+    // now is what the batch builder does anyway, at the cost of the chunk reuse it can add.
     void preempt_normalize_started(server_slot & slot) {
         if (slot.state != SLOT_STATE_STARTED || !slot.task) {
             return;
@@ -3248,10 +3185,8 @@ struct server_context_impl {
             return;
         }
 
-        // a memory that cannot remove part of a sequence (a recurrent state without rollback
-        // room for the stale suffix) aborts on a partial removal; for it the whole stale
-        // sequence goes, and the prompt is processed from the start on resume, as it would be
-        // without a usable checkpoint
+        // a memory that cannot remove part of a sequence aborts on a partial removal; for it
+        // the whole stale sequence goes and the prompt is reprocessed from the start
         const bool partial_ok = ctx_tgt_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART &&
                                 (!ctx_dft || ctx_dft_seq_rm_type == COMMON_CONTEXT_SEQ_RM_TYPE_PART);
 
@@ -3269,10 +3204,8 @@ struct server_context_impl {
         server_slot * leader    = nullptr;
         int32_t       n_running = 0;
 
-        // a slot just given a task is measured by the prefix it keeps, not by the previous
-        // request's prompt it still mirrors: measured by the mirror, a short request over a
-        // large stale cache would be the never-parked leader while the longest live
-        // conversation was parked in its place
+        // measure a just-started slot by the prefix it keeps, not by the stale prompt it
+        // mirrors, or a short request over a large stale cache becomes the leader
         for (auto & slot : slots) {
             preempt_normalize_started(slot);
         }
@@ -3288,19 +3221,16 @@ struct server_context_impl {
         }
 
         if (n_running < 2) {
-            // a single conversation that does not fit the pool on its own is a real context
-            // overflow and not a scheduling problem - leave it to the existing error path
+            // one conversation that does not fit alone is a real overflow, not a scheduling
+            // problem - leave it to the existing error path
             return nullptr;
         }
 
         server_slot * victim = nullptr;
 
         for (auto & slot : slots) {
-            // Before the batch is built every one of these is at a token boundary: a
-            // generating slot between two sampled tokens, a prompt-processing slot between
-            // two chunks of its prompt, a started slot with only a cached prefix (or
-            // nothing) in the pool. A slot holding no cells is still worth parking - it
-            // is about to ask for a whole batch of them.
+            // before the batch is built every one of these is at a token boundary. A slot
+            // holding no cells is still worth parking - it is about to ask for a batch.
             if (slot.state != SLOT_STATE_GENERATING &&
                 slot.state != SLOT_STATE_PROCESSING_PROMPT &&
                 slot.state != SLOT_STATE_STARTED) {
@@ -3332,9 +3262,8 @@ struct server_context_impl {
         return victim;
     }
 
-    // is a the better victim of the two? the smallest slot under the shipped policy: it
-    // gives up the least work and its restore is the cheapest (see the PR's simulation);
-    // the other choices exist for the comparison runs behind LLAMA_SERVER_PREEMPT_POLICY
+    // is a the better victim? the smallest under the shipped policy, since it gives up the
+    // least work; the rest exist for comparison runs behind LLAMA_SERVER_PREEMPT_POLICY
     bool preempt_better_victim(const server_slot & a, const server_slot & b) const {
         if (preempt_test_policy == "largest") {
             return a.prompt.n_tokens() > b.prompt.n_tokens();
@@ -3351,10 +3280,8 @@ struct server_context_impl {
         return a.prompt.n_tokens() < b.prompt.n_tokens();
     }
 
-    // called once per update_slots(), before the batch is built: at that point every slot is
-    // at a token boundary, prompt.tokens is exactly what the cache holds for it, and no
-    // draft is in flight, so a slot can be removed from the picture without unpicking a
-    // half-decoded batch
+    // called once per update_slots(), before the batch is built: every slot is then at a token
+    // boundary with no draft in flight, so one can be removed without unpicking a batch
     void update_preemption() {
         if (!params_base.kv_unified || slots.size() < 2) {
             return; // with a cache per slot, no slot can take another one's cells
@@ -3370,10 +3297,7 @@ struct server_context_impl {
 
         const int32_t n_cells = n_ctx;
 
-        // Put back what fits, in the order preempt_resume_head_of_line() describes: by default
-        // the slot parked longest, and only that one until it fits; under
-        // LLAMA_SERVER_PREEMPT_RESUME=pass the most-preempted slot first, then the one parked
-        // longest, and a smaller slot may pass a head that does not fit.
+        // put back what fits, in the order preempt_resume_head_of_line() describes
         const bool head_of_line = preempt_resume_head_of_line();
 
         for (;;) {
@@ -3403,12 +3327,9 @@ struct server_context_impl {
 
             server_slot * best = nullptr;
 
-            // A parked slot whose sequence plus its next step would not fit an empty pool can
-            // never be restored, and would otherwise sit at the head of the line for ever
-            // without a restore ever being attempted: a prompt within n_ctx that was parked
-            // before it took any cells, but too close to n_ctx to leave room for its first
-            // batch. That is the single-conversation overflow the KV-full path reports, so
-            // report it the same way and rescan the line without it.
+            // a parked slot that would not fit an empty pool can never be restored and would
+            // sit at the head of the line for ever. That is the single-conversation overflow
+            // the KV-full path reports, so report it the same way and rescan without it.
             {
                 server_slot * impossible = nullptr;
 
@@ -3428,12 +3349,10 @@ struct server_context_impl {
                 }
             }
 
-            // Room for the sequence AND for the next step of everything already running,
-            // so that a resume cannot immediately trigger the preemption of someone else.
-            // The margin is headroom for the others; with nothing resident there is nobody
-            // to keep it for, so a sequence that fits the pool exactly is let back in.
-            // A cached prompt on an idle slot is worth less than a conversation waiting to
-            // continue, so give those cells up first - same call the KV-full path makes.
+            // room for the sequence AND for the next step of everything running, so a resume
+            // cannot immediately preempt someone else. The margin is headroom for the others,
+            // so with nothing resident a sequence that fits exactly is let back in. Cached
+            // prompts on idle slots are worth less than a waiting conversation, so go first.
             for (;;) {
                 const int32_t occupied = preempt_kv_used() + preempt_kv_reserve();
                 const int32_t margin   = occupied == 0 ? 0 : preempt_n_margin();
@@ -3450,20 +3369,15 @@ struct server_context_impl {
                 }
             }
 
-            // Nothing fits. A resident that has reached the pool's limit and is cycling
-            // through context shifts holds the room for as long as it likes to generate,
-            // and the head behind it would wait for ever. After the head has waited its
-            // turn, that resident is parked in its place: it is at a token boundary like
-            // any other park, and when it comes back it is the one waiting, so the two
-            // take turns instead of one taking everything.
+            // nothing fits. A resident cycling through context shifts holds the room for as
+            // long as it generates, so once the head has waited its turn that resident is
+            // parked in its place and the two take turns.
             if (!best) {
                 server_slot * head = parked.front();
 
                 if (ggml_time_us() - head->t_preempt_us >= PREEMPT_ROTATE_US) {
-                    // the resident whose cells let the head in, the smallest of those; failing
-                    // one that does so alone, the largest, since it makes the most room. Taking
-                    // the first shifting resident in slot order could park one too small to
-                    // matter, spend the park budget on it, and leave the head waiting anyway.
+                    // the smallest resident whose cells let the head in; failing one that does
+                    // so alone, the largest, since it makes the most room
                     const int32_t occupied = preempt_kv_used() + preempt_kv_reserve();
                     const int32_t need     = preempt_n_need(*head) + PREEMPT_N_MARGIN;
 
@@ -3520,9 +3434,8 @@ struct server_context_impl {
             const int64_t t_start = ggml_time_us();
 
             if (!best->preempt_restore()) {
-                // update_slots() runs in a tight loop while tasks are pending, so a counter
-                // alone burns its whole budget in a couple of milliseconds. Give up only on
-                // a slot that has been failing for a while, and keep the log quiet.
+                // update_slots() loops tightly, so a counter alone burns its budget in
+                // milliseconds: give up only on a slot failing for a while, and log quietly
                 if (best->n_preempt_fail % 64 == 1) {
                     SLT_WRN(*best, "resume failed (%d in a row, parked %.1f s), staying preempted\n",
                             best->n_preempt_fail, (ggml_time_us() - best->t_preempt_us) / 1e6);
@@ -3646,8 +3559,7 @@ struct server_context_impl {
             }
         }
 
-        // [TAG_PREEMPT] make the pool fit the step that is about to be built, measured after
-        // any context shift
+        // [TAG_PREEMPT] make the pool fit the step about to be built, measured after any shift
         pre_decode_shift();
         update_preemption();
 
@@ -3700,8 +3612,7 @@ struct server_context_impl {
 #endif
 
                 if (preempt_batch_abandoned) {
-                    // [TAG_PREEMPT] the rest of this batch was never decoded and the slots no
-                    // longer describe it; the next pass builds a new one
+                    // [TAG_PREEMPT] the rest of this batch never ran; the next pass rebuilds it
                     preempt_batch_abandoned = false;
                     break;
                 }
@@ -3735,8 +3646,7 @@ struct server_context_impl {
 
     // apply context-shift if needed
     // TODO: simplify and improve
-    // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift,
-    // not with the cells the shift is about to give back
+    // [TAG_PREEMPT] runs before update_preemption() so the pool is measured after the shift
     void pre_decode_shift() {
         iterate(slots, [&](server_slot & slot) {
             if (slot.state == SLOT_STATE_GENERATING && slot.prompt.n_tokens() + 1 >= slot.n_ctx) {
@@ -3949,8 +3859,7 @@ struct server_context_impl {
                     return; // batch is full, skip remaining slots
                 }
 
-                // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to
-                // batch; it takes no part in this pass until it is restored
+                // [TAG_PREEMPT] a parked slot is processing but has nothing in the cache to batch
                 if (!slot.is_processing() || slot.state == SLOT_STATE_PREEMPTED) {
                     return;
                 }
@@ -4477,14 +4386,11 @@ struct server_context_impl {
         }
     }
 
-    // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell. Upstream this is
-    // the context error for every slot in the batch. With a park budget the batch is given
-    // up instead: every resident slot is rewound to the token boundary the cache is at (a
-    // batch is applied one chunk at a time, and the chunk that failed left nothing behind),
-    // the smallest are parked until the planner's own bound holds again, and the next
-    // update_slots() rebuilds the batch from the survivors. The planner brings the parked
-    // ones back as cells free up. A multimodal prompt has no boundary the cache can name,
-    // so it keeps the old path.
+    // [TAG_PREEMPT] the retry ladder ran out: a single token found no cell, which upstream is
+    // the context error for every slot in the batch. With a park budget the batch is given up
+    // instead: resident slots are rewound to the token boundary the cache is at, the smallest
+    // are parked until the planner's bound holds, and the next update_slots() rebuilds the
+    // batch. A multimodal prompt has no boundary the cache can name, so it keeps the old path.
     bool preempt_last_resort_possible() const {
         return params_base.kv_unified && params_base.preempt_ram_mib != 0 && slots.size() >= 2 && llama_get_memory(ctx_tgt);
     }
@@ -4620,11 +4526,9 @@ struct server_context_impl {
             {
                 std::string err;
 
-                // [TAG_PREEMPT] with speculation on, a slot's sampled token and its draft have
-                // to stay in one view: a narrower view splits the group and the verify step
-                // throws for the slot whose tokens straddle it. Halving is no help there, so
-                // after the idle slots the ladder goes to its last resort straight away. With
-                // no budget to park into the ladder is what it always was.
+                // [TAG_PREEMPT] a slot's sampled token and its draft have to stay in one view,
+                // so halving would split the group and make the verify step throw: after the
+                // idle slots the ladder goes straight to its last resort
                 if (ret == 1 && n_batch > 1 && preempt_last_resort_possible() && batch_has_spec_groups()) {
                     if (try_clear_idle_slots()) {
                         SRV_WRN("%s", "failed to find free space in the KV cache, retrying after purging an idle slot\n");
@@ -4661,8 +4565,7 @@ struct server_context_impl {
                     SRV_ERR("%s off = %d, n_batch = %d, ret = %d\n", err.c_str(), off, n_batch, ret);
 
                     for (auto & slot : slots) {
-                        // [TAG_PREEMPT] a parked slot has nothing in this batch and nothing in the
-                        // cache; it is not part of this failure and comes back when there is room
+                        // [TAG_PREEMPT] a parked slot is not part of this failure
                         if (slot.is_processing() && slot.state != SLOT_STATE_PREEMPTED) {
                             send_error(slot, err);
                             slot.release();
@@ -5270,11 +5173,9 @@ std::unique_ptr<server_res_generator> server_routes::handle_completions_impl(
             task.params.oaicompat_cmpl_id = completion_id;
             task.params.oaicompat_model   = meta->model_name;
 
-            // [TAG_EXACT_CONCURRENCY] the children of an n_cmpl > 1 task are served by
-            // copying the parent's cells to another sequence id, and exact mode gives a KV
-            // page to one sequence, so there is nothing for that copy to land in. Refuse
-            // the request here, where it becomes a 400 the client can read, rather than
-            // letting it reach seq_cp with nothing to do.
+            // [TAG_EXACT_CONCURRENCY] children of an n_cmpl > 1 task are served by copying the
+            // parent's cells to another sequence id, and exact mode gives a page to a single
+            // sequence, so refuse here where it becomes a 400 rather than at seq_cp
             if (task.params.n_cmpl > 1 && server_exact_concurrency()) {
                 throw std::runtime_error(
                     "n > 1 is not supported while LLAMA_EXACT_CONCURRENCY is set: each "
diff --git a/tools/server/tests/unit/test_preempt.py b/tools/server/tests/unit/test_preempt.py
index 7170efea080..238afa98f41 100644
--- a/tools/server/tests/unit/test_preempt.py
+++ b/tools/server/tests/unit/test_preempt.py
@@ -5,10 +5,9 @@
 import pytest
 from utils import *
 
-# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked
-# (its sequence copied to host RAM, its cells released) instead of every slot being
-# terminated. Both tests need more than one slot and --kv-unified, which is the only
-# configuration where one slot can take another one's cells.
+# Preemption on a unified KV pool: when the next decode does not fit, one slot is parked (its
+# sequence copied to host RAM, its cells released) instead of every slot being terminated. Needs
+# more than one slot and --kv-unified, the only configuration where slots share cells.
 
 server = ServerPreset.tinyllama2()
 
@@ -57,9 +56,8 @@ def _complete(n_predict: int, prompt: str = "Hi how are you"):
 
 
 def test_forced_preemption_does_not_change_the_output():
-    # Park and restore the only running slot every 8 tokens. With one request the batch
-    # has the same shape at every step whether or not the slot was parked in between, so
-    # any difference in the output is the preemption's fault and nothing else's.
+    # park and restore the only running slot every 8 tokens: the batch shape is the same at every
+    # step, so any difference in the output is the preemption's fault
     global server
     server.n_ctx = 512
     server.start()
@@ -86,11 +84,9 @@ def test_forced_preemption_does_not_change_the_output():
 
 
 def test_two_slots_that_overflow_the_pool_together_both_finish():
-    # Each request alone fits in the pool: 8 prompt tokens plus 160 generated is well
-    # under 256. Together they do not, 336 against 256. Without preemption the retry
-    # ladder ends with "Context size has been exceeded" on every processing slot; with it
-    # the smaller slot is parked until the leader finishes and its cells are purged, and
-    # then it resumes from the token it was parked on.
+    # each request fits the pool alone (168 of 256 cells) but not together (336). Without
+    # preemption both end with "Context size has been exceeded"; with it the smaller is parked
+    # until the leader finishes, then resumes from the token it was parked on.
     global server
     server.n_ctx = 256
     server.start()
@@ -115,15 +111,11 @@ def test_two_slots_that_overflow_the_pool_together_both_finish():
 
 
 def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages():
-    # A pool that hands out cells in blocks gives a whole block to one sequence, so a sequence
-    # of n tokens occupies round_up(n, block) cells and holds the rest of its tail block against
-    # everybody else. The planner has to count those cells: counting tokens, it sees room the
-    # allocator cannot find, never parks anybody, and the retry ladder ends every request.
-    #
-    # llama_memory_alloc_granularity() reports the block size, and the only mode that returns
-    # more than 1 today is exact concurrency, whose paged attention kernel needs a head size this
-    # model does not have. LLAMA_SERVER_PREEMPT_GRANULARITY injects the figure instead: what is
-    # under test is the server's arithmetic, which is the same at 64 as at 256.
+    # a pool that allocates in blocks gives a whole block to one sequence, so n tokens occupy
+    # round_up(n, block) cells and the planner has to count cells: counting tokens it sees room
+    # the allocator cannot find, never parks anybody, and the retry ladder ends every request.
+    # LLAMA_SERVER_PREEMPT_GRANULARITY injects the block size, since the only mode that reports
+    # one needs a head size this model does not have; the arithmetic is the same at 64 as at 256.
     global server
     server.n_ctx = 256
     os.environ["LLAMA_SERVER_PREEMPT_GRANULARITY"] = "64"
@@ -142,9 +134,8 @@ def test_the_planner_counts_whole_pages_when_the_pool_allocates_in_pages():
     assert "preempted:" in text
     assert "resumed after" in text
 
-    # every figure the planner logs is a whole number of blocks: "kv N/256" is what the pool is
-    # holding and "(wanted N)" is that plus what the next decode reserves. Counting tokens, both
-    # land wherever the sequences happen to be.
+    # every figure the planner logs is a whole number of blocks: "kv N/256" is what the pool holds
+    # and "(wanted N)" is that plus the next decode's reservation
     held   = [int(n) for n in re.findall(r"kv (\d+)/256", text)]
     wanted = [int(n) for n in re.findall(r"\(wanted (\d+)\)", text)]
     assert held and wanted, f"the planner logged no figures:\n{text}"
@@ -184,9 +175,8 @@ def _prompt_of_about(n_tokens: int, salt: str = "") -> tuple[str, int]:
 
 
 def test_two_prompts_that_overflow_the_pool_together_both_finish():
-    # Neither slot ever generates before the pool is full: both are still processing their
-    # prompts. A prompt-processing slot is between two chunks of its prompt, which is as
-    # clean a boundary as the one between two sampled tokens, so it is parked the same way.
+    # neither slot generates before the pool is full: a slot between two chunks of its prompt is
+    # as clean a boundary as one between two sampled tokens, so it is parked the same way
     global server
     server.n_ctx = 256
     server.start()
@@ -215,20 +205,17 @@ def test_two_prompts_that_overflow_the_pool_together_both_finish():
 
 
 def test_a_generating_slot_and_a_large_prompt_both_finish():
-    # One slot is generating a long answer to a short prompt when a large prompt arrives
-    # beside it. Together they need far more than the pool has. The prompt is admitted
-    # chunk by chunk, whoever is smaller is parked when the pool fills, and both finish.
-    # This model produces a thousand tokens a second, so the second request is sent right
-    # behind the first rather than after a delay: its prompt takes several batches to
-    # process, which is enough for the two to overlap however fast the first one runs.
+    # a slot generating a long answer to a short prompt meets a large prompt arriving beside it,
+    # needing far more than the pool has: the prompt is admitted chunk by chunk, whoever is
+    # smaller is parked, and both finish. The second request follows immediately, since its
+    # prompt takes several batches and that is enough overlap however fast the first one runs.
     global server
     server.n_ctx = 256
     server.start()
     log = LogReader(server.log_path)
 
     prompt_b, n_b = _prompt_of_about(150, "Charlie")
-    # b lives long enough for the two to collide: the first run of this used 16 tokens
-    # and b was finished and purged before a had grown into it
+    # b has to live long enough for the two to collide
     n_predict_a = 230
     n_predict_b = 90
     assert 8 + n_predict_a <= 256 and n_b + n_predict_b <= 256
@@ -254,8 +241,8 @@ def _late(n_predict, prompt):
 
 
 def test_preempt_ram_zero_disables_preemption():
-    # --preempt-ram 0 is the switch back to the old behaviour: nothing is parked and the
-    # KV-full path ends the requests the way it always did.
+    # --preempt-ram 0 switches back to the old behaviour: nothing is parked and the KV-full path
+    # ends the requests the way it always did
     global server
     server.n_ctx = 256
     os.environ["LLAMA_ARG_PREEMPT_RAM"] = "0"
@@ -275,9 +262,8 @@ def test_preempt_ram_zero_disables_preemption():
 
 
 def test_metrics_and_slots_report_the_parked_state():
-    # A client that wants to tell a parked chat from a slow one reads /slots, and an
-    # operator reads /metrics. Both must show the preemption happening, and the counters
-    # must survive the requests finishing.
+    # /slots tells a parked chat from a slow one and /metrics reports it to an operator; both
+    # must show the preemption, and the counters must survive the requests finishing
     global server
     server.n_ctx = 256
     server.server_metrics = True
@@ -315,11 +301,9 @@ def test_metrics_and_slots_report_the_parked_state():
 
 
 def test_two_prompts_near_the_context_size_both_complete():
-    # Two prompts that each fit the context alone but not together. The second one is
-    # parked before it takes any cells, and it is close enough to n_ctx that its sequence
-    # plus its first batch would not leave the usual scheduling margin. It must still be
-    # restored once the first one finishes: with nothing resident there is nobody to keep
-    # the margin for. Before the fix it was parked for ever, with no restore ever tried.
+    # two prompts that each fit the context alone but not together. The second is parked before
+    # it takes any cells and is too close to n_ctx to leave the usual margin, but must still be
+    # restored once the first finishes: with nothing resident there is nobody to keep it for.
     global server
     server.n_ctx = 256
     # the whole prompt in one batch, so the parked slot's first step is the whole prompt
@@ -342,11 +326,9 @@ def test_two_prompts_near_the_context_size_both_complete():
 
 
 def test_the_last_resort_parks_instead_of_ending_everyone():
-    # With the planner off nothing is parked ahead of the decode, so two generations that
-    # fit alone but not together fill the pool until a single token finds no cell. That
-    # is where upstream ends every slot with the context error. Instead the batch is
-    # given up, the smaller slot is parked, the larger one finishes, and the parked one
-    # comes back and finishes too.
+    # with the planner off, two generations that fit alone but not together fill the pool until a
+    # single token finds no cell, where upstream ends every slot with the context error. Instead
+    # the batch is given up, the smaller slot is parked, and both finish.
     global server
     server.n_ctx = 256
     os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off"
@@ -400,10 +382,8 @@ def test_the_last_resort_works_with_an_unlimited_budget():
 
 
 def test_the_last_resort_rewinds_a_prompt_in_flight():
-    # Same, with a prompt being processed when the pool runs out: the chunk that failed
-    # is taken back off the slot's tokens and processed again after the resume, so the
-    # prompt is neither skipped nor fed twice. The prompt is far longer than a batch, so
-    # the failing chunk is a chunk of it, not its last token.
+    # same, with a prompt being processed when the pool runs out: the failed chunk comes back off
+    # the slot's tokens and is processed again after the resume, neither skipped nor fed twice
     global server
     server.n_ctx = 256
     os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off"
@@ -439,12 +419,10 @@ def _late(n_predict, prompt):
 
 
 def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_head():
-    # Two generations that each outgrow the pool on their own, with context shift on. The
-    # resident reaches the limit, shifts, keeps about half the pool and would keep going
-    # for as long as it has tokens to make, while the parked one never fits beside it.
-    # After the head has waited its turn the resident is parked in its place, and the two
-    # take turns until both finish. Long enough that the resident is still going when the
-    # head's turn comes: this model makes a couple of thousand tokens a second.
+    # two generations that each outgrow the pool, with context shift on: the resident shifts and
+    # would hold half the pool for as long as it generates, while the parked one never fits
+    # beside it. After the head has waited its turn the resident is parked in its place and the
+    # two take turns. n_predict is large enough that the resident is still going by then.
     global server
     server.n_ctx = 256
     server.enable_ctx_shift = True
@@ -468,9 +446,8 @@ def test_a_resident_cycling_through_context_shifts_takes_turns_with_a_parked_hea
 
 
 def test_the_rotation_parks_a_resident_that_lets_the_head_in():
-    # Three generations with no end in a 256-cell pool with context shift on: two residents
-    # cycle through shifts while the third waits parked. Every rotation must let the head
-    # in, so all three keep finishing their tokens and no stream ends short.
+    # three endless generations with context shift on: two residents cycle through shifts while
+    # the third waits parked, and every rotation must let the head in so no stream ends short
     global server
     server.n_slots = 3
     server.n_ctx = 384
@@ -496,11 +473,9 @@ def test_the_rotation_parks_a_resident_that_lets_the_head_in():
 
 
 def test_a_parent_and_child_that_do_not_fit_alone_get_the_context_error_and_the_server_lives():
-    # One request asking for two completions is one conversation in two slots: a parent
-    # and a child sharing the prompt. When the two together do not fit the pool there is
-    # nobody else to park, since the family is charged once and a member of it is not a
-    # victim for the other, so the request gets the context error it would get alone, and
-    # the server carries on serving.
+    # a two-completion request is one conversation in two slots, and a family member is not a
+    # victim for the other, so with nobody else to park it gets the context error it would get
+    # alone and the server carries on serving
     global server
     server.n_ctx = 256
     os.environ["LLAMA_SERVER_PREEMPT_PLANNER"] = "off"