diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu
index f0d08152..6ca696fc 100644
--- a/ggml/src/ggml-cuda/gated_delta_net.cu
+++ b/ggml/src/ggml-cuda/gated_delta_net.cu
@@ -1,13 +1,13 @@
#include "gated_delta_net.cuh"
#include "ggml-cuda/common.cuh"
-// Host and device must agree on columns/warp. Kernel uses __CUDA_ARCH__; host uses ggml_cuda_highest_compiled_arch(cc). NVIDIA Ampere+ and S_v==128 and !KDA -> 4; HIP/MUSA stay 1.
+// Host and device must agree on columns/warp. Kernel uses __CUDA_ARCH__; host uses ggml_cuda_highest_compiled_arch(cc). NVIDIA Ampere+ and S_v==128 and !KDA -> 8 (RTX 3060, 1024-token prefill: 2.1 ms vs 2.6 ms at 4); HIP/MUSA stay 1.
static constexpr __host__ __device__ int gdn_cols_per_warp(int arch, int S_v, bool KDA) {
#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA)
(void) arch; (void) S_v; (void) KDA;
return 1;
#else
- return GGML_CUDA_CC_IS_NVIDIA(arch) && arch >= GGML_CUDA_CC_AMPERE && S_v == 128 && !KDA ? 4 : 1;
+ return GGML_CUDA_CC_IS_NVIDIA(arch) && arch >= GGML_CUDA_CC_AMPERE && S_v == 128 && !KDA ? 8 : 1;
#endif
}
@@ -18,6 +18,35 @@ static __global__ void gdn_precompute_exp(const float * g, float * g_exp, int64_
}
}
+// Sums v[c] across the warp for C columns at once. The first log2(C) steps halve the columns a lane carries
+// (recursive halving), pairing lanes at the same distances as warp_reduce_sum, so every column's sum has the
+// same bits as warp_reduce_sum would give it, from fewer shuffles. Lane l returns column l / (width / C).
+template <int C, int width>
+static __device__ __forceinline__ float warp_reduce_sum_cols(const float (&v)[C], const int lane) {
+ static_assert(C >= 1 && (C & (C - 1)) == 0 && C <= width, "C must be a power of two no wider than the warp");
+ float cur[C];
+#pragma unroll
+ for (int c = 0; c < C; ++c) {
+ cur[c] = v[c];
+ }
+#pragma unroll
+ for (int half = C / 2, offset = width / 2; half > 0; half /= 2, offset /= 2) {
+ const bool upper = lane & offset;
+#pragma unroll
+ for (int i = 0; i < half; ++i) {
+ const float keep = upper ? cur[half + i] : cur[i];
+ const float send = upper ? cur[i] : cur[half + i];
+ cur[i] = keep + __shfl_xor_sync(0xffffffff, send, offset, width);
+ }
+ }
+ float x = cur[0];
+#pragma unroll
+ for (int offset = width / C / 2; offset > 0; offset /= 2) {
+ x += __shfl_xor_sync(0xffffffff, x, offset, width);
+ }
+ return x;
+}
+
// 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.
@@ -90,65 +119,141 @@ gated_delta_net_cuda(const float * q,
}
}
- for (int t = 0; t < n_tokens; t++) {
- const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1;
- const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1;
- const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1;
-
- const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1;
- const float * beta_t = beta + gb_offset;
- const float * g_t = g + gb_offset * (KDA ? S_v : 1);
-
- float beta_val = *beta_t;
- if constexpr (RAW) {
- beta_val = 1.0f / (1.0f + expf(-beta_val));
+ // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
+ // When n_tokens < K only slots 0..n_tokens-1 are written; older slots are caller-owned.
+ auto snapshot = [&](const int t) {
+ if constexpr (keep_rs_t) {
+ const int target_slot = (int) n_tokens - 1 - t;
+ if (target_slot >= 0 && target_slot < K) {
+ float * curr_state = state + target_slot * state_slot_stride;
+#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;
+ curr_state[(col + c) * S_v + i] = s_shard[c][r];
+ }
+ }
+ }
}
-
- // Cache k and q in registers
- float k_reg[rows_per_lane];
- float q_reg[rows_per_lane];
+ };
+
+ if constexpr (!KDA) {
+ static_assert(!(RAW && G_PRECOMPUTED), "exp(g) precompute is only defined for activated gates");
+
+ // token t + 1 is loaded before token t's reductions, so a step does not start with a global-memory round trip
+ float k_nx[rows_per_lane];
+ float q_nx[rows_per_lane];
+ float v_nx[cols_per_warp];
+ float beta_nx = 0.0f;
+ float g_nx = 0.0f;
+ auto fetch = [&](const int t) {
+ const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1;
+ const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1;
+ const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1;
+ const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1;
#pragma unroll
- for (int r = 0; r < rows_per_lane; r++) {
- const int i = r * warp_size + lane;
- k_reg[r] = k_t[i];
- q_reg[r] = q_t[i];
+ for (int r = 0; r < rows_per_lane; r++) {
+ const int i = r * warp_size + lane;
+ k_nx[r] = k_t[i];
+ q_nx[r] = q_t[i];
+ }
+#pragma unroll
+ for (int c = 0; c < cols_per_warp; ++c) {
+ v_nx[c] = v_t[col + c];
+ }
+ beta_nx = beta[gb_offset];
+ g_nx = g[gb_offset];
+ };
+ if (n_tokens > 0) {
+ fetch(0);
}
- if constexpr (!KDA) {
- static_assert(!(RAW && G_PRECOMPUTED), "exp(g) precompute is only defined for activated gates");
- float g0 = *g_t;
+ constexpr int lanes_per_col = warp_size / cols_per_warp;
+
+ for (int t = 0; t < n_tokens; t++) {
+ float k_reg[rows_per_lane];
+ float q_reg[rows_per_lane];
+ float v_reg[cols_per_warp];
+#pragma unroll
+ for (int r = 0; r < rows_per_lane; r++) {
+ k_reg[r] = k_nx[r];
+ q_reg[r] = q_nx[r];
+ }
+#pragma unroll
+ for (int c = 0; c < cols_per_warp; ++c) {
+ v_reg[c] = v_nx[c];
+ }
+ float beta_val = beta_nx;
+ float g0 = g_nx;
+ if (t + 1 < n_tokens) {
+ fetch(t + 1);
+ }
+
if constexpr (RAW) {
+ beta_val = 1.0f / (1.0f + expf(-beta_val));
const float x = g0 + raw_dt_bias[h_idx];
g0 = raw_a[h_idx] * ((x > 20.0f) ? x : logf(1.0f + expf(x)));
}
const float g_val = G_PRECOMPUTED ? g0 : expf(g0);
// Each warp owns one or more columns and reuses the common q/k registers.
+ float kv_part[cols_per_warp];
#pragma unroll
for (int c = 0; c < cols_per_warp; ++c) {
- float kv_shard = 0.0f;
+ kv_part[c] = 0.0f;
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
- kv_shard += s_shard[c][r] * k_reg[r];
+ kv_part[c] += s_shard[c][r] * k_reg[r];
}
- float kv_col = warp_reduce_sum<warp_size>(kv_shard);
+ }
+ const float kv_own = warp_reduce_sum_cols<cols_per_warp, warp_size>(kv_part, lane);
- float delta_col = (v_t[col + c] - g_val * kv_col) * beta_val;
+ float attn_part[cols_per_warp];
+#pragma unroll
+ for (int c = 0; c < cols_per_warp; ++c) {
+ const float kv_col = __shfl_sync(0xffffffff, kv_own, c * lanes_per_col, warp_size);
+ const float delta_col = (v_reg[c] - g_val * kv_col) * beta_val;
- float attn_partial = 0.0f;
+ attn_part[c] = 0.0f;
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
- s_shard[c][r] = g_val * s_shard[c][r] + k_reg[r] * delta_col;
- attn_partial += s_shard[c][r] * q_reg[r];
+ s_shard[c][r] = g_val * s_shard[c][r] + k_reg[r] * delta_col;
+ attn_part[c] += s_shard[c][r] * q_reg[r];
}
+ }
+ const float attn_own = warp_reduce_sum_cols<cols_per_warp, warp_size>(attn_part, lane);
+ if (lane % lanes_per_col == 0) {
+ attn_data[col + lane / lanes_per_col] = attn_own * scale;
+ }
- float attn_col = warp_reduce_sum<warp_size>(attn_partial);
+ attn_data += S_v * H;
+ snapshot(t);
+ }
+ } else {
+ for (int t = 0; t < n_tokens; t++) {
+ const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1;
+ const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1;
+ const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1;
- if (lane == 0) {
- attn_data[col + c] = attn_col * scale;
- }
+ const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1;
+ const float * g_t = g + gb_offset * S_v;
+
+ float beta_val = beta[gb_offset];
+ if constexpr (RAW) {
+ beta_val = 1.0f / (1.0f + expf(-beta_val));
}
- } else {
+
+ // Cache k and q in registers
+ float k_reg[rows_per_lane];
+ float q_reg[rows_per_lane];
+#pragma unroll
+ for (int r = 0; r < rows_per_lane; r++) {
+ const int i = r * warp_size + lane;
+ k_reg[r] = k_t[i];
+ q_reg[r] = q_t[i];
+ }
+
// kv[col] = sum_i g[i] * S[i][col] * k[i]
float kv_shard = 0.0f;
#pragma unroll
@@ -177,27 +282,10 @@ gated_delta_net_cuda(const float * q,
if (lane == 0) {
attn_data[col] = attn_col * scale;
}
- }
-
- attn_data += S_v * H;
- if constexpr (keep_rs_t) {
- // snapshot slot mapping: slot 0 = most recent state, slot s = s tokens back.
- // When n_tokens < K only slots 0..n_tokens-1 are written; older slots are caller-owned.
- const int target_slot = (int) n_tokens - 1 - t;
- if (target_slot >= 0 && target_slot < K) {
- float * curr_state = state + target_slot * state_slot_stride;
-#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;
- curr_state[(col + c) * S_v + i] = s_shard[c][r];
- }
- }
- }
+ attn_data += S_v * H;
+ snapshot(t);
}
-
}
if constexpr (!keep_rs_t) {