Skip to content

cuda: use the 4-column GDN warp layout on all Ampere+ NVIDIA, not only GB10 (+6% prefill on Ada) #216

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
bri-prism merged 2 commits into from
Sep 23, 2026
Merged

cuda: use the 4-column GDN warp layout on all Ampere+ NVIDIA, not only GB10 (+6% prefill on Ada) #216

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a02de73
cuda: use the 4-column GDN warp layout on all Ampere+ NVIDIA, not onl…
professorpalmer Sep 19, 2026
937c4d8
cuda: keep only the GDN cols-per-warp launch invariant in comments.
professorpalmer Sep 21, 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
19 changes: 15 additions & 4 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,6 +1,16 @@
#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.
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;
#endif
}

static __global__ void gdn_precompute_exp(const float * g, float * g_exp, int64_t n) {
for (int64_t i = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; i < n;
i += (int64_t) blockDim.x*gridDim.x) {
Expand Down Expand Up @@ -44,10 +54,10 @@ gated_delta_net_cuda(const float * q,
const uint32_t sequence = blockIdx.y;
// Each warp owns one or more columns, using warp-level primitives to reduce across rows.
const int lane = threadIdx.x;
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK
constexpr int cols_per_warp = S_v == 128 && !KDA ? 4 : 1;
#if defined(__CUDA_ARCH__)
constexpr int cols_per_warp = gdn_cols_per_warp(__CUDA_ARCH__, S_v, KDA);
#else
constexpr int cols_per_warp = 1;
constexpr int cols_per_warp = 1; // host pass only; never executed
#endif
const int col = (blockIdx.z * blockDim.y + threadIdx.y) * cols_per_warp;

Expand Down Expand Up @@ -217,7 +227,8 @@ static void launch_gated_delta_net(
const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
const int num_warps = 4;
const int cols_per_warp = cc == GGML_CUDA_CC_DGX_SPARK && S_v == 128 && !KDA ? 4 : 1;
// Same predicate as the kernel, fed the arch compiled for this cc.
const int cols_per_warp = GGML_CUDA_CC_IS_NVIDIA(cc) ? gdn_cols_per_warp(ggml_cuda_highest_compiled_arch(cc), S_v, KDA) : 1;
dim3 grid_dims(H, n_seqs, (S_v + num_warps * cols_per_warp - 1) / (num_warps * cols_per_warp));
dim3 block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1);

Expand Down