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 +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(kv_shard); + } + const float kv_own = warp_reduce_sum_cols(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(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(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) {