diff --git a/ggml/src/ggml-cuda/fwht.cu b/ggml/src/ggml-cuda/fwht.cu index 467a84a9..281a3d97 100644 --- a/ggml/src/ggml-cuda/fwht.cu +++ b/ggml/src/ggml-cuda/fwht.cu @@ -1,5 +1,6 @@ #include "common.cuh" #include "fwht.cuh" +#include "unary.cuh" #include @@ -13,10 +14,19 @@ __device__ __forceinline__ float fwht_load(const half value) { return __half2float(value); } -template +// the transform's input: src, or silu(src) * src2 for a SWIGLU folded into the load (same arithmetic as the GLU kernel) +template +__device__ __forceinline__ float fwht_value(const T * src, const T * src2, const int i) { + if constexpr (glu) { + return ggml_cuda_op_silu_single((float) src[i]) * (float) src2[i]; + } + return fwht_load(src[i]); +} + +template __launch_bounds__(4*ggml_cuda_get_physical_warp_size(), 1) __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, const float scale, - const float * signs, const int n_blk) { + const float * signs, const int n_blk, const T * src2) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); const int64_t r = (int64_t) blockIdx.x * blockDim.y + threadIdx.y; @@ -27,6 +37,9 @@ __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, cons src += r * N; dst += r * N; + if (glu) { + src2 += r * N; + } static constexpr int el_w = N / warp_size; float reg[el_w]; @@ -36,7 +49,7 @@ __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, cons const float * signs_row = has_signs ? signs + (r % n_blk) * N : nullptr; #pragma unroll for (int i = 0; i < el_w; ++i) { - reg[i] = fwht_load(src[i * warp_size + lane]) * scale; + reg[i] = fwht_value(src, src2, i * warp_size + lane) * scale; if (has_signs) { reg[i] *= signs_row[i * warp_size + lane]; } @@ -83,10 +96,10 @@ __global__ void fwht_cuda(const T * src, float * dst, const int64_t n_rows, cons // row than the register path, so it is used only where that path cannot go. #define FWHT_SMEM_THREADS 256 -template +template __launch_bounds__(FWHT_SMEM_THREADS, 1) __global__ void fwht_cuda_smem(const T * src, float * dst, const int64_t n_rows, const float scale, - const float * signs, const int n_blk) { + const float * signs, const int n_blk, const T * src2) { __shared__ float s[N]; const int64_t r = blockIdx.x; @@ -96,12 +109,15 @@ __global__ void fwht_cuda_smem(const T * src, float * dst, const int64_t n_rows, src += r * N; dst += r * N; + if (glu) { + src2 += r * N; + } const float * signs_row = has_signs ? signs + (r % n_blk) * N : nullptr; ggml_cuda_pdl_sync(); for (int i = threadIdx.x; i < N; i += FWHT_SMEM_THREADS) { - float v = fwht_load(src[i]) * scale; + float v = fwht_value(src, src2, i) * scale; if (has_signs) { v *= signs_row[i]; } @@ -134,10 +150,10 @@ __global__ void fwht_cuda_smem(const T * src, float * dst, const int64_t n_rows, // Both leave most of the GPU idle at these shapes. #define FWHT_BLOCK_THREADS 256 -template +template __launch_bounds__(NT, 1) __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows, const float scale, - const float * signs, const int n_blk) { + const float * signs, const int n_blk, const T * src2) { constexpr int warp_size = ggml_cuda_get_physical_warp_size(); constexpr int NE = N / NT; static_assert(NE >= 1 && N % NT == 0 && NT % warp_size == 0, "bad FWHT block shape"); @@ -151,6 +167,9 @@ __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows src += r * N; dst += r * N; + if (glu) { + src2 += r * N; + } const int tid = threadIdx.x; const int lane = tid % warp_size; @@ -161,7 +180,7 @@ __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows float reg[NE]; #pragma unroll for (int i = 0; i < NE; ++i) { - reg[i] = fwht_load(src[i * NT + tid]) * scale; + reg[i] = fwht_value(src, src2, i * NT + tid) * scale; if (has_signs) { reg[i] *= signs_row[i * NT + tid]; } @@ -220,7 +239,7 @@ __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows template static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float * dst_d, const int n, const int64_t rows, const float scale, - const float * signs, const int n_blk) { + const float * signs, const int n_blk, const T * src2_d = nullptr) { const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; const int rows_per_block = 4; const int64_t num_blocks = (rows + rows_per_block - 1) / rows_per_block; @@ -233,10 +252,12 @@ static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float switch (n) { #define FWHT_CASE(NN) \ case NN: \ - if (signs) { \ - ggml_cuda_kernel_launch(fwht_cuda, launch_params, src_d, dst_d, rows, scale, signs, n_blk); \ + if (src2_d) { \ + ggml_cuda_kernel_launch(fwht_cuda, launch_params, src_d, dst_d, rows, scale, signs, n_blk, src2_d); \ + } else if (signs) { \ + ggml_cuda_kernel_launch(fwht_cuda, launch_params, src_d, dst_d, rows, scale, signs, n_blk, nullptr); \ } else { \ - ggml_cuda_kernel_launch(fwht_cuda, launch_params, src_d, dst_d, rows, scale, nullptr, 1); \ + ggml_cuda_kernel_launch(fwht_cuda, launch_params, src_d, dst_d, rows, scale, nullptr, 1, nullptr); \ } \ return true; FWHT_CASE(64) @@ -252,10 +273,12 @@ static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float case NN: { \ const dim3 g((unsigned) rows, 1, 1), b(FWHT_SMEM_THREADS, 1, 1); \ const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(g, b, 0, stream); \ - if (signs) { \ - ggml_cuda_kernel_launch(fwht_cuda_smem, lp, src_d, dst_d, rows, scale, signs, n_blk); \ + if (src2_d) { \ + ggml_cuda_kernel_launch(fwht_cuda_smem, lp, src_d, dst_d, rows, scale, signs, n_blk, src2_d); \ + } else if (signs) { \ + ggml_cuda_kernel_launch(fwht_cuda_smem, lp, src_d, dst_d, rows, scale, signs, n_blk, nullptr); \ } else { \ - ggml_cuda_kernel_launch(fwht_cuda_smem, lp, src_d, dst_d, rows, scale, nullptr, 1); \ + ggml_cuda_kernel_launch(fwht_cuda_smem, lp, src_d, dst_d, rows, scale, nullptr, 1, nullptr); \ } \ return true; \ } @@ -263,10 +286,12 @@ static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float case NN: { \ const dim3 g((unsigned) rows, 1, 1), b(FWHT_BLOCK_THREADS, 1, 1); \ const ggml_cuda_kernel_launch_params lp = ggml_cuda_kernel_launch_params(g, b, 0, stream); \ - if (signs) { \ - ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, signs, n_blk); \ + if (src2_d) { \ + ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, signs, n_blk, src2_d); \ + } else if (signs) { \ + ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, signs, n_blk, nullptr); \ } else { \ - ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, nullptr, 1); \ + ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, nullptr, 1, nullptr); \ } \ return true; \ } @@ -337,3 +362,19 @@ bool ggml_cuda_op_fwht_signed(ggml_backend_cuda_context & ctx, const ggml_tensor const ggml_tensor * signs, ggml_tensor * dst) { return fwht_dispatch(ctx, src, dst, signs); } + +bool ggml_cuda_op_fwht_swiglu_signed(ggml_backend_cuda_context & ctx, const ggml_tensor * gate, const ggml_tensor * up, + const ggml_tensor * signs_t, ggml_tensor * dst) { + // F32 split SWIGLU only, both halves laid out like dst + if (gate->type != GGML_TYPE_F32 || up->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + !ggml_is_contiguous(gate) || !ggml_is_contiguous(up) || !ggml_is_contiguous(dst) || + ggml_nelements(gate) != ggml_nelements(dst) || ggml_nelements(up) != ggml_nelements(dst)) { + return false; + } + const int n = dst->ne[0]; + if (signs_t->type != GGML_TYPE_F32 || !ggml_is_contiguous(signs_t) || signs_t->ne[0] % n != 0) { + return false; + } + return fwht_launch(ctx, (const float *) gate->data, (float *) dst->data, n, ggml_nelements(dst) / n, + 1 / sqrtf(n), (const float *) signs_t->data, signs_t->ne[0] / n, (const float *) up->data); +} diff --git a/ggml/src/ggml-cuda/fwht.cuh b/ggml/src/ggml-cuda/fwht.cuh index 62b2f288..faf3e88b 100644 --- a/ggml/src/ggml-cuda/fwht.cuh +++ b/ggml/src/ggml-cuda/fwht.cuh @@ -4,3 +4,6 @@ bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst); bool ggml_cuda_op_fwht_signed(ggml_backend_cuda_context & ctx, const ggml_tensor * src, const ggml_tensor * signs, ggml_tensor * dst); +// the same, with a split F32 SWIGLU (silu(gate) * up) folded into the transform's load +bool ggml_cuda_op_fwht_swiglu_signed(ggml_backend_cuda_context & ctx, const ggml_tensor * gate, const ggml_tensor * up, + const ggml_tensor * signs, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 7536a5d2..69b2652b 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2971,7 +2971,18 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_count, const int * out_nodes, const int out_count, - const bool is_topk_moe = false) { + const bool is_topk_moe = false, + const bool allow_exact_alias = false) { + // An exact alias is a destination that shares the base pointer, type and + // contiguous layout of a source, i.e. element k of the destination is + // element k of the source. Fusions whose kernels are elementwise within the + // unit they parallelize over can opt into accepting those; a partial + // overlap is still rejected for everyone. + auto is_exact_alias = [&](const ggml_tensor * a, const ggml_tensor * b) { + return allow_exact_alias && a->data == b->data && a->type == b->type && + ggml_are_same_shape(a, b) && ggml_is_contiguous(a) && ggml_is_contiguous(b); + }; + auto nodes_overlap = [&](const ggml_tensor * a, const ggml_tensor * b) { const int64_t a_start = (int64_t) a->data; const int64_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); @@ -3007,7 +3018,7 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, continue; } - if (nodes_overlap(dst, src)) { + if (nodes_overlap(dst, src) && !is_exact_alias(dst, src)) { bool found = false; for (int k = node_idx; k < j; ++k) { @@ -3359,13 +3370,20 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph // Preserve the residual sum for later graph consumers while normalizing it // in the same launch. This removes a full read of the residual tensor. + // + // add_rms_norm_f32 uses no architecture-specific instructions, so this runs + // on every NVIDIA device rather than GB10 only. It is elementwise within a + // row -- each thread reads a[col] and b[col] before writing sum[col], and + // block_reduce synchronizes the block before dst[col] is written -- so the + // in-place residual add and the norm output reusing an input buffer are + // both safe, hence allow_exact_alias below. if (node->op == GGML_OP_ADD && i + 2 < cgraph->n_nodes) { ggml_tensor * rms_norm = cgraph->nodes[i + 1]; ggml_tensor * mul = cgraph->nodes[i + 2]; const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const ggml_op ops[] = { GGML_OP_ADD, GGML_OP_RMS_NORM, GGML_OP_MUL }; const int out_nodes[] = { i, i + 2 }; - if (cc == GGML_CUDA_CC_DGX_SPARK && rms_norm->op == GGML_OP_RMS_NORM && + if (GGML_CUDA_CC_IS_NVIDIA(cc) && rms_norm->op == GGML_OP_RMS_NORM && mul->op == GGML_OP_MUL && (mul->src[0] == rms_norm || mul->src[1] == rms_norm) && rms_norm->src[0] == node && node->src[0] && node->src[1] && node->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 && @@ -3375,10 +3393,12 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_is_contiguous(node->src[1]) && ggml_is_contiguous(node) && ggml_is_contiguous(rms_norm) && ggml_is_contiguous(mul) && ggml_can_fuse_subgraph(cgraph, i, 3, ops, out_nodes, 2) && - ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 2)) { + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 2, false, true)) { const ggml_tensor * weight = mul->src[0] == rms_norm ? mul->src[1] : mul->src[0]; if (weight && weight->type == GGML_TYPE_F32 && ggml_is_contiguous(weight) && - weight->ne[0] == node->ne[0] && ggml_nrows(weight) == 1) { + weight->ne[0] == node->ne[0] && ggml_nrows(weight) == 1 && + weight->data != node->data && weight->data != rms_norm->data && + weight->data != mul->data) { ggml_cuda_op_add_rms_norm_fused(*cuda_ctx, node, rms_norm, mul); return 2; } @@ -3475,6 +3495,49 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + // SWIGLU [+ reshapes] + Hadamard sign flip + reshape + FWHT-hint matmul (the FFN down projection's input, and the + // DeltaNet gated norm's output before ssm_out): the transform computes silu(a) * b as it loads, so the SWIGLU output + // never goes through memory + static const bool swiglu_fwht = getenv("GGML_CUDA_NO_SWIGLU_FWHT") == nullptr; + if (swiglu_fwht && node->op == GGML_OP_GLU && ggml_get_glu_op(node) == GGML_GLU_OP_SWIGLU) { + int k = i + 1; + while (k < cgraph->n_nodes && k < i + 4 && cgraph->nodes[k]->op == GGML_OP_RESHAPE) { + ++k; + } + const int n_ops = k - i + 3; + if (k + 2 < cgraph->n_nodes && n_ops <= 6) { + ggml_op ops[6]; + for (int j = 0; j < n_ops; ++j) { + ops[j] = cgraph->nodes[i + j]->op; + } + const int out_nodes[] = { k + 2 }; + const ggml_tensor * mul = cgraph->nodes[k]; + const ggml_tensor * reshape = cgraph->nodes[k + 1]; + ggml_tensor * mm = cgraph->nodes[k + 2]; + + bool chain_ok = true; + for (int j = i + 1; j <= k; ++j) { + chain_ok = chain_ok && cgraph->nodes[j]->src[0] == cgraph->nodes[j - 1]; + } + const ggml_tensor * signs = mul->src[1]; + + const bool pattern_ok = chain_ok && mul->op == GGML_OP_MUL && reshape->op == GGML_OP_RESHAPE && + mm->op == GGML_OP_MUL_MAT && node->src[1] != nullptr && + ggml_get_op_params_i32(node, 1) == 0 && node->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && + ggml_get_op_params_i32(mm, 1) == GGML_HINT_SRC0_IS_HADAMARD && + mm->src[1] == reshape && reshape->src[0] == mul && signs && + signs->ne[1] == 1 && signs->ne[2] == 1 && signs->ne[3] == 1 && + ggml_are_same_shape(node->src[0], node) && ggml_are_same_shape(node->src[1], node) && + ggml_nelements(mul) == ggml_nelements(node) && ggml_is_contiguous(node) && + signs->ne[0] == mul->ne[0] && signs->ne[0] % mm->src[0]->ne[0] == 0 && + ggml_can_fuse_subgraph(cgraph, i, n_ops, ops, out_nodes, 1); + + if (pattern_ok && ggml_cuda_op_fwht_swiglu_signed(*cuda_ctx, node->src[0], node->src[1], signs, mm)) { + return n_ops - 1; + } + } + } + // Hadamard sign flip + reshape + FWHT-hint matmul: multiply the sign // vector during the transform's load instead of a separate full pass if (ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL, GGML_OP_RESHAPE, GGML_OP_MUL_MAT }, { i + 2 })) {