Skip to content

CUDA: AR gated delta net improvements #20391

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ggerganov merged 13 commits into from
Mar 11, 2026
Merged

CUDA: AR gated delta net improvements #20391

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
55d2e3a
Add FastDiv to gated_delta_net_cuda
ORippler Mar 10, 2026
2352caa
Shard columns across warps
ORippler Mar 10, 2026
f62852a
Remove unneded include in gated_delta_net.cu
ORippler Mar 11, 2026
da33f8e
Improve comments
ORippler Mar 11, 2026
e1cfcf8
Apply code-formating
ORippler Mar 11, 2026
0211798
Make sharding HIP-compatible
ORippler Mar 11, 2026
e26d75b
Remove fastdiv_s64, as we can treat neqk1 and rq3 as uint32_t
ORippler Mar 11, 2026
1623bbc
Rename variables
ORippler Mar 11, 2026
2068908
Enable GDN also for prefill, move TODO for chunked_GDN
ORippler Mar 11, 2026
226a1ac
Actually remove the TODO from 206890897546bd16602c3b79394fd5ea09ef199f
ORippler Mar 11, 2026
fd54823
Get warp size at runtime
IMbackK Mar 11, 2026
2636f68
Merge pull request #1 from IMbackK/patch-1
ORippler Mar 11, 2026
c896fe7
Don't expose ggml_cuda_get_physical_warp_size on host
ORippler Mar 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
85 changes: 85 additions & 0 deletions ggml/src/ggml-cuda/common.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -821,6 +821,91 @@ __device__ __forceinline__ uint8_t ggml_cuda_float_to_fp4_e2m1(float x, float e)
return static_cast<uint8_t>(best_i | sign_bit);
}

struct fastdiv_consts_s64 {
int64_t mp; // magic number
int64_t d; // divisor
uint8_t L; // need at most 6 bits to represent L, use 7th bit to signal sign of d
};

static inline uint8_t floor_log2(uint64_t x) {
uint64_t exp;
#if defined(__GNUC__) || defined(__clang__)
exp = 63 - __builtin_clzll(x);
#elif defined(_MSC_VER)
// MSVC: _BitScanReverse64 finds the index of the MSB (0 to 63)
_BitScanReverse64(&exp, x);
#endif
return exp;
}

// Helper: Safely computes floor(2^{64 + l} * (1 + fraction) / d)
static uint64_t compute_m(uint64_t d, int l, int prec) {
// 1. Calculate q and r such that: 2^64 = (q * d) + r
// Since 2^64 overflows uint64_t, we use 0xFF...FF and adjust.
uint64_t q = (0xFFFFFFFFFFFFFFFF / d);
uint64_t r = (0xFFFFFFFFFFFFFFFF % d) + 1;
if (r >= d) {
r -= d;
q++;
}

// 2. We need (2^l * 2^64) / d => (2^l * q) + (2^l * r) / d
uint64_t base_q = (q << l);
uint64_t base_r = (r << l) / d;
uint64_t m = base_q + base_r;

// 3. For m_high, we need to add the precision term: (2^{64 + l - prec}) / d
// (2^{64 + l - prec}) / d => (2^{64} / d) >> (prec - l)
// This is safe because (prec - l) >= 0 in this algorithm
uint64_t extra = (q >> (prec - l)) + ((r >> (prec - l)) / d);
m += extra;

return m;
}

static const fastdiv_consts_s64 init_fastdiv_s64(int64_t d) {
GGML_ASSERT(d != 0);
uint64_t abs_d = d < 0 ? -d : d;

uint8_t L = floor_log2(abs_d);

if (uint64_t{ 1 } << L == abs_d) {
// signal negative divisor in L's 7th bit
d < 0 ? L |= 0x40 : 0;
// multiply with 0 to avoid branching in fastdiv_s64 kernel when d is power of 2
return { 0, d, L };
}

uint64_t mh = compute_m(abs_d, L, 63);
// signal negative divisor in L's 7th bit
d < 0 ? L |= 0x40 : 0;
return { (int64_t) mh, d, L };
}

static __device__ int64_t fastdiv_s64(int64_t n, fastdiv_consts_s64 c) {
int64_t q;
q = __mul64hi(n, c.mp);
q += n;

// Extract the sign bit
uint64_t q_sign = q >> 63;
// L's lower 6 bits are the shift amount
uint8_t shift = c.L & 0x3F;
q += q_sign & ((1ULL << shift) - (c.mp == 0));
q >>= shift;

// if divisor is negative, negate the quotient
int64_t d_sign = c.L >> 6;
d_sign ? q = -q : 0;

return q;
}

__device__ __forceinline__ int64_t fastmodulo_s64(int64_t n, fastdiv_consts_s64 c) {
int64_t q = fastdiv_s64(n, c);
return n - (q * c.d);
}

// See https://gmplib.org/~tege/divcnst-pldi94.pdf figure 4.1.
// Precompute mp (m' in the paper) and L such that division
// can be computed using a multiply (high 32b of 64b result)
Expand Down
129 changes: 76 additions & 53 deletions ggml/src/ggml-cuda/gated_delta_net.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,35 +1,36 @@
#include "gated_delta_net.cuh"
#include "ggml-cuda/common.cuh"

