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).