perf(cuda): the recurrent-state GET_ROWS folds into GATED_DELTA_NET's load, f32 or q8_0 cache (#220 plus -cts q8_0)

TORAD.md

@@ -33,6 +33,7 @@ which pins a commit of this branch as a submodule.
 | `d79f6e0` | a sequence's first KQ mask row tests the KV cells 16 at a time and masks a group past the sequence's last ubatch position on the positions alone, before any sequence bitset is read (a bit-packed mask is written a word at a time). With `--kv-unified` and another slot holding 118K tokens, building the mask went 230 → 79 µs for a verify and 251 → 73 µs for a draft step (131,328 cells). Served together with the next two rows (RTX 5070 Ti, 2 slots × 131,072, MTP n_max 2, greedy, legs new/legacy/legacy/new): empty pool +2.6 %, a slot decoding beside a resident 118K context +9.5 % (8 of 8 requests +8.9 to +9.7 %), the long slot's own decode at 118K +2.6 %; text identical 13/13 across all legs. Fuzzed against the old builder over 20,000 random caches (streams, multi-sequence cells, 2D positions, SWA, ALiBi, every mask layout): 0 byte mismatches | `LLAMA_KQ_MASK_SCAN_LEGACY=1` |
 | `6ea8df2` | a speculative rollback's `seq_rm` (every verify, on the target's and the draft's cache) walks back from the head to the last of the 0-2 cells it removes instead of scanning every cell; more than 1,024 cells keep the scan. Measured with the row above; fuzzed against the old walk over 72,709 calls: 0 mismatches | `LLAMA_KV_SEQ_RM_SCAN_LEGACY=1` |
 | `da382d9` | the batch's coupled-sequence check (n_seq_max² flags per decode, 256² with `--kv-unified`) runs only when a token is in several sequences, the only way a flag is set. Measured with the rows above | `LLAMA_BATCH_CPL_CHECK_LEGACY=1` |
+| `8fd3a28`, `b534681` | a Gated DeltaNet layer's recurrent state is read by the GATED_DELTA_NET kernel straight from its cache row (`ids[seq]`), f32 or q8_0 dequantized as it loads with get_rows' formula, and the GET_ROWS that gathered it into a temp is skipped (from PrismML-Eng/llama.cpp#220, plus the q8_0 cache); declined for more than one sequence, another reader of the temp, a node between them that writes the cache or sits on ids' memory, and the chunked prefill. RTX 5070 Ti, served flags (`-cts q8_0`, MTP n_max 2), Nsight per step at 1 / 3 rows: the dequantizing gather's 48 launches (270 µs) gone, the GDN kernel +39 µs, 12.698 → 12.338 ms / 14.774 → 14.427 ms (−2.8 / −2.3 %, A-B-B-A), perplexity identical to the last digit; served with the MTP draft head, 48 greedy eval requests, legs cancelled: +1.69 % tok/s (95 % CI +1.59 to +1.80), 48 of 48 faster, text identical 48/48, same draft acceptance | `GGML_CUDA_GDN_STATE_GATHER_LEGACY=1` |
 
 Every switch in the last column is read once per process and parses as an integer: a `*_LEGACY` switch set to `0`
 is the same as unset (the change stays on), and `=0` turns off `GGML_CUDA_LORA_RANK1_FUSE` and

ggml/src/ggml-cuda/common.cuh

@@ -1486,6 +1486,36 @@ struct ggml_cuda_stream_context {
     }
 };
 
