fix: the code review's findings across the server, the k-pool, sparse attention, the MoE kernels and the graph backend

common/speculative.cpp

@@ -1915,6 +1915,42 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         pending_pos[seq_id] = verify_pos[seq_id] + i_h;
     }
 
+    // the kept row and the position it is for, so a restored checkpoint pairs the token after it with the row the
+    // generation it belongs to produced. A sequence with no kept row has no state, and an empty one clears it: a slot
+    // whose cells are removed or replaced by another conversation's keeps a row that is no longer the one before
+    // anything (position alone does not tell the two apart, h_before)
+    bool get_state(llama_seq_id seq_id, std::vector<uint8_t> & data) const override {
+        if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq || pending_pos[seq_id] < 0) {
+            return false;
+        }
+
+        const llama_pos pos = pending_pos[seq_id];
+
+        data.resize(sizeof(llama_pos) + (size_t) n_embd * sizeof(float));
+        std::memcpy(data.data(),                     &pos,                       sizeof(llama_pos));
+        std::memcpy(data.data() + sizeof(llama_pos), pending_h[seq_id].data(), (size_t) n_embd * sizeof(float));
+        return true;
+    }
+
+    void set_state(llama_seq_id seq_id, const std::vector<uint8_t> & data) override {
+        if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) {
+            return;
+        }
+        if (data.empty()) {
+            pending_pos[seq_id] = -1;
+            return;
+        }
+        if (data.size() != sizeof(llama_pos) + (size_t) n_embd * sizeof(float)) {
+            return;
+        }
+
+        llama_pos pos = -1;
+        std::memcpy(&pos, data.data(), sizeof(llama_pos));
+
+        pending_pos[seq_id] = pos;
+        std::memcpy(pending_h[seq_id].data(), data.data() + sizeof(llama_pos), (size_t) n_embd * sizeof(float));
+    }
+
     void end(llama_seq_id seq_id) override {
         if (seq_id >= 0 && seq_id < (llama_seq_id) n_seq) {
             generating[seq_id] = 0;

ggml/include/ggml.h

@@ -2468,7 +2468,7 @@ extern "C" {
             const struct ggml_tensor * a);
 
     // Use finite mask entries as a sparse K/V set. Set 0 to disable.
-    // n_kv_max must bound the number of finite entries in every mask row.
+    // n_kv_max must bound the number of finite entries in every mask row: CUDA traps on a row over it rather than drop cells.
     GGML_API void ggml_flash_attn_ext_set_n_kv_max(
             struct ggml_tensor * a,
             int32_t              n_kv_max);
@@ -2606,7 +2606,9 @@ extern "C" {
     // tensor shapes (S_k == S_v, H_v % H_k == 0):
     //   q, k  : [S_k, H_k, n_tokens, n_seqs]
     //   v     : [S_v, H_v, n_tokens, n_seqs]
-    //   g     : [1, H_v, n_tokens, n_seqs] (scalar gate) or [S_v, H_v, n_tokens, n_seqs] (KDA)
+    //   g     : [1, H_v, n_tokens, n_seqs] (scalar gate) or [S_v, H_v, n_tokens, n_seqs] (KDA), a log-decay: KDA's must be
+    //           <= 0 (a decay exp(g) of at most 1), since CUDA's chunked path stages exp of its sums in fp16 and traps on
+    //           one above 0; a raw KDA gate (ggml_gated_delta_net_set_raw_kda_gates) with lower_bound > 0 runs recurrent
     //   beta  : [1, H_v, n_tokens, n_seqs]
     //   state : [S_v, S_v, H_v, n_seqs] -- initial recurrent state s0
     //

ggml/src/ggml-cuda/chunk_gated_delta_net.cu

@@ -311,6 +311,7 @@ __launch_bounds__(128, 4) __global__ void cgdr_kda_fwdsub_intra_kernel(
 
     // Step 0: q (scaled), k, the activated log-decay and beta into SMEM; tokens past seq_len are zeros, which neither
     // decay the state nor update it
+    bool g_above_0 = false;
     for (int i = tid; i < CS * BK; i += 128) {
         const int t = i / BK, c = i % BK;
         float     qv = 0.f, kv = 0.f, gv = 0.f;
@@ -320,6 +321,8 @@ __launch_bounds__(128, 4) __global__ void cgdr_kda_fwdsub_intra_kernel(
             gv = g_chunk[t * sb2 * BK + c];
             if constexpr (RAW) {
                 gv = raw_lb * (1.0f / (1.0f + expf(gv * raw_a[h])));
+            } else {
+                g_above_0 |= gv > 0.0f;
             }
         }
         s_q[t * sk + c] = qv;
@@ -336,7 +339,15 @@ __launch_bounds__(128, 4) __global__ void cgdr_kda_fwdsub_intra_kernel(
         }
         s_beta[i] = beta_val;
     }
-    __syncthreads();
+    // stage 3's fp16 operands hold exp of the gate's sums, at most 1 for a log-decay at or under 0 (ggml_gated_delta_net);
+    // a gate above it saturates them to inf, so the launch traps. A raw gate is under 0 by eligibility (raw_lb <= 0)
+    if constexpr (RAW) {
+        __syncthreads();
+    } else if (__syncthreads_or(g_above_0) && tid == 0) {
+        printf("ERROR: chunked KDA: a log-decay above 0 in chunk %d of head %d; ggml_gated_delta_net's g must be <= 0 "
+               "(GGML_CUDA_KDA_CHUNKED_LEGACY=1 takes the recurrent kernel)\n", pid_chunk, h);
+        __trap();
+    }
 
     // Step 1: G, each channel's prefix sum along the chunk
     {

ggml/src/ggml-cuda/dsv4-hc.cu

@@ -938,7 +938,9 @@ static __device__ __forceinline__ void dsv4_hc_gram_pre(const float * mix, const
         }
         sumsq += pre[a]*row;
     }
-    rms = rsqrtf(sumsq/pr.n_embd + pr.eps_norm);
+    // pre' G pre is a sum of squares, but rounded it can fall under 0 when the streams cancel, and rsqrtf of a negative
+    // makes the whole token NaN
+    rms = rsqrtf(fmaxf(sumsq, 0.0f)/pr.n_embd + pr.eps_norm);
 }
 
 // token it's pre, post and comb weights into weights_out, by one warp: the comb in registers (dsv4_hc_comb_regs), lane

ggml/src/ggml-cuda/fattn.cu

@@ -6,8 +6,19 @@
 #include "fattn.cuh"
 
 #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA)
+// A row with more finite cells than n_kv_max breaks ggml_flash_attn_ext_set_n_kv_max's contract, and its first n_kv_max
+// would attend to part of the row without a sign: the launch traps instead, and the next sync fails
+static __device__ __forceinline__ void flash_attn_sparse_check_count(
+        const int count, const int n_kv_max, const int query, const int sequence) {
+    if (count > n_kv_max && threadIdx.x == 0) {
+        printf("ERROR: sparse flash attention: mask row %d of sequence %d has %d finite cells, over n_kv_max %d\n",
+               query, sequence, count, n_kv_max);
+        __trap();
+    }
+}
+
 // A block a mask row: its finite entries' cells, ascending, into the row's n_kv_max indices, -1 past its count (a count
-// over n_kv_max keeps the first n_kv_max: the bound is the graph's to keep).
+// over n_kv_max writes the first n_kv_max and traps: the bound is the graph's to keep).
 // Upstream's scan (8e93a9773): 256 threads, 2048 columns a round of scalar loads. A decode's one row is one block, so
 // a round is a memory round trip: 11 us at 32K columns. GGML_CUDA_FATTN_SPARSE_SCAN_LEGACY=1, or a mask row not
 // 16-byte aligned.
@@ -84,6 +95,7 @@ static __global__ void flash_attn_mask_to_sparse_indices_legacy(
     }
 
     const int count = row_count;
+    flash_attn_sparse_check_count(count, n_kv_max, query, sequence);
     for (int i = count + tid; i < n_kv_max; i += blockDim.x) {
         indices[i] = -1;
     }
@@ -173,6 +185,7 @@ static __global__ void flash_attn_mask_to_sparse_indices(
         __syncthreads(); // the next chunk rewrites warp_sums
     }
 
+    flash_attn_sparse_check_count(row_count, n_kv_max, query, sequence);
     for (int i = row_count + tid; i < n_kv_max; i += sparse_scan_threads) {
         indices[i] = -1;
     }
@@ -255,7 +268,9 @@ void ggml_cuda_flash_attn_ext_gather_k_q8_0(
 #endif // defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
 }
 
-// GGML_CUDA_FATTN_SPARSE_LEGACY=1: the hint is ignored, the kernel reads the whole cache under the mask.
+// GGML_CUDA_FATTN_SPARSE_LEGACY=1: the hint is ignored, and the dense kernel runs as it did before the gather: over the
+// KV steps a row of its Q tile sees (launch_fattn's kv_live, 2f0ca785a), the whole cache under the mask only with
+// GGML_CUDA_FATTN_LIVE_TILES_LEGACY=1 as well.
 // The gather reads n_kv_max cells a query; the dense kernel reads the cache once for up to 64/ncols2 queries. The gather
 // wins where the cache is at least those queries' cells (GLM-5.3's DSA shape, 32 heads on the latent, RTX 5080 and
 // 5070 Ti: at 0.25-0.99x the dense time wherever K >= n_gather, at 1.13-1.27x wherever it is under; upstream's 2x margin

ggml/src/ggml-cuda/gated_delta_net.cu

@@ -428,7 +428,10 @@ bool ggml_cuda_gdn_chunked_shape_eligible(const ggml_tensor * dst) {
     // - q/k/v rows contiguous with any head/token/seq stride (the views qwen35 takes of the conv output are
     //   read in place; q and k share strides), g/beta/state contiguous
     // - n_tokens >= 128; with K > 1 snapshot slots the last K-1 tokens go to the recurrent kernel
-    return (!kda || kda_chunked) && dst->src[6] == nullptr
+    // - KDA's raw gate, lower_bound * sigmoid, at or under 0: the chunked kernels stage exp of its sums in fp16, which a
+    //   gate above 0 saturates (an activated gate above 0 traps there instead, cgdr_kda_fwdsub_intra_kernel)
+    const bool kda_raw_ok = !kda || ggml_get_op_params_i32(dst, 1) != 2 || ggml_get_op_params_f32(dst, 2) <= 0.0f;
+    return (!kda || kda_chunked) && kda_raw_ok && dst->src[6] == nullptr
         && dst->type == GGML_TYPE_F32 && src_q->type == GGML_TYPE_F32 && src_k->type == GGML_TYPE_F32
         && src_v->type == GGML_TYPE_F32 && src_g->type == GGML_TYPE_F32 && src_beta->type == GGML_TYPE_F32
         && src_state->type == GGML_TYPE_F32

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -142,10 +142,7 @@ void ggml_cuda_set_device(int device) {
 // attribution taken from a trace is unreliable at this scale -- including the one that put 66 % of a decode token's
 // 444.6 us device-idle gap on this call. Off (the default) this costs one test of a static flag.
 static bool ggml_cuda_time_launch() {
-    static const bool enabled = [] {
-        const char * env = getenv("GGML_CUDA_TIME_LAUNCH");
-        return env != nullptr && atoi(env) == 1;
-    }();
+    static const bool enabled = ggml_env_switch("GGML_CUDA_TIME_LAUNCH");
     return enabled;
 }
 
@@ -168,10 +165,7 @@ static void ggml_cuda_time_launch_record(int device, double us) {
 // GGML_CUDA_COUNT_INFLIGHT=1 prints, before each replay launch, the device's earlier replays not yet finished, so the last
 // line before a hang tells a host that ran far ahead (many) from one launch that cannot enqueue (0 or 1).
 static bool ggml_cuda_count_inflight() {
-    static const bool enabled = [] {
-        const char * env = getenv("GGML_CUDA_COUNT_INFLIGHT");
-        return env != nullptr && atoi(env) == 1;
-    }();
+    static const bool enabled = ggml_env_switch("GGML_CUDA_COUNT_INFLIGHT");
     return enabled;
 }
 
@@ -2073,7 +2067,10 @@ static bool ggml_cuda_pq2_mma_group_enabled() {
 
 // A MUL_MAT that ggml_cuda_mul_mat sends to ggml_cuda_mul_mat_vec_q and that sends it, unfused, to the tensor-core kernel
 static bool ggml_cuda_pq2_mma_group_member(ggml_backend_cuda_context & ctx, const ggml_tensor * mm) {
-    return mm->op == GGML_OP_MUL_MAT && ggml_get_op_params_i32(mm, 1) != GGML_HINT_SRC0_IS_HADAMARD &&
+    // a node the graph does not compute (GGML_TENSOR_FLAG_COMPUTE unset, ggml_build_forward_select) is no member: the node
+    // loop skips it, and a group would compute it
+    return mm->op == GGML_OP_MUL_MAT && (mm->flags & GGML_TENSOR_FLAG_COMPUTE) && !ggml_is_empty(mm) &&
+        ggml_get_op_params_i32(mm, 1) != GGML_HINT_SRC0_IS_HADAMARD &&
         mm->src[1]->nb[0] == sizeof(float) && mm->nb[0] == sizeof(float) &&
         ggml_cuda_should_use_mmvq(GGML_TYPE_PQ2_0, ggml_cuda_info().devices[ctx.device].cc, mm->src[1]->ne[1]) &&
         ggml_cuda_pq2_mma_fuses(ctx, mm, nullptr, nullptr);
@@ -3600,7 +3597,7 @@ static ggml_cuda_graph_key ggml_cuda_graph_get_key(ggml_cgraph * cgraph) {
     static const bool legacy = ggml_env_switch("GGML_CUDA_GRAPH_KEY_LEGACY");
 
     ggml_cuda_graph_key key;
-    key.first_node = cgraph->nodes[0];
+    key.first_node = cgraph->n_nodes > 0 ? cgraph->nodes[0] : nullptr; // an empty graph's slot 0 is no node of it
     if (legacy) {
         return key;
     }
@@ -4590,13 +4587,26 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
     };
 
     bool is_ok = true;
-    // exception for topk-moe, which reads all its rows before it writes when they fit one block
-    // (ggml_cuda_topk_moe_reads_before_writes). GGML_CUDA_TOPK_MOE_ALIAS_LEGACY=1: at one row only
-    static const bool topk_moe_alias_legacy = ggml_env_switch("GGML_CUDA_TOPK_MOE_ALIAS_LEGACY");
-    const int64_t     topk_moe_rows         = ggml_nrows(cgraph->nodes[node_idx]);
-    if (is_topk_moe && (topk_moe_alias_legacy ? topk_moe_rows == 1 : ggml_cuda_topk_moe_reads_before_writes(topk_moe_rows))) {
+    // exception for topk-moe at one row: its one warp reads every src before it writes
+    const int64_t topk_moe_rows = ggml_nrows(cgraph->nodes[node_idx]);
+    if (is_topk_moe && topk_moe_rows == 1) {
         return true;
     }
+    // at the rows that fit one block, its warps meet at a barrier once the logits (the first node's src) are read, so the
+    // outputs may overlap the logits themselves, or gpt-oss's reshape of them; the bias is read past the barrier, so it
+    // is checked (ggml_cuda_topk_moe_reads_before_writes). GGML_CUDA_TOPK_MOE_ALIAS_LEGACY=1: at one row only
+    static const bool   topk_moe_alias_legacy = ggml_env_switch("GGML_CUDA_TOPK_MOE_ALIAS_LEGACY");
+    const ggml_tensor * read_first            = is_topk_moe && !topk_moe_alias_legacy &&
+        ggml_cuda_topk_moe_reads_before_writes(topk_moe_rows) ? cgraph->nodes[node_idx]->src[0] : nullptr;
+    // the logits and a reshape of them, which share both, and nothing else: a src merely CONTAINED in the logits is a
+    // strict sub-view, and the two that can be one are both read past the barrier, not before it. The bias is one (its
+    // every instance here is a model weight, so GGML_OP_NONE above already skips it, and a weight shares no buffer with
+    // the logits); the other is the tail of a logits tensor with ne[2] > 1, which ggml_nrows counts and the launch's
+    // logits->ne[1] does not read.
+    auto reads_first_whole = [&](const ggml_tensor * src) {
+        return read_first != nullptr && src->buffer == read_first->buffer && src->data == read_first->data &&
+            ggml_nbytes(src) == ggml_nbytes(read_first);
+    };
 
     for (int i = 0; i < out_count; ++i) {
         const ggml_tensor * dst = cgraph->nodes[out_nodes[i]];
@@ -4609,7 +4619,7 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph,
             for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) {
                 const ggml_tensor * src = cgraph->nodes[j]->src[src_idx];
 
-                if (!src || src->op == GGML_OP_NONE || src == staged_src) {
+                if (!src || src->op == GGML_OP_NONE || src == staged_src || reads_first_whole(src)) {
                     continue;
                 }
 

ggml/src/ggml-cuda/mmq.cu

@@ -320,7 +320,9 @@ int64_t ggml_cuda_mmq_moe_ncols() {
 }
 
 // One block: each pass scans the column tiles of blockDim.x experts, ceil(columns/J) each, on top of the earlier
-// passes' total, and every expert writes its own entries at its exclusive prefix.
+// passes' total, and every expert writes its own entries at its exclusive prefix. A list longer than the launch's
+// column tiles (ntiles_max, whose first term counts an expert at most once a token) would leave an expert's last
+// columns unlaunched and their rows of dst unwritten, with no sign: the launch traps instead, and the next sync fails.
 static __global__ void mmq_moe_tiles_kernel(const int32_t * __restrict__ expert_bounds, const int n_expert, const int J,
                                             const int ntiles_max, int32_t * __restrict__ moe_tiles) {
     __shared__ int warp_incl[WARP_SIZE];
@@ -365,7 +367,14 @@ static __global__ void mmq_moe_tiles_kernel(const int32_t * __restrict__ expert_
         __syncthreads();
 
         const int first = done + (warp > 0 ? warp_incl[warp - 1] : 0) + incl - n;
-        for (int k = 0; k < n && first + k < ntiles_max; ++k) {
+        if (first + n > ntiles_max) {
+            // first and n, not a range: an expert that takes none still trips this once the running offset is past the
+            // bound, and "tiles 64 to 63" names no tile
+            printf("ERROR: MUL_MAT_ID: expert %d takes %d column tiles from %d, over the launch's %d (an expert named "
+                   "twice for one token?)\n", e, n, first, ntiles_max);
+            __trap();
+        }
+        for (int k = 0; k < n; ++k) {
             moe_tiles[first + k] = (k << 16) | e;
         }
         __syncthreads();

ggml/src/ggml-cuda/mmq.cuh

@@ -1607,7 +1607,9 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
 
     // MUL_MAT_ID: the grid's column tiles are the experts' non-empty ones back to back, one channel. An expert of c
     // columns fills ceil(c/J) <= ntx of them and an empty one none, so with n columns over n_expert experts there are
-    // at most (n + min(n_expert, n)*(J - 1))/J.
+    // at most (n + min(n_expert, n)*(J - 1))/J. The first term holds while ids names an expert at most once for a
+    // token, which every graph here builds (a top-k of distinct experts); mmq_moe_tiles_kernel traps on a list past
+    // this bound rather than drop an expert's last tiles.
     int64_t nchannels_grid = args.nchannels_y;
     ggml_cuda_pool_alloc<int32_t> moe_tiles(ctx.pool(id));
     if (args.ids_dst && !ggml_cuda_mmq_moe_tiles_legacy()) {

ggml/src/ggml-cuda/mmvq-moe.cu

@@ -201,7 +201,12 @@ static __device__ __forceinline__ void mmvq_moe_quantize_y(const mmvq_moe_dev_ar
     constexpr int NT    = MMVQ_MOE_NW*32;
     const int     nf4   = ncols / 4;
     const int     total = a.ntokens*a.n_used*nf4; // a whole number of q8_1 blocks: 8 lanes are all in or all out
-    for (int j0 = threadIdx.x; j0 < total; j0 += MMVQ_MOE_QY_CHUNK*NT) {
+    // the bound is rounded up to a whole warp, so a warp's lanes are all in or all out: the shuffles below name all 32
+    // lanes, and a warp the bound split would wait on lanes that never reach them. Today total is a multiple of 32
+    // only because ggml_cuda_mmvq_moe_y_bytes admits a copy of a whole number of 16-byte units, which holds the rows to
+    // 128 columns; the j < total guards cover the load and the store
+    const int bound = (total + 31)/32*32;
+    for (int j0 = threadIdx.x; j0 < bound; j0 += MMVQ_MOE_QY_CHUNK*NT) {
         float4 v[MMVQ_MOE_QY_CHUNK];
 #pragma unroll
         for (int c = 0; c < MMVQ_MOE_QY_CHUNK; ++c) {