From fa2898d8cf3f7ed47f61d9a650089d125dff973d Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Tue, 28 Jul 2026 10:27:37 -0700 Subject: [PATCH] CUDA: fused chunked gated_delta_net kernel (RDNA3.5) GATED_DELTA_NET is 18% of prefill GPU time on Qwen3.5-35B-A3B / gfx1151. The existing kernel is a token-by-token scan with one warp per state column, which leaves the machine latency bound: k/q are re-fetched by all 128 warps of a head on every token, a 256x read amplification that puts MemUnitBusy at 94.5% while moving only 17.9 GB/s. Add a chunked delta rule (WY representation / UT transform) as a single fused kernel, one block per (head, seq), state resident in registers. Keeping a whole chunk in 64 KB of LDS needs two things: - CS = 32 rather than 64. At CS = 64 the chunk's U + W + kq alone are 80 KB. - Never materialising U and W. Instead of v_new = U - W S with U = Tinv(V beta) and W = Tinv(K beta exp(g_cs)), use the identical v_new = Tinv (V beta - (K beta exp(g_cs)) S), so one [CS][S] buffer holds V*beta, becomes Y, then becomes v_new in place. The in-place triangular product needs no second buffer because rows CS/2..CS-1 are accumulated into registers before any row is overwritten. The two triangular products (Gram and kq) are packed across all 1024 threads with a triangle-to-rectangle fold: 496 + 496 + 32 is exactly the block size. The natural (tid/CS, tid%CS) mapping makes a wave issue a full 128-MAC dot with most of its lanes masked off. Other things that mattered: __launch_bounds__(1024, 16) caps VGPRs at 96, which is what lets two 32-wave workgroups share a WGP; folding the lane-uniform exp(g_cs[i]) into the value before the butterfly halves the number of values to reduce; two state columns per thread so each staged K/Q read feeds two FMAs. That occupancy target is only reachable on wave32, where 1024 threads are 32 waves and two workgroups per WGP give 16 waves per SIMD. On wave64 the same workgroup is 16 waves and tops out at 4, which -Wpass-failed turns into a build error, so the request is conditioned on the physical warp size. The dispatch is also gated on wave32: the lane mapping, the 16-lane butterflies and the LDS padding are all tuned for RDNA3.5. Verified to build clean with -Werror -Wpass-failed for gfx908/90a/942/950, gfx1030, gfx1100-1103, gfx1150-1153 and gfx1200/1201. GDN op at pp2048: 220.9 -> 116.7 ms. End to end against 0f0db6292, interleaved: +5.7% at pp1024, +10.5% at pp2048, +20.0% at pp4096 with -ub 4096. Note the op only ever sees n_ubatch tokens, so at the default -ub 512 the gain is ~+2.5%. The larger figures need a larger ubatch - which is itself the point: with the sequential kernel -ub 4096 is a pessimisation at pp4096 (1520 vs 1656 t/s at -ub 512) because the scan's serial length grows with n_tokens, while with this kernel it becomes the best configuration (1825 t/s). Scalar gate and K == 1 only, S_v == 128, n_tokens >= 128; everything else falls back to the existing kernel, which is untouched. Off by default; GGML_CUDA_GDN_CHUNKED=1 enables it, =2 forces it for any multi-chunk shape so the tests can reach it. Also add five test-backend-ops cases: no existing GDN case combined head_size=128 with more than 64 tokens, and none covered GQA (v_repeat > 1), which this model actually uses (H_k=16, H_v=32). Assisted-by: Claude Opus 4.7 --- ggml/src/ggml-cuda/gated_delta_net.cu | 26 ++ ggml/src/ggml-cuda/gated_delta_net_chunked.cu | 383 ++++++++++++++++++ .../src/ggml-cuda/gated_delta_net_chunked.cuh | 34 ++ tests/test-backend-ops.cpp | 8 + 4 files changed, 451 insertions(+) create mode 100644 ggml/src/ggml-cuda/gated_delta_net_chunked.cu create mode 100644 ggml/src/ggml-cuda/gated_delta_net_chunked.cuh diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 1b431a724d72..c22762578a69 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -1,4 +1,5 @@ #include "gated_delta_net.cuh" +#include "gated_delta_net_chunked.cuh" #include "ggml-cuda/common.cuh" template @@ -294,6 +295,31 @@ static void ggml_cuda_op_gated_delta_net_impl( state_slot_stride = cache->slot_stride; } + if (ggml_cuda_gdn_chunked_supported(kda, keep_rs, S_v, n_tokens)) { + ggml_cuda_gdn_chunked_args args = {}; + args.q = q_d; + args.k = k_d; + args.v = v_d; + args.g = g_d; + args.beta = b_d; + args.state_in = s_d; + args.dst = dst_d; + args.state_out = state_d; + args.S_v = S_v; + args.H = H; + args.n_tokens = n_tokens; + args.n_seqs = n_seqs; + args.sq1 = sq1; args.sq2 = sq2; args.sq3 = sq3; + args.sv1 = sv1; args.sv2 = sv2; args.sv3 = sv3; + args.sb1 = sb1; args.sb2 = sb2; args.sb3 = sb3; + args.neqk1 = neqk1; + args.rq3 = rq3; + args.scale = scale; + + ggml_cuda_gdn_chunked(ctx, args); + return; + } + if (kda) { if (keep_rs) { launch_gated_delta_net(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, diff --git a/ggml/src/ggml-cuda/gated_delta_net_chunked.cu b/ggml/src/ggml-cuda/gated_delta_net_chunked.cu new file mode 100644 index 000000000000..8d2f51f806c7 --- /dev/null +++ b/ggml/src/ggml-cuda/gated_delta_net_chunked.cu @@ -0,0 +1,383 @@ +#include "gated_delta_net_chunked.cuh" +#include "ggml-cuda/common.cuh" + +// Chunked delta rule (WY representation / UT transform), scalar-gate only, fully fused. +// +// The sequential kernel walks one token at a time, which leaves the machine latency bound: k/q are +// re-read by every warp of a head on every token. Here the sequence is cut into chunks of CS tokens +// and the CS rank-1 state updates of a chunk are folded into (I+A)^-1, so the inner loops become +// dense matmuls out of LDS. +// +// A two-kernel split (per-chunk precompute writing U/W/kq to scratch, then a recurrence reading +// them back) was measured at 194.7 ms/layer-set with 61% of the precompute kernel's time in global +// traffic - 360 KB per (chunk, head) round-tripping 82 MB of scratch per layer. Fusing removes all +// of it: one block per (head, seq) walks the chunks and never spills U/W/kq out of LDS, leaving +// 64 KB per (chunk, head) of unavoidable K/Q/V reads and output writes. +// +// Fitting that in 64 KB of LDS needs two things: CS = 32 rather than 64, and never materialising +// U and W separately. Instead of +// v_new = U - W S with U = Tinv (V beta), W = Tinv (K beta exp(g_cs)) +// use the equivalent +// v_new = Tinv (V beta - (K beta exp(g_cs)) S) +// so one [CS][S] buffer holds V*beta, becomes Y, then becomes v_new in place. +// +// Per head h and chunk c, rows = tokens, math indices (row, col): +// g_cs = inclusive cumsum of the log decay within the chunk, g_last = g_cs[CS-1] +// A[i][j] = beta[i] * dot(K_i, K_j) * exp(g_cs[i] - g_cs[j]) for j < i, else 0 +// Tinv = (I + A)^-1 unit lower triangular, no division needed +// kq[i][j] = dot(Q_i, K_j) * exp(g_cs[i] - g_cs[j]) for j <= i, else 0 +// Y[r][j] = V[r][j]*beta[r] - sum_a K[r][a]*beta[r]*exp(g_cs[r]) * S[a][j] +// v_new[t][j] = sum_{r<=t} Tinv[t][r] * Y[r][j] +// out[i][j] = sum_a S[a][j]*Q[i][a]*exp(g_cs[i]) + sum_t v_new[t][j]*kq[i][t] +// S[a][j] = S[a][j]*exp(g_last) + sum_t K[t][a]*exp(g_last-g_cs[t]) * v_new[t][j] +// Everything after Tinv is independent per state column j, which is what allows the state to live +// in registers distributed over the block. + +#define GDN_CH_CS 32 // chunk size; 64 does not leave room to keep the chunk resident +#define GDN_CH_S 128 // head dim (only size supported by this path) +#define GDN_CH_NTHR 1024 +#define GDN_CH_SROW (GDN_CH_S + 1) // +1 breaks the power-of-two bank stride (64 banks) +#define GDN_CH_CROW (GDN_CH_CS + 1) +#define GDN_CH_NSG 16 // lanes cooperating on one column pair +#define GDN_CH_NU (GDN_CH_S / GDN_CH_NSG) // state rows per thread, per column + +// GGML_CUDA_GDN_CHUNKED: 0 / unset = off, 1 = on above the measured crossover, 2 = on for every +// multi-chunk shape (used by test-backend-ops, whose cases are all far shorter than the crossover). +bool ggml_cuda_gdn_chunked_supported(bool kda, bool keep_rs, int64_t S_v, int64_t n_tokens) { + static const int mode = []() { + const char * env = getenv("GGML_CUDA_GDN_CHUNKED"); + return env == nullptr ? 0 : std::atoi(env); + }(); + + if (mode == 0 || kda || keep_rs || S_v != GDN_CH_S) { + return false; + } + + // The lane mapping, the 16-lane butterflies and the LDS padding all assume wave32. It compiles + // and is correct on wave64, but every constant here is tuned for RDNA3.5, so do not take it. + if (ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size != 32) { + return false; + } + + // Measured on gfx1151 (Qwen3.5-35B-A3B), interleaved: +1.9% at 128, +1.1% at 256, +3.3% at 512, + // +5.3% at 1024, +12.6% at 2048, +16.3% at 4096. It wins wherever it applies; the gain is + // smaller at short prompts only because GDN is 8% of prefill there rather than 18%. + return n_tokens >= (mode >= 2 ? GDN_CH_CS + 1 : 128); +} + +// 1024 threads is 32 waves on wave32, and two workgroups per WGP make that 16 waves per SIMD; +// asking for it caps the VGPR allocation at 1536/16 = 96, which is what keeps both workgroups +// resident. On wave64 the same workgroup is only 16 waves, so 16 per SIMD is unreachable and +// -Wpass-failed turns the missed target into a build error. Ask for what the target can give. +#define GDN_CH_MIN_WAVES (ggml_cuda_get_physical_warp_size() == 32 ? 16 : 1) + +__global__ void __launch_bounds__(GDN_CH_NTHR, GDN_CH_MIN_WAVES) +gdn_ch_fused(const float * __restrict__ q, + const float * __restrict__ k, + const float * __restrict__ v, + const float * __restrict__ g, + const float * __restrict__ beta, + const float * __restrict__ state_in, + float * __restrict__ dst, + float * __restrict__ state_out, + int64_t n_tokens, + int64_t H, + 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, + uint3 neqk1_magic, + uint3 rq3_magic, + float scale) { + constexpr int CS = GDN_CH_CS; + constexpr int S = GDN_CH_S; + constexpr int NSG = GDN_CH_NSG; + constexpr int NU = GDN_CH_NU; + + const int h_idx = blockIdx.x; + const int sequence = blockIdx.y; + const int tid = threadIdx.x; + + const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic); + const uint32_t iq3 = fastdiv(sequence, rq3_magic); + + // Each thread owns an adjacent pair of state columns, so every staged K/Q element read from + // LDS feeds two FMAs. Its 16 partners hold 8 rows each and all sit inside one wave32, so the + // reductions are 4-step shuffle butterflies with no LDS traffic and no barrier. One column per + // thread with an 8-lane butterfly (3 shuffles instead of 8) was measured 8% slower: the doubled + // LDS reads cost more than the extra cross-lane steps. + const int js = tid / NSG; + const int sg = tid - js * NSG; + const int j0 = js * 2; + const int j1 = j0 + 1; + + __shared__ float sK [CS][GDN_CH_SROW]; + __shared__ float sQ [CS][GDN_CH_SROW]; + __shared__ float sY [CS][GDN_CH_SROW]; // V*beta -> Y -> v_new, in place + __shared__ float sT [CS][GDN_CH_CROW]; // A -> Tinv + __shared__ float sKQ[CS][GDN_CH_CROW]; + __shared__ float sg_cs[CS]; + __shared__ float sbeta[CS]; + __shared__ float skbg [CS]; // beta[r]*exp(g_cs[r]) + __shared__ float seg [CS]; // exp(g_cs[i]) + __shared__ float sdg [CS]; // exp(g_last - g_cs[t]) + + float st0[NU]; + float st1[NU]; + + const int64_t st_in_off = (int64_t) sequence * H * S * S + (int64_t) h_idx * S * S; +#pragma unroll + for (int u = 0; u < NU; u++) { + st0[u] = state_in[st_in_off + (int64_t) j0 * S + u * NSG + sg]; + st1[u] = state_in[st_in_off + (int64_t) j1 * S + u * NSG + sg]; + } + + // Only the lower triangles are ever written: the packed Gram/kq loop touches j <= i and the + // inversion writes zeros above the diagonal, so clearing once here is enough for every chunk. + for (int idx = tid; idx < CS * GDN_CH_CROW; idx += GDN_CH_NTHR) { + sT [idx / GDN_CH_CROW][idx % GDN_CH_CROW] = 0.0f; + sKQ[idx / GDN_CH_CROW][idx % GDN_CH_CROW] = 0.0f; + } + + const int64_t n_chunks = (n_tokens + CS - 1) / CS; + + for (int64_t chunk = 0; chunk < n_chunks; chunk++) { + const int64_t tok0 = chunk * CS; + const int n_tok = (int) (n_tokens - tok0 < CS ? n_tokens - tok0 : CS); + + __syncthreads(); + + // ---- stage the chunk ------------------------------------------------------------------- + if (tid < CS) { + const int64_t gb = sequence * sb3 + (tok0 + tid) * sb2 + h_idx * sb1; + sg_cs[tid] = tid < n_tok ? g[gb] : 0.0f; + sbeta[tid] = tid < n_tok ? beta[gb] : 0.0f; + } + __syncthreads(); + + // CS == WARP_SIZE, so the inclusive scan is a single shuffle scan in wave 0. It runs while + // the other waves stream K/Q/V in. + if (tid < WARP_SIZE) { + float a = sg_cs[tid]; +#pragma unroll + for (int off = 1; off < CS; off <<= 1) { + const float y = __shfl_up_sync(0xffffffff, a, off, WARP_SIZE); + if (tid >= off) { + a += y; + } + } + // wave 0 still holds the scan in registers, so the per-token exp tables cost no extra + // barrier here + const float gl = __shfl_sync(0xffffffff, a, CS - 1, WARP_SIZE); + sg_cs[tid] = a; + skbg[tid] = sbeta[tid] * expf(a); + seg [tid] = expf(a); + sdg [tid] = expf(gl - a); + } + + for (int idx = tid; idx < CS * S; idx += GDN_CH_NTHR) { + const int t = idx / S; + const int s = idx - t * S; + const bool ok = t < n_tok; + sK[t][s] = ok ? k[iq3 * sq3 + (tok0 + t) * sq2 + iq1 * sq1 + s] : 0.0f; + sQ[t][s] = ok ? q[iq3 * sq3 + (tok0 + t) * sq2 + iq1 * sq1 + s] * scale : 0.0f; + sY[t][s] = ok ? v[sequence * sv3 + (tok0 + t) * sv2 + h_idx * sv1 + s] : 0.0f; + } + __syncthreads(); + + const float g_last = sg_cs[CS - 1]; + + // ---- A = strictly-lower(beta[i] * K K^T * decay), kq = lower-incl(Q K^T * decay) -------- + // The natural (i,j) = (tid/CS, tid%CS) mapping wastes half the machine: a wave covers one + // row i, so it issues the full 128-MAC dot even though only the j < i lanes are unmasked. + // Pack instead - 496 strictly-lower pairs for A, 496 + 32 for kq, which is exactly NTHR - + // using the standard triangle-to-rectangle fold, so every lane does useful work. + { + constexpr int NTRI = CS * (CS - 1) / 2; // 496 + + int i, j; + bool is_kq; + if (tid < 2 * NTRI) { + const int m = tid < NTRI ? tid : tid - NTRI; + is_kq = tid >= NTRI; + // fold the (CS-1) x (CS/2) rectangle onto { 0 <= j < i < CS } + const int a = m / (CS / 2); + const int b = m - a * (CS / 2); + if (b <= a) { i = a + 1; j = b; } + else { i = CS - 1 - a; j = CS - 1 - b; } + } else { + i = j = tid - 2 * NTRI; // kq diagonal + is_kq = true; + } + + float acc0 = 0.0f; + float acc1 = 0.0f; + const float * __restrict__ row = is_kq ? &sQ[i][0] : &sK[i][0]; + for (int s = 0; s < S; s += 2) { + acc0 += row[s] * sK[j][s]; + acc1 += row[s + 1] * sK[j][s + 1]; + } + const float d = (acc0 + acc1) * expf(sg_cs[i] - sg_cs[j]); + + if (is_kq) { + sKQ[i][j] = d; + } else { + sT[i][j] = sbeta[i] * d; + } + } + __syncthreads(); + + // ---- Tinv = (I+A)^-1, right-looking so no cross-lane reduction is needed --------------- + // One column per wave: 32 waves, CS = 32 columns, lane r owns row r. + { + const int lane = tid & (WARP_SIZE - 1); + const int c = tid / WARP_SIZE; + + float x = lane == c ? 1.0f : 0.0f; + for (int i = c; i < CS; i++) { + const float xi = __shfl_sync(0xffffffff, x, i, WARP_SIZE); + if (lane > i) { + x -= sT[lane][i] * xi; + } + } + __syncthreads(); // every wave has finished reading A + sT[lane][c] = x; + } + __syncthreads(); + + // ---- Y[r][j] = V[r][j]*beta[r] - sum_a Kbg[r][a] * S[a][j] ----------------------------- + for (int r = 0; r < CS; r++) { + float p0 = 0.0f; + float p1 = 0.0f; +#pragma unroll + for (int u = 0; u < NU; u++) { + const float kv = sK[r][u * NSG + sg]; // one read, two FMAs + p0 += kv * st0[u]; + p1 += kv * st1[u]; + } +#pragma unroll + for (int m = 1; m < NSG; m <<= 1) { + p0 += __shfl_xor_sync(0xffffffff, p0, m, WARP_SIZE); + p1 += __shfl_xor_sync(0xffffffff, p1, m, WARP_SIZE); + } + if (sg == 0) { + const float bs = sbeta[r]; + const float kb = skbg[r]; + sY[r][j0] = sY[r][j0] * bs - kb * p0; + sY[r][j1] = sY[r][j1] * bs - kb * p1; + } + } + __syncthreads(); + + // ---- v_new = Tinv * Y, in place --------------------------------------------------------- + // Tinv is lower triangular so row t only needs Y[0..t]. Doing the upper half first, into + // registers, then the lower half means no row is read after it has been overwritten and no + // second [CS][S] buffer is needed. + { + const int t_hi = CS / 2 + sg; + float h0 = 0.0f; + float h1 = 0.0f; + for (int r = 0; r <= t_hi; r++) { + const float tr = sT[t_hi][r]; + h0 += tr * sY[r][j0]; + h1 += tr * sY[r][j1]; + } + __syncthreads(); + sY[t_hi][j0] = h0; + sY[t_hi][j1] = h1; + + const int t_lo = sg; + float l0 = 0.0f; + float l1 = 0.0f; + for (int r = 0; r <= t_lo; r++) { + const float tr = sT[t_lo][r]; + l0 += tr * sY[r][j0]; + l1 += tr * sY[r][j1]; + } + __syncthreads(); + sY[t_lo][j0] = l0; + sY[t_lo][j1] = l1; + } + __syncthreads(); + + // ---- out[i][j] = sum_a S[a][j]*Q[i][a]*exp(g_cs[i]) + sum_t v_new[t][j]*kq[i][t] ------- + for (int i = 0; i < CS; i++) { + float a0 = 0.0f; + float a1 = 0.0f; +#pragma unroll + for (int u = 0; u < NU; u++) { + const float qv = sQ[i][u * NSG + sg]; + a0 += qv * st0[u]; + a1 += qv * st1[u]; + } + + float b0 = 0.0f; + float b1 = 0.0f; +#pragma unroll + for (int w = 0; w < CS / NSG; w++) { + const int t = sg * (CS / NSG) + w; + const float kv = sKQ[i][t]; + b0 += sY[t][j0] * kv; + b1 += sY[t][j1] * kv; + } + + // seg[i] is lane-uniform and the reduction is linear, so folding it in before the + // butterfly halves the number of values to reduce. + float c0 = a0 * seg[i] + b0; + float c1 = a1 * seg[i] + b1; +#pragma unroll + for (int m = 1; m < NSG; m <<= 1) { + c0 += __shfl_xor_sync(0xffffffff, c0, m, WARP_SIZE); + c1 += __shfl_xor_sync(0xffffffff, c1, m, WARP_SIZE); + } + + if (sg == 0 && i < n_tok) { + const int64_t o = ((int64_t) sequence * n_tokens + tok0 + i) * H * S + (int64_t) h_idx * S; + *(float2 *) &dst[o + j0] = make_float2(c0, c1); // j0 even, S even -> 8B aligned + } + } + + // ---- S[a][j] = S[a][j]*exp(g_last) + sum_t K[t][a]*exp(g_last-g_cs[t]) * v_new[t][j] ---- + const float eg_last = expf(g_last); +#pragma unroll + for (int u = 0; u < NU; u++) { + st0[u] *= eg_last; + st1[u] *= eg_last; + } + + for (int t = 0; t < CS; t++) { + const float d = sdg[t]; + const float dv0 = d * sY[t][j0]; + const float dv1 = d * sY[t][j1]; +#pragma unroll + for (int u = 0; u < NU; u++) { + const float kv = sK[t][u * NSG + sg]; + st0[u] += kv * dv0; + st1[u] += kv * dv1; + } + } + } + + const int64_t st_out_off = ((int64_t) sequence * H + h_idx) * S * S; +#pragma unroll + for (int u = 0; u < NU; u++) { + state_out[st_out_off + (int64_t) j0 * S + u * NSG + sg] = st0[u]; + state_out[st_out_off + (int64_t) j1 * S + u * NSG + sg] = st1[u]; + } +} + +void ggml_cuda_gdn_chunked(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_chunked_args & a) { + GGML_ASSERT(a.S_v == GDN_CH_S); + + const uint3 neqk1_magic = init_fastdiv_values(a.neqk1); + const uint3 rq3_magic = init_fastdiv_values(a.rq3); + + const dim3 grid(a.H, a.n_seqs, 1); + const dim3 block(GDN_CH_NTHR, 1, 1); + const ggml_cuda_kernel_launch_params p(grid, block, 0, ctx.stream()); + + ggml_cuda_kernel_launch(gdn_ch_fused, p, + a.q, a.k, a.v, a.g, a.beta, a.state_in, a.dst, a.state_out, + a.n_tokens, a.H, + a.sq1, a.sq2, a.sq3, a.sv1, a.sv2, a.sv3, a.sb1, a.sb2, a.sb3, + neqk1_magic, rq3_magic, a.scale); +} diff --git a/ggml/src/ggml-cuda/gated_delta_net_chunked.cuh b/ggml/src/ggml-cuda/gated_delta_net_chunked.cuh new file mode 100644 index 000000000000..27b54f94c201 --- /dev/null +++ b/ggml/src/ggml-cuda/gated_delta_net_chunked.cuh @@ -0,0 +1,34 @@ +#include "common.cuh" +#include "ggml.h" + +// Chunked (WY / UT transform) gated delta net. Replaces the token-by-token recurrence with a +// per-chunk formulation whose inner loops are dense matmuls out of LDS. Covers the scalar-gate +// case only; callers must check ggml_cuda_gdn_chunked_supported() and fall back otherwise. +struct ggml_cuda_gdn_chunked_args { + const float * q; + const float * k; + const float * v; + const float * g; + const float * beta; + const float * state_in; + float * dst; + float * state_out; + + int64_t S_v; + int64_t H; + int64_t n_tokens; + int64_t n_seqs; + + int64_t sq1, sq2, sq3; + int64_t sv1, sv2, sv3; + int64_t sb1, sb2, sb3; + + int64_t neqk1; + int64_t rq3; + + float scale; +}; + +bool ggml_cuda_gdn_chunked_supported(bool kda, bool keep_rs, int64_t S_v, int64_t n_tokens); + +void ggml_cuda_gdn_chunked(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_chunked_args & args); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index cac4ac6b4140..90b5139e8f94 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9415,6 +9415,14 @@ static std::vector> 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)); + // head_size 128 multi-chunk, incl. GQA (v_repeat > 1), which nothing above covers. + // head_count=16 v_repeat=2 is the Qwen3.5 shape: H_k=16, H_v=32. + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 256, 1)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 520, 1)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 256, 1, 2)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 130, 1, 2)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 192, 3, 2, true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 2080, 1, 2)); // 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).