diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 4ae9a06b..0037db19 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -1584,6 +1584,18 @@ struct ggml_backend_cuda_context { ggml_cuda_pool & pool() { return pool(device); } + + // MMQ src1 tiles written by fused FWHTs for their PTQ1_0 consumers, innermost last (the pool is stack-like). + // Host-side bookkeeping for one graph evaluation or capture; empty between graphs. + struct fwht_q8_tile { + const void * data; // the transform's output, which the consumers' src1 views + int64_t ne[4]; // the shape of that src1 + std::unique_ptr> q8; + int last_node; // graph index of the last consumer + int remaining; // consumers still to run + bool f32; // whether the f32 output was written too + }; + std::vector fwht_q8; }; struct ggml_cuda_mm_fusion_args_host { diff --git a/ggml/src/ggml-cuda/fwht.cu b/ggml/src/ggml-cuda/fwht.cu index 281a3d97..8cdc33b3 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 "mmq.cuh" #include "unary.cuh" #include @@ -145,6 +146,56 @@ __global__ void fwht_cuda_smem(const T * src, float * dst, const int64_t n_rows, } +// Quantizes a transformed row (reg[j] holds value j*NT + tid) into the MMQ src1 tile, with the exact arithmetic of +// quantize_mmq_q8_1 for MMQ_Q8_1_DS_LAYOUT_D1: the row is restaged in s so that each warp owns 128 consecutive values, +// 4 per lane. Row r is block r % n_blk of column r / n_blk. +template +static __device__ __forceinline__ void fwht_store_q8(const float * reg, float * s, const int64_t r, const int n_blk, + const fwht_q8_out q8) { + constexpr int NE = N / NT; + static_assert(N % (4 * QK8_1_MMQ) == 0, "FWHT rows must hold whole q8 warps"); + +#pragma unroll + for (int j = 0; j < NE; ++j) { + s[j * NT + threadIdx.x] = reg[j]; + } + __syncthreads(); + + const int64_t col = r / n_blk; + const int64_t channel = col / q8.ne11; + const int64_t ib0 = channel * q8.ne11 * n_blk * (N / QK8_1_MMQ) + col % q8.ne11 + + (int64_t) (r % n_blk) * (N / QK8_1_MMQ) * q8.ne11; + block_q8_1_mmq * y = (block_q8_1_mmq *) q8.y; + + for (int u = threadIdx.x; u < N / 4; u += NT) { + const float4 xi = ((const float4 *) s)[u]; + + float amax = fabsf(xi.x); + amax = fmaxf(amax, fabsf(xi.y)); + amax = fmaxf(amax, fabsf(xi.z)); + amax = fmaxf(amax, fabsf(xi.w)); +#pragma unroll + for (int offset = QK8_1_MMQ/8; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE)); + } + + const float d_inv = 127.0f / amax; + char4 q; + q.x = roundf(xi.x*d_inv); + q.y = roundf(xi.y*d_inv); + q.z = roundf(xi.z*d_inv); + q.w = roundf(xi.w*d_inv); + const float d = 1.0f / d_inv; + + block_q8_1_mmq & yb = y[ib0 + (int64_t) (u / (QK8_1_MMQ/4)) * q8.ne11]; + const int iqs = (u % (QK8_1_MMQ/4)) * 4; + ((char4 *) yb.qs)[iqs/4] = q; + if (iqs % 32 == 0) { + yb.d4[iqs/32] = d; + } + } +} + // Wide rows at small row counts (decode): one row per block instead of per warp. // The warp kernel serialises every stage on one warp, and the shared-memory kernel above synchronises on each stage. // Both leave most of the GPU idle at these shapes. @@ -153,7 +204,7 @@ __global__ void fwht_cuda_smem(const T * src, float * dst, const int64_t n_rows, 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 T * src2) { + const float * signs, const int n_blk, const T * src2, const fwht_q8_out q8) { 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"); @@ -166,7 +217,9 @@ __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows } src += r * N; - dst += r * N; + if (dst) { + dst += r * N; + } if (glu) { src2 += r * N; } @@ -230,16 +283,35 @@ __global__ void fwht_cuda_block(const T * src, float * dst, const int64_t n_rows } } + if (dst) { #pragma unroll - for (int i = 0; i < NE; ++i) { - dst[i * NT + tid] = reg[i]; + for (int i = 0; i < NE; ++i) { + dst[i * NT + tid] = reg[i]; + } + } + if (q8.y) { + fwht_store_q8(reg, s, r, n_blk, q8); } } +static bool fwht_legacy() { + static const bool legacy = getenv("GGML_CUDA_FWHT_LEGACY") != nullptr; + return legacy; +} + +bool ggml_cuda_fwht_can_emit_q8(const int n) { + // the block kernel's sizes; the per-128 tile is PTQ1_0's MMQ_Q8_1_DS_LAYOUT_D1 + return GGML_CUDA_PTQ1_0_MMQ_Y128 && !fwht_legacy() && n >= 512 && n <= 8192 && (n & (n - 1)) == 0; +} + 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 T * src2_d = nullptr) { + const float * signs, const int n_blk, const T * src2_d = nullptr, + const fwht_q8_out & q8 = {}) { + if (q8.y && (!signs || !ggml_cuda_fwht_can_emit_q8(n))) { + return false; + } 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; @@ -287,16 +359,16 @@ static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float 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 (src2_d) { \ - ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, signs, n_blk, src2_d); \ + ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, f32_d, rows, scale, signs, n_blk, src2_d, q8); \ } else if (signs) { \ - ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, signs, n_blk, nullptr); \ + ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, f32_d, rows, scale, signs, n_blk, nullptr, q8); \ } else { \ - ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, nullptr, 1, nullptr); \ + ggml_cuda_kernel_launch(fwht_cuda_block, lp, src_d, dst_d, rows, scale, nullptr, 1, nullptr, q8); \ } \ return true; \ } - static const bool legacy = getenv("GGML_CUDA_FWHT_LEGACY") != nullptr; - if (legacy) { + float * f32_d = q8.skip_f32 ? nullptr : dst_d; + if (fwht_legacy()) { switch (n) { FWHT_CASE(512) FWHT_CASE(1024) @@ -322,7 +394,7 @@ static bool fwht_launch(ggml_backend_cuda_context & ctx, const T * src_d, float } static bool fwht_dispatch(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst, - const ggml_tensor * signs_t) { + const ggml_tensor * signs_t, const fwht_q8_out & q8 = {}) { GGML_ASSERT(ggml_nelements(src) == ggml_nelements(dst)); if (!ggml_is_contiguous(src) || !ggml_is_contiguous(dst)) { return false; @@ -348,9 +420,9 @@ static bool fwht_dispatch(ggml_backend_cuda_context & ctx, const ggml_tensor * s const float scale = 1 / sqrtf(n); if (src->type == GGML_TYPE_F32) { - return fwht_launch(ctx, (const float *) src->data, dst_d, n, rows, scale, signs, n_blk); + return fwht_launch(ctx, (const float *) src->data, dst_d, n, rows, scale, signs, n_blk, nullptr, q8); } - return fwht_launch(ctx, (const half *) src->data, dst_d, n, rows, scale, signs, n_blk); + return fwht_launch(ctx, (const half *) src->data, dst_d, n, rows, scale, signs, n_blk, nullptr, q8); } bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, ggml_tensor * dst) { @@ -359,12 +431,12 @@ bool ggml_cuda_op_fwht(ggml_backend_cuda_context & ctx, const ggml_tensor * src, } bool ggml_cuda_op_fwht_signed(ggml_backend_cuda_context & ctx, const ggml_tensor * src, - const ggml_tensor * signs, ggml_tensor * dst) { - return fwht_dispatch(ctx, src, dst, signs); + const ggml_tensor * signs, ggml_tensor * dst, const fwht_q8_out & q8) { + return fwht_dispatch(ctx, src, dst, signs, q8); } 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) { + const ggml_tensor * signs_t, ggml_tensor * dst, const fwht_q8_out & q8) { // 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) || @@ -376,5 +448,5 @@ bool ggml_cuda_op_fwht_swiglu_signed(ggml_backend_cuda_context & ctx, const ggml 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); + 1 / sqrtf(n), (const float *) signs_t->data, signs_t->ne[0] / n, (const float *) up->data, q8); } diff --git a/ggml/src/ggml-cuda/fwht.cuh b/ggml/src/ggml-cuda/fwht.cuh index faf3e88b..8a1f0079 100644 --- a/ggml/src/ggml-cuda/fwht.cuh +++ b/ggml/src/ggml-cuda/fwht.cuh @@ -1,9 +1,20 @@ #include "common.cuh" +// The MMQ src1 tile (block_q8_1_mmq, one scale per 128 values) a fused FWHT can write next to, or +// instead of, its f32 output, for PTQ1_0 matmuls that read the transform through a [K, ne11, ...] view. +struct fwht_q8_out { + void * y = nullptr; + int64_t ne11 = 0; // columns per channel of the consumers' src1 + bool skip_f32 = false; // no one reads the f32 output +}; + +// Returns whether a fused FWHT of rows of n values can write fwht_q8_out. +bool ggml_cuda_fwht_can_emit_q8(int n); + // Returns whether the Fast Walsh-Hadamard transform could be used. 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); + const ggml_tensor * signs, ggml_tensor * dst, const fwht_q8_out & q8 = {}); // 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); + const ggml_tensor * signs, ggml_tensor * dst, const fwht_q8_out & q8 = {}); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 69b2652b..2f5ad147 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -1815,13 +1815,20 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { bool ggml_cuda_mul_mat_q1_hopper(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); -static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - GGML_TENSOR_BINARY_OP_LOCALS +enum ggml_cuda_mul_mat_path { + GGML_CUDA_MUL_MAT_CUBLAS, + GGML_CUDA_MUL_MAT_VEC_F, + GGML_CUDA_MUL_MAT_VEC_F_TRANSPOSED, // a transposed vector (ne01 == 1) through the F32 vector kernel + GGML_CUDA_MUL_MAT_F, + GGML_CUDA_MUL_MAT_VEC_Q, + GGML_CUDA_MUL_MAT_Q1_HOPPER, // the opt-in Hopper wgmma path, which falls back to MMQ when unsupported + GGML_CUDA_MUL_MAT_Q, +}; - const int32_t hint = ggml_get_op_params_i32(dst, 1); - if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { - return; - } +// Returns the kernel ggml_cuda_mul_mat runs for a matmul without the Hadamard hint. +static ggml_cuda_mul_mat_path ggml_cuda_mul_mat_choose(const int device, const ggml_tensor * src0, + const ggml_tensor * src1, const ggml_tensor * dst) { + GGML_TENSOR_BINARY_OP_LOCALS // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. @@ -1829,52 +1836,116 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; if (bad_padding_clear || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); - return; + return GGML_CUDA_MUL_MAT_CUBLAS; } - const int cc = ggml_cuda_info().devices[ctx.device].cc; - const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; + const int cc = ggml_cuda_info().devices[device].cc; + const int warp_size = ggml_cuda_info().devices[device].warp_size; if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) { // The custom F16 vector kernel can be used over batched cuBLAS GEMM. // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) - ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MUL_MAT_VEC_F; } // A transposed vector can still use MMVQ (i.e. ne01 == 1) if (ne01 == 1 && ne11 > MMVF_MAX_BATCH_SIZE && ne2 == 1 && ne3 == 1 && src0->type == GGML_TYPE_F32 && ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && ggml_cuda_should_use_mmvf(src1->type, cc, src1->ne, src1->nb, /*ne11 =*/ 1)) { - ggml_tensor dst_vec = *dst; - dst_vec.ne[0] = ne11; - dst_vec.ne[1] = 1; - dst_vec.nb[1] = dst_vec.nb[0]*ne11; - dst_vec.nb[2] = dst_vec.nb[1]; - dst_vec.nb[3] = dst_vec.nb[1]; - ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); - return; + return GGML_CUDA_MUL_MAT_VEC_F_TRANSPOSED; } if (ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, ne11, /*mul_mat_id =*/ false)) { - ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); - return; + return GGML_CUDA_MUL_MAT_F; } if (ggml_cuda_should_use_mmvq(src0->type, cc, ne11)) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); - return; - } - if ((src0->type == GGML_TYPE_Q1_0 || src0->type == GGML_TYPE_PQ2_0) && ne11 >= 128 - && ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0) - && ggml_cuda_mul_mat_q1_hopper(ctx, src0, src1, dst)) { - // handled by the opt-in Hopper wgmma path (returns false to fall through when unsupported) - return; + return GGML_CUDA_MUL_MAT_VEC_Q; } if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { - ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return (src0->type == GGML_TYPE_Q1_0 || src0->type == GGML_TYPE_PQ2_0) && ne11 >= 128 ? + GGML_CUDA_MUL_MAT_Q1_HOPPER : GGML_CUDA_MUL_MAT_Q; + } + return GGML_CUDA_MUL_MAT_CUBLAS; +} + +// Returns whether node is a matmul that reads src1 through MMQ in the per-128 layout a fused FWHT can write. +static bool ggml_cuda_mul_mat_reads_fwht_q8(const int device, const ggml_tensor * node) { + if (node->op != GGML_OP_MUL_MAT || ggml_get_op_params_i32(node, 1) == GGML_HINT_SRC0_IS_HADAMARD) { + return false; + } + const ggml_tensor * src0 = node->src[0]; + const ggml_tensor * src1 = node->src[1]; + return src0->type == GGML_TYPE_PTQ1_0 && ggml_is_contiguous(src1) && src1->ne[0] % MATRIX_ROW_PADDING == 0 && + ggml_cuda_mul_mat_choose(device, src0, src1, node) == GGML_CUDA_MUL_MAT_Q; +} + +// Returns the FWHT tile for src1, or nullptr when there is none. A matmul that will not read the tile may only +// read src1 when the transform also wrote its f32 output. +static ggml_backend_cuda_context::fwht_q8_tile * ggml_cuda_fwht_q8_find(ggml_backend_cuda_context & ctx, + const ggml_tensor * src1, const bool reads_q8) { + for (auto & tile : ctx.fwht_q8) { + if (tile.data != src1->data) { + continue; + } + const bool same_shape = std::equal(std::begin(tile.ne), std::end(tile.ne), src1->ne); + GGML_ASSERT((reads_q8 && same_shape) || tile.f32); + return reads_q8 && same_shape ? &tile : nullptr; + } + return nullptr; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_TENSOR_BINARY_OP_LOCALS + + const int32_t hint = ggml_get_op_params_i32(dst, 1); + if (hint == GGML_HINT_SRC0_IS_HADAMARD && ggml_cuda_op_fwht(ctx, src1, dst)) { return; } - ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + + if (!ctx.fwht_q8.empty()) { + if (auto * tile = ggml_cuda_fwht_q8_find(ctx, src1, ggml_cuda_mul_mat_reads_fwht_q8(ctx.device, dst))) { + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst, nullptr, nullptr, nullptr, tile->q8->get(), + /*quantize_external =*/ false); + if (--tile->remaining == 0) { + // the innermost tile: every tile pushed after this one ended before its last consumer + GGML_ASSERT(tile == &ctx.fwht_q8.back()); + ctx.fwht_q8.pop_back(); + } + return; + } + } + + switch (ggml_cuda_mul_mat_choose(ctx.device, src0, src1, dst)) { + case GGML_CUDA_MUL_MAT_CUBLAS: + ggml_cuda_mul_mat_cublas(ctx, src0, src1, dst); + return; + case GGML_CUDA_MUL_MAT_VEC_F: + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MUL_MAT_VEC_F_TRANSPOSED: { + ggml_tensor dst_vec = *dst; + dst_vec.ne[0] = ne11; + dst_vec.ne[1] = 1; + dst_vec.nb[1] = dst_vec.nb[0]*ne11; + dst_vec.nb[2] = dst_vec.nb[1]; + dst_vec.nb[3] = dst_vec.nb[1]; + ggml_cuda_mul_mat_vec_f(ctx, src1, src0, nullptr, &dst_vec); + return; + } + case GGML_CUDA_MUL_MAT_F: + ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MUL_MAT_VEC_Q: + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); + return; + case GGML_CUDA_MUL_MAT_Q1_HOPPER: + if (ggml_cuda_mul_mat_q1_hopper(ctx, src0, src1, dst)) { + return; + } + [[fallthrough]]; // unsupported here + case GGML_CUDA_MUL_MAT_Q: + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + return; + } } // returns true when ggml_cuda_mul_mat_id takes the fallback path that requires stream synchronization @@ -3292,6 +3363,71 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, return false; } +// Plans the MMQ tile a fused FWHT writes for the PTQ1_0 matmuls that read its output mm (graph node i_mm, rows of +// K-wide activations): every reader after i_mm that ggml_cuda_mul_mat will run through MMQ in the per-128 layout takes +// the tile instead of quantizing mm itself, and the f32 output is skipped when nothing else reads it. Returns an empty +// fwht_q8_out when no tile is worth writing; otherwise the tile is on cuda_ctx->fwht_q8 until its last reader ran. +static fwht_q8_out ggml_cuda_fwht_q8_plan(ggml_backend_cuda_context * cuda_ctx, const ggml_cgraph * cgraph, + const int i_mm, const int64_t K) { + static const bool enabled = getenv("GGML_CUDA_NO_FWHT_Q8") == nullptr; + const ggml_tensor * mm = cgraph->nodes[i_mm]; + // one stream only: the tile lives in stream 0's pool and is read in stream order + if (!enabled || cuda_ctx->curr_stream_no != 0 || !cuda_ctx->stream_context().concurrent_events.empty() || + mm->type != GGML_TYPE_F32 || mm->view_src || !ggml_is_contiguous(mm) || + !ggml_cuda_fwht_can_emit_q8(mm->ne[0]) || K % mm->ne[0] != 0) { + return {}; + } + + bool f32 = (mm->flags & GGML_TENSOR_FLAG_OUTPUT) != 0; + const ggml_tensor * shape = nullptr; // src1 of the first tile reader + int readers = 0; + int last = -1; + size_t size = 0; + // uses of mm and its views not seen yet; the counts cover the whole graph, so a use in another split stays pending + int pending = ggml_node_get_use_count(cgraph, i_mm); + for (int j = i_mm + 1; j < cgraph->n_nodes && pending > 0; ++j) { + const ggml_tensor * node = cgraph->nodes[j]; + const bool view = ggml_cuda_is_view_or_noop(node); + if (view && node->view_src == mm) { + pending += ggml_node_get_use_count(cgraph, j); + f32 = f32 || (node->flags & GGML_TENSOR_FLAG_OUTPUT); + } + for (int s = 0; s < GGML_MAX_SRC; ++s) { + const ggml_tensor * src = node->src[s]; + if (!src || (src != mm && src->view_src != mm)) { + continue; + } + --pending; + if (view) { + continue; + } + const bool tile = s == 1 && (node->flags & GGML_TENSOR_FLAG_COMPUTE) && src->data == mm->data && + src->ne[0] == K && ggml_nelements(src) == ggml_nelements(mm) && + (!shape || ggml_are_same_shape(src, shape)) && + ggml_cuda_mul_mat_reads_fwht_q8(cuda_ctx->device, node); + if (!tile) { + f32 = true; + continue; + } + shape = src; + ++readers; + last = j; + size = std::max(size, ggml_cuda_mul_mat_q_q8_size(node->src[0], src)); + } + } + f32 = f32 || pending > 0; + + // the pool is stack-like: a tile must end before the one under it + if (readers == 0 || (!cuda_ctx->fwht_q8.empty() && last >= cuda_ctx->fwht_q8.back().last_node)) { + return {}; + } + auto q8 = std::make_unique>(cuda_ctx->pool(), size); + const fwht_q8_out out = { q8->get(), shape->ne[1], !f32 }; + cuda_ctx->fwht_q8.push_back({ mm->data, { shape->ne[0], shape->ne[1], shape->ne[2], shape->ne[3] }, std::move(q8), + last, readers, f32 }); + return out; +} + // try and fuse nodes and return the number of nodes to skip static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { @@ -3532,8 +3668,14 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph 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; + if (pattern_ok) { + const fwht_q8_out q8 = ggml_cuda_fwht_q8_plan(cuda_ctx, cgraph, k + 2, signs->ne[0]); + if (ggml_cuda_op_fwht_swiglu_signed(*cuda_ctx, node->src[0], node->src[1], signs, mm, q8)) { + return n_ops - 1; + } + if (q8.y) { + cuda_ctx->fwht_q8.pop_back(); + } } } } @@ -3558,8 +3700,14 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_is_contiguous(x) && ggml_is_contiguous(signs) && signs->ne[0] == x->ne[0] && signs->ne[0] % mm->src[0]->ne[0] == 0; - if (pattern_ok && ggml_cuda_op_fwht_signed(*cuda_ctx, x, signs, mm)) { - return 2; + if (pattern_ok) { + const fwht_q8_out q8 = ggml_cuda_fwht_q8_plan(cuda_ctx, cgraph, i + 2, signs->ne[0]); + if (ggml_cuda_op_fwht_signed(*cuda_ctx, x, signs, mm, q8)) { + return 2; + } + if (q8.y) { + cuda_ctx->fwht_q8.pop_back(); + } } } @@ -4530,6 +4678,8 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud } } + GGML_ASSERT(cuda_ctx->fwht_q8.empty()); + #ifdef USE_CUDA_GRAPH ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 2e610278..c93cacf5 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -4751,6 +4751,65 @@ struct test_fwht_signed : public test_case { } }; +// sign flip [+ SWIGLU] + FWHT read by two PTQ1_0 matmuls: through MMQ the transform writes their quantized src1 +// itself; f32_reader adds an F32 matmul on the same activation, so the transform's f32 output is needed too +struct test_fwht_q8_mul_mat : public test_fwht_signed { + const int64_t m; + const bool swiglu; + const bool f32_reader; + + test_fwht_q8_mul_mat(int64_t blk, int64_t width, int64_t n_tokens, int64_t m, bool swiglu, bool f32_reader) + : test_fwht_signed(blk, width, n_tokens), m(m), swiglu(swiglu), f32_reader(f32_reader) {} + + std::string vars() override { + return VARS_TO_STR6(blk, width, n_tokens, m, swiglu, f32_reader) + ",type_w=ptq1_0"; + } + + std::string op_desc(ggml_tensor * t) override { + GGML_UNUSED(t); + return "MUL_MAT"; + } + + double max_nmse_err() override { + return 5e-4; + } + + ggml_tensor * build_graph(ggml_context * ctx) override { + ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, blk, blk); + ggml_set_name(a, "a"); + ggml_tensor * x = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, n_tokens); + if (swiglu) { + x = ggml_swiglu_split(ctx, x, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, n_tokens)); + } + ggml_tensor * s = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, width); + ggml_set_name(s, "s"); + + ggml_tensor * cur = ggml_mul(ctx, x, s); + cur = ggml_reshape_2d(ctx, cur, blk, width / blk * n_tokens); + cur = ggml_mul_mat(ctx, a, cur); + ggml_mul_mat_set_hint(cur, GGML_HINT_SRC0_IS_HADAMARD); + cur = ggml_reshape_2d(ctx, cur, width, n_tokens); + + ggml_tensor * w0 = ggml_new_tensor_2d(ctx, GGML_TYPE_PTQ1_0, width, m); + ggml_tensor * w1 = ggml_new_tensor_2d(ctx, GGML_TYPE_PTQ1_0, width, m); + ggml_tensor * out = ggml_add(ctx, ggml_mul_mat(ctx, w0, cur), ggml_mul_mat(ctx, w1, cur)); + if (f32_reader) { + out = ggml_add(ctx, out, ggml_mul_mat(ctx, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, width, m), cur)); + } + ggml_set_name(out, "out"); + return out; + } + + void initialize_tensors(ggml_context * ctx) override { + test_fwht_signed::initialize_tensors(ctx); + for (ggml_tensor * t = ggml_get_first_tensor(ctx); t != NULL; t = ggml_get_next_tensor(ctx, t)) { + if (t->type == GGML_TYPE_PTQ1_0) { + init_tensor_uniform(t); + } + } + } +}; + static void init_mul_mat_id_tensors(ggml_context * ctx, int n_mats) { std::random_device rd; std::default_random_engine rng(rd()); @@ -9311,6 +9370,14 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_fwht_signed(1024, 5120, 32)); test_cases.emplace_back(new test_fwht_signed(1024, 6144, 7, GGML_TYPE_F16)); test_cases.emplace_back(new test_fwht_signed(1024, 17408, 3)); + for (const int64_t n_tokens : { 3, 7, 64, 512 }) { + for (const bool f32_reader : { false, true }) { + test_cases.emplace_back(new test_fwht_q8_mul_mat(1024, 5120, n_tokens, 256, false, f32_reader)); + test_cases.emplace_back(new test_fwht_q8_mul_mat(1024, 17408, n_tokens, 128, true, f32_reader)); + } + test_cases.emplace_back(new test_fwht_q8_mul_mat(512, 2048, n_tokens, 128, false, false)); + test_cases.emplace_back(new test_fwht_q8_mul_mat(8192, 8192, n_tokens, 128, true, false)); + } // Block widths above the register path's reach, plus a couple below it as controls. 4096 and // 8192 exercise the shared-memory kernel; before it existed the CUDA backend declined them and // the op fell back, which cost both speed and (measurably) a little accuracy.