+// Fused recurrent-state gather for GATED_DELTA_NET (PrismML-Eng/llama.cpp#220, plus a q8_0 cache): build_rs
+// materialises GET_ROWS(cache, s_copy) into a temp that only the GDN kernel reads; the graph evaluator skips that
+// GET_ROWS (ggml_cuda_try_gdn_gather_skip) and records the gather here, so the kernel reads each sequence's state
+// row ids[seq] out of the cache itself.
+struct ggml_cuda_gated_delta_net_gather {
+    const void *    base       = nullptr; // cache rows, f32 or q8_0
+    const int32_t * ids        = nullptr; // per-seq row index
+    int64_t         row_stride = 0;       // between rows, in elements
+    bool            q8_0       = false;   // a q8_0 cache (-cts q8_0): the kernel dequantizes as it loads
+};
+
+// Owned by the backend context that evaluates the graph: registrations are keyed by node pointer, so they only mean
+// something for the evaluation that made them. Cleared at the start of every graph evaluation/capture.
+struct ggml_cuda_gdn_gather_context {
+    std::unordered_map<const ggml_tensor *, ggml_cuda_gated_delta_net_gather> gathers;
+
+    void reset() {
+        gathers.clear();
+    }
+
+    void set(const ggml_tensor * gdn, const ggml_cuda_gated_delta_net_gather & gather) {
+        gathers[gdn] = gather;
+    }
+
+    const ggml_cuda_gated_delta_net_gather * find(const ggml_tensor * gdn) const {
+        const auto it = gathers.find(gdn);
+        return it == gathers.end() ? nullptr : &it->second;
+    }
+};
+
 // Names one ggml graph across computes, to find its CUDA graph. The first node's address alone names the memory a
 // graph is built in, and graphs of different shapes built there in turn (a speculative draft context's catch-up batch
 // and its one-row draft step) would share one CUDA graph, each finding the other's capture stale on every compute and
@@ -1589,6 +1619,7 @@ struct ggml_backend_cuda_context {
     }
 
     ggml_cuda_stream_context concurrent_stream_context;
+    ggml_cuda_gdn_gather_context gdn_gather_context;
 
     ~ggml_backend_cuda_context();
 
