Commit 24e41838e for llama.cpp
commit 24e41838e0c834f67db214e6d280af51f266056e
Author: SXX <song_xiaoxi@126.com>
Date: Thu Oct 8 14:07:34 2026 +0800
ggml-cuda: assign four GDN state columns per warp (#30087)
* ggml-cuda: assign two GDN state columns per warp
* ggml-cuda: use 4 GDN state columns per warp at S_v=128
* ggml-cuda: default cols_per_warp=4
* ggml-cuda: address GDN review nits
diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu
index 1b431a724..b58257067 100644
--- a/ggml/src/ggml-cuda/gated_delta_net.cu
+++ b/ggml/src/ggml-cuda/gated_delta_net.cu
@@ -1,7 +1,9 @@
#include "gated_delta_net.cuh"
#include "ggml-cuda/common.cuh"
-template <int S_v, bool KDA, bool keep_rs_t>
+constexpr int gdn_cols_per_warp = 4;
+
+template <int S_v, bool KDA, bool keep_rs_t, int cols_per_warp = gdn_cols_per_warp>
__global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2)
gated_delta_net_cuda(const float * q,
const float * k,
@@ -30,9 +32,19 @@ gated_delta_net_cuda(const float * q,
int K) {
const uint32_t h_idx = blockIdx.x;
const uint32_t sequence = blockIdx.y;
- // each warp owns one column, using warp-level primitives to reduce across rows
- const int lane = threadIdx.x;
- const int col = blockIdx.z * blockDim.y + threadIdx.y;
+
+ constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v;
+ static_assert(S_v % warp_size == 0, "S_v must be a multiple of warp_size");
+ // the warp is split into cols_per_warp segments of lanes_per_col lanes; each segment owns
+ // one state column and reduces within itself
+ constexpr int lanes_per_col = warp_size / cols_per_warp;
+ constexpr int rows_per_lane = S_v / lanes_per_col;
+ static_assert(S_v % lanes_per_col == 0, "S_v must be a multiple of lanes_per_col");
+
+ const int lane = threadIdx.x;
+ const int col_in_warp = lane / lanes_per_col; // column slot within the warp
+ const int lane_in_col = lane - col_in_warp * lanes_per_col; // lane within the column's reduction segment
+ const int col = (blockIdx.z * blockDim.y + threadIdx.y) * cols_per_warp + col_in_warp;
const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic);
const uint32_t iq3 = fastdiv(sequence, rq3_magic);
@@ -47,16 +59,13 @@ gated_delta_net_cuda(const float * q,
curr_state += state_in_offset + col * S_v;
attn_data += (sequence * n_tokens * H + h_idx) * S_v;
- constexpr int warp_size = ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v;
- static_assert(S_v % warp_size == 0, "S_v must be a multiple of warp_size");
- constexpr int rows_per_lane = (S_v + warp_size - 1) / warp_size;
float s_shard[rows_per_lane];
// state is stored transposed: M[col][i] = S[i][col], row col is contiguous
ggml_cuda_pdl_sync();
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
- const int i = r * warp_size + lane;
+ const int i = r * lanes_per_col + lane_in_col;
s_shard[r] = curr_state[i];
}
@@ -76,7 +85,7 @@ gated_delta_net_cuda(const float * q,
float q_reg[rows_per_lane];
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
- const int i = r * warp_size + lane;
+ const int i = r * lanes_per_col + lane_in_col;
k_reg[r] = k_t[i];
q_reg[r] = q_t[i];
}
@@ -90,7 +99,7 @@ gated_delta_net_cuda(const float * q,
for (int r = 0; r < rows_per_lane; r++) {
kv_shard += s_shard[r] * k_reg[r];
}
- float kv_col = warp_reduce_sum<warp_size>(kv_shard);
+ float kv_col = warp_reduce_sum<lanes_per_col>(kv_shard);
// delta[col] = (v[col] - g * kv[col]) * beta
float delta_col = (v_t[col] - g_val * kv_col) * beta_val;
@@ -104,9 +113,9 @@ gated_delta_net_cuda(const float * q,
attn_partial += s_shard[r] * q_reg[r];
}
- float attn_col = warp_reduce_sum<warp_size>(attn_partial);
+ float attn_col = warp_reduce_sum<lanes_per_col>(attn_partial);
- if (lane == 0) {
+ if (lane_in_col == 0) {
attn_data[col] = attn_col * scale;
}
} else {
@@ -114,11 +123,11 @@ gated_delta_net_cuda(const float * q,
float kv_shard = 0.0f;
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
- const int i = r * warp_size + lane;
+ const int i = r * lanes_per_col + lane_in_col;
kv_shard += expf(g_t[i]) * s_shard[r] * k_reg[r];
}
- float kv_col = warp_reduce_sum<warp_size>(kv_shard);
+ float kv_col = warp_reduce_sum<lanes_per_col>(kv_shard);
// delta[col] = (v[col] - kv[col]) * beta
float delta_col = (v_t[col] - kv_col) * beta_val;
@@ -128,14 +137,14 @@ gated_delta_net_cuda(const float * q,
float attn_partial = 0.0f;
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
- const int i = r * warp_size + lane;
+ const int i = r * lanes_per_col + lane_in_col;
s_shard[r] = expf(g_t[i]) * s_shard[r] + k_reg[r] * delta_col;
attn_partial += s_shard[r] * q_reg[r];
}
- float attn_col = warp_reduce_sum<warp_size>(attn_partial);
+ float attn_col = warp_reduce_sum<lanes_per_col>(attn_partial);
- if (lane == 0) {
+ if (lane_in_col == 0) {
attn_data[col] = attn_col * scale;
}
}
@@ -150,7 +159,7 @@ gated_delta_net_cuda(const float * q,
float * curr_state = state + target_slot * state_slot_stride;
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
- const int i = r * warp_size + lane;
+ const int i = r * lanes_per_col + lane_in_col;
curr_state[col * S_v + i] = s_shard[r];
}
}
@@ -160,7 +169,7 @@ gated_delta_net_cuda(const float * q,
if constexpr (!keep_rs_t) {
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
- const int i = r * warp_size + lane;
+ const int i = r * lanes_per_col + lane_in_col;
state[col * S_v + i] = s_shard[r];
}
}
@@ -179,8 +188,16 @@ static void launch_gated_delta_net(
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);
+ // four columns per warp (see the kernel); shrink the CTA when the wider CTA would leave
+ // SMs without a CTA, so small head counts keep the device filled
+ const int nsm = ggml_cuda_info().devices[ggml_cuda_get_device()].nsm;
+ const int cols_per_warp = gdn_cols_per_warp;
+ int num_warps = 4;
+ while (num_warps > 1 && H*n_seqs*(S_v / (cols_per_warp * num_warps)) < nsm) {
+ num_warps /= 2;
+ }
+ // one CTA covers cols_per_warp*num_warps columns (see the kernel)
+ dim3 grid_dims(H, n_seqs, (S_v + cols_per_warp * num_warps - 1) / (cols_per_warp * num_warps));
dim3 block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1);
const uint3 neqk1_magic = init_fastdiv_values(neqk1);