From 5bb19725e4452edf33cd35e861edaf87c6cec012 Mon Sep 17 00:00:00 2001 From: Robert Esclapez Garcia Date: Thu, 16 Jul 2026 05:44:29 -0700 Subject: [PATCH] CUDA: gated_delta_net - share per-token k/q via shared memory (RDNA3.5) RDNA3.5-only optimisation. The serial GATED_DELTA_NET kernel reloaded the per-token k/q from global once per warp, even though k/q are identical for every state column in a block. On RDNA3.5 (gfx115x) stage them in shared once per token (one block-wide load) for long, memory-bound sequences. Gated at compile time (#ifdef RDNA3_5) and at runtime (GGML_CUDA_CC_IS_RDNA3_5(cc) && n_tokens >= 1536) so no other arch pays the extra LDS or the barrier. Short sequences stay latency-bound (hidden by inter-warp overlap) and use the direct global load. Measured on gfx1151 (d=128, 32 heads, median of 3): ~1.05x at 1536, ~1.20x at 2048, ~1.60x at 4096; neutral below the threshold. Also removes the super-linear slowdown of the serial scan at long sequences. --- ggml/src/ggml-cuda/gated_delta_net.cu | 63 ++++++++++++++++++++++----- tests/test-backend-ops.cpp | 3 ++ 2 files changed, 55 insertions(+), 11 deletions(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 1b431a724d72..d23e828e7d34 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -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 @@ -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++) { @@ -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) { @@ -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); @@ -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: diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4edaba125079..ac7c8203d259 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -9734,6 +9734,9 @@ static std::vector> 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