CUDA: Support of GDN chunked kernel for prefill - #26001
Conversation
github-actions
Bot
added
testing
|
Hi @BLSharda, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
On RTX 6000 Pro MaxQ, PR 24561 needed MIN_BLOCKS_PER_SM=1 to get the new kernel to launch. And with larger ubatch sizes (≥512) PR 24561 gives 1.9–3.9% prefill performance, against this current PR delivers about a ~10% speedup for ubatch ≥512.
Chunked GDN requires at least 128 tokens, so MTP steps 1/2/3 are very short and stay on the recurrent path. Also, speculative decoding fails the K==1 check and therefore falls back to the recurrent kernel. |
How much effort would it be to make this PR use
How much effort would it be to get the same support for speculative decoding than we have for MTP? Seems like we would need to fail only based on the tokens, not the K==1 check |
Sure, I can give it a try to switch to using
This kernel is specifically designed for use in the prefill phase only and does not need to support MTP. Also, supporting K>1 is substantial work. |
|
I found this "deep dive" post from moonshot talking about flash kda https://github.com/MoonshotAI/FlashKDA/blob/master/docs/20260420-flashkda-v1-deep-dive.md They seem to be using 16 token chunks. Also they seem to link to a proof that f16 does not need to be clamped, but not sure how exact this applies here. |
I’ve now added ggml_cuda_mma support in this kernel, but it regresses performance by about 1–2%, so I’ve retained WMMA for NVIDIA GPUs and using |
|
hip has rocwmma with the equivalent interface to cuda wmma, but we should avoid using it if possible and not have two implementations. I can check the kernel on some hip devices. |
Yeah, if you can also try the ggml_cuda_mma path with HIP/MUSA that would be really helpful. Right now I’ve explicitly disabled AMD via runtime checks in ggml_cuda_gdn_op_is_chunked. |
I see similar KLD score when comparing CUDA recurrent path against Vulkan backend running same model checkpoint, so it seems this is expected due to very small difference in the values. |
There was a problem hiding this comment.
We need some more correctness tests apart from wikitext. Perhaps take a look at llama-eval and see if running a benchmark there matches what master produces. Also my guess is that this is in-fact not ready for review? Since @gaugarg-nv @ORippler have not approved it, so maybe you can move it back to draft
| return is_nvidia | ||
| && cc_dev >= GGML_CUDA_CC_AMPERE | ||
| && !chunk_disabled | ||
| && !kda && K == 1 | ||
| && neq0 == 128 && S_v == 128 && nev1 % neq1 == 0 | ||
| && src_k->ne[1] == neq1 | ||
| && n_tokens >= 128 | ||
| && ggml_is_contiguous(src_q) && ggml_is_contiguous(src_k) && ggml_is_contiguous(src_g) | ||
| && src_v->nb[0] == ggml_type_size(src_v->type) && src_v->nb[1] == (size_t)S_v * ggml_type_size(src_v->type) | ||
| && src_v->nb[3] == (size_t) n_tokens * src_v->nb[2] | ||
| && ggml_is_contiguous(src_beta) && ggml_is_contiguous(src_state); |
There was a problem hiding this comment.
@ggml-org/amd are you guys willing to test this on some hardware and post results? You will need modify this condition and we enable it for more hardware
| const int v_loc = idx % BV; | ||
| oi_regs[j] = s_result[t_idx * BV + v_loc] * __expf(fminf(s_gcum[t_idx], 88.72f)); |
|
Validated the Hardware: 2× MI210 (gfx90a, wave64), ROCm 7.14, Ubuntu 24.04. Branch at Starting pointWidening Every failure matched the chunked predicate exactly ( The four defects1. 2. 3. Plain 4. The accumulator's data layout is wrong on CDNA — this was the real one. All three GEMM helpers declare static __device__ __forceinline__ int get_i(const int l) {
return tile<I_, J_, T, DATA_LAYOUT_I_MAJOR>::get_j(l);
}So every lane reports transposed coordinates when writing its accumulator elements back to shared memory, and the entire tile is scattered. Symptom is #if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
# define CGDR_C_DL ggml_cuda_mma::DATA_LAYOUT_J_MAJOR
#else
# define CGDR_C_DL ggml_cuda_mma::DATA_LAYOUT_I_MAJOR
#endifFor what it's worth I also checked and ruled out a suspicion of mine: Results after the fixesCorrectness — Performance —
Token-level — For context on the ceiling: I measured the recurrent GDN op standalone at ~1.31 TFLOPS on this hardware while the same box does 75–80 TFLOPS on this model's own GEMMs, and attributed the op at ~19% of prefill wall clock. That caps any chunked implementation at ~1.24× end-to-end here, so +12% is roughly half the theoretical headroom. One honest caveatThe single remaining failure is the longest shape: That is a precision margin, not a correctness failure — NMSE ~3e-7 is ~5.5e-4 RMS relative error, right at fp16 epsilon, and 2048 tokens is 128 chunks at Not covered
Patch53 insertions / 18 deletions across |
The patch (53 insertions / 18 deletions, NVIDIA paths unchanged)diff --git a/ggml/src/ggml-cuda/chunk_gated_delta_net.cu b/ggml/src/ggml-cuda/chunk_gated_delta_net.cu
index 0dd6f5947..23e7ec5fe 100644
--- a/ggml/src/ggml-cuda/chunk_gated_delta_net.cu
+++ b/ggml/src/ggml-cuda/chunk_gated_delta_net.cu
@@ -15,6 +15,27 @@
#define GDN_TC_MMA 0
#endif
+// GDN_ACC_DL: the 16x16 f32 accumulator fragment layout is J_MAJOR on CDNA/RDNA4 and
+// I_MAJOR on NVIDIA. get_i/get_j report where each lane's accumulator element lands, so
+// declaring the wrong layout transposes the mapping and scatters the whole tile to the
+// wrong shared-memory slots. Matches PR #24561, validated 48/48 on MI250X/gfx90a.
+#if defined(AMD_MFMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA4))
+# define CGDR_C_DL ggml_cuda_mma::DATA_LAYOUT_J_MAJOR
+#else
+# define CGDR_C_DL ggml_cuda_mma::DATA_LAYOUT_I_MAJOR
+#endif
+
+// GDN_MFMA_LOAD: plain load_ldmatrix in mma.cuh has TURING and AMD_WMMA branches but
+// NO AMD_MFMA branch -- on CDNA it falls through to NO_DEVICE_CODE and traps at runtime
+// ("unspecified launch failure"). load_generic is the portable element-wise loader and is
+// defined in terms of the same get_i/get_j, so it is correct for MFMA tile layouts.
+// load_ldmatrix_trans is left alone: it does have an AMD_MFMA branch.
+#if defined(AMD_MFMA_AVAILABLE)
+# define CGDR_LOAD ggml_cuda_mma::load_generic
+#else
+# define CGDR_LOAD ggml_cuda_mma::load_ldmatrix
+#endif
+
// Check if tensor-core kernels are supported on this architecture; otherwise, fallback or no-op.
#if GDN_TC_MMA
# if defined(TURING_MMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
@@ -40,12 +61,12 @@ template <int BK>
__device__ __forceinline__ void cgdr_gemm_ABt_16(const __half * s_a, const __half * s_b, float * s_c, int ldc, int c_col)
{
#if GDN_TC_MMA
- ggml_cuda_mma::tile<16, 16, float> acc;
+ ggml_cuda_mma::tile<16, 16, float, CGDR_C_DL> acc;
#pragma unroll
for (int kt = 0; kt < BK / 16; kt++) {
ggml_cuda_mma::tile<16, 8, half2> ta, tb;
- ggml_cuda_mma::load_ldmatrix(ta, (const half2 *) s_a + kt * 8, BK / 2);
- ggml_cuda_mma::load_ldmatrix(tb, (const half2 *) s_b + kt * 8, BK / 2);
+ CGDR_LOAD(ta, (const half2 *) s_a + kt * 8, BK / 2);
+ CGDR_LOAD(tb, (const half2 *) s_b + kt * 8, BK / 2);
ggml_cuda_mma::mma(acc, ta, tb);
}
#pragma unroll
@@ -79,7 +100,7 @@ __device__ __forceinline__ void cgdr_gemm_ktv(const __half * s_vnew, const __hal
for (int nk = 0; nk < BK; nk += 16) {
ggml_cuda_mma::tile<16, 8, half2> y_kch;
ggml_cuda_mma::load_ldmatrix_trans(y_kch, (const half2 *) (s_kch + nk), BK / 2);
- ggml_cuda_mma::tile<16, 16, float> acc;
+ ggml_cuda_mma::tile<16, 16, float, CGDR_C_DL> acc;
ggml_cuda_mma::mma(acc, x_vnew, y_kch);
#pragma unroll
for (int l = 0; l < acc.ne; l++) {
@@ -107,9 +128,9 @@ __device__ __forceinline__ void cgdr_gemm_qkv(const __half * s_qk, const __half
{
#if GDN_TC_MMA
ggml_cuda_mma::tile<16, 8, half2> x_qk, y_vnew;
- ggml_cuda_mma::load_ldmatrix(x_qk, (const half2 *) s_qk, 16 / 2);
+ CGDR_LOAD(x_qk, (const half2 *) s_qk, 16 / 2);
ggml_cuda_mma::load_ldmatrix_trans(y_vnew, (const half2 *) s_vnew + n_off / 2, BV / 2);
- ggml_cuda_mma::tile<16, 16, float> acc;
+ ggml_cuda_mma::tile<16, 16, float, CGDR_C_DL> acc;
ggml_cuda_mma::mma(acc, x_qk, y_vnew);
#pragma unroll
for (int l = 0; l < acc.ne; l++) {
@@ -283,9 +304,13 @@ __launch_bounds__(128, 4) __global__ void cgdr_fwdsub_intra_kernel(
// Masked Q@K^T on tensor cores (fp16 WMMA, one warp per block):
// qk_buf[i,j] = (Q_ch . K_ch[j]) * exp(g_cum[i] - g_cum[j]) for j <= i, else 0.
-// Grid (B*H, num_chunks); 32 threads. Requires CS==16, BK%16==0.
+// Grid (B*H, num_chunks); one wavefront per block (32 on NVIDIA, 64 on CDNA).
+// Requires CS==16, BK%16==0.
template <int CS, int BK>
-__launch_bounds__(32, 8) __global__ void cgdr_precompute_qk_wmma_kernel(const float * __restrict__ Q_raw,
+// GDN_WAVE64: launch bounds must match the launch, which is one full wavefront
+// (64 on CDNA). Declaring 32 here made the 64-thread launch fail outright.
+// Same pattern as mmid.cu.
+__launch_bounds__(ggml_cuda_get_physical_warp_size(), 8) __global__ void cgdr_precompute_qk_wmma_kernel(const float * __restrict__ Q_raw,
const float * __restrict__ K_raw,
const float * __restrict__ g_cum,
float * __restrict__ qk_buf,
@@ -307,7 +332,7 @@ __launch_bounds__(32, 8) __global__ void cgdr_precompute_qk_wmma_kernel(const fl
const int bh = blockIdx.x;
const int c = blockIdx.y;
- const int tid = threadIdx.x; // 0..31
+ const int tid = threadIdx.x; // 0..warp_size-1 (GDN_WAVE64)
const int b_idx = bh / H;
const int h_idx = bh % H; // v-head
@@ -326,7 +351,7 @@ __launch_bounds__(32, 8) __global__ void cgdr_precompute_qk_wmma_kernel(const fl
// seq_len is not a multiple of CS -- zero-fill rows past valid_cs to avoid out-of-bounds reads.
// Q*scale and K are small (unit-length vectors), so fp16 is safe.
const int valid_cs = min(CS, seq_len - t_off);
- for (int i = tid; i < CS * BK; i += 32) {
+ for (int i = tid; i < CS * BK; i += blockDim.x) { // GDN_WAVE64: was hardcoded 32
const int row = i / BK, col = i % BK;
const float qv = (row < valid_cs) ? Q_chunk[(long long) row * HK + col] : 0.f;
const float kv = (row < valid_cs) ? K_chunk[(long long) row * HK + col] : 0.f;
@@ -340,11 +365,11 @@ __launch_bounds__(32, 8) __global__ void cgdr_precompute_qk_wmma_kernel(const fl
__syncthreads();
// Causal mask + cumulative-decay scaling, then write qk_buf.
- float * O_base = qk_buf + (long long) (bh * num_chunks + c) * CS * CS;
- constexpr int EPT = CS * CS / 32;
- #pragma unroll
- for (int e = 0; e < EPT; e++) {
- const int flat = tid + e * 32;
+ float * O_base = qk_buf + (long long) (bh * num_chunks + c) * CS * CS;
+ // GDN_WAVE64: stride by the actual block size. A 64-lane wavefront owns CS*CS
+ // in half as many steps as a 32-lane warp, so the old constexpr trip count
+ // double-wrote on AMD.
+ for (int flat = tid; flat < CS * CS; flat += blockDim.x) {
const int row = flat / CS;
const int col = flat % CS;
O_base[flat] = (col <= row) ? s_acc[flat] * __expf(s_gcum[row] - s_gcum[col]) : 0.f;
@@ -704,7 +729,15 @@ static void ggml_cuda_op_gated_delta_net_chunked_impl(ggml_backend_cuda_context
{
const size_t qk_smem = cgdr_smem_preqk_wmma(CS, K_dim); // 9.1 KB < 48 KB -> no opt-in needed
const dim3 qk_grid(B * H, num_chunks, 1);
- cgdr_precompute_qk_wmma_kernel<CS, 128><<<qk_grid, 32, qk_smem, stream>>>(
+ // GDN_WAVE64: one full wavefront. ggml_cuda_mma tiles span all 64 lanes on
+ // CDNA (mma.cuh: ne = I*J/64), so a 32-thread block half-populates the
+ // accumulator. warp_size is 32 on NVIDIA, so that path is unchanged.
+#if GDN_TC_MMA
+ const int qk_block = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size;
+#else
+ const int qk_block = 32;
+#endif
+ cgdr_precompute_qk_wmma_kernel<CS, 128><<<qk_grid, qk_block, qk_smem, stream>>>(
q_in, k_in, g_cum_buf.get(), qk_buf.get(), num_chunks, scale, H, num_k_heads, T);
}
CUDA_CHECK(cudaGetLastError());
diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu
index b3977e7ff..6fe4d66b2 100644
--- a/ggml/src/ggml-cuda/gated_delta_net.cu
+++ b/ggml/src/ggml-cuda/gated_delta_net.cu
@@ -248,15 +248,18 @@ bool ggml_cuda_gdn_op_is_chunked(const ggml_tensor * dst) {
return s && s[0] && !(s[0] == '0' && s[1] == '\0');
}();
const int cc_dev = ggml_cuda_info().devices[ggml_cuda_get_device()].cc;
- // NVIDIA-only for now. The HIP/MUSA ggml_cuda_mma backend intentionally not dispatched until validated.
- const bool is_nvidia = GGML_CUDA_CC_IS_NVIDIA(cc_dev);
+ // GDN_CDNA: admit CDNA to the ggml_cuda_mma path. CDNA has the fp16 matrix
+ // cores this kernel wants (v_mfma_f32_16x16x16f16). NVIDIA's condition is
+ // unchanged. Deliberately not all of AMD -- RDNA's WMMA path is unvalidated.
+ const bool arch_ok = (GGML_CUDA_CC_IS_NVIDIA(cc_dev) && cc_dev >= GGML_CUDA_CC_AMPERE)
+ || GGML_CUDA_CC_IS_CDNA(cc_dev);
// - NVIDIA Ampere+ (fp16 WMMA); not KDA; K == 1 (final state only)
// - Q/K/G/beta/state must be contiguous
// (nb[0]/nb[1] packed) with arbitrary token stride (fused QKV view)
// - V is packed per token (nb[2]) and across sequences (nb[3] == n_tokens*nb[2]).
// - 128-wide heads, GQA-aligned head counts, n_tokens >= 128
- return is_nvidia
+ return arch_ok
&& cc_dev >= GGML_CUDA_CC_AMPERE
&& !chunk_disabled
&& !kda && K == 1
|
Thanks, @davetha — this is really helpful. I’ll take the warp_size-related changes that I can test on NVIDIA GPUs for cuda_mma. Once this PR is merged, could you open a follow-up PR with the remaining changes? Since I’m not testing on AMD, I’d prefer not to enable it in this PR.
|
|
Sure! |
|
I used commit 1e1885f from this PR in my final Qwen3.8-27B runtime. My model is a 17.1 GB, 5.01 BPW iMatrix/NVFP4 hybrid with an embedded MTP layer. The target model, MTP, recurrent state, CUDA graphs and 256K Q4_0 KV cache run on one RTX PRO 4000 Blackwell SFF: 24,467 MiB VRAM, sm120a, 432 GB/s rated bandwidth and a 70 W board limit. An RTX 2000 Ada holds the optional F16 vision projector. I tested with CUDA 12.9.86, driver 610.57.04, Debian 13, batch 512, ubatch 256 and four recurrent checkpoints. I did not isolate #26001 from the other two CUDA patches in the final runtime sweep, so I cannot honestly assign it an individual percentage. The #26001 + #26048 + #26705 bundle moved deterministic MTP decode from 45.422 to 45.866 tok/s, or +0.98%. The more relevant result for this prefill patch was the real long-context test: 261,500 input tokens were processed at 226.750 tok/s, followed by 256 generated tokens, without truncation or OOM. Full process and exact arguments: Released GGUF: |
The GDN prefill kernel processed tokens serially (see the TODO in gated_delta_net.cu). Adds a chunked kernel following the scheme of the Metal Mamba-2 chunked prefill; related CUDA/Vulkan efforts: ggml-org#26001, ggml-org#20377. Off by default; passes test-backend-ops GATED_DELTA_NET.
The GDN prefill kernel processed tokens serially (see the TODO in gated_delta_net.cu). Adds a chunked kernel following the scheme of the Metal Mamba-2 chunked prefill; related CUDA/Vulkan efforts: ggml-org#26001, ggml-org#20377. Passes test-backend-ops GATED_DELTA_NET. Opt-in on RDNA3/RDNA4 via GGML_HIP_GDN_CHUNK=1: arch_ok ties that variable to the RDNA branch only, so the kernel is default-active on NVIDIA Ampere+ and CDNA. It is not active in the published gfx1151 numbers, which were measured without the variable set.
The GDN prefill kernel processed tokens serially (see the TODO in gated_delta_net.cu). Adds a chunked kernel following the scheme of the Metal Mamba-2 chunked prefill; related CUDA/Vulkan efforts: ggml-org#26001, ggml-org#20377. Passes test-backend-ops GATED_DELTA_NET. Opt-in on RDNA3/RDNA4 via GGML_HIP_GDN_CHUNK=1: arch_ok ties that variable to the RDNA branch only, so the kernel is default-active on NVIDIA Ampere+ and CDNA. It is not active in the published gfx1151 numbers, which were measured without the variable set.
1e1885f to
10c1e46
Compare
|
@BLSharda have you benchmarked this kernel w.r.t to FLA? at ISL 8192 there is a huge diff between vLLM and llama.cpp and it is entirely because of this kernel. It would be good to close the gap (it is like 3x worse rather than a few percent) |
|
@am17an I ran that comparison. On an RTX 3090, the chunked kernel in this PR is 2.3–2.7x To be upfront about provenance: the kernel was generated by an LLM-driven Code based on the current head of this pr:
Measurements: RTX 3090, CUDA 12.0. µs/op via
End to end, Qwen3.8-27B Q4_K:
Correctness: Limitations: tuned on GA102. The grid is n_v_heads x 4 blocks per sequence, which @BLSharda happy to open this as a PR against |
I think this #26001 (comment) comment addresses your concern (increases kernel perf 2.6+ times). |
Thanks for this work, let me try to merge my change and you can raise a separate PR. |
That's an entirely different kernel from what I see. So should we try merging some version of that? |
Fixed review comment Co-authored-by: Aman Gupta <amangupta052@gmail.com>
On correctness beyond wikitext, ran the full MMLU-Pro through lm-eval against llama-server, comparing the chunked path to a recurrent baseline — i.e. similar what master produces. Tested Qwen3.6-35B-A3B UD-Q4_K_M, all 12,032 questions per arm, 5-shot, greedy, n_ctx 8192.
@gaugarg-nv and @ORippler, the review comments are addressed, can you help review this? |
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
Assisted-by: OpenAI Codex
9 participants
Overview
This PR adds a chunked mode to the Gated Delta Net (GDN) CUDA operator, significantly speeding up prefill for >=128 tokens to the previous recurrent (token-by-token) kernel. All changes are implemented within the CUDA backend as a single GDN operator update.
The new chunked execution path utilizes a three-stage pipeline, inspired by the (vLLM) Triton/FLA GDN kernel, and leverages a mix of BF16, FP16, and FP32 to achieve throughput improvements via 16-bit tensor cores. This mode is enabled on supported hardware (NVIDIA Ampere+ with BF16 tensor core), otherwise original recurrent kernel is used as a fallback. The chunked kernel is GQA-aware and reads the fused QK and strided V tensor, similarly to the recurrent kernel.
Implementation in
chunk_gated_delta_net.cuFP32 forward substitution (WY inverse)
Computes intra-chunk corrections:
V_corr,K_cumdecay, andg_cumBF16 WMMA masked attention
Performs
Q @ K^Twith causal masking and cumulative decayWMMA state update + fused output
k^T · vandqk · v_newAccuracy
qwen_3_6_35b_a3b_q4_k_mmodel:Addtional Memory
The buffers are pool-allocated and freed per op, so they're reused across all GDN layers and stays independent of model quant (always FP32).
Performance
Below are the measured performance improvements for Qwen 3.6 35B MoE and 27B models; performance gains apply broadly to all supported NVIDIA Ampere+ GPUs
-ub perf-sweep
pp=16384 with qwen_3_6_35b_a3b_q4_k_m
RTX 6000 Pro MaxQ
DGX Spark
Requirements
Known limitations / future work