@@ -1604,6 +1635,8 @@ struct ggml_backend_cuda_context {
 
     ggml_cuda_stream_context & stream_context() { return concurrent_stream_context; }
 
+    ggml_cuda_gdn_gather_context & gdn_gathers() { return gdn_gather_context; }
+
     cublasHandle_t cublas_handle() {
         if (cublas_handles[device][curr_stream_no] == nullptr) {
             ggml_cuda_set_device(device);

ggml/src/ggml-cuda/gated_delta_net.cu

@@ -36,6 +36,19 @@ static __device__ __forceinline__ void gdn_store_state(void * dst, const int64_t
     }
 }
 
+// Reads element e of a state source: f32, or a q8_0 cache row (the fused gather, ggml_cuda_try_gdn_gather_skip). q8_0
+// is dequantized with get_rows' formula (dequantize_q8_0), so the kernel sees the values the skipped GET_ROWS wrote.
+static __device__ __forceinline__ float gdn_load_state(const void * src, const int64_t e, const bool q8) {
+    if (q8) {
+        const block_q8_0 * b = (const block_q8_0 *) src + e / QK8_0;
+        const float        d = b->d;
+        float              v = b->qs[e % QK8_0];
+        v *= d;
+        return v;
+    }
+    return ((const float *) src)[e];
+}
+
 // RAW: beta and g arrive pre-activation (ggml_gated_delta_net_set_raw_gates); the kernel applies
 // sigmoid(beta) and raw_a[h] * softplus(g + raw_dt_bias[h]) with the unary kernels' formulas.
 // G_PRECOMPUTED: g already holds exp(g) (GB10 long-prompt path); only used with RAW == false.
@@ -48,7 +61,7 @@ gated_delta_net_cuda(const float * q,
                                      const float * beta,
                                      const float * raw_dt_bias,
                                      const float * raw_a,
-                                     const float * curr_state,
+                                     const void *  curr_state,
                                      float *       dst,
                                      void *        state,
                                      int64_t       H,
@@ -68,7 +81,10 @@ gated_delta_net_cuda(const float * q,
                                      float         scale,
                                      int64_t       state_slot_stride,
                                      int           K,
-                                     int64_t       attn_seq_stride) {
+                                     int64_t       attn_seq_stride,
+                                     const int32_t * s_ids,
+                                     int64_t       s_row_stride,
+                                     bool          s_q8) {
     const uint32_t h_idx    = blockIdx.x;
     const uint32_t sequence = blockIdx.y;
     // Each warp owns one or more columns, using warp-level primitives to reduce across rows.
@@ -85,11 +101,10 @@ gated_delta_net_cuda(const float * q,
 
     float *       attn_data        = dst;
 
-    // input state holds s0 only: [S_v, S_v, H, n_seqs] — seq stride is D = H * S_v * S_v.
+    // input state holds s0 only: [S_v, S_v, H, n_seqs] — seq stride is D = H * S_v * S_v; with the fused gather
+    // (s_ids) it is the cache itself, this sequence's live state at row s_ids[sequence] (read after the PDL wait).
     // output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before.
-    const int64_t state_in_offset      = sequence * H * S_v * S_v + h_idx * S_v * S_v;
     const int64_t state_out_offset     = (sequence * H + h_idx) * S_v * S_v; // in elements of the destination
-    curr_state += state_in_offset;
     // attention rows of one sequence are attn_seq_stride apart: n_tokens * H * S_v, unless this launch
     // covers only the tail of each sequence (the chunked prefill path hands the last K-1 tokens here)
     attn_data += sequence * attn_seq_stride + h_idx * S_v;
@@ -101,12 +116,14 @@ gated_delta_net_cuda(const float * q,
     // state is stored transposed: M[col][i] = S[i][col], row col is contiguous
 
     ggml_cuda_pdl_sync();
+    const int64_t state_in_offset = (s_ids ? (int64_t) s_ids[sequence] * s_row_stride : (int64_t) sequence * H * S_v * S_v)
+                                    + h_idx * S_v * S_v;
 #pragma unroll
     for (int c = 0; c < cols_per_warp; ++c) {
 #pragma unroll
         for (int r = 0; r < rows_per_lane; r++) {
             const int i = r * warp_size + lane;
-            s_shard[c][r] = curr_state[(col + c) * S_v + i];
+            s_shard[c][r] = gdn_load_state(curr_state, state_in_offset + (col + c) * S_v + i, s_q8);
         }
     }
 
@@ -235,14 +252,15 @@ gated_delta_net_cuda(const float * q,
 template <bool KDA, bool keep_rs_t, bool RAW, bool G_PRECOMPUTED, bool STATE_Q8>
 static void launch_gated_delta_net(
         const float * q_d, const float * k_d, const float * v_d,
-        const float * g_d, const float * b_d, const float * rb_d, const float * ra_d, const float * s_d,
+        const float * g_d, const float * b_d, const float * rb_d, const float * ra_d, const void * s_d,
         float * dst_d, void * state_d,
         int64_t S_v,   int64_t H, int64_t n_tokens, int64_t n_seqs,
         int64_t sq1,   int64_t sq2, int64_t sq3,
         int64_t sv1,   int64_t sv2, int64_t sv3,
         int64_t sb1,   int64_t sb2, int64_t sb3,
         int64_t neqk1, int64_t rq3,
-        float scale, int64_t state_slot_stride, int K, int64_t attn_seq_stride, cudaStream_t stream) {
+        float scale, int64_t state_slot_stride, int K, int64_t attn_seq_stride,
+        const int32_t * s_ids, int64_t s_row_stride, bool s_q8, cudaStream_t stream) {
     const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
     const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
     const int num_warps = 4;
@@ -260,28 +278,28 @@ static void launch_gated_delta_net(
                 ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t, RAW, G_PRECOMPUTED, false>, launch_params,
                     q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H,
                     n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                    sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride);
+                    sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_q8);
                 break;
             }
             GGML_ABORT("a q8_0 recurrent state needs S_v >= 32");
         case 32:
             ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t, RAW, G_PRECOMPUTED, STATE_Q8>, launch_params,
                 q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_q8);
             break;
         case 64: {
             ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t, RAW, G_PRECOMPUTED, STATE_Q8>, launch_params,
                 q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_q8);
             break;
         }
         case 128: {
             ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t, RAW, G_PRECOMPUTED, STATE_Q8>, launch_params,
                 q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, H,
                 n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
-                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride);
+                sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, attn_seq_stride, s_ids, s_row_stride, s_q8);
             break;
         }
         default:
@@ -383,6 +401,14 @@ static void ggml_cuda_op_gated_delta_net_impl(
     const float * s_d   = (const float *) src_state->data;
     float *       dst_d = (float *) dst->data;
 
+    // the fused state gather, registered for this node by this context's graph evaluator (ggml_cuda_try_gdn_gather_skip):
+    // the state comes from the cache rows, not from src_state (the skipped GET_ROWS's temp, never written)
+    const ggml_cuda_gated_delta_net_gather * gather = ctx.gdn_gathers().find(dst);
+    const void *    s_in         = gather ? gather->base : (const void *) s_d;
+    const int32_t * s_ids        = gather ? gather->ids : nullptr;
+    const int64_t   s_row_stride = gather ? gather->row_stride : 0;
+    const bool      s_q8         = gather != nullptr && gather->q8_0;
+
     GGML_ASSERT(ggml_is_contiguous_rows(src_q));
     GGML_ASSERT(ggml_is_contiguous_rows(src_k));
     GGML_ASSERT(ggml_is_contiguous_rows(src_v));
@@ -434,6 +460,7 @@ static void ggml_cuda_op_gated_delta_net_impl(
     // the fused cache).
     if (ggml_cuda_should_use_chunked_gdn(dst)) {
         GGML_ASSERT(!state_q8); // the chunk pipeline writes f32 (ggml_cuda_try_gdn_cache_fusion keeps q8_0 off it)
+        GGML_ASSERT(!gather);   // and reads a gathered s0 (ggml_cuda_try_gdn_gather_skip keeps the GET_ROWS for it)
         const int64_t n_tail      = K - 1;
         const int64_t n_chunked   = n_tokens - n_tail;
         float *       chunk_state = (float *) state_d + n_tail * state_slot_stride;
@@ -469,7 +496,7 @@ static void ggml_cuda_op_gated_delta_net_impl(
             launch_gated_delta_net<false, true, RAW_, false, false>(q_d + t0 * sq2, k_d + t0 * sq2, v_d + t0 * sv2, \
                 g_d + t0 * sb2, b_d + t0 * sb2, rb_d, ra_d, chunk_state, dst_d + t0 * S_v * H, state_d,        \
                 S_v, H, n_tail, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,                                           \
-                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, S_v * H * n_tokens, stream)
+                sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, S_v * H * n_tokens, nullptr, 0, false, stream)
             if (raw) { GDN_TAIL_LAUNCH(true); } else { GDN_TAIL_LAUNCH(false); }
 #undef GDN_TAIL_LAUNCH
         }
@@ -492,9 +519,10 @@ static void ggml_cuda_op_gated_delta_net_impl(
     }
 
 #define GDN_LAUNCH(KDA_, KEEP_, RAW_, PRE_, Q8_)                                                  \
-    launch_gated_delta_net<KDA_, KEEP_, RAW_, PRE_, Q8_>(q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_d, dst_d, state_d, \
+    launch_gated_delta_net<KDA_, KEEP_, RAW_, PRE_, Q8_>(q_d, k_d, v_d, g_d, b_d, rb_d, ra_d, s_in, dst_d, state_d, \
         S_v, H, n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,                                    \
-        sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, S_v * H * n_tokens, stream)
+        sb1, sb2, sb3, neqk1, rq3, scale, state_slot_stride, K, S_v * H * n_tokens,                \
+        s_ids, s_row_stride, s_q8, stream)
 
     // a q8_0 cache is fused for the scalar gate only (ggml_cuda_try_gdn_cache_fusion)
     GGML_ASSERT(!(state_q8 && kda));