Commit 62df1815b for llama.cpp
commit 62df1815bc053066761f4eac58e288f23778e164
Author: Aman Gupta <amangupta052@gmail.com>
Date: Wed Oct 14 19:30:58 2026 +0530
cuda, hip: GDN chunked kernel (#29353)
* cuda, hip: GDN chunked kernel
Ported-from: https://github.com/tetherto/qvac-fabric-llm.cpp/pull/257
Source-commit: 53f0027d8070715fe9ad2131f3201bb0001d4eed
Co-authored-by: Aman Gupta <aman.gupta@tether.io>
* fix older builds
* review fixes
* cuda: enable GDN MMA on tested Ampere and Ada GPUs
Enable SM86 and SM89 after RTX 3090 and RTX 4090 correctness and prefill benchmarks. Keep other untested architectures on the existing fallback.
Assisted-by: Codex
* cuda : use refactored swizzle API in GDN MMA
* ggml : GDN only guarantees the final state when n_tokens > K
* cuda : add short GDN MMA phase comments
---------
Co-authored-by: gianni-cor <gianfranco.cordella@tether.io>
Co-authored-by: Aman Gupta <aman.gupta@tether.io>
diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h
index 23ea6437e..c37ba9471 100644
--- a/ggml/include/ggml.h
+++ b/ggml/include/ggml.h
@@ -2654,6 +2654,7 @@ extern "C" {
// the output packs the attention scores [S_v, H_v, n_tokens, n_seqs] followed by K state
// snapshots, most-recent first (slot 0 = final state, slot s = state s tokens back). K == 1
// keeps only the final state; when n_tokens < K only slots 0..n_tokens-1 are written.
+ // when n_tokens > K only slot 0 is guaranteed; backends may leave slots 1..K-1 unwritten.
GGML_API struct ggml_tensor * ggml_gated_delta_net(
struct ggml_context * ctx,
struct ggml_tensor * q,
diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh
index c6dc9dd62..fea1cf9c4 100644
--- a/ggml/src/ggml-cuda/common.cuh
+++ b/ggml/src/ggml-cuda/common.cuh
@@ -52,6 +52,7 @@
#define GGML_CUDA_CC_VOLTA 700
#define GGML_CUDA_CC_TURING 750
#define GGML_CUDA_CC_AMPERE 800
+#define GGML_CUDA_CC_AMPERE_GA10X 860
#define GGML_CUDA_CC_ORIN 870
#define GGML_CUDA_CC_ADA_LOVELACE 890
#define GGML_CUDA_CC_HOPPER 900
diff --git a/ggml/src/ggml-cuda/gated-delta-net-mma.cu b/ggml/src/ggml-cuda/gated-delta-net-mma.cu
new file mode 100644
index 000000000..8a9e926de
--- /dev/null
+++ b/ggml/src/ggml-cuda/gated-delta-net-mma.cu
@@ -0,0 +1,442 @@
+#include "gated-delta-net-mma.cuh"
+#include "mma.cuh"
+
+#include <climits>
+#include <mutex>
+
+namespace {
+constexpr int D = 128;
+using bf16 = nv_bfloat16;
+
+template <int R, int C> struct matrix {
+ bf16 hi[R * C], lo[R * C];
+
+ static __device__ __forceinline__ int index(int r, int c) {
+#ifdef GGML_USE_HIP
+ return r * C + c;
+#else
+ return 2 * ggml_cuda_mma::swizzle<C / 2, nv_bfloat162>(r * (C / 2) + c / 2, r) + c % 2;
+#endif // GGML_USE_HIP
+ }
+
+ __device__ __forceinline__ void store2(int r, int c, float2 x) {
+#ifdef GGML_USE_MUSA
+ store(r, c, x.x);
+ store(r, c + 1, x.y);
+#else
+ const int i = index(r, c);
+ const auto h = __float22bfloat162_rn(x);
+ const float2 rounded = __bfloat1622float2(h);
+ *reinterpret_cast<nv_bfloat162 *>(hi + i) = h;
+ *reinterpret_cast<nv_bfloat162 *>(lo + i) =
+ __float22bfloat162_rn(make_float2(x.x - rounded.x, x.y - rounded.y));
+#endif // GGML_USE_MUSA
+ }
+
+ __device__ __forceinline__ void store(int r, int c, float x) {
+ const int i = index(r, c);
+ const bf16 h = __float2bfloat16(x);
+ hi[i] = h;
+ lo[i] = __float2bfloat16(x - __bfloat162float(h));
+ }
+};
+
+template <int C, int V> struct alignas(128) shared {
+ matrix<C, D> q, k;
+ matrix<C, C> p, inverse;
+ matrix<C, V> delta;
+
+ union {
+ matrix<D, V> state;
+ matrix<C, V> solved;
+ } scratch;
+
+ float lower[C * C], prefix[C], beta[C];
+
+ static __device__ __forceinline__ int lower_index(int r, int c) { return r * C + c; }
+};
+
+#if defined(AMPERE_MMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA3))
+using namespace ggml_cuda_mma;
+#ifdef GGML_USE_HIP
+using acc = tile<16, 16, float, DATA_LAYOUT_J_MAJOR>;
+using a_tile = tile<16, 8, nv_bfloat162, DATA_LAYOUT_I_MAJOR_MIRRORED>;
+using b_tile = a_tile;
+#else
+using acc = tile<16, 8, float>;
+using a_tile = tile<16, 8, nv_bfloat162>;
+using b_tile = tile<8, 8, nv_bfloat162>;
+#endif // GGML_USE_HIP
+
+template <class tile_t, bool TRANS, int R, int C>
+__device__ __forceinline__ tile_t load(const matrix<R, C> & m, int row, int col, bool low) {
+ tile_t t;
+ const bf16 * data = low ? m.lo : m.hi;
+#ifdef GGML_USE_HIP
+#pragma unroll
+ for (int i = 0; i < tile_t::ne; ++i) {
+ const int r = row + tile_t::get_i(i), c = col + 2 * tile_t::get_j(i);
+ const bf16 x = data[TRANS ? m.index(c, r) : m.index(r, c)];
+ const bf16 y = data[TRANS ? m.index(c + 1, r) : m.index(r, c + 1)];
+ t.x[i] = __float22bfloat162_rn(make_float2(__bfloat162float(x), __bfloat162float(y)));
+ }
+#else
+ const auto * packed = reinterpret_cast<const nv_bfloat162 *>(data);
+ if constexpr (TRANS) {
+ load_ldmatrix_trans_swizzled<C / 2>(t, packed, col * (C / 2) + row / 2);
+ } else {
+ load_ldmatrix_swizzled<C / 2>(t, packed, row * (C / 2) + col / 2);
+ }
+#endif // GGML_USE_HIP
+ return t;
+}
+
+// B is addressed as B-transpose. Preserve both BF16 first-order residual terms.
+template <int K, bool TA = false, bool TB = false, int AR, int AC, int BR, int BC>
+__device__ __forceinline__ void product(acc & c, const matrix<AR, AC> & a, const matrix<BR, BC> & b, int row, int col) {
+#pragma unroll
+ for (int k = 0; k < K; k += 16) {
+ const auto ah = load<a_tile, TA>(a, row, k, false);
+ const auto al = load<a_tile, TA>(a, row, k, true);
+ const auto bh = load<b_tile, TB>(b, col, k, false);
+ const auto bl = load<b_tile, TB>(b, col, k, true);
+ mma(c, al, bh);
+ mma(c, ah, bl);
+ mma(c, ah, bh);
+ }
+}
+
+template <int R, int C> __device__ __forceinline__ void store(matrix<R, C> & m, const acc & x, int r, int c) {
+#pragma unroll
+#ifdef GGML_USE_HIP
+ for (int i = 0; i < acc::ne; ++i) {
+ m.store(r + acc::get_i(i), c + acc::get_j(i), x.x[i]);
+ }
+#else
+ for (int i = 0; i < acc::ne; i += 2) {
+ m.store2(r + acc::get_i(i), c + acc::get_j(i), make_float2(x.x[i], x.x[i + 1]));
+ }
+#endif // GGML_USE_HIP
+}
+#endif // defined(AMPERE_MMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA3))
+
+// Each block owns V independent value columns and advances by C tokens.
+// Keep recurrent state in FP32; pack high/residual BF16 operands inside the block.
+template <int C, int V, int BLOCKS = 1>
+static __global__ __launch_bounds__(256, BLOCKS) void gdn_single(ggml_cuda_gdn_mma_args a) {
+#if defined(AMPERE_MMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA3))
+ static_assert(ggml_cuda_get_physical_warp_size() == 32, "GDN MMA requires 32-lane warps");
+ constexpr int WARPS = 8;
+#ifdef GGML_USE_HIP
+ constexpr int N = 16;
+#else
+ constexpr int N = 8;
+#endif // GGML_USE_HIP
+ static_assert(C >= 16 && (C & (C - 1)) == 0 && V % N == 0 && D % V == 0, "invalid GDN tile");
+ extern __shared__ __align__(128) unsigned char bytes[];
+ auto & s = *reinterpret_cast<shared<C, V> *>(bytes);
+ const int lane = threadIdx.x, warp = threadIdx.y, tid = warp * 32 + lane;
+ const int slice = blockIdx.x % (D / V);
+ const int h = (blockIdx.x / (D / V)) % a.H;
+ const int seq = blockIdx.x / ((D / V) * a.H);
+ const int v0 = slice * V;
+ const int64_t soff = ((int64_t) seq * a.H + h) * D * D;
+ const int64_t qoff = (seq / a.rq3) * a.sq3 + (h % a.H_k) * a.sq1;
+ const int64_t voff = seq * a.sv3 + h * a.sv1 + v0;
+ const int64_t goff = seq * a.sb3 + h * a.sb1;
+ float * out = a.dst + ((int64_t) seq * a.n_tokens * a.H + h) * D + v0;
+ constexpr int STATE_TILES = (D / 16) * (V / N);
+ constexpr int PER_WARP = STATE_TILES / WARPS;
+ static_assert(STATE_TILES % WARPS == 0, "state must divide across warps");
+ acc state[PER_WARP];
+#pragma unroll
+ for (int j = 0; j < PER_WARP; ++j) {
+ const int tile_index = warp + j * WARPS, r = (tile_index / (V / N)) * 16, c = (tile_index % (V / N)) * N;
+#pragma unroll
+ for (int i = 0; i < acc::ne; ++i) {
+ state[j].x[i] = a.state[soff + (v0 + c + acc::get_j(i)) * D + r + acc::get_i(i)];
+ }
+ }
+ auto read_qk = [&](int64_t t, float4 & q, float4 & k) {
+ t += warp;
+ q = k = make_float4(0, 0, 0, 0);
+ if (t < a.n_tokens) {
+ const int64_t off = qoff + t * a.sq2 + 4 * lane;
+ if ((((uintptr_t) (a.q + off) | (uintptr_t) (a.k + off)) & 15) == 0) {
+ q = *reinterpret_cast<const float4 *>(a.q + off);
+ k = *reinterpret_cast<const float4 *>(a.k + off);
+ } else {
+ q = make_float4(a.q[off], a.q[off + 1], a.q[off + 2], a.q[off + 3]);
+ k = make_float4(a.k[off], a.k[off + 1], a.k[off + 2], a.k[off + 3]);
+ }
+ }
+ };
+ // Carry the first eight input rows across chunks to hide global-load latency.
+ constexpr bool PREFETCH = C == 16 && V == 32;
+ float4 next_q, next_k;
+ if constexpr (PREFETCH) {
+ read_qk(0, next_q, next_k);
+ }
+ for (int64_t t0 = 0; t0 < a.n_tokens; t0 += C) {
+ const int valid = min((int64_t) C, a.n_tokens - t0);
+#pragma unroll
+ for (int j = 0; j < C / 8; ++j) {
+ const int t = warp + 8 * j, d = 4 * lane;
+ float4 q, k;
+ if (PREFETCH && j == 0) {
+ q = next_q;
+ k = next_k;
+ } else {
+ read_qk(t0 + 8 * j, q, k);
+ }
+ s.q.store2(t, d, make_float2(q.x * a.scale, q.y * a.scale));
+ s.q.store2(t, d + 2, make_float2(q.z * a.scale, q.w * a.scale));
+ s.k.store2(t, d, make_float2(k.x, k.y));
+ s.k.store2(t, d + 2, make_float2(k.z, k.w));
+ }
+ // Load gates in parallel before the FP32 prefix scan.
+ if (warp == 0) {
+ float carry = 0.f;
+ for (int base = 0; base < C; base += 32) {
+ const int t = base + lane;
+ float g = t < valid ? a.g[goff + (t0 + t) * a.sb2] : 0.f;
+ const float beta = t < valid ? a.beta[goff + (t0 + t) * a.sb2] : 0.f;
+#pragma unroll
+ for (int offset = 1; offset < 32; offset *= 2) {
+ const float previous = __shfl_up_sync(0xffffffff, g, offset, 32);
+ if (lane >= offset) {
+ g += previous;
+ }
+ }
+ g += carry;
+ if (t < C) {
+ s.prefix[t] = g;
+ s.beta[t] = beta;
+ }
+ carry = __shfl_sync(0xffffffff, g, 31, 32);
+ }
+ }
+#pragma unroll
+ for (int j = 0; j < PER_WARP; ++j) {
+ const int ti = warp + j * WARPS;
+ store(s.scratch.state, state[j], (ti / (V / N)) * 16, (ti % (V / N)) * N);
+ }
+ __syncthreads();
+ if constexpr (PREFETCH) {
+ if (t0 + C < a.n_tokens) {
+ read_qk(t0 + C, next_q, next_k);
+ }
+ }
+ // V32 overlaps Gram work with one output tile per projection warp.
+ constexpr bool SPLIT_PROJECTION = C == 16 && V == 32;
+ acc base_output[((C / 16) * (V / N) + WARPS - 1) / WARPS];
+ // delta = beta * (v - decay * K S), base output = decay * Q S
+ auto project = [&]() {
+ if constexpr (V == 32) {
+ if (warp < 4) {
+ return;
+ }
+ }
+ for (int ti = (V == 32 ? warp - 4 : warp); ti < (C / 16) * (V / N); ti += (V == 32 ? 4 : WARPS)) {
+ const int r = (ti / (V / N)) * 16, c = (ti % (V / N)) * N;
+ acc qs, ks;
+ product<D, false, true>(qs, s.q, s.scratch.state, r, c);
+ product<D, false, true>(ks, s.k, s.scratch.state, r, c);
+#pragma unroll
+ for (int i = 0; i < acc::ne; ++i) {
+ const int t = r + acc::get_i(i), v = c + acc::get_j(i);
+ const float decay = expf(s.prefix[t]);
+ const float vv = t < valid ? a.v[voff + (t0 + t) * a.sv2 + v] : 0.f;
+ s.delta.store(t, v, s.beta[t] * (vv - decay * ks.x[i]));
+ base_output[ti / WARPS].x[i] = decay * qs.x[i];
+ }
+ }
+ };
+ // A = beta * strict_tril(K K^T * decay), P = tril(Q K^T * decay)
+ for (int ti = warp; ti < 2 * (C / 16) * (C / N) && (V != 32 || warp < 4); ti += (V == 32 ? 4 : WARPS)) {
+ const bool gram = ti < (C / 16) * (C / N);
+ const int t = ti % ((C / 16) * (C / N)), r = (t / (C / N)) * 16, c = (t % (C / N)) * N;
+ acc x;
+ if (gram) {
+ product<D>(x, s.k, s.k, r, c);
+ } else {
+ product<D>(x, s.q, s.k, r, c);
+ }
+#pragma unroll
+ for (int i = 0; i < acc::ne; ++i) {
+ const int rr = r + acc::get_i(i), cc = c + acc::get_j(i);
+ const float v = rr < valid && cc <= rr ? x.x[i] * expf(s.prefix[rr] - s.prefix[cc]) : 0.f;
+ if (gram) {
+ if (cc < rr) {
+ s.lower[s.lower_index(rr, cc)] = v * s.beta[rr];
+ }
+ } else {
+ s.p.store(rr, cc, v);
+ }
+ }
+ }
+ if constexpr (V == 32) {
+ project();
+ }
+ __syncthreads();
+ // inverse = (I + A)^-1 by forward substitution
+ for (int row = warp; row < C; row += WARPS) {
+ float x[(C + 31) / 32];
+#pragma unroll
+ for (int j = 0; j < (C + 31) / 32; ++j) {
+ const int col = lane + 32 * j;
+ x[j] = col == row ? 1.f : (col < row ? -s.lower[s.lower_index(row, col)] : 0.f);
+ }
+ for (int pivot = row - 1; pivot > 0; --pivot) {
+ const float xp = __shfl_sync(0xffffffff, x[pivot / 32], pivot % 32, 32);
+#pragma unroll
+ for (int j = 0; j < (C + 31) / 32; ++j) {
+ const int col = lane + 32 * j;
+ if (col < pivot) {
+ x[j] -= xp * s.lower[s.lower_index(pivot, col)];
+ }
+ }
+ }
+#pragma unroll
+ for (int j = 0; j < (C + 31) / 32; ++j) {
+ if (lane + 32 * j < C) {
+ s.inverse.store(row, lane + 32 * j, x[j]);
+ }
+ }
+ }
+ if constexpr (V != 32) {
+ project();
+ }
+ __syncthreads();
+ // solved = inverse * delta
+ for (int ti = warp; ti < (C / 16) * (V / N); ti += WARPS) {
+ const int r = (ti / (V / N)) * 16, c = (ti % (V / N)) * N;
+ acc x;
+ product<C, false, true>(x, s.inverse, s.delta, r, c);
+ store(s.scratch.solved, x, r, c);
+ }
+ __syncthreads();
+ // out = base output + P * solved
+ for (int ti = (SPLIT_PROJECTION ? warp - 4 : warp); ti < (C / 16) * (V / N) && (!SPLIT_PROJECTION || warp >= 4);
+ ti += WARPS) {
+ const int r = (ti / (V / N)) * 16, c = (ti % (V / N)) * N;
+ acc x;
+ product<C, false, true>(x, s.p, s.scratch.solved, r, c);
+#pragma unroll
+ for (int i = 0; i < acc::ne; ++i) {
+ const int t = r + acc::get_i(i), v = c + acc::get_j(i);
+ if (t < valid) {
+ out[(t0 + t) * a.H * D + v] = base_output[ti / WARPS].x[i] + x.x[i];
+ }
+ }
+ }
+ // S = decay(chunk) * S + K^T * (solved * decay to chunk end)
+ for (int i = tid; i < C * V; i += 256) {
+ const int t = i / V, v = i % V, j = s.scratch.solved.index(t, v);
+ const float x = __bfloat162float(s.scratch.solved.hi[j]) + __bfloat162float(s.scratch.solved.lo[j]);
+ s.delta.store(t, v, x * expf(s.prefix[valid - 1] - s.prefix[t]));
+ }
+ __syncthreads();
+ const float decay = expf(s.prefix[valid - 1]);
+#pragma unroll
+ for (int j = 0; j < PER_WARP; ++j) {
+ const int ti = warp + j * WARPS, r = (ti / (V / N)) * 16, c = (ti % (V / N)) * N;
+#pragma unroll
+ for (int i = 0; i < acc::ne; ++i) {
+ state[j].x[i] *= decay;
+ }
+ product<C, true, true>(state[j], s.k, s.delta, r, c);
+ }
+ __syncthreads();
+ }
+#pragma unroll
+ for (int j = 0; j < PER_WARP; ++j) {
+ const int ti = warp + j * WARPS, r = (ti / (V / N)) * 16, c = (ti % (V / N)) * N;
+#pragma unroll
+ for (int i = 0; i < acc::ne; ++i) {
+ a.state_out[soff + (v0 + c + acc::get_j(i)) * D + r + acc::get_i(i)] = state[j].x[i];
+ }
+ }
+#else
+ GGML_UNUSED_VARS(a);
+ NO_DEVICE_CODE;
+#endif // defined(AMPERE_MMA_AVAILABLE) || (defined(AMD_WMMA_AVAILABLE) && defined(RDNA3))
+}
+
+template <int C, int V, int BLOCKS = 1> bool init(int device) {
+ static std::once_flag once[GGML_CUDA_MAX_DEVICES];
+ static bool available[GGML_CUDA_MAX_DEVICES] = {};
+ std::call_once(once[device], [device] {
+ ggml_cuda_set_device(device);
+ if (sizeof(shared<C, V>) > ggml_cuda_info().devices[device].smpbo) {
+ return;
+ }
+ const auto status = cudaFuncSetAttribute((const void *) gdn_single<C, V, BLOCKS>,
+ cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(shared<C, V>));
+ available[device] = status == cudaSuccess;
+ if (status != cudaSuccess) {
+ (void) cudaGetLastError();
+ }
+ });
+ return available[device];
+}
+
+enum class gdn_path { ar, slice32_one_block, slice32_two_blocks, slice64, whole_head };
+
+static gdn_path select_gdn_mma_path(int device, const ggml_cuda_gdn_mma_args & a) {
+ GGML_ASSERT(device >= 0 && device < GGML_CUDA_MAX_DEVICES);
+ if (!a.eligible || a.H_k != 16 || (a.H != 16 && a.H != 32 && a.H != 48 && a.H != 64) || a.n_tokens < 64 ||
+ a.n_tokens > INT64_MAX - 32 || a.rq3 <= 0 || a.n_seqs <= 0 || a.n_seqs > INT_MAX / (a.H * 4)) {
+ return gdn_path::ar;
+ }
+ const auto & info = ggml_cuda_info().devices[device];
+#ifdef GGML_USE_HIP
+ if (!GGML_CUDA_CC_IS_RDNA3_5(info.cc) || a.n_tokens < 2048) {
+ return gdn_path::ar;
+ }
+ return init<16, 64>(device) ? gdn_path::slice64 : gdn_path::ar;
+#elif defined(GGML_USE_MUSA)
+ GGML_UNUSED_VARS(info);
+ return gdn_path::ar;
+#else
+ // Enable only architectures benchmarked for this schedule. MMA support alone does not imply a speedup.
+ if (!ampere_mma_available(info.cc) ||
+ (info.cc != GGML_CUDA_CC_AMPERE_GA10X && info.cc != GGML_CUDA_CC_ADA_LOVELACE &&
+ info.cc != GGML_CUDA_CC_BLACKWELL && info.cc != GGML_CUDA_CC_DGX_SPARK)) {
+ return gdn_path::ar;
+ }
+ // Value splitting pays off when whole heads leave at least half the SMs idle.
+ if (a.H * a.n_seqs * 2 <= info.nsm) {
+ // A grid that already fits in one wave benefits from a larger register budget.
+ if (a.H * a.n_seqs * 4 <= info.nsm) {
+ return init<16, 32, 1>(device) ? gdn_path::slice32_one_block : gdn_path::ar;
+ }
+ return init<16, 32, 2>(device) ? gdn_path::slice32_two_blocks : gdn_path::ar;
+ }
+ return init<16, 128>(device) ? gdn_path::whole_head : gdn_path::ar;
+#endif // GGML_USE_HIP
+}
+} // namespace
+
+bool ggml_cuda_gdn_mma_launch(int device, const ggml_cuda_gdn_mma_args & a, cudaStream_t stream) {
+ if (!a.state_out) {
+ return false;
+ }
+ const auto path = select_gdn_mma_path(device, a);
+ if (path == gdn_path::ar) {
+ return false;
+ }
+#ifdef GGML_USE_HIP
+ gdn_single<16, 64><<<a.H * a.n_seqs * 2, dim3(32, 8), sizeof(shared<16, 64>), stream>>>(a);
+#else
+ if (path == gdn_path::slice32_two_blocks) {
+ gdn_single<16, 32, 2><<<a.H * a.n_seqs * 4, dim3(32, 8), sizeof(shared<16, 32>), stream>>>(a);
+ } else if (path == gdn_path::slice32_one_block) {
+ gdn_single<16, 32, 1><<<a.H * a.n_seqs * 4, dim3(32, 8), sizeof(shared<16, 32>), stream>>>(a);
+ } else {
+ gdn_single<16, 128><<<a.H * a.n_seqs, dim3(32, 8), sizeof(shared<16, 128>), stream>>>(a);
+ }
+#endif // GGML_USE_HIP
+ CUDA_CHECK(cudaGetLastError());
+ return true;
+}
diff --git a/ggml/src/ggml-cuda/gated-delta-net-mma.cuh b/ggml/src/ggml-cuda/gated-delta-net-mma.cuh
new file mode 100644
index 000000000..8eb7c4f33
--- /dev/null
+++ b/ggml/src/ggml-cuda/gated-delta-net-mma.cuh
@@ -0,0 +1,35 @@
+#pragma once
+
+#include "common.cuh"
+
+#include <cstdint>
+
+struct ggml_cuda_gdn_mma_args {
+ const float * q;
+ const float * k;
+ const float * v;
+ const float * g;
+ const float * beta;
+ const float * state;
+ float * dst;
+ float * state_out;
+
+ int64_t H;
+ int64_t H_k;
+ int64_t n_tokens;
+ int64_t n_seqs;
+ int64_t rq3;
+ 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;
+ float scale;
+ bool eligible;
+};
+
+bool ggml_cuda_gdn_mma_launch(int device, const ggml_cuda_gdn_mma_args & args, cudaStream_t stream);
diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu
index a703e78be..ae2374cbc 100644
--- a/ggml/src/ggml-cuda/gated_delta_net.cu
+++ b/ggml/src/ggml-cuda/gated_delta_net.cu
@@ -1,6 +1,9 @@
#include "gated_delta_net.cuh"
+#include "gated-delta-net-mma.cuh"
#include "ggml-cuda/common.cuh"
+#include <cstdint>
+
constexpr int gdn_cols_per_warp = 4;
template <int S_v, bool KDA, bool keep_rs_t, int cols_per_warp = gdn_cols_per_warp>
@@ -236,6 +239,44 @@ static void launch_gated_delta_net(
}
}
+static ggml_cuda_gdn_mma_args ggml_cuda_gdn_mma_args_from_tensor(const ggml_tensor * dst) {
+ const ggml_tensor * q = dst->src[0];
+ const ggml_tensor * k = dst->src[1];
+ const ggml_tensor * v = dst->src[2];
+ const ggml_tensor * g = dst->src[3];
+ const ggml_tensor * beta = dst->src[4];
+ const ggml_tensor * state = dst->src[5];
+
+ const int64_t S_v = v->ne[0];
+ const int64_t H = v->ne[1];
+ return {
+ (const float *) q->data,
+ (const float *) k->data,
+ (const float *) v->data,
+ (const float *) g->data,
+ (const float *) beta->data,
+ (const float *) state->data,
+ (float *) dst->data,
+ nullptr,
+ H,
+ q->ne[1],
+ v->ne[2],
+ v->ne[3],
+ v->ne[3] / q->ne[3],
+ (int64_t) (q->nb[1] / sizeof(float)),
+ (int64_t) (q->nb[2] / sizeof(float)),
+ (int64_t) (q->nb[3] / sizeof(float)),
+ (int64_t) (v->nb[1] / sizeof(float)),
+ (int64_t) (v->nb[2] / sizeof(float)),
+ (int64_t) (v->nb[3] / sizeof(float)),
+ (int64_t) (beta->nb[1] / sizeof(float)),
+ (int64_t) (beta->nb[2] / sizeof(float)),
+ (int64_t) (beta->nb[3] / sizeof(float)),
+ 1.0f / sqrtf((float) S_v),
+ g->ne[0] == 1 && ggml_get_op_params_i32(dst, 0) < v->ne[2] && S_v == 128 && q->ne[0] == 128,
+ };
+}
+
static void ggml_cuda_op_gated_delta_net_impl(
ggml_backend_cuda_context & ctx, ggml_tensor * dst, const ggml_cuda_gated_delta_net_fused_cache * cache) {
ggml_tensor * src_q = dst->src[0];
@@ -263,6 +304,8 @@ static void ggml_cuda_op_gated_delta_net_impl(
GGML_ASSERT(neq1 == nek1);
const int64_t neqk1 = neq1;
+ GGML_ASSERT(neq3 > 0 && nev3 >= neq3 && nev3 % neq3 == 0);
+ GGML_ASSERT(nek3 == neq3);
const int64_t rq3 = nev3 / neq3;
const float * q_d = (const float *) src_q->data;
@@ -310,6 +353,13 @@ static void ggml_cuda_op_gated_delta_net_impl(
state_slot_stride = cache->slot_stride;
}
+ {
+ ggml_cuda_gdn_mma_args args = ggml_cuda_gdn_mma_args_from_tensor(dst);
+ args.state_out = state_d;
+ if (ggml_cuda_gdn_mma_launch(ctx.device, args, stream)) {
+ return;
+ }
+ }
if (kda) {
if (keep_rs) {
launch_gated_delta_net<true, true>(q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d,
diff --git a/ggml/src/ggml-cuda/mma.cuh b/ggml/src/ggml-cuda/mma.cuh
index 93a51aa0b..1cec98b06 100644
--- a/ggml/src/ggml-cuda/mma.cuh
+++ b/ggml/src/ggml-cuda/mma.cuh
@@ -949,16 +949,22 @@ namespace ggml_cuda_mma {
static __device__ __forceinline__ void load_ldmatrix_swizzled(
tile<I, J, T, dl> & t, const T * __restrict__ xs0, const int offset) {
#if defined(TURING_MMA_AVAILABLE)
- static_assert(I == 16, "bad tile width");
+ static_assert(I == 8 || I == 16, "bad tile width");
static_assert(J == 8, "bad tile height");
const int i = threadIdx.x % t.I;
- const int j = (threadIdx.x / t.I) * (t.J / 2);
+ const int j = ((threadIdx.x / t.I) * (t.J / 2)) % t.J;
int offset_ij = offset + i * stride + j;
offset_ij = swizzle<stride, T>(offset_ij, i);
int * xi = (int *) t.x;
- asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
- : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3])
- : "l"(xs0 + offset_ij));
+ if constexpr (I == 16) {
+ asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];"
+ : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3])
+ : "l"(xs0 + offset_ij));
+ } else {
+ asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];"
+ : "=r"(xi[0]), "=r"(xi[1])
+ : "l"(xs0 + offset_ij));
+ }
#elif defined(VOLTA_MMA_AVAILABLE)
#pragma unroll
for (int o = 0; o < t.ne; o += 4) {
@@ -1011,16 +1017,22 @@ namespace ggml_cuda_mma {
static __device__ __forceinline__ void load_ldmatrix_trans_swizzled(
tile<I, 8, T, dl> & t, const T * __restrict__ xs0, const int offset) {
#if defined(TURING_MMA_AVAILABLE)
- static_assert(I == 16, "bad tile width");
+ static_assert(I == 8 || I == 16, "bad tile width");
static_assert(dl == DATA_LAYOUT_I_MAJOR, "bad data layout");
- const int i = threadIdx.x % t.I;
- const int j = (threadIdx.x / t.I) * (t.J / 2);
+ const int i = threadIdx.x % 16;
+ const int j = I == 16 ? (threadIdx.x / 16) * (t.J / 2) : 0;
int offset_ij = offset + i * stride + j;
offset_ij = swizzle<stride, T>(offset_ij, i);
int * xi = (int *) t.x;
- asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];"
- : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3])
- : "l"(xs0 + offset_ij));
+ if constexpr (I == 16) {
+ asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];"
+ : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3])
+ : "l"(xs0 + offset_ij));
+ } else {
+ asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.b16 {%0, %1}, [%2];"
+ : "=r"(xi[0]), "=r"(xi[1])
+ : "l"(xs0 + offset_ij));
+ }
#elif defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)
static_assert(dl == DATA_LAYOUT_I_MAJOR || dl == DATA_LAYOUT_I_MAJOR_MIRRORED, "bad data layout");
if constexpr (I == 32) {
diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp
index 1e95bf50c..0d7a7ed81 100644
--- a/tests/test-backend-ops.cpp
+++ b/tests/test-backend-ops.cpp
@@ -11950,6 +11950,18 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 33, 1, 1, false, true));
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 100, 1, 1, false, true));
+ // MMA-eligible GDN shapes: whole heads, value slices, tails, and the HIP prefill threshold.
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 65, 1, 1));
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 65, 3, 4));
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 64, 1, 4));
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 257, 1, 3));
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 512, 1, 2));
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 513, 1, 3));
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 513, 2, 2, true));
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 2048, 1, 2));
+ // Keep the WebGPU dispatch below 65535 workgroups (32 per token for this shape).
+ test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 16, 128, 2047, 2, 3));
+
// K > 1: output keeps the last min(n_tokens, K) per-token snapshots, ordered most-recent-first
// (slot 0 = final state, slot s = state s tokens back).
// exact-match cases (K == n_seq_tokens):
@@ -11959,9 +11971,6 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 128, 4, 1, 1, false, false, /*K=*/4));
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 4, 2, 1, false, true, /*K=*/4));
test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 8, 32, 4, 2, 2, false, true, /*K=*/4));
- // overflow: n_tokens > K — only the last K snapshots kept.
- test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 32, 8, 1, 1, false, false, /*K=*/3));
- test_cases.emplace_back(new test_gated_delta_net(GGML_TYPE_F32, 4, 64, 16, 2, 1, false, false, /*K=*/4));
// gdn + cache cpy fusion (K > 1)
test_cases.emplace_back(new test_gated_delta_net_cache_fusion(GGML_TYPE_F32, 4, 32, 2, 1, 2));