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 +#include + +namespace { +constexpr int D = 128; +using bf16 = nv_bfloat16; + +template 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(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(hi + i) = h; + *reinterpret_cast(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 struct alignas(128) shared { + matrix q, k; + matrix p, inverse; + matrix delta; + + union { + matrix state; + matrix 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 +__device__ __forceinline__ tile_t load(const matrix & 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(data); + if constexpr (TRANS) { + load_ldmatrix_trans_swizzled(t, packed, col * (C / 2) + row / 2); + } else { + load_ldmatrix_swizzled(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 +__device__ __forceinline__ void product(acc & c, const matrix & a, const matrix & b, int row, int col) { +#pragma unroll + for (int k = 0; k < K; k += 16) { + const auto ah = load(a, row, k, false); + const auto al = load(a, row, k, true); + const auto bh = load(b, col, k, false); + const auto bl = load(b, col, k, true); + mma(c, al, bh); + mma(c, ah, bl); + mma(c, ah, bh); + } +} + +template __device__ __forceinline__ void store(matrix & 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 +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 *>(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(a.q + off); + k = *reinterpret_cast(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(qs, s.q, s.scratch.state, r, c); + product(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(x, s.k, s.k, r, c); + } else { + product(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(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(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(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 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) > ggml_cuda_info().devices[device].smpbo) { + return; + } + const auto status = cudaFuncSetAttribute((const void *) gdn_single, + cudaFuncAttributeMaxDynamicSharedMemorySize, sizeof(shared)); + 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><<), stream>>>(a); +#else + if (path == gdn_path::slice32_two_blocks) { + gdn_single<16, 32, 2><<), stream>>>(a); + } else if (path == gdn_path::slice32_one_block) { + gdn_single<16, 32, 1><<), stream>>>(a); + } else { + gdn_single<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 + +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 + constexpr int gdn_cols_per_warp = 4; template @@ -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(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 & 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(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 & 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(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> 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> 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));