template <int S_v, bool KDA>
__global__ void gated_delta_net_cuda(const float * q,
const float * k,
const float * v,
const float * g,
const float * beta,
const float * curr_state,
float * dst,
int64_t H,
int64_t n_tokens,
int64_t n_seqs,
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,
int64_t neqk1,
int64_t rq3,
float scale) {
__global__ void gated_delta_net_cuda(const float * q,
const float * k,
const float * v,
const float * g,
const float * beta,
const float * curr_state,
float * dst,
int64_t H,
int64_t n_tokens,
int64_t n_seqs,
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,
const fastdiv_consts_s64 neqk1_magic,
const fastdiv_consts_s64 rq3_magic,
float scale) {
const int64_t h_idx = blockIdx.x;
const int64_t sequence = blockIdx.y;
const int col = threadIdx.x; // each thread owns one column
// 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;

const int64_t iq1 = h_idx % neqk1;
const int64_t iq3 = sequence / rq3;
const int64_t iq1 = fastmodulo_s64(h_idx, neqk1_magic);
const int64_t iq3 = fastdiv_s64(sequence, rq3_magic);

Comment thread
ORippler marked this conversation as resolved.
Outdated
const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs;
float * attn_data = dst;
Expand All @@ -40,11 +41,13 @@ __global__ void gated_delta_net_cuda(const float * q,
curr_state += state_offset;
attn_data += (sequence * n_tokens * H + h_idx) * S_v;

// Load state column into registers
float s[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;
float s_shard[ROWS_PER_LANE];
#pragma unroll
for (int i = 0; i < S_v; i++) {
s[i] = curr_state[i * S_v + col];
for (int r = 0; r < ROWS_PER_LANE; r++) {
const int i = r * WARP_SIZE + lane;
s_shard[r] = curr_state[i * S_v + col];
}

for (int t = 0; t < n_tokens; t++) {
Expand All @@ -62,55 +65,71 @@ __global__ void gated_delta_net_cuda(const float * q,
const float g_val = expf(*g_t);

// kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i]
float kv_col = 0.0f;
float kv_shard = 0.0f;
#pragma unroll
for (int i = 0; i < S_v; i++) {
kv_col += s[i] * k_t[i];
for (int r = 0; r < ROWS_PER_LANE; r++) {
const int i = r * WARP_SIZE + lane;
kv_shard += s_shard[r] * k_t[i];
}
float kv_col = warp_reduce_sum(kv_shard);

// delta[col] = (v[col] - g * kv[col]) * beta
float delta_col = (v_t[col] - g_val * kv_col) * beta_val;

// fused: S[i][col] = g * S[i][col] + k[i] * delta[col]
// attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i]
float attn_col = 0.0f;
float attn_partial = 0.0f;
#pragma unroll
for (int i = 0; i < S_v; i++) {
s[i] = g_val * s[i] + k_t[i] * delta_col;
attn_col += s[i] * q_t[i];
for (int r = 0; r < ROWS_PER_LANE; r++) {
const int i = r * WARP_SIZE + lane;
s_shard[r] = g_val * s_shard[r] + k_t[i] * delta_col;
attn_partial += s_shard[r] * q_t[i];
}

attn_data[col] = attn_col * scale;
float attn_col = warp_reduce_sum(attn_partial);

if (lane == 0) {
attn_data[col] = attn_col * scale;
}
} else {
// kv[col] = sum_i g[i] * S[i][col] * k[i]
float kv_col = 0.0f;
float kv_shard = 0.0f;
#pragma unroll
for (int i = 0; i < S_v; i++) {
kv_col += expf(g_t[i]) * s[i] * k_t[i];
for (int r = 0; r < ROWS_PER_LANE; r++) {
const int i = r * WARP_SIZE + lane;
kv_shard += expf(g_t[i]) * s_shard[r] * k_t[i];
}

float kv_col = warp_reduce_sum(kv_shard);

// delta[col] = (v[col] - kv[col]) * beta
float delta_col = (v_t[col] - kv_col) * beta_val;

// fused: S[i][col] = g[i] * S[i][col] + k[i] * delta[col]
// attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i]
float attn_col = 0.0f;
float attn_partial = 0.0f;
#pragma unroll
for (int i = 0; i < S_v; i++) {
s[i] = expf(g_t[i]) * s[i] + k_t[i] * delta_col;
attn_col += s[i] * q_t[i];
for (int r = 0; r < ROWS_PER_LANE; r++) {
const int i = r * WARP_SIZE + lane;
s_shard[r] = expf(g_t[i]) * s_shard[r] + k_t[i] * delta_col;
attn_partial += s_shard[r] * q_t[i];
}

attn_data[col] = attn_col * scale;
float attn_col = warp_reduce_sum(attn_partial);

if (lane == 0) {
attn_data[col] = attn_col * scale;
}
}

attn_data += S_v * H;
}

// Write state back to global memory
#pragma unroll
for (int i = 0; i < S_v; i++) {
state[i * S_v + col] = s[i];
for (int r = 0; r < ROWS_PER_LANE; r++) {
const int i = r * WARP_SIZE + lane;
state[i * S_v + col] = s_shard[r];
}
}

Expand All @@ -126,27 +145,31 @@ static void launch_gated_delta_net(
int64_t neqk1, int64_t rq3,
float scale, cudaStream_t stream) {

dim3 grid_dims(H, n_seqs, 1);
dim3 block_dims(S_v, 1, 1);
const int num_warps = 4;
dim3 grid_dims(H, n_seqs, (S_v + num_warps - 1) / num_warps);
dim3 block_dims(WARP_SIZE, num_warps, 1);

const fastdiv_consts_s64 neqk1_magic = init_fastdiv_s64(neqk1);
const fastdiv_consts_s64 rq3_magic = init_fastdiv_s64(rq3);

switch (S_v) {
case 32:
gated_delta_net_cuda<32, KDA><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1, rq3, scale);
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
break;
case 64:
gated_delta_net_cuda<64, KDA><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1, rq3, scale);
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
break;
case 128:
gated_delta_net_cuda<128, KDA><<<grid_dims, block_dims, 0, stream>>>(
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1, rq3, scale);
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale);
break;
default:
GGML_ABORT("fatal error");
Expand Down
Loading