CUDA: Support of GDN chunked kernel for prefill

ggml/src/ggml-backend.cpp

@@ -2101,6 +2101,7 @@ ggml_backend_t ggml_backend_sched_get_tensor_backend(ggml_backend_sched_t sched,
 bool ggml_op_alloc_size_may_expand(enum ggml_op op) {
     switch (op) {
         case GGML_OP_FLASH_ATTN_EXT:
+        case GGML_OP_GATED_DELTA_NET:
         case GGML_OP_MUL_MAT:
         case GGML_OP_MUL_MAT_ID:
         case GGML_OP_CUMSUM:

ggml/src/ggml-cuda/chunk_gated_delta_net.cuh

@@ -0,0 +1,29 @@
+#pragma once
+
+#include "common.cuh"
+#include "gated_delta_net.cuh"
+#include "ggml.h"
+
+void ggml_cuda_op_gated_delta_net_chunked(ggml_backend_cuda_context & ctx, ggml_tensor * dst,
+                                          const ggml_cuda_gated_delta_net_fused_cache * cache);
+
+// Chunked-GDN scratch, carved out of the tail of dst's own allocation. Two reasons:
+//  1. The graph allocator sizes every tensor through ggml_backend_buft_get_alloc_size, so scratch
+//     that lives in dst's buffer is visible to llama_params_fit (--fit). A context-owned cudaMalloc
+//     (or the old ggml_cuda_pool_alloc) is invisible to that projection, which made --fit
+//     under-estimate VRAM by the full scratch size.
+//  2. The address is a fixed offset from dst->data, so it is stable across CUDA-graph capture and
+//     replay without needing a separate persistent allocation to be pre-sized before capture.
+struct ggml_cuda_gdn_chunked_scratch {
+    float *   v_corr;
+    float *   k_cumdecay;
+    float *   g_cum;
+    float *   qk;
+    uintptr_t end;
+};
+
+// Pure function of the tensor graph, so allocation time and execution time cannot disagree.
+ggml_cuda_gdn_chunked_scratch ggml_cuda_gdn_get_chunked_scratch(const ggml_tensor * dst);
+
+// ggml_nbytes(dst) plus the scratch above, or just ggml_nbytes(dst) when the shape is ineligible.
+size_t ggml_cuda_gdn_get_alloc_size(const ggml_tensor * dst);

ggml/src/ggml-cuda/gated_delta_net.cu

@@ -1,4 +1,5 @@
 #include "gated_delta_net.cuh"
+#include "chunk_gated_delta_net.cuh"
 #include "ggml-cuda/common.cuh"
 
 template <int S_v, bool KDA, bool keep_rs_t>
@@ -39,8 +40,8 @@ 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.
-    // output state layout (per-slot D * n_seqs) — same per-(seq,head) offset as before.
+    // input state holds s0 only: [S_v, S_v, H, n_seqs] - seq stride is D = H * S_v * S_v.
+    // 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;
     state += state_out_offset;
@@ -177,7 +178,6 @@ static void launch_gated_delta_net(
         int64_t sb1,   int64_t sb2, int64_t sb3,
         int64_t neqk1, int64_t rq3,
         float scale, int64_t state_slot_stride, int K, cudaStream_t stream) {
-    //TODO: Add chunked kernel for even faster pre-fill
     const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
     const int num_warps = 4;
     dim3      grid_dims(H, n_seqs, (S_v + num_warps - 1) / num_warps);
@@ -220,6 +220,63 @@ static void launch_gated_delta_net(
     }
 }
 
+// Shape-only half of the chunked predicate, split out because the buffer-type get_alloc_size hook
+// has to size the chunked scratch. That sizing must NOT depend on which device is current: deciding
+// from shape alone can only over-allocate on a device that ends up ineligible, never under-allocate
+// (which would corrupt memory). See ggml_cuda_gdn_get_alloc_size.
+bool ggml_cuda_gdn_chunked_shape_eligible(const ggml_tensor * dst) {
+    if (dst->op != GGML_OP_GATED_DELTA_NET) {
+        return false;
+    }
+
+    const ggml_tensor * src_q     = dst->src[0];
+    const ggml_tensor * src_k     = dst->src[1];
+    const ggml_tensor * src_v     = dst->src[2];
+    const ggml_tensor * src_g     = dst->src[3];
+    const ggml_tensor * src_beta  = dst->src[4];
+    const ggml_tensor * src_state = dst->src[5];
+
+    const int64_t S_v      = src_v->ne[0];
+    const int64_t n_tokens = src_v->ne[2];
+    const int64_t neq0     = src_q->ne[0];
+    const int64_t neq1     = src_q->ne[1];   // q head count
+    const int64_t nev1     = src_v->ne[1];   // v head count
+    const bool    kda      = (src_g->ne[0] == S_v);
+    const int     K        = ggml_get_op_params_i32(dst, 0);
+
+    // - not KDA; K == 1 (final state only)
+    // - Q/K/G/beta/state must be contiguous; V must be contiguous within each token (nb[0]/nb[1]
+    //   packed) and packed across sequences (nb[3] == n_tokens*nb[2]), with an arbitrary per-token
+    //   stride nb[2] (fused QKV view). The nb[3] check matches the chunked-entry assert: without it a
+    //   view with inter-sequence padding would pass dispatch and then read the wrong batch slice.
+    // - 128-wide heads, GQA-aligned head counts, n_tokens >= 128
+    return !kda && K == 1
+        && neq0 == 128 && S_v == 128 && nev1 % neq1 == 0
+        && src_k->ne[1] == neq1
+        && n_tokens >= 128
+        && ggml_is_contiguous(src_q) && ggml_is_contiguous(src_k) && ggml_is_contiguous(src_g)
+        && src_v->nb[0] == ggml_type_size(src_v->type) && src_v->nb[1] == (size_t)S_v * ggml_type_size(src_v->type)
+        && src_v->nb[3] == (size_t) n_tokens * src_v->nb[2]
+        && ggml_is_contiguous(src_beta) && ggml_is_contiguous(src_state);
+}
+
+bool ggml_cuda_should_use_chunked_gdn(const ggml_tensor * dst) {
+#ifdef GGML_CUDA_NO_GDN_CHUNK
+    // Recurrent-only baseline build (see tools/bench_ab_*): everything routes to the recurrent
+    // kernel, which also re-enables CUDA graphs for the op.
+    GGML_UNUSED(dst);
+    return false;
+#else
+    if (!ggml_cuda_gdn_chunked_shape_eligible(dst)) {
+        return false;
+    }
+    // NVIDIA Ampere+ only (fp16 WMMA). The HIP/MUSA ggml_cuda_mma backend is intentionally not
+    // dispatched until validated.
+    const int cc_dev = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
+    return GGML_CUDA_CC_IS_NVIDIA(cc_dev) && cc_dev >= GGML_CUDA_CC_AMPERE;
+#endif // GGML_CUDA_NO_GDN_CHUNK
+}
+
 static void ggml_cuda_op_gated_delta_net_impl(
         ggml_backend_cuda_context & ctx, ggml_tensor * dst, const ggml_cuda_gated_delta_net_fused_cache * cache) {
     ggml_tensor * src_q     = dst->src[0];
@@ -286,6 +343,15 @@ static void ggml_cuda_op_gated_delta_net_impl(
     const int K = ggml_get_op_params_i32(dst, 0);
     const bool keep_rs = K > 1;
 
+    // Route to the chunked prefill kernel when eligible.
+    // Passes cache so the kernel can write the final state directly to the fused destination
+    // (cache->data). Scratch lives in dst's own allocation, so its address is stable across CUDA
+    // graph capture and replay.
+    if (ggml_cuda_should_use_chunked_gdn(dst)) {
+        ggml_cuda_op_gated_delta_net_chunked(ctx, dst, cache);
+        return;
+    }
+
     // recurrent state -> gdn_out tail (after attention scores), or the cache when fusing
     float * state_d           = dst_d + S_v * H * n_tokens * n_seqs;
     int64_t state_slot_stride = S_v * S_v * H * n_seqs;

ggml/src/ggml-cuda/gated_delta_net.cuh

@@ -1,3 +1,4 @@
+#pragma once
 #include "common.cuh"
 #include "ggml.h"
 
@@ -12,3 +13,10 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor *
 // same op, but writes the snapshot(s) into the cache instead of dst (see ggml_cuda_try_gdn_cache_fusion)
 void ggml_cuda_op_gated_delta_net_fused_cache(ggml_backend_cuda_context & ctx, ggml_tensor * dst,
                                               ggml_cuda_gated_delta_net_fused_cache cache);
+
+// Returns true if chunked prefill can be used; false for recurrent kernel
+bool ggml_cuda_should_use_chunked_gdn(const ggml_tensor * dst);
+
+// Shape-only part of the above, with no dependence on the current device. Used to size the chunked
+// scratch at allocation time, where the eventual execution device may not be current yet.
+bool ggml_cuda_gdn_chunked_shape_eligible(const ggml_tensor * dst);

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

@@ -59,6 +59,7 @@
 #include "ggml-cuda/upscale.cuh"
 #include "ggml-cuda/wkv.cuh"
 #include "ggml-cuda/gla.cuh"
+#include "ggml-cuda/chunk_gated_delta_net.cuh"
 #include "ggml-cuda/gated_delta_net.cuh"
 #include "ggml-cuda/dsv4-hc.cuh"
 #include "ggml-cuda/set.cuh"
@@ -908,9 +909,12 @@ static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_ty
 static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) {
     ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *) buft->context;
 
-    size_t size = tensor->op == GGML_OP_FLASH_ATTN_EXT
-        ? ggml_cuda_flash_attn_ext_get_alloc_size(buft_ctx->device, tensor)
-        : ggml_nbytes(tensor);
+    size_t size = ggml_nbytes(tensor);
+    if (tensor->op == GGML_OP_FLASH_ATTN_EXT) {
+        size = ggml_cuda_flash_attn_ext_get_alloc_size(buft_ctx->device, tensor);
+    } else if (tensor->op == GGML_OP_GATED_DELTA_NET) {
+        size = ggml_cuda_gdn_get_alloc_size(tensor);
+    }
     int64_t ne0 = tensor->ne[0];
 
     // [TAG_ALLOC_SIZE_EXPAND]
@@ -2752,6 +2756,12 @@ static int ggml_cuda_try_gdn_cache_fusion(
         return 0;
     }
 
+    // Chunked prefill writes state to dst; it cannot scatter into the snapshot cache, so skip fusion.
+    // Decode and fallback recurrent (T < 128) do fuse.
+    if (ggml_cuda_should_use_chunked_gdn(gdn)) {
+        return 0;
+    }
+
     const ggml_tensor * src_v     = gdn->src[2];
     const int64_t       S_v       = src_v->ne[0];
     const int64_t       H         = src_v->ne[1];

tests/test-backend-ops.cpp

@@ -4580,16 +4580,23 @@ struct test_gated_delta_net : public test_case {
     const bool    permuted;
     const bool    kda;
     const int64_t K; // snapshot slot count: 1 = final-only, >1 = last K states
+    const bool    strided_v; // V is a view into a fused QKV buffer (model path)
 
     std::string vars() override {
-        return VARS_TO_STR9(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K);
+        return VARS_TO_STR10(type, head_count, head_size, n_seq_tokens, n_seqs, v_repeat, permuted, kda, K, strided_v);
+    }
+
+    // Dominant cost: B * H * T * head_size * (head_size * v_repeat) mul-adds, counted twice.
+    uint64_t op_flops(ggml_tensor * t) override {
+        GGML_UNUSED(t);
+        return (uint64_t)2 * n_seqs * head_count * n_seq_tokens * head_size * head_size * v_repeat;
     }
 
     test_gated_delta_net(ggml_type type = GGML_TYPE_F32,
             int64_t head_count = 4, int64_t head_size = 16, int64_t n_seq_tokens = 1, int64_t n_seqs = 1,
-            int v_repeat = 1, bool permuted = false, bool kda = false, int64_t K = 1)
+            int v_repeat = 1, bool permuted = false, bool kda = false, int64_t K = 1, bool strided_v = false)
         : type(type), head_count(head_count), head_size(head_size), n_seq_tokens(n_seq_tokens), n_seqs(n_seqs),
-          v_repeat(v_repeat), permuted(permuted), kda(kda), K(K) {}
+          v_repeat(v_repeat), permuted(permuted), kda(kda), K(K), strided_v(strided_v) {}
 
     ggml_tensor * build_graph(ggml_context * ctx) override {
         ggml_tensor * q;
@@ -4600,6 +4607,19 @@ struct test_gated_delta_net : public test_case {
             q = ggml_permute(ctx, ggml_new_tensor_4d(ctx, type, head_size, n_seq_tokens, head_count, n_seqs), 0, 2, 1, 3);
             k = ggml_permute(ctx, ggml_new_tensor_4d(ctx, type, head_size, n_seq_tokens, head_count, n_seqs), 0, 2, 1, 3);
             v = ggml_permute(ctx, ggml_new_tensor_4d(ctx, type, head_size, n_seq_tokens, head_count * v_repeat, n_seqs), 0, 2, 1, 3);
+        } else if (strided_v) {
+            // Fused QKV layout as in qwen35: V is a view with larger token stride; Q/K stay contiguous
+            // (as after L2-norm in the model graph).
+            const int64_t n_v_heads = head_count * v_repeat;
+            const int64_t qkv_dim   = head_size * (2 * head_count + n_v_heads);
+            ggml_tensor * qkv = ggml_new_tensor_3d(ctx, type, qkv_dim, n_seq_tokens, n_seqs);
+            ggml_set_name(qkv, "v_qkv");
+            const size_t nb1_qkv  = ggml_row_size(type, qkv_dim);
+            const size_t v_offset = ggml_row_size(type, 2 * head_size * head_count);
+            q = ggml_new_tensor_4d(ctx, type, head_size, head_count, n_seq_tokens, n_seqs);
+            k = ggml_new_tensor_4d(ctx, type, head_size, head_count, n_seq_tokens, n_seqs);
+            v = ggml_view_4d(ctx, qkv, head_size, n_v_heads, n_seq_tokens, n_seqs,
+                    ggml_row_size(type, head_size), nb1_qkv, nb1_qkv * n_seq_tokens, v_offset);
         } else {
             q = ggml_new_tensor_4d(ctx, type, head_size, head_count, n_seq_tokens, n_seqs);
             k = ggml_new_tensor_4d(ctx, type, head_size, head_count, n_seq_tokens, n_seqs);
@@ -4629,13 +4649,18 @@ struct test_gated_delta_net : public test_case {
                 init_tensor_uniform(t, -20.0f, -1e-4f);
             } else if (strcmp(t->name, "beta") == 0) {
                 init_tensor_uniform(t, 0.0f, 1.0f);
-            } else if (strcmp(t->name, "v") == 0) {
+            } else if (strcmp(t->name, "v") == 0 || strcmp(t->name, "v_qkv") == 0) {
                 init_tensor_uniform(t, -0.3f, 5.0f);
             } else {
                 init_tensor_uniform(t);
             }
         }
     }
+
+    double max_nmse_err() override {
+        // CUDA GDN chunked requires a slightly higher threshold than the default (~1e-7)
+        return 2e-7;
+    }
 };
 
 // GGML_OP_GATED_DELTA_NET + GGML_OP_CPY (recurrent cache fusion)
@@ -10832,7 +10857,7 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 32, 4, 2, 2, false, true));
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, true,  true));
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 16, 4, 2, 1, true,  true));
-    // chunked path: multi-chunk and non-multiple-of-chunk-size (chunk_size=64 GDN, 16 KDA)
+    // Recurrent path: multi-chunk and non-multiple-of-chunk-size (chunk_size=64 GDN, 16 KDA)
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64,  64, 1));
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 127, 1));
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 256, 1));
@@ -10843,6 +10868,26 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64,  64, 1, 1, false, true));
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64,  33, 1, 1, false, true));
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 100, 1, 1, false, true));
+    // Chunked GDN path: K_dim==V_dim==128, !kda, K==1, n_tokens>=128
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 128, 1));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 512, 1));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 2048, 1));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128,  256, 4));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32,  4, 128,  256, 4));
+    // GQA (v_repeat>1)
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 256, 1, /*v_repeat=*/2));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 512, 1, /*v_repeat=*/3));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 256, 4, /*v_repeat=*/2));
+    // Partial final chunk (T not multiple of CS=16)
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 129, 1));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 143, 1));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 200, 1));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 200, 1, /*v_repeat=*/2));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 143, 1, /*v_repeat=*/3));
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 300, 4));
+    // Strided V from fused QKV buffer (qwen35 model path without ggml_cont_4d on V)
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 256, 1, /*v_repeat=*/2,
+            /*permuted=*/false, /*kda=*/false, /*K=*/1, /*strided_v=*/true));
 
     // K > 1: output keeps the last min(n_tokens, K) per-token snapshots, ordered most-recent-first
     // (slot 0 = final state, slot s = state s tokens back).
@@ -11308,6 +11353,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 512, 1));  // 4h PP-512
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 1024, 1)); // 4h PP-1024
     test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 64, 1, 1, false, true)); // KDA PP-64
+    // Long PP and B=4
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 2048, 1)); // PP-2048
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 4096, 1)); // PP-4096
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 2048, 4)); // B4 PP-2048
+    test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 4096, 4)); // B4 PP-4096
 
     // lightning_indexer
     for (int kv : { 256, 4096, 65536 }) {