From 0d3ee65244b6a195593aa3b0023cd042b7be97ee Mon Sep 17 00:00:00 2001 From: Paul Flynn Date: Wed, 11 Mar 2026 21:28:09 -0400 Subject: [PATCH 01/12] ggml : transpose fused GDN state access for coalesced memory reads (#20436) The fused Gated Delta Net kernel accessed the [S_v, S_v] state matrix column-wise on row-major storage, causing strided reads (stride S_v = 128 floats = 512 bytes) that waste GPU cache bandwidth. This produced a 39% regression on Qwen3.5-9B (Metal, M4 Max) compared to the unfused path. Transpose the state indexing so threads read contiguously: - Metal: s_ptr[is*S_v] -> s_ptr[is] (stride 1 vs S_v) - CUDA: curr_state[i*S_v+col] -> curr_state[col*S_v+i] (coalesced) - CPU: restructured loops for row-wise transposed access Also add --fused-gdn [on|off|auto] CLI flag (mirrors --flash-attn) so users can control fused GDN independently of auto-detection. All GATED_DELTA_NET backend-ops tests pass. Co-Authored-By: Claude Opus 4.6 --- common/arg.cpp | 15 ++++++++++ common/common.cpp | 1 + common/common.h | 1 + ggml/src/ggml-cpu/ops.cpp | 42 +++++++++++++++++---------- ggml/src/ggml-cuda/gated_delta_net.cu | 7 +++-- ggml/src/ggml-metal/ggml-metal.metal | 9 +++--- include/llama.h | 9 ++++++ src/llama-context.cpp | 8 +++-- src/llama.cpp | 12 ++++++++ 9 files changed, 79 insertions(+), 25 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 41da8563d632..0f7caf1c8f47 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1339,6 +1339,21 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str())); } }).set_env("LLAMA_ARG_FLASH_ATTN")); + add_opt(common_arg({ "-fgdn", "--fused-gdn" }, "[on|off|auto]", + string_format("set fused Gated Delta Net kernel use ('on', 'off', or 'auto', default: '%s')", + llama_fused_gdn_type_name(params.fused_gdn_type)), + [](common_params & params, const std::string & value) { + if (is_truthy(value)) { + params.fused_gdn_type = LLAMA_FUSED_GDN_TYPE_ENABLED; + } else if (is_falsey(value)) { + params.fused_gdn_type = LLAMA_FUSED_GDN_TYPE_DISABLED; + } else if (is_autoy(value)) { + params.fused_gdn_type = LLAMA_FUSED_GDN_TYPE_AUTO; + } else { + throw std::runtime_error( + string_format("error: unknown value for --fused-gdn: '%s'\n", value.c_str())); + } + }).set_env("LLAMA_ARG_FUSED_GDN")); add_opt(common_arg( {"-p", "--prompt"}, "PROMPT", "prompt to start generation with; for system message, use -sys", diff --git a/common/common.cpp b/common/common.cpp index cc423d3439fc..2e7ae7b0f312 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1369,6 +1369,7 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.pooling_type = params.pooling_type; cparams.attention_type = params.attention_type; cparams.flash_attn_type = params.flash_attn_type; + cparams.fused_gdn_type = params.fused_gdn_type; cparams.cb_eval = params.cb_eval; cparams.cb_eval_user_data = params.cb_eval_user_data; cparams.offload_kqv = !params.no_kv_offload; diff --git a/common/common.h b/common/common.h index c5645bba460e..0b30f0c8aab1 100644 --- a/common/common.h +++ b/common/common.h @@ -411,6 +411,7 @@ struct common_params { enum llama_pooling_type pooling_type = LLAMA_POOLING_TYPE_UNSPECIFIED; // pooling type for embeddings enum llama_attention_type attention_type = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings enum llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention + enum llama_fused_gdn_type fused_gdn_type = LLAMA_FUSED_GDN_TYPE_AUTO; // whether to use fused Gated Delta Net kernel struct common_params_sampling sampling; struct common_params_speculative speculative; diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index fa9d27046b5e..13ca087faf27 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -10477,34 +10477,46 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( const float beta_val = *(const float *)((const char *)src_beta->data + iv3 * nbb3 + t * nbb2 + iv1 * nbb1); const float * g_d = (const float *)((const char *)src_g->data + iv3 * nbg3 + t * nbg2 + iv1 * nbg1); + // state is stored transposed: s_out[j*S_v + i] = S[i][j] + // so row j of s_out = column j of S (contiguous access) + if (kda) { + // precompute exp(g) into delta scratch (reused below) for (int64_t i = 0; i < S_v; ++i) { - ggml_vec_scale_f32(S_v, &s_out[i * S_v], expf(g_d[i])); + delta[i] = expf(g_d[i]); + } + // S[i][:] *= exp(g[i]) => for each row j of M: M[j][i] *= exp(g[i]) + for (int64_t j = 0; j < S_v; ++j) { + ggml_vec_mul_f32(S_v, &s_out[j * S_v], &s_out[j * S_v], delta); } } else { ggml_vec_scale_f32(S_v * S_v, s_out, expf(g_d[0])); } - // delta[j] = sum_i S[j][i] * k[i] - memset(delta, 0, S_v * sizeof(float)); - for (int64_t i = 0; i < S_v; ++i) { - ggml_vec_mad_f32(S_v, delta, &s_out[i * S_v], k_d[i]); - } + // delta[j] = sum_i S[i][j] * k[i] = dot(row j of M, k) for (int64_t j = 0; j < S_v; ++j) { - delta[j] = (v_d[j] - delta[j]) * beta_val; + const float * row_j = s_out + j * S_v; + float sum = 0.0f; + for (int64_t i = 0; i < S_v; ++i) { + sum += row_j[i] * k_d[i]; + } + delta[j] = (v_d[j] - sum) * beta_val; } - // outer product: S[j][i] += k[i] * delta[j] - for (int64_t i = 0; i < S_v; ++i) { - ggml_vec_mad_f32(S_v, &s_out[i * S_v], delta, k_d[i]); + // outer product: S[i][j] += k[i] * delta[j] => M[j][i] += delta[j] * k[i] + for (int64_t j = 0; j < S_v; ++j) { + ggml_vec_mad_f32(S_v, &s_out[j * S_v], k_d, delta[j]); } - // attn_out[j] = sum_i S[j][i] * q[i] - memset(attn_data, 0, S_v * sizeof(float)); - for (int64_t i = 0; i < S_v; ++i) { - ggml_vec_mad_f32(S_v, attn_data, &s_out[i * S_v], q_d[i]); + // attn_out[j] = sum_i S[i][j] * q[i] = dot(row j of M, q) + for (int64_t j = 0; j < S_v; ++j) { + const float * row_j = s_out + j * S_v; + float sum = 0.0f; + for (int64_t i = 0; i < S_v; ++i) { + sum += row_j[i] * q_d[i]; + } + attn_data[j] = sum * scale; } - ggml_vec_scale_f32(S_v, attn_data, scale); attn_data += S_v * H; // advance to next token } diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 5f0fa8e58dfe..eb38588202a2 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -45,10 +45,11 @@ __global__ void gated_delta_net_cuda(const float * q, 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 #pragma unroll 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]; + s_shard[r] = curr_state[col * S_v + i]; } for (int t = 0; t < n_tokens; t++) { @@ -126,11 +127,11 @@ __global__ void gated_delta_net_cuda(const float * q, attn_data += S_v * H; } - // Write state back to global memory + // Write state back to global memory (transposed layout) #pragma unroll for (int r = 0; r < rows_per_lane; r++) { const int i = r * warp_size + lane; - state[i * S_v + col] = s_shard[r]; + state[col * S_v + i] = s_shard[r]; } } diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 0b77d5349b86..e06f266b5337 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -2466,13 +2466,14 @@ kernel void kernel_gated_delta_net_impl( const float scale = 1.0f / sqrt((float)S_v); - device const float * s_ptr = (device const float *) (s) + (i23*args.ne21 + i21)*S_v*S_v + i20; + // state is stored transposed: M[i20][is] = S[is][i20], so row i20 is contiguous + device const float * s_ptr = (device const float *) (s) + (i23*args.ne21 + i21)*S_v*S_v + i20*S_v; float ls[NSG]; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; - ls[j] = s_ptr[is*S_v]; + ls[j] = s_ptr[is]; } device float * dst_attn = (device float *) (dst) + (i23*args.ne22*args.ne21 + i21)*S_v + i20; @@ -2533,11 +2534,11 @@ kernel void kernel_gated_delta_net_impl( g_ptr += args.ne21*G; } - device float * dst_state = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20; + device float * dst_state = (device float *) (dst) + args.ne23*args.ne22*args.ne21*S_v + (i23*args.ne21 + i21)*S_v*S_v + i20*S_v; FOR_UNROLL (short j = 0; j < NSG; j++) { const short is = tx*NSG + j; - dst_state[is*S_v] = ls[j]; + dst_state[is] = ls[j]; } #undef S_v diff --git a/include/llama.h b/include/llama.h index c6e102abe519..05e5152d0dbf 100644 --- a/include/llama.h +++ b/include/llama.h @@ -190,6 +190,14 @@ extern "C" { LLAMA_API const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_type); + enum llama_fused_gdn_type { + LLAMA_FUSED_GDN_TYPE_AUTO = -1, + LLAMA_FUSED_GDN_TYPE_DISABLED = 0, + LLAMA_FUSED_GDN_TYPE_ENABLED = 1, + }; + + LLAMA_API const char * llama_fused_gdn_type_name(enum llama_fused_gdn_type fused_gdn_type); + enum llama_split_mode { LLAMA_SPLIT_MODE_NONE = 0, // single GPU LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs @@ -338,6 +346,7 @@ extern "C" { enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id enum llama_attention_type attention_type; // attention type to use for embeddings enum llama_flash_attn_type flash_attn_type; // when to enable Flash Attention + enum llama_fused_gdn_type fused_gdn_type; // when to enable fused Gated Delta Net kernel // ref: https://github.com/ggml-org/llama.cpp/pull/2054 float rope_freq_base; // RoPE base frequency, 0 = from model diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0be949391025..5deb81f8b174 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -150,9 +150,9 @@ llama_context::llama_context( cparams.flash_attn = params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED; cparams.auto_fa = params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO; - cparams.fused_gdn_ar = true; - cparams.fused_gdn_ch = true; - cparams.auto_fgdn = true; + cparams.fused_gdn_ar = params.fused_gdn_type != LLAMA_FUSED_GDN_TYPE_DISABLED; + cparams.fused_gdn_ch = params.fused_gdn_type != LLAMA_FUSED_GDN_TYPE_DISABLED; + cparams.auto_fgdn = params.fused_gdn_type == LLAMA_FUSED_GDN_TYPE_AUTO; // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; @@ -200,6 +200,7 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); + LLAMA_LOG_INFO("%s: fused_gdn = %s\n", __func__, llama_fused_gdn_type_name(params.fused_gdn_type)); LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); @@ -2869,6 +2870,7 @@ llama_context_params llama_context_default_params() { /*.pooling_type =*/ LLAMA_POOLING_TYPE_UNSPECIFIED, /*.attention_type =*/ LLAMA_ATTENTION_TYPE_UNSPECIFIED, /*.flash_attn_type =*/ LLAMA_FLASH_ATTN_TYPE_AUTO, + /*.fused_gdn_type =*/ LLAMA_FUSED_GDN_TYPE_AUTO, /*.rope_freq_base =*/ 0.0f, /*.rope_freq_scale =*/ 0.0f, /*.yarn_ext_factor =*/ -1.0f, diff --git a/src/llama.cpp b/src/llama.cpp index 872e659edcaf..4bc996bbee95 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -45,6 +45,18 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty GGML_ABORT("fatal error"); } +const char * llama_fused_gdn_type_name(enum llama_fused_gdn_type fused_gdn_type) { + switch (fused_gdn_type) { + case LLAMA_FUSED_GDN_TYPE_AUTO: + return "auto"; + case LLAMA_FUSED_GDN_TYPE_DISABLED: + return "disabled"; + case LLAMA_FUSED_GDN_TYPE_ENABLED: + return "enabled"; + } + GGML_ABORT("fatal error"); +} + struct llama_device_memory_data { int64_t total; int64_t free; From 9ec08b3d2bfb082dc2f928f6df0d86a40afb1930 Mon Sep 17 00:00:00 2001 From: Paul Flynn Date: Wed, 11 Mar 2026 22:53:50 -0400 Subject: [PATCH 02/12] ggml : use SIMD dot products in CPU GDN kernel, couple AR/chunked fused flags - Replace scalar inner loops with ggml_vec_dot_f32 for SIMD-optimized dot products in the CPU fused GDN kernel (delta and attention output) - Couple fused_gdn_ar and fused_gdn_ch flags in auto-detection: if one path lacks device support, disable both to prevent state layout mismatch between transposed (fused) and non-transposed (unfused) formats Co-Authored-By: Claude Opus 4.6 --- ggml/src/ggml-cpu/ops.cpp | 10 ++-------- src/llama-context.cpp | 8 ++++++++ 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 13ca087faf27..e027325cd0a6 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -10495,11 +10495,8 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( // delta[j] = sum_i S[i][j] * k[i] = dot(row j of M, k) for (int64_t j = 0; j < S_v; ++j) { - const float * row_j = s_out + j * S_v; float sum = 0.0f; - for (int64_t i = 0; i < S_v; ++i) { - sum += row_j[i] * k_d[i]; - } + ggml_vec_dot_f32(S_v, &sum, 0, &s_out[j * S_v], 0, k_d, 0, 1); delta[j] = (v_d[j] - sum) * beta_val; } @@ -10510,11 +10507,8 @@ static void ggml_compute_forward_gated_delta_net_one_chunk( // attn_out[j] = sum_i S[i][j] * q[i] = dot(row j of M, q) for (int64_t j = 0; j < S_v; ++j) { - const float * row_j = s_out + j * S_v; float sum = 0.0f; - for (int64_t i = 0; i < S_v; ++i) { - sum += row_j[i] * q_d[i]; - } + ggml_vec_dot_f32(S_v, &sum, 0, &s_out[j * S_v], 0, q_d, 0, 1); attn_data[j] = sum * scale; } diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5deb81f8b174..4586bd07192d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -538,6 +538,14 @@ void llama_context::sched_reserve() { } } + // the fused kernel uses a transposed state layout, so both paths must agree + // to avoid a state layout mismatch when switching between AR and chunked + if (cparams.fused_gdn_ar != cparams.fused_gdn_ch) { + cparams.fused_gdn_ar = false; + cparams.fused_gdn_ch = false; + LLAMA_LOG_WARN("%s: fused Gated Delta Net AR/chunked support mismatch, disabling both\n", __func__); + } + cparams.auto_fgdn = false; } From 1b82571e4f5f71d9a989f1b076bbe444e22d78c0 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 12 Mar 2026 08:29:28 +0200 Subject: [PATCH 03/12] llama : rever fgdn argument changes --- common/arg.cpp | 15 --------------- common/common.cpp | 1 - common/common.h | 1 - include/llama.h | 9 --------- src/llama-context.cpp | 16 +++------------- src/llama.cpp | 12 ------------ 6 files changed, 3 insertions(+), 51 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index 0f7caf1c8f47..41da8563d632 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1339,21 +1339,6 @@ common_params_context common_params_parser_init(common_params & params, llama_ex string_format("error: unknown value for --flash-attn: '%s'\n", value.c_str())); } }).set_env("LLAMA_ARG_FLASH_ATTN")); - add_opt(common_arg({ "-fgdn", "--fused-gdn" }, "[on|off|auto]", - string_format("set fused Gated Delta Net kernel use ('on', 'off', or 'auto', default: '%s')", - llama_fused_gdn_type_name(params.fused_gdn_type)), - [](common_params & params, const std::string & value) { - if (is_truthy(value)) { - params.fused_gdn_type = LLAMA_FUSED_GDN_TYPE_ENABLED; - } else if (is_falsey(value)) { - params.fused_gdn_type = LLAMA_FUSED_GDN_TYPE_DISABLED; - } else if (is_autoy(value)) { - params.fused_gdn_type = LLAMA_FUSED_GDN_TYPE_AUTO; - } else { - throw std::runtime_error( - string_format("error: unknown value for --fused-gdn: '%s'\n", value.c_str())); - } - }).set_env("LLAMA_ARG_FUSED_GDN")); add_opt(common_arg( {"-p", "--prompt"}, "PROMPT", "prompt to start generation with; for system message, use -sys", diff --git a/common/common.cpp b/common/common.cpp index 2e7ae7b0f312..cc423d3439fc 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1369,7 +1369,6 @@ struct llama_context_params common_context_params_to_llama(const common_params & cparams.pooling_type = params.pooling_type; cparams.attention_type = params.attention_type; cparams.flash_attn_type = params.flash_attn_type; - cparams.fused_gdn_type = params.fused_gdn_type; cparams.cb_eval = params.cb_eval; cparams.cb_eval_user_data = params.cb_eval_user_data; cparams.offload_kqv = !params.no_kv_offload; diff --git a/common/common.h b/common/common.h index 0b30f0c8aab1..c5645bba460e 100644 --- a/common/common.h +++ b/common/common.h @@ -411,7 +411,6 @@ struct common_params { enum llama_pooling_type pooling_type = LLAMA_POOLING_TYPE_UNSPECIFIED; // pooling type for embeddings enum llama_attention_type attention_type = LLAMA_ATTENTION_TYPE_UNSPECIFIED; // attention type for embeddings enum llama_flash_attn_type flash_attn_type = LLAMA_FLASH_ATTN_TYPE_AUTO; // whether to use Flash Attention - enum llama_fused_gdn_type fused_gdn_type = LLAMA_FUSED_GDN_TYPE_AUTO; // whether to use fused Gated Delta Net kernel struct common_params_sampling sampling; struct common_params_speculative speculative; diff --git a/include/llama.h b/include/llama.h index 05e5152d0dbf..c6e102abe519 100644 --- a/include/llama.h +++ b/include/llama.h @@ -190,14 +190,6 @@ extern "C" { LLAMA_API const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_type); - enum llama_fused_gdn_type { - LLAMA_FUSED_GDN_TYPE_AUTO = -1, - LLAMA_FUSED_GDN_TYPE_DISABLED = 0, - LLAMA_FUSED_GDN_TYPE_ENABLED = 1, - }; - - LLAMA_API const char * llama_fused_gdn_type_name(enum llama_fused_gdn_type fused_gdn_type); - enum llama_split_mode { LLAMA_SPLIT_MODE_NONE = 0, // single GPU LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs @@ -346,7 +338,6 @@ extern "C" { enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id enum llama_attention_type attention_type; // attention type to use for embeddings enum llama_flash_attn_type flash_attn_type; // when to enable Flash Attention - enum llama_fused_gdn_type fused_gdn_type; // when to enable fused Gated Delta Net kernel // ref: https://github.com/ggml-org/llama.cpp/pull/2054 float rope_freq_base; // RoPE base frequency, 0 = from model diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4586bd07192d..0be949391025 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -150,9 +150,9 @@ llama_context::llama_context( cparams.flash_attn = params.flash_attn_type != LLAMA_FLASH_ATTN_TYPE_DISABLED; cparams.auto_fa = params.flash_attn_type == LLAMA_FLASH_ATTN_TYPE_AUTO; - cparams.fused_gdn_ar = params.fused_gdn_type != LLAMA_FUSED_GDN_TYPE_DISABLED; - cparams.fused_gdn_ch = params.fused_gdn_type != LLAMA_FUSED_GDN_TYPE_DISABLED; - cparams.auto_fgdn = params.fused_gdn_type == LLAMA_FUSED_GDN_TYPE_AUTO; + cparams.fused_gdn_ar = true; + cparams.fused_gdn_ch = true; + cparams.auto_fgdn = true; // with causal attention, the batch size is limited by the context size cparams.n_batch = cparams.causal_attn ? std::min(cparams.n_ctx, params.n_batch) : params.n_batch; @@ -200,7 +200,6 @@ llama_context::llama_context( LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); - LLAMA_LOG_INFO("%s: fused_gdn = %s\n", __func__, llama_fused_gdn_type_name(params.fused_gdn_type)); LLAMA_LOG_INFO("%s: kv_unified = %s\n", __func__, cparams.kv_unified ? "true" : "false"); LLAMA_LOG_INFO("%s: freq_base = %.1f\n", __func__, cparams.rope_freq_base); LLAMA_LOG_INFO("%s: freq_scale = %g\n", __func__, cparams.rope_freq_scale); @@ -538,14 +537,6 @@ void llama_context::sched_reserve() { } } - // the fused kernel uses a transposed state layout, so both paths must agree - // to avoid a state layout mismatch when switching between AR and chunked - if (cparams.fused_gdn_ar != cparams.fused_gdn_ch) { - cparams.fused_gdn_ar = false; - cparams.fused_gdn_ch = false; - LLAMA_LOG_WARN("%s: fused Gated Delta Net AR/chunked support mismatch, disabling both\n", __func__); - } - cparams.auto_fgdn = false; } @@ -2878,7 +2869,6 @@ llama_context_params llama_context_default_params() { /*.pooling_type =*/ LLAMA_POOLING_TYPE_UNSPECIFIED, /*.attention_type =*/ LLAMA_ATTENTION_TYPE_UNSPECIFIED, /*.flash_attn_type =*/ LLAMA_FLASH_ATTN_TYPE_AUTO, - /*.fused_gdn_type =*/ LLAMA_FUSED_GDN_TYPE_AUTO, /*.rope_freq_base =*/ 0.0f, /*.rope_freq_scale =*/ 0.0f, /*.yarn_ext_factor =*/ -1.0f, diff --git a/src/llama.cpp b/src/llama.cpp index 4bc996bbee95..872e659edcaf 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -45,18 +45,6 @@ const char * llama_flash_attn_type_name(enum llama_flash_attn_type flash_attn_ty GGML_ABORT("fatal error"); } -const char * llama_fused_gdn_type_name(enum llama_fused_gdn_type fused_gdn_type) { - switch (fused_gdn_type) { - case LLAMA_FUSED_GDN_TYPE_AUTO: - return "auto"; - case LLAMA_FUSED_GDN_TYPE_DISABLED: - return "disabled"; - case LLAMA_FUSED_GDN_TYPE_ENABLED: - return "enabled"; - } - GGML_ABORT("fatal error"); -} - struct llama_device_memory_data { int64_t total; int64_t free; From 7ea6ee4fa578c69df262bdc6ff5b9c1ded995f0f Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Thu, 12 Mar 2026 08:30:40 +0200 Subject: [PATCH 04/12] graph : remove GDN state transposes --- src/models/delta-net-base.cpp | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index a62dbc15dd05..6bc989c95099 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -225,9 +225,8 @@ std::pair llm_build_delta_net_base::build_delta_ne ggml_tensor * kg_t = ggml_cont(ctx0, ggml_transpose(ctx0, kg)); cb(kg_t, "key_gdiff_t", il); - ggml_tensor * s_t = ggml_transpose(ctx0, s); - s_t = ggml_cont_4d(ctx0, s_t, S_v, S_v, 1, H_v * n_seqs); - cb(s_t, "dnet_add_ch_state", il); + s = ggml_reshape_4d(ctx0, s, S_v, S_v, 1, H_v * n_seqs); + cb(s, "dnet_add_ch_state", il); // [CS, S_v, n_chunks, H_v * n_seqs] ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, v)); @@ -240,7 +239,7 @@ std::pair llm_build_delta_net_base::build_delta_ne ggml_tensor * ch_kg_t = get_slice_2d(ctx0, kg_t, chunk); // [ CS, S_k, 1, H_v * n_seqs] // [CS, S_v, 1, H_v * n_seqs] - ggml_tensor * v_t_p = ggml_mul_mat(ctx0, ch_k_cd, s_t); + ggml_tensor * v_t_p = ggml_mul_mat(ctx0, ch_k_cd, s); cb(v_t_p, "v_prime", il); // [CS, S_v, 1, H_v * n_seqs] @@ -252,7 +251,7 @@ std::pair llm_build_delta_net_base::build_delta_ne cb(v_attn, "v_attn", il); // [S_v, CS, 1, H_v * n_seqs] - ggml_tensor * attn_inter = ggml_mul_mat(ctx0, s_t, ch_q_g_exp); + ggml_tensor * attn_inter = ggml_mul_mat(ctx0, s, ch_q_g_exp); cb(attn_inter, "attn_inter", il); // [S_v, CS, 1, H_v * n_seqs] @@ -268,13 +267,11 @@ std::pair llm_build_delta_net_base::build_delta_ne // last_recurrent_state = last_recurrent_state * g_last + kgdmulvnew ggml_tensor * ch_g_last_exp_t = get_slice_2d(ctx0, g_last_exp_t, chunk); - s_t = ggml_mul(ctx0, s_t, ch_g_last_exp_t); - s_t = ggml_add(ctx0, s_t, kgv); - cb(s_t, "dnet_add_ch_state", il); + s = ggml_mul(ctx0, s, ch_g_last_exp_t); + s = ggml_add(ctx0, s, kgv); + cb(s, "dnet_add_ch_state", il); } - s_t = ggml_reshape_4d(ctx0, s_t, S_v, S_v, H_v, n_seqs); - // truncate padded tokens ggml_tensor * o = ggml_view_4d(ctx0, v, S_v, n_tokens, H_v, n_seqs, @@ -282,7 +279,7 @@ std::pair llm_build_delta_net_base::build_delta_ne ggml_row_size(v->type, S_v * CS * n_chunks), ggml_row_size(v->type, S_v * CS * n_chunks * H_v), 0); o = ggml_permute (ctx0, o, 0, 2, 1, 3); // [S_v, H_v, n_tokens, n_seqs] - s = ggml_transpose(ctx0, s_t); + s = ggml_reshape_4d(ctx0, s, S_v, S_v, H_v, n_seqs); cb(s, "output_state", il); return {o, s}; @@ -341,11 +338,9 @@ std::pair llm_build_delta_net_base::build_delta_ne g = ggml_exp(ctx0, g); s = ggml_mul(ctx0, s, g); - ggml_tensor * s_t = ggml_cont(ctx0, ggml_transpose(ctx0, s)); - // [1, S_v, H_v, n_seqs] ggml_tensor * sk; - sk = ggml_mul (ctx0, s_t, k); + sk = ggml_mul (ctx0, s, k); sk = ggml_sum_rows(ctx0, sk); // [S_v, 1, H_v, n_seqs] @@ -362,15 +357,14 @@ std::pair llm_build_delta_net_base::build_delta_ne k = ggml_repeat(ctx0, k, s); kd = ggml_mul (ctx0, k, d_t); - s_t = ggml_add(ctx0, s_t, kd); + s = ggml_add(ctx0, s, kd); - cb(s_t, "dnet_add_ar_state", il); + cb(s, "dnet_add_ar_state", il); - ggml_tensor * s_q = ggml_mul (ctx0, s_t, q); + ggml_tensor * s_q = ggml_mul (ctx0, s, q); ggml_tensor * o = ggml_sum_rows(ctx0, s_q); o = ggml_permute (ctx0, o, 2, 0, 1, 3); // [S_v, H_v, n_tokens, n_seqs] - s = ggml_transpose(ctx0, s_t); // [S_v, S_v, H_v, n_seqs] return {o, s}; } From 86d393cc0cefe3068eeeec84b8aae5f2773e1365 Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Thu, 12 Mar 2026 11:23:50 +0100 Subject: [PATCH 05/12] WIP showing that data-access is the limitation in PP --- ggml/src/ggml-cuda/gated_delta_net.cu | 165 ++++++++++++++------------ 1 file changed, 91 insertions(+), 74 deletions(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index eb38588202a2..d0669e4f3ae7 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -1,6 +1,6 @@ #include "gated_delta_net.cuh" -template +template __global__ void gated_delta_net_cuda(const float * q, const float * k, const float * v, @@ -45,86 +45,88 @@ __global__ void gated_delta_net_cuda(const float * q, 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]; + int idx[rows_per_lane]; // state is stored transposed: M[col][i] = S[i][col], row col is contiguous #pragma unroll for (int r = 0; r < rows_per_lane; r++) { const int i = r * warp_size + lane; s_shard[r] = curr_state[col * S_v + i]; + idx[r] = i; } - for (int t = 0; t < n_tokens; t++) { - const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; - const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; - const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; + for (int ts = 0; ts < n_tokens; ts += n_tokens_per_loop) { +#pragma unroll + for (int i = 0; i < n_tokens_per_loop; i++) { + const int t = ts + i; + const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; - const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1; - const float * beta_t = beta + gb_offset; - const float * g_t = g + gb_offset * (KDA ? S_v : 1); + const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1; + const float * beta_t = beta + gb_offset; + const float * g_t = g + gb_offset * (KDA ? S_v : 1); - const float beta_val = *beta_t; + const float beta_val = *beta_t; - if constexpr (!KDA) { - const float g_val = expf(*g_t); + if constexpr (!KDA) { + const float g_val = expf(*g_t); - // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] - float kv_shard = 0.0f; + // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] + float kv_shard = 0.0f; #pragma unroll - 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); + for (int r = 0; r < rows_per_lane; r++) { + kv_shard += s_shard[r] * k_t[idx[r]]; + } + 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; + // 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_partial = 0.0f; + // 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_partial = 0.0f; #pragma unroll - 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]; - } - - 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_shard = 0.0f; + for (int r = 0; r < rows_per_lane; r++) { + s_shard[r] = g_val * s_shard[r] + k_t[idx[r]] * delta_col; + attn_partial += s_shard[r] * q_t[idx[r]]; + } + + 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_shard = 0.0f; #pragma unroll - 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]; - } + for (int r = 0; r < rows_per_lane; r++) { + kv_shard += expf(g_t[idx[r]]) * s_shard[r] * k_t[idx[r]]; + } - float kv_col = warp_reduce_sum(kv_shard); + 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; + // 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_partial = 0.0f; + // 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_partial = 0.0f; #pragma unroll - 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]; - } + for (int r = 0; r < rows_per_lane; r++) { + s_shard[r] = expf(g_t[idx[r]]) * s_shard[r] + k_t[idx[r]] * delta_col; + attn_partial += s_shard[r] * q_t[idx[r]]; + } - float attn_col = warp_reduce_sum(attn_partial); + float attn_col = warp_reduce_sum(attn_partial); - if (lane == 0) { - attn_data[col] = attn_col * scale; + if (lane == 0) { + attn_data[col] = attn_col * scale; + } } - } - attn_data += S_v * H; + attn_data += S_v * H; + } } // Write state back to global memory (transposed layout) @@ -135,15 +137,6 @@ __global__ void gated_delta_net_cuda(const float * q, } } -static size_t calculate_smem(const int sv, int cc) -{ - size_t smem = 0; - if ((GGML_CUDA_CC_IS_AMD(cc) && !GGML_CUDA_CC_IS_RDNA3(cc) && !GGML_CUDA_CC_IS_RDNA4(cc)) || GGML_CUDA_CC_IS_MTHREADS(cc)) { - smem = sv * sv * sizeof(float); - } - return smem; -} - template static void launch_gated_delta_net( const float * q_d, const float * k_d, const float * v_d, @@ -168,33 +161,57 @@ static void launch_gated_delta_net( switch (S_v) { case 16: - gated_delta_net_cuda<16, KDA><<>>( + if (n_tokens >= 4){ + gated_delta_net_cuda<16, KDA, 4><<>>( 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_magic, rq3_magic, scale); + } else{ + gated_delta_net_cuda<16, KDA, 1><<>>( + 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_magic, rq3_magic, scale); + } break; case 32: - gated_delta_net_cuda<32, KDA><<>>( + if (n_tokens >= 4){ + gated_delta_net_cuda<32, KDA, 4><<>>( 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_magic, rq3_magic, scale); + } else{ + gated_delta_net_cuda<32, KDA, 1><<>>( + 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_magic, rq3_magic, scale); + } break; case 64: { - constexpr int sv = 64; - size_t smem = calculate_smem(sv, cc); - gated_delta_net_cuda<<>>( + if (n_tokens >= 4){ + gated_delta_net_cuda<64, KDA, 4><<>>( + 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_magic, rq3_magic, scale); + } else{ + gated_delta_net_cuda<64, KDA, 1><<>>( 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_magic, rq3_magic, scale); + } break; } case 128: { - constexpr int sv = 128; - size_t smem = calculate_smem(sv, cc); - gated_delta_net_cuda<<>>( + if (n_tokens >= 4){ + gated_delta_net_cuda<128, KDA, 4><<>>( 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_magic, rq3_magic, scale); + } else{ + gated_delta_net_cuda<128, KDA, 1><<>>( + 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_magic, rq3_magic, scale); + } break; } default: From 2b211be01256040d7753ea0cc221c464f42a7497 Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Thu, 12 Mar 2026 12:09:38 +0100 Subject: [PATCH 06/12] Explicitly prefetch so the compiler knows what is expected of him --- ggml/src/ggml-cuda/gated_delta_net.cu | 72 +++++++++++++++++---------- 1 file changed, 45 insertions(+), 27 deletions(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index d0669e4f3ae7..d561d73a308f 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -45,6 +45,12 @@ __global__ void gated_delta_net_cuda(const float * q, 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]; + float qs[n_tokens_per_loop][rows_per_lane]; + float ks[n_tokens_per_loop][rows_per_lane]; + float vs[n_tokens_per_loop]; + float beta_ts[n_tokens_per_loop]; + // non-KDA has one g per column, KDA has one g per row + float g_vals[n_tokens_per_loop][rows_per_lane]; int idx[rows_per_lane]; // state is stored transposed: M[col][i] = S[i][col], row col is contiguous #pragma unroll @@ -54,41 +60,54 @@ __global__ void gated_delta_net_cuda(const float * q, idx[r] = i; } - for (int ts = 0; ts < n_tokens; ts += n_tokens_per_loop) { + for (int ts = 0; ts < n_tokens; ts+= n_tokens_per_loop) { #pragma unroll for (int i = 0; i < n_tokens_per_loop; i++) { - const int t = ts + i; - const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; - const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; - const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; - + const int t = ts + i; + const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1; const float * beta_t = beta + gb_offset; const float * g_t = g + gb_offset * (KDA ? S_v : 1); - const float beta_val = *beta_t; - + beta_ts[i] = *beta_t; + vs[i] = v_t[col]; + if constexpr (!KDA) { + g_vals[i][0] = expf(*g_t); + } +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + qs[i][r] = q_t[idx[r]]; + ks[i][r] = k_t[idx[r]]; + if constexpr (KDA) { + g_vals[i][r] = expf(g_t[idx[r]]); + } + } + } +#pragma unroll + for (int i = 0; i < n_tokens_per_loop; i++) { if constexpr (!KDA) { - const float g_val = expf(*g_t); + const float g_val = g_vals[i][0]; // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] float kv_shard = 0.0f; #pragma unroll for (int r = 0; r < rows_per_lane; r++) { - kv_shard += s_shard[r] * k_t[idx[r]]; + kv_shard += s_shard[r] * ks[i][r]; } 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; + float delta_col = (vs[i] - g_val * kv_col) * beta_ts[i]; // 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_partial = 0.0f; #pragma unroll for (int r = 0; r < rows_per_lane; r++) { - s_shard[r] = g_val * s_shard[r] + k_t[idx[r]] * delta_col; - attn_partial += s_shard[r] * q_t[idx[r]]; + s_shard[r] = g_val * s_shard[r] + ks[i][r] * delta_col; + attn_partial += s_shard[r] * qs[i][r]; } float attn_col = warp_reduce_sum(attn_partial); @@ -101,21 +120,21 @@ __global__ void gated_delta_net_cuda(const float * q, float kv_shard = 0.0f; #pragma unroll for (int r = 0; r < rows_per_lane; r++) { - kv_shard += expf(g_t[idx[r]]) * s_shard[r] * k_t[idx[r]]; + kv_shard += g_vals[i][r] * s_shard[r] * ks[i][r]; } 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; + float delta_col = (vs[i] - kv_col) * beta_ts[i]; // 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_partial = 0.0f; #pragma unroll for (int r = 0; r < rows_per_lane; r++) { - s_shard[r] = expf(g_t[idx[r]]) * s_shard[r] + k_t[idx[r]] * delta_col; - attn_partial += s_shard[r] * q_t[idx[r]]; + s_shard[r] = g_vals[i][r] * s_shard[r] + ks[i][r] * delta_col; + attn_partial += s_shard[r] * qs[i][r]; } float attn_col = warp_reduce_sum(attn_partial); @@ -132,8 +151,7 @@ __global__ void gated_delta_net_cuda(const float * q, // Write state back to global memory (transposed layout) #pragma unroll for (int r = 0; r < rows_per_lane; r++) { - const int i = r * warp_size + lane; - state[col * S_v + i] = s_shard[r]; + state[col * S_v + idx[r]] = s_shard[r]; } } @@ -161,8 +179,8 @@ static void launch_gated_delta_net( switch (S_v) { case 16: - if (n_tokens >= 4){ - gated_delta_net_cuda<16, KDA, 4><<>>( + if (n_tokens >= 24){ + gated_delta_net_cuda<16, KDA, 24><<>>( 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_magic, rq3_magic, scale); @@ -174,8 +192,8 @@ static void launch_gated_delta_net( } break; case 32: - if (n_tokens >= 4){ - gated_delta_net_cuda<32, KDA, 4><<>>( + if (n_tokens >= 24){ + gated_delta_net_cuda<32, KDA, 24><<>>( 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_magic, rq3_magic, scale); @@ -187,8 +205,8 @@ static void launch_gated_delta_net( } break; case 64: { - if (n_tokens >= 4){ - gated_delta_net_cuda<64, KDA, 4><<>>( + if (n_tokens >= 24){ + gated_delta_net_cuda<64, KDA, 24><<>>( 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_magic, rq3_magic, scale); @@ -201,8 +219,8 @@ static void launch_gated_delta_net( break; } case 128: { - if (n_tokens >= 4){ - gated_delta_net_cuda<128, KDA, 4><<>>( + if (n_tokens >= 24){ + gated_delta_net_cuda<128, KDA, 24><<>>( 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_magic, rq3_magic, scale); From f2b24fd939d7c73f3516980f31bcde10bf5850be Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Thu, 12 Mar 2026 12:25:06 +0100 Subject: [PATCH 07/12] Add unit-tests for n_tokens % n_tokens_per_loop != 0 --- tests/test-backend-ops.cpp | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index a821655d10c7..e21eaeeee444 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -8456,6 +8456,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 32, 4, 2, 2)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 1, 1, true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 33, 1, 1, true)); // KDA (vector gate) test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 1, 1, 1, false, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 1, 2, 1, false, true)); @@ -8464,6 +8465,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 32, 4, 2, 2, false, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, true, true)); + test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 33, 1, 1, true, true)); test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 16, 4, 2, 1, true, true)); #if 0 From e60c68f9ccf8a7b70d0b405bd5f9ae8ec3057ff3 Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Thu, 12 Mar 2026 15:03:43 +0100 Subject: [PATCH 08/12] Handle case where n_tokens % n_tokens_per_loop != 0 --- ggml/src/ggml-cuda/gated_delta_net.cu | 117 ++++++++++++++++++++++++-- 1 file changed, 108 insertions(+), 9 deletions(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index d561d73a308f..33c4bd64f787 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -60,7 +60,9 @@ __global__ void gated_delta_net_cuda(const float * q, idx[r] = i; } - for (int ts = 0; ts < n_tokens; ts+= n_tokens_per_loop) { + // Iterate over all valid chunks first, then handle the tail case if n_tokens is not divisible by n_tokens_per_loop + const int tokens_to_process = n_tokens - (n_tokens % n_tokens_per_loop); + for (int ts = 0; ts < tokens_to_process; ts += n_tokens_per_loop) { #pragma unroll for (int i = 0; i < n_tokens_per_loop; i++) { const int t = ts + i; @@ -148,6 +150,103 @@ __global__ void gated_delta_net_cuda(const float * q, } } + if (tokens_to_process != n_tokens) { + // handle tail case +// handle tail +#pragma unroll + for (int i = 0; i < n_tokens_per_loop; i++) { + const int t = tokens_to_process + i; + if (t >= n_tokens) { + break; + } + const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; + const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1; + const float * beta_t = beta + gb_offset; + const float * g_t = g + gb_offset * (KDA ? S_v : 1); + + beta_ts[i] = *beta_t; + vs[i] = v_t[col]; + if constexpr (!KDA) { + g_vals[i][0] = expf(*g_t); + } +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + qs[i][r] = q_t[idx[r]]; + ks[i][r] = k_t[idx[r]]; + if constexpr (KDA) { + g_vals[i][r] = expf(g_t[idx[r]]); + } + } + } +#pragma unroll + for (int i = 0; i < n_tokens_per_loop; i++) { + const int t = tokens_to_process + i; + if (t >= n_tokens) { + break; + } + if constexpr (!KDA) { + const float g_val = g_vals[i][0]; + + // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] + float kv_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + kv_shard += s_shard[r] * ks[i][r]; + } + float kv_col = warp_reduce_sum(kv_shard); + + // delta[col] = (v[col] - g * kv[col]) * beta + float delta_col = (vs[i] - g_val * kv_col) * beta_ts[i]; + + // 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_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + s_shard[r] = g_val * s_shard[r] + ks[i][r] * delta_col; + attn_partial += s_shard[r] * qs[i][r]; + } + + 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_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + kv_shard += g_vals[i][r] * s_shard[r] * ks[i][r]; + } + + float kv_col = warp_reduce_sum(kv_shard); + + // delta[col] = (v[col] - kv[col]) * beta + float delta_col = (vs[i] - kv_col) * beta_ts[i]; + + // 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_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + s_shard[r] = g_vals[i][r] * s_shard[r] + ks[i][r] * delta_col; + attn_partial += s_shard[r] * qs[i][r]; + } + + 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 (transposed layout) #pragma unroll for (int r = 0; r < rows_per_lane; r++) { @@ -179,8 +278,8 @@ static void launch_gated_delta_net( switch (S_v) { case 16: - if (n_tokens >= 24){ - gated_delta_net_cuda<16, KDA, 24><<>>( + if (n_tokens >= 20){ + gated_delta_net_cuda<16, KDA, 20><<>>( 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_magic, rq3_magic, scale); @@ -192,8 +291,8 @@ static void launch_gated_delta_net( } break; case 32: - if (n_tokens >= 24){ - gated_delta_net_cuda<32, KDA, 24><<>>( + if (n_tokens >= 20){ + gated_delta_net_cuda<32, KDA, 20><<>>( 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_magic, rq3_magic, scale); @@ -205,8 +304,8 @@ static void launch_gated_delta_net( } break; case 64: { - if (n_tokens >= 24){ - gated_delta_net_cuda<64, KDA, 24><<>>( + if (n_tokens >= 20){ + gated_delta_net_cuda<64, KDA, 20><<>>( 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_magic, rq3_magic, scale); @@ -219,8 +318,8 @@ static void launch_gated_delta_net( break; } case 128: { - if (n_tokens >= 24){ - gated_delta_net_cuda<128, KDA, 24><<>>( + if (n_tokens >= 20){ + gated_delta_net_cuda<128, KDA, 20><<>>( 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_magic, rq3_magic, scale); From c86d89a30ec53815758b3dac8a7e1f18b33231ec Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Thu, 12 Mar 2026 15:33:55 +0100 Subject: [PATCH 09/12] Account for register pressure of KDA due to per-row g --- ggml/src/ggml-cuda/gated_delta_net.cu | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 33c4bd64f787..49f33cf69bde 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -276,10 +276,12 @@ static void launch_gated_delta_net( int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + // for KDA we have to store one g per row, so we have higher register pressure + constexpr int nt_thresh = KDA ? 14 : 20; switch (S_v) { case 16: - if (n_tokens >= 20){ - gated_delta_net_cuda<16, KDA, 20><<>>( + if (n_tokens >= nt_thresh){ + gated_delta_net_cuda<16, KDA, nt_thresh><<>>( 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_magic, rq3_magic, scale); @@ -291,8 +293,8 @@ static void launch_gated_delta_net( } break; case 32: - if (n_tokens >= 20){ - gated_delta_net_cuda<32, KDA, 20><<>>( + if (n_tokens >= nt_thresh){ + gated_delta_net_cuda<32, KDA, nt_thresh><<>>( 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_magic, rq3_magic, scale); @@ -304,8 +306,8 @@ static void launch_gated_delta_net( } break; case 64: { - if (n_tokens >= 20){ - gated_delta_net_cuda<64, KDA, 20><<>>( + if (n_tokens >= nt_thresh){ + gated_delta_net_cuda<64, KDA, nt_thresh><<>>( 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_magic, rq3_magic, scale); @@ -318,8 +320,8 @@ static void launch_gated_delta_net( break; } case 128: { - if (n_tokens >= 20){ - gated_delta_net_cuda<128, KDA, 20><<>>( + if (n_tokens >= nt_thresh){ + gated_delta_net_cuda<128, KDA, nt_thresh><<>>( 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_magic, rq3_magic, scale); From 02ed849246a57acd619ec14f7d93907e148ddbfc Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Thu, 12 Mar 2026 16:57:24 +0100 Subject: [PATCH 10/12] Code cleanup --- ggml/src/ggml-cuda/gated_delta_net.cu | 1 - 1 file changed, 1 deletion(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 49f33cf69bde..48e5e8506850 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -152,7 +152,6 @@ __global__ void gated_delta_net_cuda(const float * q, if (tokens_to_process != n_tokens) { // handle tail case -// handle tail #pragma unroll for (int i = 0; i < n_tokens_per_loop; i++) { const int t = tokens_to_process + i; From f0629e274a08071b0d31797d08a0a01e002d0b3c Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Thu, 12 Mar 2026 17:37:36 +0100 Subject: [PATCH 11/12] Try adding untailed kernel for speed-ups in KDA setting --- ggml/src/ggml-cuda/gated_delta_net.cu | 305 ++++++++++++++++++++------ 1 file changed, 235 insertions(+), 70 deletions(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 48e5e8506850..415c23161b7c 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -1,28 +1,182 @@ #include "gated_delta_net.cuh" -template -__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 uint3 neqk1_magic, - const uint3 rq3_magic, - float scale) { +template +__global__ void gated_delta_net_cuda_no_tail(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 uint3 neqk1_magic, + const uint3 rq3_magic, + float scale) { + 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; + + const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic); + const uint32_t iq3 = fastdiv(sequence, rq3_magic); + + const int64_t attn_score_elems = S_v * H * n_tokens * n_seqs; + float * attn_data = dst; + float * state = dst + attn_score_elems; + + const int64_t state_offset = (sequence * H + h_idx) * S_v * S_v; + state += state_offset; + curr_state += state_offset; + 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]; + float qs[n_tokens_per_loop][rows_per_lane]; + float ks[n_tokens_per_loop][rows_per_lane]; + float vs[n_tokens_per_loop]; + float beta_ts[n_tokens_per_loop]; + // non-KDA has one g per column, KDA has one g per row + float g_vals[n_tokens_per_loop][rows_per_lane]; + int idx[rows_per_lane]; + // state is stored transposed: M[col][i] = S[i][col], row col is contiguous +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + s_shard[r] = curr_state[col * S_v + i]; + idx[r] = i; + } + + for (int ts = 0; ts < n_tokens; ts += n_tokens_per_loop) { +#pragma unroll + for (int i = 0; i < n_tokens_per_loop; i++) { + const int t = ts + i; + const float * q_t = q + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * k_t = k + iq3 * sq3 + t * sq2 + iq1 * sq1; + const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; + const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1; + const float * beta_t = beta + gb_offset; + const float * g_t = g + gb_offset * (KDA ? S_v : 1); + + beta_ts[i] = *beta_t; + vs[i] = v_t[col]; + if constexpr (!KDA) { + g_vals[i][0] = expf(*g_t); + } +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + qs[i][r] = q_t[idx[r]]; + ks[i][r] = k_t[idx[r]]; + if constexpr (KDA) { + g_vals[i][r] = expf(g_t[idx[r]]); + } + } + } +#pragma unroll + for (int i = 0; i < n_tokens_per_loop; i++) { + if constexpr (!KDA) { + const float g_val = g_vals[i][0]; + + // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] + float kv_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + kv_shard += s_shard[r] * ks[i][r]; + } + float kv_col = warp_reduce_sum(kv_shard); + + // delta[col] = (v[col] - g * kv[col]) * beta + float delta_col = (vs[i] - g_val * kv_col) * beta_ts[i]; + + // 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_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + s_shard[r] = g_val * s_shard[r] + ks[i][r] * delta_col; + attn_partial += s_shard[r] * qs[i][r]; + } + + 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_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + kv_shard += g_vals[i][r] * s_shard[r] * ks[i][r]; + } + + float kv_col = warp_reduce_sum(kv_shard); + + // delta[col] = (v[col] - kv[col]) * beta + float delta_col = (vs[i] - kv_col) * beta_ts[i]; + + // 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_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + s_shard[r] = g_vals[i][r] * s_shard[r] + ks[i][r] * delta_col; + attn_partial += s_shard[r] * qs[i][r]; + } + + 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 (transposed layout) +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + state[col * S_v + idx[r]] = s_shard[r]; + } +} + +template +__global__ void gated_delta_net_cuda_tail(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 uint3 neqk1_magic, + const uint3 rq3_magic, + float scale) { 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 @@ -276,62 +430,73 @@ static void launch_gated_delta_net( int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; // for KDA we have to store one g per row, so we have higher register pressure - constexpr int nt_thresh = KDA ? 14 : 20; + const int no_tail = (n_tokens % 32) == 0 && (n_tokens != 1); + constexpr int tail_thresh = KDA ? 14 : 20; switch (S_v) { case 16: - if (n_tokens >= nt_thresh){ - gated_delta_net_cuda<16, KDA, nt_thresh><<>>( - 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_magic, rq3_magic, scale); - } else{ - gated_delta_net_cuda<16, KDA, 1><<>>( - 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_magic, rq3_magic, scale); + if (no_tail) { + gated_delta_net_cuda_no_tail<16, KDA, 32><<>>( + 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_magic, rq3_magic, scale); + } else if (n_tokens >= tail_thresh) { + gated_delta_net_cuda_tail<16, KDA, tail_thresh><<>>( + 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_magic, rq3_magic, scale); + } else { + gated_delta_net_cuda_no_tail<16, KDA, 1><<>>( + 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_magic, rq3_magic, scale); } break; case 32: - if (n_tokens >= nt_thresh){ - gated_delta_net_cuda<32, KDA, nt_thresh><<>>( - 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_magic, rq3_magic, scale); - } else{ - gated_delta_net_cuda<32, KDA, 1><<>>( - 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_magic, rq3_magic, scale); + if (no_tail) { + gated_delta_net_cuda_no_tail<32, KDA, 32><<>>( + 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_magic, rq3_magic, scale); + } else if (n_tokens >= tail_thresh) { + gated_delta_net_cuda_tail<32, KDA, tail_thresh><<>>( + 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_magic, rq3_magic, scale); + } else { + gated_delta_net_cuda_no_tail<32, KDA, 1><<>>( + 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_magic, rq3_magic, scale); } break; - case 64: { - if (n_tokens >= nt_thresh){ - gated_delta_net_cuda<64, KDA, nt_thresh><<>>( - 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_magic, rq3_magic, scale); - } else{ - gated_delta_net_cuda<64, KDA, 1><<>>( - 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_magic, rq3_magic, scale); + case 64: + { + if (no_tail) { + gated_delta_net_cuda_no_tail<64, KDA, 32><<>>( + 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_magic, rq3_magic, scale); + } else if (n_tokens >= tail_thresh) { + gated_delta_net_cuda_tail<64, KDA, tail_thresh><<>>( + 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_magic, rq3_magic, scale); + } else { + gated_delta_net_cuda_no_tail<64, KDA, 1><<>>( + 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_magic, rq3_magic, scale); + } + break; } - break; - } - case 128: { - if (n_tokens >= nt_thresh){ - gated_delta_net_cuda<128, KDA, nt_thresh><<>>( - 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_magic, rq3_magic, scale); - } else{ - gated_delta_net_cuda<128, KDA, 1><<>>( - 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_magic, rq3_magic, scale); + case 128: + { + if (no_tail) { + gated_delta_net_cuda_no_tail<128, KDA, 32><<>>( + 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_magic, rq3_magic, scale); + } else if (n_tokens >= tail_thresh) { + gated_delta_net_cuda_tail<128, KDA, tail_thresh><<>>( + 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_magic, rq3_magic, scale); + } else { + gated_delta_net_cuda_no_tail<128, KDA, 1><<>>( + 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_magic, rq3_magic, scale); + } + break; } - break; - } default: GGML_ABORT("fatal error"); break; From 33ee71b3ba90e0c9c98d5d1a5cbb8158f695e4bd Mon Sep 17 00:00:00 2001 From: Oliver Simons Date: Thu, 12 Mar 2026 19:20:52 +0100 Subject: [PATCH 12/12] Pull attn_writes out of the compute loop for no_tail --- ggml/src/ggml-cuda/gated_delta_net.cu | 25 ++++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 415c23161b7c..706f1097fe16 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -52,6 +52,8 @@ __global__ void gated_delta_net_cuda_no_tail(const float * q, // non-KDA has one g per column, KDA has one g per row float g_vals[n_tokens_per_loop][rows_per_lane]; int idx[rows_per_lane]; + float attn_col_values[n_tokens_per_loop]; + // state is stored transposed: M[col][i] = S[i][col], row col is contiguous #pragma unroll for (int r = 0; r < rows_per_lane; r++) { @@ -112,9 +114,7 @@ __global__ void gated_delta_net_cuda_no_tail(const float * q, float attn_col = warp_reduce_sum(attn_partial); - if (lane == 0) { - attn_data[col] = attn_col * scale; - } + attn_col_values[i] = attn_col * scale; } else { // kv[col] = sum_i g[i] * S[i][col] * k[i] float kv_shard = 0.0f; @@ -139,13 +139,24 @@ __global__ void gated_delta_net_cuda_no_tail(const float * q, float attn_col = warp_reduce_sum(attn_partial); - if (lane == 0) { - attn_data[col] = attn_col * scale; - } + attn_col_values[i] = attn_col * scale; } + } - attn_data += S_v * H; + // Pulling the writes out of the loop is compiler-friendly + if constexpr (warp_size < n_tokens_per_loop) { + for (int i = lane; i < n_tokens_per_loop; i += warp_size) { + if (i < n_tokens_per_loop) { + attn_data[i * S_v * H + col] = attn_col_values[i]; + } + } + } else { + if (lane < n_tokens_per_loop) { + attn_data[lane * S_v * H + col] = attn_col_values[lane]; + } } + + attn_data += n_tokens_per_loop * S_v * H; } // Write state back to global memory (transposed layout) #pragma unroll