Skip to content

CUDA: gated_delta_net - share per-token k/q via shared memory for long prefills #54

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

Draft
roberteg16 wants to merge 1 commit into
base: gfx11
Choose a base branch
from
Draft

CUDA: gated_delta_net - share per-token k/q via shared memory for long prefills #54

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
63 changes: 52 additions & 11 deletions ggml/src/ggml-cuda/gated_delta_net.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@ gated_delta_net_cuda(const float * q,
const uint3 rq3_magic,
float scale,
int64_t state_slot_stride,
int K) {
int K,
bool share_kq) {
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
Expand All @@ -53,6 +54,18 @@ gated_delta_net_cuda(const float * q,
float s_shard[rows_per_lane];
// state is stored transposed: M[col][i] = S[i][col], row col is contiguous

// RDNA3.5-only optimisation (tuned/measured on gfx1151): k_t/q_t are identical for every
// column (warp) in the block, so stage them in shared once per token and let the block's
// warps share a single global load instead of one load each. Other archs use the direct load.
#ifdef RDNA3_5
__shared__ float ksh[S_v];
__shared__ float qsh[S_v];
const int nwarps = blockDim.y;
const int tid = threadIdx.y * warp_size + lane;
#else
GGML_UNUSED(share_kq);
#endif

ggml_cuda_pdl_sync();
#pragma unroll
for (int r = 0; r < rows_per_lane; r++) {
Expand All @@ -71,14 +84,36 @@ gated_delta_net_cuda(const float * q,

const float beta_val = *beta_t;

// Cache k and q in registers
// Cache k and q in registers.
float k_reg[rows_per_lane];
float q_reg[rows_per_lane];
// On RDNA3.5, long sequences are memory-bound and k_t/q_t are identical for every column,
// so staging them in shared (one block-wide load instead of one per warp) wins. Short
// sequences are latency-bound and hidden by inter-warp overlap, which the barrier would
// break, so load straight from global there. Other archs always use the direct load.
#ifdef RDNA3_5
if (share_kq) {
for (int i = tid; i < S_v; i += nwarps * warp_size) {
ksh[i] = k_t[i];
qsh[i] = q_t[i];
}
__syncthreads();
#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_reg[r] = ksh[i];
q_reg[r] = qsh[i];
}
__syncthreads();
} else
#endif
{
#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];
}
}

if constexpr (!KDA) {
Expand Down Expand Up @@ -177,9 +212,15 @@ static void launch_gated_delta_net(
int64_t sb1, int64_t sb2, int64_t sb3,
int64_t neqk1, int64_t rq3,
float scale, int64_t state_slot_stride, int K, cudaStream_t stream) {
//TODO: Add chunked kernel for even faster pre-fill
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
const int num_warps = 4;
// RDNA3.5 only: long sequences are memory-bound and the per-token k/q are identical across
// columns, so staging them in shared lets the block's warps share one load. Short sequences
// are latency-bound (hidden by inter-warp overlap) and the extra barrier would hurt, so it is
// gated on a token threshold. Measured crossover on gfx1151 (d=128): ~1.05x at 1536, ~1.20x
// at 2048, ~1.60x at 4096; below 1536 the barrier costs more than it saves.
const bool share_kq = n_tokens >= 1536 && GGML_CUDA_CC_IS_RDNA3_5(cc);
dim3 grid_dims(H, n_seqs, (S_v + num_warps - 1) / num_warps);
dim3 block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1);

Expand All @@ -192,26 +233,26 @@ static void launch_gated_delta_net(
ggml_cuda_kernel_launch(gated_delta_net_cuda<16, KDA, keep_rs_t>, launch_params,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, share_kq);
break;
case 32:
ggml_cuda_kernel_launch(gated_delta_net_cuda<32, KDA, keep_rs_t>, launch_params,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, share_kq);
break;
case 64: {
ggml_cuda_kernel_launch(gated_delta_net_cuda<64, KDA, keep_rs_t>, launch_params,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, share_kq);
break;
}
case 128: {
ggml_cuda_kernel_launch(gated_delta_net_cuda<128, KDA, keep_rs_t>, launch_params,
q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H,
n_tokens, n_seqs, sq1, sq2, sq3, sv1, sv2, sv3,
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K);
sb1, sb2, sb3, neqk1_magic, rq3_magic, scale, state_slot_stride, K, share_kq);
break;
}
default:
Expand Down
3 changes: 3 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9734,6 +9734,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 256, 1)); // PP-256
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 512, 1)); // PP-512
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 1024, 1)); // PP-1024
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 1536, 1)); // PP-1536
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 2048, 1)); // PP-2048
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 32, 128, 4096, 1)); // PP-4096
// Small model configs (fewer heads = less GPU occupancy for autoregressive)
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 64, 1)); // 4h PP-64
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 128, 256, 1)); // 4h PP-256
Expand Down
Loading